[https://nvbugs/6223556][fix] Propagate gen-first ctx usage via aux buffer to postproc - #15246
Conversation
📝 WalkthroughWalkthroughThis PR adds context-worker usage propagation for gen-first disaggregated scheduling. A new ChangesGen-first disaggregated context-usage propagation
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
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
🤖 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/llm_request.py`:
- Around line 569-574: The LlmResult.__init__ currently sets self.ctx_usage =
None without a type annotation which can be inferred as NoneType; update
LlmResult.__init__ to declare the attribute with the proper optional dict type,
e.g. self.ctx_usage: Optional[Dict[str, Any]] = None, and ensure you import
Optional and Dict/Any in the module if not already present so type checkers
recognize the intended type used elsewhere (ctx_usage.get(...)).
🪄 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: e5ab7d4b-f062-44be-8a06-69ac992fa7e2
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/executor/result.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yamltests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yamltests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
/bot run --disable-fail-fast |
|
PR_Github #53496 [ run ] triggered by Bot. Commit: |
|
PR_Github #53496 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53584 [ run ] triggered by Bot. Commit: |
|
PR_Github #53584 [ run ] completed with state
|
a175926 to
bb645c2
Compare
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #53737 [ run ] triggered by Bot. Commit: |
|
PR_Github #53739 [ run ] triggered by Bot. Commit: |
|
PR_Github #53737 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53739 [ run ] completed with state
|
|
PR_Github #53852 [ run ] triggered by Bot. Commit: |
|
PR_Github #53852 [ run ] completed with state
|
e615062 to
1bca389
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54205 [ run ] triggered by Bot. Commit: |
1bca389 to
6b6953d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54271 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #54299 [ run ] triggered by Bot. Commit: |
|
PR_Github #54271 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #54396 [ run ] triggered by Bot. Commit: |
|
PR_Github #54299 [ run ] completed with state |
…uffer to postproc In gen-first disaggregated scheduling the context and generation requests run concurrently, so the context worker's usage cannot be injected into the request as in the context-first path. It is instead delivered to the generation worker through the KV-transfer aux buffer (RxSession.unpack_aux), which sets py_disaggregated_params.ctx_usage. That value was never surfaced onto the response, so the postprocessor fell back to the generation worker's local accounting, which treats the entire transferred prompt as cached and over-reports cached_tokens (e.g. the whole prompt length 17 instead of the reused prompt_tokens-1 = 16). Carry ctx_usage on LlmResult, set it on the response alongside cached_tokens in the executor, and surface it onto the output's disaggregated_params so the existing postprocessor path (_ctx_usage_from_outputs -> rewrite_usage_info_from_ctx) adopts the context-side accounting. This keeps the orchestrator a thin proxy and fixes both streaming and non-streaming. Enable block/partial reuse in the overlap_gen_first test configs so the cache-reuse usage check is meaningful, and remove the related waives. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…ttribute access Both py_disaggregated_params and ctx_usage are unconditionally initialized in their respective __init__ methods, so getattr fallbacks are unnecessary. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
response_result in _handle_response can be the raw C++ bindings.executor.Result (non-disagg / benchmark path), which has no ctx_usage attribute -- only the Python LlmResult wrapper sets it. The previous direct attribute access therefore raised AttributeError on the non-disagg path (e.g. trtllm-bench). Fall back to None via getattr, the same pattern the adjacent cached_tokens line already uses. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
6b6953d to
aa306ca
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54464 [ run ] triggered by Bot. Commit: |
|
PR_Github #54396 [ run ] completed with state |
|
PR_Github #54464 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54580 [ run ] triggered by Bot. Commit: |
|
PR_Github #54580 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54603 [ run ] triggered by Bot. Commit: |
|
PR_Github #54603 [ run ] completed with state |
…uffer to postproc (NVIDIA#15246) Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
…uffer to postproc (NVIDIA#15246) Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
6 participants
Summary
Root cause: In gen-first disaggregated scheduling the context and generation requests run concurrently, so the context worker's usage accounting cannot be injected into the request as it is in the context-first path. It is instead delivered to the generation worker through the KV-transfer aux buffer (
RxSession.unpack_aux), which setspy_disaggregated_params.ctx_usage. However, that value was never surfaced onto the response, so the postprocessor fell back to the generation worker's local accounting — which treats the entire transferred prompt as cached and over-reportscached_tokens(the full prompt length 17 instead of the reusedprompt_tokens-1 = 16). This causedtest_disaggregated_overlap_gen_firstto fail its cache-reuse usage check.Fix: Complete the existing aux-buffer transmission pipeline (introduced in #14177) instead of rewriting usage at the orchestrator:
ctx_usageonLlmResult.cached_tokensinpy_executor(all three response-creation sites: generation, fast-transfer, and streaming emit).disaggregated_paramsinresult.pyso the existing postprocessor path (_ctx_usage_from_outputs->rewrite_usage_info_from_ctx) adopts the context-side accounting.This keeps the disaggregated orchestrator a thin proxy (consistent with #14177) and fixes both streaming and non-streaming gen-first usage, since the postprocessor already handles both uniformly.
The
overlap_gen_firsttest configs are updated to enable block/partial reuse so the cache-reuse usage check is meaningful, and the corresponding waives are removed.Test plan
test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0]passes (B200). Before fix: every request reportedcached_tokens == prompt_tokens(17/17). After fix: first (cold) request reports the context-side count (e.g.cached_tokens: 3), second (reuse) request reportsprompt_tokens - 1(16 for completions, 30 for chat).ctx_pp4variant — verified by CI (locally blocked by GPU availability; logic is identical and independent of context-side pipeline parallelism).Links
Summary by CodeRabbit
Bug Fixes
Improvements