[https://nvbugs/5961736][fix] Prebuild disagg ctx response to avoid ctx_request_id race#12466

peihu-nv · 2026-03-23T23:13:30Z

Summary by CodeRabbit

Bug Fixes
- Fixed duplicate termination behavior in request handling for context-only scenarios, preventing resource conflicts.
Performance
- Optimized response creation timing to reduce latency in asynchronous request processing.

Description

Fixes a race in disaggregated serving where the CTX response can reach the orchestrator with ctx_request_id=None, causing a 500 error.

Root cause: _check_disagg_ctx_cache_transfer_status(0) in _send_kv_async() can advance a ctx-only request from DISAGG_CONTEXT_TRANS_IN_PROGRESS to DISAGG_CONTEXT_COMPLETE before _handle_responses() runs. Once that happens, C++ createResult() returns nullopt for contextPhaseParams, so the CTX response no longer carries ctx_request_id.

Fix: call create_response() immediately after respond_and_send_async() while the request is still guaranteed to be in the transfer-in-progress state, store the result as _prebuilt_ctx_response, and consume it later in _handle_responses() without calling create_response() again.

Test Coverage

N/A

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
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.

…capture The ctx response (containing ctx_request_id and opaque KV state) was only accessible while the request was in DISAGG_CONTEXT_TRANS_IN_PROGRESS. _check_disagg_ctx_cache_transfer_status(0) could advance the state to DISAGG_CONTEXT_COMPLETE before _handle_responses ran, causing C++ createResult() to return nullopt and leaving ctx_request_id=None. Fix: eagerly call create_response() immediately after respond_and_send_async() in _send_kv_async, while TRANS_IN_PROGRESS is guaranteed. The response is latched on the request (_eager_ctx_response) and consumed directly by _handle_responses, which skips create_response() for those requests. Safety of calling releaseState() here: CacheFormatter::format() (the CTX-side KV sender) reads only mRequestId, mSamplingConfig, getContextProgress(), and KV block data — it does NOT read mContextPhaseParams. The opaque state is only needed by the GEN-side unformat(), which runs in a separate process. Lifecycle is preserved: _handle_responses removes the request from active_requests (does not re-add to new_active_requests) without calling _terminate_request, leaving termination to _check_disagg_ctx_cache_transfer_status in both the in-flight (Case A) and already-complete (Case B) cases. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

coderabbitai · 2026-03-23T23:17:04Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: baacd9fe-cd94-49aa-b43e-adcc68bb291f

📥 Commits

Reviewing files that changed from the base of the PR and between 90c1cb7 and 1ce2186.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/py_executor.py

📝 Walkthrough

Walkthrough

Modified response handling in the PyExecutor to capture and store context-only responses upfront during asynchronous KV transfer, then reuse the prebuilt response during response processing while adjusting termination logic to prevent duplicate request termination.

Changes

Cohort / File(s)	Summary
Response Lifecycle Management `tensorrt_llm/_torch/pyexecutor/py_executor.py`	Modified `_send_kv_async` to capture and cache the response object in `req._prebuilt_ctx_response` immediately after starting async KV transfer. Updated `_handle_responses` to conditionally use prebuilt responses when available, bypassing standard response creation and properly merging cached tokens. Adjusted termination logic to prevent re-adding requests that already consumed prebuilt responses, eliminating duplicate termination scenarios.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: prebuilding the disagg ctx response to fix a race condition involving ctx_request_id, matching the core fix described in the summary.
Description check	✅ Passed	The description includes a clear explanation of the issue, root cause, and solution, plus the completed PR checklist. However, 'Test Coverage' section states 'N/A' without justifying why no tests are provided.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

peihu-nv · 2026-03-24T00:06:00Z

/bot run

peihu-nv · 2026-03-24T16:03:04Z

/bot run

tensorrt-cicd · 2026-03-24T16:09:07Z

PR_Github #40146 [ run ] triggered by Bot. Commit: 1ce2186 Link to invocation

tensorrt-cicd · 2026-03-24T17:16:08Z

PR_Github #40146 [ run ] completed with state FAILURE. Commit: 1ce2186
/LLM/main/L0_MergeRequest_PR pipeline #31292 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

Link to invocation

peihu-nv · 2026-03-24T17:34:26Z

/bot run

tensorrt-cicd · 2026-03-24T17:40:51Z

PR_Github #40156 [ run ] triggered by Bot. Commit: fdb2bdb Link to invocation

tensorrt-cicd · 2026-03-24T20:03:04Z

PR_Github #40156 [ run ] completed with state SUCCESS. Commit: fdb2bdb
/LLM/main/L0_MergeRequest_PR pipeline #31302 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

Link to invocation

peihu-nv · 2026-03-24T20:37:44Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-03-24T20:43:41Z

PR_Github #40168 [ run ] triggered by Bot. Commit: fdb2bdb Link to invocation

tensorrt-cicd · 2026-03-25T06:17:29Z

PR_Github #40168 [ run ] completed with state SUCCESS. Commit: fdb2bdb
/LLM/main/L0_MergeRequest_PR pipeline #31313 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

Link to invocation

…eateResult fix When the KV transfer poll in _send_kv_async advanced request state to DISAGG_CONTEXT_COMPLETE before _handle_responses ran, createResult() returned nullopt (its disagg-ctx branch only fired for TRANS_IN_PROGRESS), causing ctx_request_id=None and the error "Invalid disaggregated params in context phase response." Fix: - isFinished(): extend to cover DISAGG_CONTEXT_COMPLETE so createResult() passes the early-return guard for that state. - createResult(): broaden the disagg-ctx branch condition from isDisaggContextTransmissionState() to also match isDisaggContextCompleteState(), so the context phase params are included in the response regardless of which state the race resolves to. mContextPhaseParams is still valid at this point since releaseState() is called by createResult() itself. - py_executor.py _handle_responses(): guard the requests_to_terminate logic to skip DISAGG_CONTEXT_COMPLETE requests (already terminated by _check_disagg_ctx_cache_transfer_status), preventing double-free. Track them in requests_finished_by_transfer for stats visibility. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

…lt with DISAGG_CONTEXT_COMPLETE Adds a unit test that constructs a context-only request in kDISAGG_CONTEXT_COMPLETE state and verifies createResult() returns a valid response with contextPhaseParams populated. Uses nullptr for the opaque state — sufficient to prove the fix (the bug was createResult() returning nullopt, not state pointer corruption). A follow-up could strengthen this with a non-null sentinel state and firstGenTokens check. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

peihu-nv · 2026-03-31T21:51:24Z

/bot run

tensorrt-cicd · 2026-03-31T21:58:01Z

PR_Github #41027 [ run ] triggered by Bot. Commit: a311df9 Link to invocation

tensorrt-cicd · 2026-04-01T00:09:37Z

PR_Github #41027 [ run ] completed with state FAILURE. Commit: a311df9
/LLM/main/L0_MergeRequest_PR pipeline #32007 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

Link to invocation

peihu-nv · 2026-04-01T16:56:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-01T17:01:57Z

PR_Github #41232 [ run ] triggered by Bot. Commit: 957218f Link to invocation

tensorrt-cicd · 2026-04-02T02:48:25Z

PR_Github #41232 [ run ] completed with state FAILURE. Commit: 957218f
/LLM/main/L0_MergeRequest_PR pipeline #32196 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

Link to invocation

peihu-nv · 2026-04-02T03:24:51Z

/bot run

peihu-nv · 2026-04-02T04:26:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-02T04:32:15Z

PR_Github #41336 [ run ] triggered by Bot. Commit: 957218f Link to invocation

tensorrt-cicd · 2026-04-02T04:36:54Z

PR_Github #41336 [ run ] completed with state FAILURE. Commit: 957218f

Link to invocation

pcastonguay · 2026-04-07T17:32:57Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-07T17:38:47Z

PR_Github #42171 [ run ] triggered by Bot. Commit: 957218f Link to invocation

tensorrt-cicd · 2026-04-07T21:02:02Z

PR_Github #42171 [ run ] completed with state SUCCESS. Commit: 957218f
/LLM/main/L0_MergeRequest_PR pipeline #32998 completed with status: 'SUCCESS'

CI Report

Link to invocation

…tx_request_id race (NVIDIA#12466) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

peihu-nv requested a review from a team as a code owner March 23, 2026 23:13

peihu-nv requested a review from achartier March 23, 2026 23:13

github-actions Bot assigned peihu-nv Mar 23, 2026

Merge branch 'main' into fix/nvbug-5961736-disagg-ctx-eager-response

fdb2bdb

pcastonguay reviewed Mar 25, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

pcastonguay reviewed Mar 25, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

peihu-nv added 2 commits March 30, 2026 12:19

Merge branch 'main' into fix/nvbug-5961736-disagg-ctx-eager-response

957218f

pcastonguay approved these changes Apr 7, 2026

View reviewed changes

pcastonguay merged commit 3bd37ae into NVIDIA:main Apr 8, 2026
5 checks passed

suyoggupta pushed a commit to nv-auto-deploy/TensorRT-LLM that referenced this pull request Apr 8, 2026

[https://nvbugs/5961736][fix] Prebuild disagg ctx response to avoid c…

cedaf8d

…tx_request_id race (NVIDIA#12466) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>

peihu-nv mentioned this pull request Jul 16, 2026

[https://nvbugs/6120535][chore] Unwaive DeepSeek V3.2 disaggregated serving test #16515

Merged

1 task

Uh oh!

Conversation

peihu-nv commented Mar 23, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Mar 23, 2026

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

peihu-nv commented Mar 24, 2026

Uh oh!

peihu-nv commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

peihu-nv commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

peihu-nv commented Mar 24, 2026

Uh oh!

tensorrt-cicd commented Mar 24, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Mar 25, 2026

Uh oh!

Uh oh!

Uh oh!

peihu-nv commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Mar 31, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

peihu-nv commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 1, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

peihu-nv commented Apr 2, 2026

Uh oh!

peihu-nv commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

tensorrt-cicd commented Apr 2, 2026

Uh oh!

pcastonguay commented Apr 7, 2026

Uh oh!

tensorrt-cicd commented Apr 7, 2026

Uh oh!

tensorrt-cicd commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

peihu-nv commented Mar 23, 2026 •

edited by coderabbitai Bot

Loading