[https://nvbugs/6156492][fix] Fix disaggregated usage propagation#14177
Conversation
📝 WalkthroughWalkthroughThis PR extends disaggregated serving to propagate exact token usage accounting from the context phase to generation phase. It introduces ChangesContext Response Usage Propagation Through Disaggregated Serving
🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
tensorrt_llm/disaggregated_params.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA copyright/license header to this modified source file.
This modified file is missing the required NVIDIA Apache-2.0 copyright header.
As per coding guidelines "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification, using the Apache 2.0 license text".
🤖 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/disaggregated_params.py` at line 1, This file is missing the required NVIDIA Apache-2.0 copyright header; add the standard NVIDIA Apache-2.0 header block (with the year of latest modification) at the very top of the source before the first import (the line "from dataclasses import dataclass"), ensuring it includes the copyright statement and the Apache-2.0 license text/SPDX identifier exactly as used across other TensorRT-LLM source files.tensorrt_llm/serve/postprocess_handlers.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA copyright/license header to this modified source file.
This modified source file lacks the required NVIDIA Apache-2.0 header.
As per coding guidelines "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification, using the Apache 2.0 license text".
🤖 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/serve/postprocess_handlers.py` at line 1, Add the required NVIDIA Apache-2.0 copyright header at the top of this modified source file (tensorrt_llm.serve.postprocess_handlers.py) so the file contains the standard NVIDIA header with the year of latest modification and the Apache-2.0 license text; place it above the first import (the line starting with "from dataclasses import dataclass, field") and ensure the header matches the project's canonical NVIDIA/Apache-2.0 header format exactly.tensorrt_llm/serve/openai_protocol.py (1)
1-2:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA copyright/license header to this modified source file.
This modified source file is missing the mandatory NVIDIA Apache-2.0 header.
As per coding guidelines "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification, using the Apache 2.0 license text".
🤖 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/serve/openai_protocol.py` around lines 1 - 2, Add the required NVIDIA Apache-2.0 copyright/license header to the top of openai_protocol.py: insert the official NVIDIA Apache-2.0 header block with the year of latest modification and the Apache-2.0 license notice above the existing file comments (before the "Adapted from" lines) so the file includes the mandatory NVIDIA copyright and license text.tensorrt_llm/_torch/disaggregation/native/transfer.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA copyright/license header to this modified source file.
This modified source file is missing the mandated NVIDIA Apache-2.0 header.
As per coding guidelines "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification, using the Apache 2.0 license text".
🤖 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/disaggregation/native/transfer.py` at line 1, This file (module tensorrt_llm._torch.disaggregation.native.transfer, file transfer.py) is missing the required NVIDIA Apache-2.0 copyright header; add the standard NVIDIA copyright/license block at the top of transfer.py containing the NVIDIA copyright line with the year of latest modification and the full Apache License, Version 2.0 notice (or the project's canonical header template), ensuring it appears before any imports (before "from __future__ import annotations").tensorrt_llm/_torch/disaggregation/native/auxiliary.py (1)
1-1:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd the required NVIDIA copyright/license header to this modified source file.
This file is modified but does not include the required NVIDIA Apache-2.0 header block.
As per coding guidelines "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification, using the Apache 2.0 license text".
🤖 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/disaggregation/native/auxiliary.py` at line 1, The file is missing the required NVIDIA Apache-2.0 copyright header; add the standard NVIDIA copyright/license header block (with the year of latest modification) at the very top of the file above the existing import(s) (i.e., before the line containing "from abc import ABC, abstractmethod"), ensuring it matches the project's required Apache-2.0 template and includes the NVIDIA copyright notice and license text.
🧹 Nitpick comments (4)
tests/unittest/disaggregated/test_openai_disagg_service.py (2)
303-392: QA list update not needed for this PR scope.These changes are unit-test-only under
tests/unittest/..., so QA integration test-list updates are unnecessary in this PR.As per coding guidelines "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."
🤖 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 `@tests/unittest/disaggregated/test_openai_disagg_service.py` around lines 303 - 392, This PR only modifies unit tests in tests/unittest/disaggregated/test_openai_disagg_service.py (e.g., test_generation_postprocessor_rewrites_usage_from_disaggregated_params and test_generation_chat_postprocessor_rewrites_usage_from_output_disaggregated_params), so update the PR description or checklist to state explicitly that QA integration list updates are unnecessary for this change; mention the file and test-only scope and that no QA list modifications are required.
310-329: ⚡ Quick winIsolate usage source in this completion postprocessor test.
On Line 313 and Line 326,
ctx_response_usageis populated in both request and output disaggregated params. That can mask regressions because either path can satisfy the same assertions. Keep one source per test (and add a separate test for the other source/precedence).♻️ Suggested adjustment to isolate this test path
request = CompletionRequest( model="test-model", prompt=[1, 2, 3], disaggregated_params=DisaggregatedParams( request_type="generation_only", - ctx_response_usage=ctx_usage, ), )As per coding guidelines "Coverage expectations: Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix."
🤖 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 `@tests/unittest/disaggregated/test_openai_disagg_service.py` around lines 310 - 329, The test currently sets ctx_response_usage in both the request DisaggregatedParams and the output LlmDisaggregatedParams, which hides regressions; modify the test to keep only one source of usage data (e.g., keep ctx_response_usage on the request side and remove it from the output LlmDisaggregatedParams) so the CompletionPostprocArgs.from_request / postprocessor behavior is exercised deterministically; update the SimpleNamespace outputs block to omit ctx_response_usage from LlmDisaggregatedParams (or alternatively remove it from the request and keep it in the output) and add a separate test to cover the other precedence path if needed.tests/unittest/disaggregated/region/test_aux.py (2)
90-118: ⚡ Quick winAdd a test for
ctx_response_usageprecedence in_get_ctx_usage.This test only exercises the fallback branch (
prompt_len/cached_tokens). Add a case wherepy_disaggregated_params.ctx_response_usageis present to lock precedence behavior.As per coding guidelines "Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix".
🤖 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 `@tests/unittest/disaggregated/region/test_aux.py` around lines 90 - 118, Add a new assertion path to the test_fill_slot_get_slot_tokens_round_trip that supplies py_disaggregated_params with a ctx_response_usage dict on the mock_request to ensure _get_ctx_usage prefers ctx_response_usage over prompt_len/cached_tokens; specifically, set mock_request.py_disaggregated_params = SimpleNamespace(ctx_response_usage={"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}) (or equivalent) before calling buf.fill_slot(slot.id, mock_request), then call buf.get_slot_data(slot.id) and assert the returned usage equals the provided ctx_response_usage structure (and that first/draft tokens still match) so the precedence behavior is tested for AuxBuffer.fill_slot / AuxBuffer._get_ctx_usage / AuxBuffer.get_slot_data.
1-129: QA test-list updates look unnecessary here since this PR touches unit-test scope and notests/integration/defs/scheduling surface.As per coding guidelines "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional".
🤖 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 `@tests/unittest/disaggregated/region/test_aux.py` around lines 1 - 129, Add an explicit note in the PR description saying QA test-list updates are unnecessary (or optional) because this change only touches unit-test scope — specifically the test module containing test_aux_buffer_meta_construction, test_aux_buffer_alloc_and_free_slot, test_fill_slot_get_slot_tokens_round_trip, and related tests — so QA scheduling changes under tests/integration are not required.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/disaggregation/native/auxiliary.py`:
- Line 1: The file is missing the required NVIDIA Apache-2.0 copyright header;
add the standard NVIDIA copyright/license header block (with the year of latest
modification) at the very top of the file above the existing import(s) (i.e.,
before the line containing "from abc import ABC, abstractmethod"), ensuring it
matches the project's required Apache-2.0 template and includes the NVIDIA
copyright notice and license text.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Line 1: This file (module tensorrt_llm._torch.disaggregation.native.transfer,
file transfer.py) is missing the required NVIDIA Apache-2.0 copyright header;
add the standard NVIDIA copyright/license block at the top of transfer.py
containing the NVIDIA copyright line with the year of latest modification and
the full Apache License, Version 2.0 notice (or the project's canonical header
template), ensuring it appears before any imports (before "from __future__
import annotations").
In `@tensorrt_llm/disaggregated_params.py`:
- Line 1: This file is missing the required NVIDIA Apache-2.0 copyright header;
add the standard NVIDIA Apache-2.0 header block (with the year of latest
modification) at the very top of the source before the first import (the line
"from dataclasses import dataclass"), ensuring it includes the copyright
statement and the Apache-2.0 license text/SPDX identifier exactly as used across
other TensorRT-LLM source files.
In `@tensorrt_llm/serve/openai_protocol.py`:
- Around line 1-2: Add the required NVIDIA Apache-2.0 copyright/license header
to the top of openai_protocol.py: insert the official NVIDIA Apache-2.0 header
block with the year of latest modification and the Apache-2.0 license notice
above the existing file comments (before the "Adapted from" lines) so the file
includes the mandatory NVIDIA copyright and license text.
In `@tensorrt_llm/serve/postprocess_handlers.py`:
- Line 1: Add the required NVIDIA Apache-2.0 copyright header at the top of this
modified source file (tensorrt_llm.serve.postprocess_handlers.py) so the file
contains the standard NVIDIA header with the year of latest modification and the
Apache-2.0 license text; place it above the first import (the line starting with
"from dataclasses import dataclass, field") and ensure the header matches the
project's canonical NVIDIA/Apache-2.0 header format exactly.
---
Nitpick comments:
In `@tests/unittest/disaggregated/region/test_aux.py`:
- Around line 90-118: Add a new assertion path to the
test_fill_slot_get_slot_tokens_round_trip that supplies py_disaggregated_params
with a ctx_response_usage dict on the mock_request to ensure _get_ctx_usage
prefers ctx_response_usage over prompt_len/cached_tokens; specifically, set
mock_request.py_disaggregated_params =
SimpleNamespace(ctx_response_usage={"prompt_tokens": 5, "completion_tokens": 2,
"total_tokens": 7}) (or equivalent) before calling buf.fill_slot(slot.id,
mock_request), then call buf.get_slot_data(slot.id) and assert the returned
usage equals the provided ctx_response_usage structure (and that first/draft
tokens still match) so the precedence behavior is tested for AuxBuffer.fill_slot
/ AuxBuffer._get_ctx_usage / AuxBuffer.get_slot_data.
- Around line 1-129: Add an explicit note in the PR description saying QA
test-list updates are unnecessary (or optional) because this change only touches
unit-test scope — specifically the test module containing
test_aux_buffer_meta_construction, test_aux_buffer_alloc_and_free_slot,
test_fill_slot_get_slot_tokens_round_trip, and related tests — so QA scheduling
changes under tests/integration are not required.
In `@tests/unittest/disaggregated/test_openai_disagg_service.py`:
- Around line 303-392: This PR only modifies unit tests in
tests/unittest/disaggregated/test_openai_disagg_service.py (e.g.,
test_generation_postprocessor_rewrites_usage_from_disaggregated_params and
test_generation_chat_postprocessor_rewrites_usage_from_output_disaggregated_params),
so update the PR description or checklist to state explicitly that QA
integration list updates are unnecessary for this change; mention the file and
test-only scope and that no QA list modifications are required.
- Around line 310-329: The test currently sets ctx_response_usage in both the
request DisaggregatedParams and the output LlmDisaggregatedParams, which hides
regressions; modify the test to keep only one source of usage data (e.g., keep
ctx_response_usage on the request side and remove it from the output
LlmDisaggregatedParams) so the CompletionPostprocArgs.from_request /
postprocessor behavior is exercised deterministically; update the
SimpleNamespace outputs block to omit ctx_response_usage from
LlmDisaggregatedParams (or alternatively remove it from the request and keep it
in the output) and add a separate test to cover the other precedence path if
needed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0b295800-b3d7-4c4e-bacf-f9dd5c9c21ad
📒 Files selected for processing (9)
tensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/disaggregated_params.pytensorrt_llm/serve/openai_disagg_service.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/postprocess_handlers.pytests/unittest/disaggregated/region/test_aux.pytests/unittest/disaggregated/test_disaggregated_params.pytests/unittest/disaggregated/test_openai_disagg_service.py
reasonsolo
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #48569 [ run ] triggered by Bot. Commit: |
|
PR_Github #48569 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48666 [ run ] triggered by Bot. Commit: |
|
PR_Github #48666 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48796 [ run ] triggered by Bot. Commit: |
|
PR_Github #48796 [ run ] completed with state |
|
It works great! |
|
PR_Github #49706 [ run ] triggered by Bot. Commit: |
|
PR_Github #49706 [ run ] completed with state
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #49790 [ run ] triggered by Bot. Commit: |
|
PR_Github #49790 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49883 [ run ] triggered by Bot. Commit: |
|
PR_Github #49883 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49903 [ run ] triggered by Bot. Commit: |
|
PR_Github #49903 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49915 [ run ] triggered by Bot. Commit: |
|
PR_Github #49915 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49932 [ run ] triggered by Bot. Commit: |
|
PR_Github #49932 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49940 [ run ] triggered by Bot. Commit: |
|
PR_Github #49940 [ run ] completed with state |
…IDIA#14177) Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…IDIA#14177) Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
8 participants
Summary by CodeRabbit
New Features
Tests
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.