[None][fix] Align GPTOSS router tokenization and disagg draft scheduling#15605
[None][fix] Align GPTOSS router tokenization and disagg draft scheduling#15605SimengLiu-nv wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughPyExecutor now synchronizes disaggregated draft-token fields before batch scheduling and uses the synchronized count in scheduled-token estimation. Router tokenization can select Harmony from ChangesDraft token synchronization
Harmony tokenization routing
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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)
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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/serve/router.py (1)
772-785: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winForward all worker-side chat-template inputs into router hashing.
The router now mirrors
tools, but the worker also renders withrequest.documentsandrequest.chat_template. Requests that differ only by RAG documents or a custom template can still produce identical router hashes while generating different worker prompts.Suggested fix
chat_template_kwargs = dict(request.chat_template_kwargs or {}) chat_template_kwargs["tools"] = self._tool_dicts(request) +chat_template_kwargs["documents"] = request.documents +if request.chat_template is not None: + chat_template_kwargs["chat_template"] = request.chat_template result = tokenizer.apply_chat_template(🤖 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/router.py` around lines 772 - 785, Forward the remaining worker-side chat-template inputs into the router hash computation in router.py’s tokenizer.apply_chat_template path: the current hash only mirrors tools via chat_template_kwargs, but it also needs request.documents and request.chat_template so router hashing matches the worker-rendered prompt. Update the chat_template_kwargs assembly in the same block that builds result to include these request fields before calling apply_chat_template, ensuring requests that differ by RAG documents or a custom template no longer collide.
🧹 Nitpick comments (3)
tests/unittest/disaggregated/test_router.py (2)
1118-1181: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the disabled-Harmony selection path.
The selection tests cover explicit
use_harmony=Trueandconfig.jsoninference, but not the server-compatible disabled path. Add amonkeypatch.setenv("DISABLE_HARMONY_ADAPTER", "1")case intests/unittest/disaggregated/test_router.pyasserting GPT-OSS requests use the tokenizer path and do not call Harmony. As per path instructions, coverage needs follow-up in this PR.🤖 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_router.py` around lines 1118 - 1181, Add a test case in test_use_harmony_flag_for_alias_model or a nearby test in KvCacheAwareRouter coverage that sets DISABLE_HARMONY_ADAPTER via monkeypatch before calling _tokenize; verify the disabled-Harmony path uses tokenizer.apply_chat_template, does not invoke harmony_adapter.get_harmony_adapter or openai_to_harmony_tokens, and leaves req.prompt_token_ids populated as expected for GPT-OSS/server-compatible requests. Use the existing test helpers and symbols like KvCacheAwareRouter, _tokenize, and ChatCompletionRequest to keep the new case aligned with the current Harmony selection tests.Source: Path instructions
1037-1117: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd parity cases for
documentsand customchat_template.Coverage is insufficient for the remaining worker-side template inputs. Extend
tests/unittest/disaggregated/test_router.pywith parametrized router/server parity cases whererequest.documentsandrequest.chat_templateare set, so cache-hash tokenization stays aligned with serving tokenization. As per path instructions, this names the concrete test file and coverage gap.🤖 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_router.py` around lines 1037 - 1117, Add router/server parity coverage for the remaining worker-side template inputs by extending the existing `test_gpt_oss_router_and_server_create_same_tokenized_input` in `test_router.py`. Parametrize new cases where `ChatCompletionRequest.documents` is populated and where a custom `chat_template` is set, then verify `KvCacheAwareRouter._tokenize` and `OpenAIServer.chat_harmony` produce identical tokenized inputs. Reuse the current `fake_harmony_tokens`, `harmony.openai_to_harmony_tokens`, and `generate_async` assertions so the new cases exercise the same cache-hash/tokenization path.Source: Path instructions
tests/unittest/_torch/executor/test_py_executor.py (1)
378-408: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage is insufficient for the executor integration paths.
tests/unittest/_torch/executor/test_py_executor.pynow covers the static helper directly, but should also cover: (1) trans-completedraft_tokens=Nonewithenable_spec_decode=True/max_total_draft_tokens > 0, and (2) the PP scheduling path that bypasses_prepare_and_schedule_batch. As per path instructions, tests undertests/**should be reviewed for actionable coverage sufficiency.🤖 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/_torch/executor/test_py_executor.py` around lines 378 - 408, The executor tests only cover the helper directly and miss key integration coverage. Add a test around trans-complete requests with draft_tokens=None when enable_spec_decode is true and max_total_draft_tokens is positive, and add a PP scheduling test that exercises the path bypassing _prepare_and_schedule_batch. Use the existing PyExecutor and _sync_disagg_generation_trans_complete_draft_tokens / _compute_scheduled_tokens helpers to locate the relevant flow, but assert behavior through the higher-level scheduling integration rather than only the static helper.Source: Path instructions
🤖 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 2617-2633: The draft-token sync in
_sync_disagg_generation_trans_complete_draft_tokens is using a different draft
count rule than _prepare_disagg_gen_transmission_complete, which can undercount
capacity when spec decode is enabled. Update the transmission-complete sync to
use the same draft-token resolver/logic as
_prepare_disagg_gen_transmission_complete, and move this sync so it runs after
the spec-decode decision is applied. Keep the behavior aligned for empty or None
context drafts and ensure py_draft_tokens, draft_tokens, and
py_draft_pages_allocated are derived from the shared resolution path.
- Around line 2651-2652: The PP executor path is missing the same disagg
generation transition sync that `_prepare_and_schedule_batch` already performs,
so `pp_size > 1` can schedule requests with stale draft-token state. Add the
`_sync_disagg_generation_trans_complete_draft_tokens` call in
`_executor_loop_pp` before it invokes `_pp_schedule_and_propagate()` or
`_schedule()`, mirroring the existing behavior in `_prepare_and_schedule_batch`
and using `self.active_requests` there as the source.
In `@tensorrt_llm/serve/router.py`:
- Around line 736-740: The Harmony tokenization decision in
`_uses_harmony_tokenization` is still defaulting GPT-OSS models to Harmony even
when the server has disabled it. Update `KvCacheAwareRouter` to use the same
resolved `use_harmony` value that `OpenAIServer` derives from
`DISABLE_HARMONY_ADAPTER`, or otherwise apply that env gate here, so router KV
hashing and worker chat template rendering stay aligned.
---
Outside diff comments:
In `@tensorrt_llm/serve/router.py`:
- Around line 772-785: Forward the remaining worker-side chat-template inputs
into the router hash computation in router.py’s tokenizer.apply_chat_template
path: the current hash only mirrors tools via chat_template_kwargs, but it also
needs request.documents and request.chat_template so router hashing matches the
worker-rendered prompt. Update the chat_template_kwargs assembly in the same
block that builds result to include these request fields before calling
apply_chat_template, ensuring requests that differ by RAG documents or a custom
template no longer collide.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 378-408: The executor tests only cover the helper directly and
miss key integration coverage. Add a test around trans-complete requests with
draft_tokens=None when enable_spec_decode is true and max_total_draft_tokens is
positive, and add a PP scheduling test that exercises the path bypassing
_prepare_and_schedule_batch. Use the existing PyExecutor and
_sync_disagg_generation_trans_complete_draft_tokens / _compute_scheduled_tokens
helpers to locate the relevant flow, but assert behavior through the
higher-level scheduling integration rather than only the static helper.
In `@tests/unittest/disaggregated/test_router.py`:
- Around line 1118-1181: Add a test case in
test_use_harmony_flag_for_alias_model or a nearby test in KvCacheAwareRouter
coverage that sets DISABLE_HARMONY_ADAPTER via monkeypatch before calling
_tokenize; verify the disabled-Harmony path uses tokenizer.apply_chat_template,
does not invoke harmony_adapter.get_harmony_adapter or openai_to_harmony_tokens,
and leaves req.prompt_token_ids populated as expected for
GPT-OSS/server-compatible requests. Use the existing test helpers and symbols
like KvCacheAwareRouter, _tokenize, and ChatCompletionRequest to keep the new
case aligned with the current Harmony selection tests.
- Around line 1037-1117: Add router/server parity coverage for the remaining
worker-side template inputs by extending the existing
`test_gpt_oss_router_and_server_create_same_tokenized_input` in
`test_router.py`. Parametrize new cases where `ChatCompletionRequest.documents`
is populated and where a custom `chat_template` is set, then verify
`KvCacheAwareRouter._tokenize` and `OpenAIServer.chat_harmony` produce identical
tokenized inputs. Reuse the current `fake_harmony_tokens`,
`harmony.openai_to_harmony_tokens`, and `generate_async` assertions so the new
cases exercise the same cache-hash/tokenization path.
🪄 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: 1001c8fd-15d1-46f1-9553-f45f1bea19cf
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/serve/router.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/disaggregated/test_router.py
1f3dde4 to
f9eafa9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56482 [ run ] triggered by Bot. Commit: |
|
PR_Github #56482 [ run ] completed with state
|
f9eafa9 to
13f8391
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57828 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #57837 [ run ] triggered by Bot. Commit: |
|
PR_Github #57828 [ run ] completed with state |
|
PR_Github #57837 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-PyTorch-Post-Merge-*" |
|
PR_Github #58515 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #59549 [ run ] triggered by Bot. Commit: |
Follow-up: the PP scheduling path still misses the disagg draft-token syncRe: the earlier CodeRabbit thread "Mirror this sync in the PP scheduling path", which was marked Resolved in new commit — on the current head ( What's in the code today (
Impact — performance only, not correctness. The synced draft-token count feeds Suggestion: either genuinely mirror the sync into the PP path (add the |
|
@SimengLiu-nv Thank you for the update! I noticed that modular tests have already been added as part of this change, which is a great practice. Given our current GPU resource constraints, would it be possible to reuse and extend existing E2E test coverage for this scenario instead of adding a new E2E test? Thank you very much for your consideration! |
|
PR_Github #59549 [ run ] completed with state
|
Use tool-aware chat template and Harmony tokenization for KV-cache-aware router hashes without mutating forwarded OpenAI requests. Sync disaggregated generation draft tokens from context-phase params before scheduling so batch capacity accounting sees transferred draft tokens. Add unit coverage for router/server tokenization parity and disagg draft-token scheduler accounting. Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
6fe5e03 to
cdcdad1
Compare
Thanks for the detailed review. Changes in the new commit. |
|
/bot run --disable-fail-fast |
|
PR_Github #60459 [ run ] triggered by Bot. Commit: |
schetlur-nv
left a comment
schetlur-nv
left a comment
There was a problem hiding this comment.
approving for multi-gpu tests
|
PR_Github #60459 [ run ] completed with state
|
| self.kv_cache_manager.commit_scheduled_kv_cache_stats( | ||
| scheduled_batch) | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
The draft-token sync itself is small, but it adds ~10 helpers, two tracking dicts, and 6 touch points in py_executor.py (our most safety-critical file) mainly to avoid an active_requests scan. Could check_gen_transfer_status return the completed request objects (as it already does for cancelled ones) so we can set the draft tokens inline on completion? That would let us drop the _disagg_generation_trans_in_progress_requests tracking layer (and the cancel-path leak it introduces) and keep the footprint much smaller. cc @chuangz0
| return | ||
|
|
||
| for req in new_gen_reqs: | ||
| self._track_disagg_generation_trans_in_progress([req]) |
There was a problem hiding this comment.
Entries added here are only popped via check_gen_transfer_status, but a user cancel of an in-flight gen transfer deletes the recv session synchronously (transceiver.cancel_request), so the rid never resurfaces and this entry (plus the LlmRequest it holds) leaks for the process lifetime. Could we pop it on the terminate/cancel path? (The timeout path may need the same.) Thanks!
| tool.model_dump() if hasattr(tool, "model_dump") else tool for tool in request.tools | ||
| ] | ||
|
|
||
| def _tokenize_harmony_chat(self, request: ChatCompletionRequest) -> list[list[int]]: |
There was a problem hiding this comment.
Could we unify this implementation with openai_server.py:2053 added by @reasonsolo ?
|
|
||
| def _mock_tokenizer(token_ids=None): | ||
| """Return a mock tokenizer with a recorded apply_chat_template. | ||
| def _mock_tokenizer(token_ids: list[int] | None = None) -> mock.MagicMock: |
There was a problem hiding this comment.
_get_weather_tool and _mock_tokenizer are each defined twice at module scope. Could we keep just one of each? Thanks!
|
Close the PR to break into smaller pieces and improve the implementation. |
12 participants
Use tool-aware chat template and Harmony tokenization for KV-cache-aware router hashes without mutating forwarded OpenAI requests.
Sync disaggregated generation draft tokens from context-phase params before scheduling so batch capacity accounting sees transferred draft tokens.
Add unit coverage for router/server tokenization parity and disagg draft-token scheduler accounting.
Summary by CodeRabbit
New Features
Bug Fixes
Description
Test Coverage
tests/unittest/_torch/executor/test_py_executor.py
tests/unittest/disaggregated/test_router.py
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.