[None][fix] Fix disaggregated draft token accounting by SimengLiu-nv · Pull Request #16805 · NVIDIA/TensorRT-LLM

SimengLiu-nv · 2026-07-23T18:04:27Z

Generation-only requests created from disaggregated context-phase handoff can carry draft tokens alongside first-generation tokens. Adopt those draft tokens into LlmRequest so decoder capacity and Python-visible request state account for every generation token handed off by prefill.

Add C++ and Python unit coverage for constructor-time and late context-phase parameter handoff paths.

Tests added:

LlmRequestTest.generationOnlyRequestAdoptsContextPhaseDraftTokens
tests/unittest/_torch/executor/test_request_utils.py::test_executor_request_to_llm_request_adopts_context_phase_draft_tokens
tests/unittest/bindings/test_bindings_ut.py::test_generation_only_llm_request_adopts_draft_tokens

Dev Engineer Review

Updated GenericLlmRequest to adopt context-phase draft tokens during construction and when context-phase parameters are set later.
Added getNumContextPhaseGenerationTokens() and integrated it into decoder sequence-length accounting.
The implementation and API changes are consistent with the stated draft-token accounting fix. No configuration or test-list changes were identified.
C++ and Python unit tests cover both constructor-time and late context-phase handoff paths.

QA Engineer Review

Added generationOnlyRequestAdoptsContextPhaseDraftTokens in cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp.
Added test_generation_only_llm_request_adopts_draft_tokens in tests/unittest/bindings/test_bindings_ut.py.
Added coverage in tests/unittest/_torch/executor/test_request_utils.py for executor_request_to_llm_request.
No corresponding tests/integration/test_lists/, test-db/, or qa/ entries were found for these unit tests, so CI/manual-QA listing coverage is not explicitly recorded.
Verdict: needs follow-up

Description

Test Coverage

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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.

Generation-only requests created from disaggregated context-phase handoff can carry draft tokens alongside first-generation tokens. Adopt those draft tokens into LlmRequest so decoder capacity and Python-visible request state account for every generation token handed off by prefill. Add C++ and Python unit coverage for constructor-time and late context-phase parameter handoff paths. Tests added: - LlmRequestTest.generationOnlyRequestAdoptsContextPhaseDraftTokens - tests/unittest/_torch/executor/test_request_utils.py::test_executor_request_to_llm_request_adopts_context_phase_draft_tokens - tests/unittest/bindings/test_bindings_ut.py::test_generation_only_llm_request_adopts_draft_tokens Signed-off-by: Simeng Liu <simengl@nvidia.com>

coderabbitai · 2026-07-23T18:11:28Z

Walkthrough

Generic LLM requests now adopt draft tokens from context-phase parameters during construction or later updates, expose combined context-phase generation-token counts, and use that count in decoder sequence-length calculations. C++ and Python tests cover these paths.

Changes

Context-phase draft tokens

Layer / File(s)	Summary
Request token adoption `cpp/include/tensorrt_llm/batch_manager/llmRequest.h`	`GenericLlmRequest` constructors and `setContextPhaseParams()` adopt non-empty context-phase draft tokens, while `getNumContextPhaseGenerationTokens()` reports first-generation and draft-token totals.
Decoder sequence-length integration `cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp`	Sequence-length calculation uses `getNumContextPhaseGenerationTokens()` directly.
Draft-token adoption validation `cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp`, `tests/unittest/_torch/executor/test_request_utils.py`, `tests/unittest/bindings/test_bindings_ut.py`	Tests cover constructor-time adoption, late context-parameter assignment, and executor-to-LLM request propagation of draft tokens.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ExecutorRequest
  participant GenericLlmRequest
  participant ContextPhaseParams
  participant DecoderRequestCreation
  ExecutorRequest->>GenericLlmRequest: construct with context-phase parameters
  GenericLlmRequest->>ContextPhaseParams: read draft tokens
  ContextPhaseParams-->>GenericLlmRequest: return draft tokens
  GenericLlmRequest->>GenericLlmRequest: adopt draft tokens
  DecoderRequestCreation->>GenericLlmRequest: get context-phase generation-token count
  GenericLlmRequest-->>DecoderRequestCreation: return first-generation plus draft-token count

Suggested reviewers: lfr-0531

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required format and clearly summarizes the main fix.
Description check	✅ Passed	The description explains the issue, solution, tests, and checklist items, though it is not formatted exactly as the template.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

🧹 Nitpick comments (1)

cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp (1)
793-826: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Exercise decoder sequence-length initialization.

This validates adoption and counting, but never runs decoder setup. A regression in decoder sequence-length accounting can still pass. Add coverage asserting prompt + generated + first-gen + draft is written to decoder sequence lengths.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp` around lines 793 -
826, Extend generationOnlyRequestAdoptsContextPhaseDraftTokens to exercise
decoder setup after adopting the context-phase tokens, then assert the decoder
sequence length equals prompt tokens plus generated tokens plus first-generation
tokens plus draft tokens. Reuse the existing token vectors and request symbols,
and verify the recorded decoder sequence lengths through the relevant
LlmRequest/decoder setup API.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp`:
- Around line 793-826: Extend generationOnlyRequestAdoptsContextPhaseDraftTokens
to exercise decoder setup after adopting the context-phase tokens, then assert
the decoder sequence length equals prompt tokens plus generated tokens plus
first-generation tokens plus draft tokens. Reuse the existing token vectors and
request symbols, and verify the recorded decoder sequence lengths through the
relevant LlmRequest/decoder setup API.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: adca7be2-268b-402f-80ef-be390c225a2c

