[https://nvbugs/6482606][fix] Enforce output length at disaggregated generation first-token handoff#16691
[https://nvbugs/6482606][fix] Enforce output length at disaggregated generation first-token handoff#16691xwang233 wants to merge 4 commits into
Conversation
…generation first-token handoff In disaggregated serving, the context server produces the first output token and hands it to the generation server, which injects it and then schedules a generation step. The injection path in _prepare_disagg_gen_transmission_complete appended this first token without applying any stop criteria, so a request whose first token already satisfies a stop condition still ran one more generation step and emitted an extra token. The visible symptom is a request with max_tokens=1 (OSL=1) receiving 2 output tokens. Aggregated serving does not hit this because the sampler applies the stop check to the context step's first token and finishes the request before any generation step is scheduled. The gap only manifests when the first token alone satisfies a stop condition (max_tokens=1, or the first token is the end id / completes a stop word); for max_tokens > 1 the per-token stop check already stops at exactly the requested length. The path is shared by disaggregated generation with and without speculative decoding. Apply the same stop criteria the aggregated sampler applies (TorchSampler._handle_stop_criteria: end id, max-token/length, stop words) to the injected first token, before the generation step runs, and mark the request finished (per beam; complete only when all beams stop) so the sampler skips it. Also skip the redundant first-token response for a request already finished at injection, mirroring _emit_first_token_responses; otherwise it would emit two final responses and double-count the token. Add tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py covering the stop-criteria cases and the finished-request response skip. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughDisaggregated generation now evaluates stop criteria on the context-side first token, completes requests only after all beams finish, and skips duplicate first-token responses for requests already completed during injection. Unit tests cover length, end-token, stop-word, beam, and response-handling cases. ChangesDisaggregated first-token completion
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ContextInjection
participant PyExecutor
participant TorchSampler
participant ResponseQueue
ContextInjection->>PyExecutor: inject first token
PyExecutor->>TorchSampler: evaluate stop criteria per beam
TorchSampler-->>PyExecutor: return finish reasons
PyExecutor->>ResponseQueue: enqueue response only when request remains unfinished
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
5873-5874: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new Python functions.
tensorrt_llm/_torch/pyexecutor/py_executor.py#L5873-L5874: add precise types forreqandfirst_gen_tokens.tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py#L39-L186: annotate helper arguments/returns and test functions with-> None.🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5873 - 5874, Annotate _apply_disagg_gen_first_token_stop_criteria in tensorrt_llm/_torch/pyexecutor/py_executor.py with precise types for req and first_gen_tokens, preserving the existing beam_width annotation. In tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py, annotate all helper parameters and return types, and add -> None to each test function; apply no other changes.Source: Coding guidelines
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5902-5907: The beam completion logic in py_executor.py must record
each non-None finish reason immediately, while transitioning to
GENERATION_COMPLETE only when all beams have finished; update the loop around
beam_finish_reasons accordingly. In
tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py lines 141-153, update
the assertion to verify beam 0 receives FinishReason.END_ID while the request
remains in progress.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5873-5874: Annotate _apply_disagg_gen_first_token_stop_criteria in
tensorrt_llm/_torch/pyexecutor/py_executor.py with precise types for req and
first_gen_tokens, preserving the existing beam_width annotation. In
tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py, annotate all helper
parameters and return types, and add -> None to each test function; apply no
other changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c01fb033-750b-4585-9bec-f7898657dbfe
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_disagg_gen_osl_stop.py
| # Complete the request only when every beam has stopped, matching the agg | ||
| # semantics in TorchSampler._handle_finish_reasons_impl. | ||
| if all(reason is not None for reason in beam_finish_reasons): | ||
| req.state = LlmRequestState.GENERATION_COMPLETE | ||
| for beam, reason in enumerate(beam_finish_reasons): | ||
| req.set_finished_reason(reason, beam) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Record finish reasons for each stopped beam.
A beam that hits END_ID or a stop word must be marked finished even while another beam continues. Deferring set_finished_reason() until all beams stop allows the stopped beam to keep decoding; the test currently codifies that incorrect behavior.
tensorrt_llm/_torch/pyexecutor/py_executor.py#L5902-L5907: callset_finished_reason(reason, beam)for every non-Nonereason, and only transition the request state toGENERATION_COMPLETEafter all beams finish.tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py#L141-L153: assert that beam 0 receivesFinishReason.END_IDwhile the request remains in progress.
📍 Affects 2 files
tensorrt_llm/_torch/pyexecutor/py_executor.py#L5902-L5907(this comment)tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py#L141-L153
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5902 - 5907, The
beam completion logic in py_executor.py must record each non-None finish reason
immediately, while transitioning to GENERATION_COMPLETE only when all beams have
finished; update the loop around beam_finish_reasons accordingly. In
tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py lines 141-153, update
the assertion to verify beam 0 receives FinishReason.END_ID while the request
remains in progress.
|
PR_Github #60807 [ run ] triggered by Bot. Commit: |
|
PR_Github #60807 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61011 [ run ] triggered by Bot. Commit: |
|
PR_Github #61011 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61025 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #61032 [ run ] triggered by Bot. Commit: |
|
PR_Github #61025 [ run ] completed with state |
|
PR_Github #61032 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61060 [ run ] triggered by Bot. Commit: |
|
PR_Github #61060 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61183 [ run ] triggered by Bot. Commit: |
|
PR_Github #61183 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #61366 [ run ] triggered by Bot. Commit: |
|
PR_Github #61366 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #61441 [ run ] triggered by Bot. Commit: |
|
PR_Github #61441 [ run ] completed with state
|
2 participants
Dev Engineer Review
GENERATION_COMPLETEwhen all beams satisfy stop conditions, preventing premature completion ordering issues and avoiding extra tokens (e.g.,max_tokens=1)._handle_first_token_responseto skip emitting a first-token response for disagg-gen requests that were already finished at injection time, preventing a redundant first-token/final-token double-count (nvbugs/6482606).TorchSamplerimport for stop-criteria evaluation; no public API/config changes.QA Engineer Review
Added tests in
tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py:test_first_token_completes_request_when_max_tokens_is_onetest_first_token_does_not_complete_when_budget_remainstest_first_token_end_id_completes_requesttest_first_token_stop_word_completes_requesttest_none_max_seq_len_still_enforces_output_lengthtest_all_beams_complete_together_for_lengthtest_request_incomplete_until_all_beams_stoptest_handle_first_token_response_skips_finished_requestCoverage in
tests/integration/test_lists/(test-db/qa):Verdict: needs follow-up — CBTS coverage/list status is unavailable.
Description
In disaggregated serving, the context server produces the first output token and hands it to the generation server, which injects it and then schedules a generation step. The injection path in
PyExecutor._prepare_disagg_gen_transmission_completeappended this first token without applying any stop criteria, so a request whose first token already satisfies a stop condition still ran one more generation step and emitted an extra token.The visible symptom is a request with
max_tokens=1(OSL=1) receiving 2 output tokens. Aggregated serving does not have this problem because the sampler applies the stop check to the context step's first token and finishes the request before any generation step is scheduled.The gap only manifests when the first token alone satisfies a stop condition (
max_tokens=1, or the first token is the end id / completes a stop word); formax_tokens > 1the per-token stop check already stops generation at exactly the requested length. The affected path is shared by disaggregated generation with and without speculative decoding.Fix
TorchSampler._handle_stop_criteria: end id, max-token / length, stop words) to the injected first token, before the generation step runs, and finish the request (per beam; complete only when all beams have stopped) so the sampler, which skips finished requests, does not append an extra token._emit_first_token_responses; otherwise a finished-at-injection request would emit two final responses and double-count the token.Test Coverage
tests/unittest/_torch/executor/test_disagg_gen_osl_stop.py(GPU-free unit tests):max_tokens=1finishes at the injected first token (FinishReason.LENGTH);max_tokens=2does not.max_seq_len=Nonestill enforces the per-request output-length budget._handle_first_token_responseskips a request already finished at injection.PR Checklist