[https://nvbugs/6683840][fix] Join the encoder launch before the decoder forward

[pranav-nvidia]

Dev Engineer Review

• _submit_encoder_step now joins the encoder launch future for tp_size == 1.
• The completed encoder result is stored in PendingEncoderStep.
• CUDA event synchronization remains deferred to _poll_encoder_steps.
• Encoder GPU work can still overlap decoder work.
• The tp_size > 1 path remains unchanged.
• The change prevents concurrent torch.fx tracing and torch.compile activity across threads.
• No public API or configuration changes were made.
• The updated lifecycle test matches the new future-joining behavior.
• No correctness or regression issues were identified from the reviewed changes.

QA Engineer Review

• Modified test: asynchronous encoder lifecycle test in tests/unittest/_torch/executor/test_py_executor.py.
• The test verifies that the launch future resolves during _submit_encoder_step and readiness remains deferred to _poll_encoder_steps.
• No matching entry was identified in tests/integration/test_lists/.
• Verdict: needs follow-up.

Description

At tp_size == 1, _submit_encoder_step returned as soon as the encoder step was submitted, so the encoder forward ran on the encoder-launch worker while the executor thread ran the decoder forward. Both call LayerNorm.forward, which carries @maybe_compile(dynamic=True).

torch.fx's tracing state is process-global, not thread-local: Tracer.trace sets _is_fx_tracing_flag and patches torch.nn.Module.__call__ for the whole process, and an ordinary torch.compile reaches it through AOTAutograd's make_fx. A fresh compile on either thread therefore captured the other thread's module calls, and dynamo raised Detected that you are using FX to symbolically trace a dynamo-optimized function, surfaced as RequestError. When the executor thread was the victim the exception ended the event loop and failed every in-flight request.

This joins the encoder future before the decoder forward runs, so the two Python launches no longer overlap. ready_event is deliberately left unsynchronized and publication stays in _poll_encoder_steps, so encoder GPU work still overlaps decoder work. The tp_size > 1 path is unchanged.

Throughput is unchanged within noise on whisper-large-v3 and whisper-tiny (fp16, SM120, batch 1/8/32, encoder CUDA graphs on and off), except large-v3 / no encoder graphs / batch 1 at about 3%. The GIL already interleaved the two threads' Python, so only the GPU-stream overlap was load-bearing and that is preserved.

modeling_bart.py and modeling_parakeet.py share the exposure: the worker is created from is_encoder_decoder alone and both use the compiled LayerNorm. modeling_t5.py uses RMSNorm, which is not compiled.

For review: this changes the fire-and-forget contract added in #16706, which that PR pinned in test_async_encoder_step_lifecycle.

Test Coverage

tests/unittest/_torch/executor/test_py_executor.py::test_async_encoder_step_lifecycle — updated to the new contract: the launch is joined inline, while publication still waits on ready_event rather than on the future.

tests/unittest/_torch/executor/ — 1769 passed, 1 skipped, 3 xfailed, with the same 10 pre-existing test_overlap_scheduler.py errors as an unpatched control run on the same tree.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2c897d9e-c9bb-4ad4-ab3f-e1dcd4f5f461

📥 Commits

Reviewing files that changed from the base of the PR and between 30316da and 7f870dd.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/unittest/_torch/executor/test_py_executor.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The non-TP encoder submission path now resolves the launch future before decoder execution continues. It stores the completed encoder result for later polling while deferring CUDA readiness synchronization. Tests now verify this lifecycle.

Changes

Encoder lifecycle

Layer / File(s)	Summary
Submission and lifecycle validation tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_py_executor.py	_submit_encoder_step resolves the launch future, routes failures through encoder error handling, and stores the completed result in PendingEncoderStep. _poll_encoder_steps remains responsible for readiness checks. Tests verify inline future resolution and deferred polling.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 7f870

The encoder launch now completes before decoder execution to prevent concurrent tracing failures while preserving asynchronous device work. This localized change has targeted lifecycle coverage, and no actionable merge-blocking risk remains.

Suggested reviewers: juney-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required NVBugs and type format and clearly describes the main change: joining the encoder launch before decoder execution.
Description check	✅ Passed	The description clearly explains the failure mode, solution, preserved GPU overlap, affected paths, test coverage, and contract change. The required sections are present and the checklist is completed…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description clearly explains the failure mode, solution, preserved GPU overlap, affected paths, test coverage, and contract change. The required sections are present and the checklist is completed.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[pranav-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #71102 [ run ] triggered by Bot. Commit: 7f870dd Link to invocation

[cascade812]

LGTM, thanks for the fix!

[tensorrt-cicd]

PR_Github #71102 [ run ] completed with state FAILURE. Commit: 7f870dd
/LLM/main/L0_MergeRequest_PR pipeline #58248 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #71288 [ run ] triggered by Bot. Commit: 7f870dd Link to invocation

[tensorrt-cicd]

PR_Github #71288 [ run ] completed with state SUCCESS. Commit: 7f870dd
/LLM/main/L0_MergeRequest_PR pipeline #58418 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #71319 [ run ] triggered by Bot. Commit: 7f870dd Link to invocation

[tensorrt-cicd]

PR_Github #71319 [ run ] completed with state SUCCESS. Commit: 7f870dd
/LLM/main/L0_MergeRequest_PR pipeline #58445 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #71442 [ run ] triggered by Bot. Commit: 7f870dd Link to invocation

[tensorrt-cicd]

PR_Github #71442 [ run ] completed with state FAILURE. Commit: 7f870dd
/LLM/main/L0_MergeRequest_PR pipeline #58550 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #71560 [ run ] triggered by Bot. Commit: 7f870dd Link to invocation

[tensorrt-cicd]

PR_Github #71560 [ run ] completed with state SUCCESS. Commit: 7f870dd
/LLM/main/L0_MergeRequest_PR pipeline #58648 completed with status: 'SUCCESS'

CI Report

Link to invocation