📥 Commits

Reviewing files that changed from the base of the PR and between 5a69240 and ed1631c.

📒 Files selected for processing (5)

cpp/include/tensorrt_llm/batch_manager/llmRequest.h
cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp
cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp
tests/unittest/_torch/executor/test_request_utils.py
tests/unittest/bindings/test_bindings_ut.py

SimengLiu-nv · 2026-07-23T18:14:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T18:21:18Z

PR_Github #61360 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

SimengLiu-nv · 2026-07-23T19:33:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T19:43:32Z

PR_Github #61379 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

tensorrt-cicd · 2026-07-23T19:45:30Z

PR_Github #61360 [ run ] completed with state ABORTED. Commit: ed1631c

Link to invocation

tensorrt-cicd · 2026-07-23T23:59:25Z

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

zhaoyangwang-nvidia

LGTM

SimengLiu-nv · 2026-07-24T02:19:28Z

/bot run --disable-fail-fast

schetlur-nv

Approving for pending groups.

tensorrt-cicd · 2026-07-24T02:25:35Z

PR_Github #61466 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

tensorrt-cicd · 2026-07-24T09:28:30Z

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

SimengLiu-nv · 2026-07-24T15:37:44Z

/bot run

tensorrt-cicd · 2026-07-24T15:43:34Z

PR_Github #61592 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

tensorrt-cicd · 2026-07-24T20:55:07Z

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

SimengLiu-nv · 2026-07-24T22:00:12Z

The previous CI failed at DGX_H100-4_GPUs-PyTorch-GptOss-1 / DGX_H100-4_GPUs-PyTorch-GptOss-1.accuracy.test_llm_api_pytorch.TestGPTOSS.test_w4_4gpus[v2_kv_cache-dp4-triton-auto] which I cannot reproduce on computelab machines.

============================================================================================================= slowest durations ==============================================================================================================
889.26s call     accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-triton-auto]
1.52s setup    accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-triton-auto]
0.56s teardown accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-triton-auto]
================================================================================================= 1 passed, 3 warnings in 901.32s (0:15:01) ==================================================================================================

SimengLiu-nv · 2026-07-24T22:00:23Z

/bot run

tensorrt-cicd · 2026-07-24T22:07:26Z

PR_Github #61652 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

tensorrt-cicd · 2026-07-24T23:04:43Z

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

SimengLiu-nv · 2026-07-24T23:07:05Z

/bot run

tensorrt-cicd · 2026-07-24T23:12:25Z

PR_Github #61656 [ run ] triggered by Bot. Commit: ed1631c Link to invocation

tensorrt-cicd · 2026-07-25T00:03:32Z

PR_Github #61656 [ run ] completed with state SUCCESS. Commit: ed1631c
/LLM/main/L0_MergeRequest_PR pipeline #49862 completed with status: 'SUCCESS'

CI Report

Link to invocation

SimengLiu-nv requested review from a team as code owners July 23, 2026 18:04

SimengLiu-nv requested review from JunyiXu-nv, QiJune, brnguyen2 and zhaoyangwang-nvidia July 23, 2026 18:04

github-actions Bot assigned SimengLiu-nv Jul 23, 2026

brnguyen2 removed their request for review July 23, 2026 18:06

schetlur-nv requested review from athena-nv, pcastonguay and zheyuf July 23, 2026 18:10

coderabbitai Bot reviewed Jul 23, 2026

View reviewed changes

SimengLiu-nv mentioned this pull request Jul 23, 2026

[None][fix] Align GPTOSS router tokenization and disagg draft scheduling #15605

Closed

1 task

zheyuf approved these changes Jul 23, 2026

View reviewed changes

pcastonguay approved these changes Jul 23, 2026

View reviewed changes

zhaoyangwang-nvidia approved these changes Jul 24, 2026

View reviewed changes

schetlur-nv approved these changes Jul 24, 2026

View reviewed changes

JunyiXu-nv approved these changes Jul 24, 2026

View reviewed changes

athena-nv approved these changes Jul 24, 2026

View reviewed changes

Uh oh!

Conversation

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

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 23, 2026

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

SimengLiu-nv commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

SimengLiu-nv commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

zhaoyangwang-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

SimengLiu-nv commented Jul 24, 2026

Uh oh!

schetlur-nv left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

SimengLiu-nv commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

SimengLiu-nv commented Jul 24, 2026

Uh oh!

SimengLiu-nv commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

SimengLiu-nv commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

SimengLiu-nv commented Jul 23, 2026 •

edited by coderabbitai Bot

Loading