[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3 - #17334
[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3#17334brnguyen2 wants to merge 14 commits into
Conversation
|
The new peer check wants both sides to have the same layer set, but each rank only holds its own pipeline stage, and the transfer code below it takes the overlap on purpose. So this rejects hybrid models with pipeline parallelism, including a Qwen3-Next test already in pre-merge. Also, the function that seeds the replay caches on the generation side is never called. |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_kda_mamba_transfer.py (1)
575-583: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an explicit hardware marker to
test_kda_transfer.This test constructs
MixedMambaHybridCacheManagerinstances and runs a NIXL loopback transfer. Mark it with the repository’s GPU/NIXL marker so CPU collection runs skip it instead of failing.🤖 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_kda_mamba_transfer.py` around lines 575 - 583, Add the repository’s existing GPU/NIXL pytest marker to test_kda_transfer, alongside its timeout and parameterization decorators. Keep the test body and parameter sets unchanged so CPU-only collection skips this MixedMambaHybridCacheManager and NIXL loopback test.
🤖 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 `@examples/disaggregated/slurm/benchmark/start_server.sh`:
- Around line 11-16: Strip one leading and trailing literal single quote from
TRTLLM_PATH_PREPEND and TRTLLM_PYTHONPATH_PREPEND before using them in the
prepend export logic. Apply this change in
examples/disaggregated/slurm/benchmark/start_server.sh:11-16 and
examples/disaggregated/slurm/benchmark/start_worker.sh:44-49, preserving the
existing conditional exports and PATH/PYTHONPATH composition.
In `@examples/disaggregated/slurm/benchmark/start_worker.sh`:
- Around line 85-91: Add the standard NVIDIA copyright header for 2026 at the
top of start_worker.sh, preserving the existing trtllm_serve_cmd selection and
launch commands unchanged.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 166-170: Update SpeculativeConfig validation in
TorchLlmArgs.validate_speculative_config to reject speculative decoding for Kimi
K3 when disaggregated serving is enabled. Add the guard before selecting the
Kimi K3 disaggregated manager, while preserving existing behavior for
non-disaggregated configurations.
In `@tests/integration/defs/kimi_k3_disagg_parity.py`:
- Around line 123-128: Update _served_model to catch connection-level failures
alongside HTTPError, including urllib.error.URLError and relevant
OSError/timeout cases, so unavailable endpoints return None instead of
propagating exceptions. Report these failures with the existing parity note
format while preserving the current HTTP status reporting.
- Around line 540-544: Update the model-name handling around the parity endpoint
construction and the NOTE: either pass the resolved candidate model name to the
candidate Endpoint while retaining the reference name for the reference
Endpoint, or change the message to accurately state that the reference model
name is used for both and that --model overrides it. Ensure the log and actual
request behavior remain consistent.
- Around line 89-93: Resolve the missing kimi_k3_sa_harness dependency used by
kimi_k3_disagg_parity.py: either add that module with PROMPTS_AND_CHECKS,
_compare_logits_parity, and _parity_prompts, or update the import to reference
their actual existing definitions. Ensure the parity test imports successfully
without changing its required symbols.
In `@tests/unittest/disaggregated/region/test_aux.py`:
- Around line 138-187: Add tests/unittest/disaggregated/region/test_aux.py,
including test_aux_buffer_zero_max_draft_len_round_trip and
test_aux_transfer_layout_ctx_no_spec_gen_sa, to the appropriate CI test lists
under tests/integration/test_lists/test-db/ and
tests/integration/test_lists/qa/. Ensure the configured test commands discover
and execute these tests.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_kda_mamba_transfer.py`:
- Around line 575-583: Add the repository’s existing GPU/NIXL pytest marker to
test_kda_transfer, alongside its timeout and parameterization decorators. Keep
the test body and parameter sets unchanged so CPU-only collection skips this
MixedMambaHybridCacheManager and NIXL loopback test.
🪄 Autofix
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: dc28e7de-971c-4e71-9024-972508653da5
📒 Files selected for processing (23)
examples/disaggregated/slurm/benchmark/run_benchmark.shexamples/disaggregated/slurm/benchmark/start_server.shexamples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yamlexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyexamples/kimi_k3/disagg/README.mdexamples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yamlexamples/kimi_k3/disagg/ctx_config.yamlexamples/kimi_k3/disagg/disagg_proxy_config.yamlexamples/kimi_k3/disagg/gen_config_no_sa.yamltensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/core.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/_util.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/kimi_k3_disagg_parity.pytests/unittest/disaggregated/region/test_aux.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
🚧 Files skipped from review as they are similar to previous changes (16)
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
- examples/kimi_k3/disagg/README.md
- tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml
- tensorrt_llm/_torch/disaggregation/native/bounce/core.py
- tensorrt_llm/_torch/disaggregation/transceiver.py
- examples/kimi_k3/disagg/ctx_config.yaml
- examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml
- examples/disaggregated/slurm/benchmark/run_benchmark.sh
- tensorrt_llm/_torch/disaggregation/native/peer.py
- examples/kimi_k3/disagg/disagg_proxy_config.yaml
- examples/kimi_k3/disagg/gen_config_no_sa.yaml
- examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml
- tensorrt_llm/_torch/disaggregation/native/bounce/config.py
- tests/unittest/disaggregated/test_bounce.py
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tensorrt_llm/_torch/disaggregation/native/bounce/impl.py
|
PR_Github #64435 [ run ] triggered by Bot. Commit: |
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The code paths shown look internally consistent, but this feature ships with no CI-runnable coverage and two author-acknowledged unresolved correctness follow-ups, so it should not merge to main as-is (mergeable_state is also blocked).
Concerns
-
[MAJOR]
tests/integration/defs/kimi_k3_disagg_parity.py:1(and missing test-list wiring) - feature ships with no CI coverage- What is wrong: Substantial production code is added (bounce byte-gate in
bounce/config.py/impl.py,MambaPolicy.validate_peer_compatible/payload_bytes, transceiver hybrid routing,_util.pydisagg enablement) plus new unit tests and a parity harness, but per the description there are no changes undertests/integration/test_lists/,test-db/, orqa/, and the parity harness is not registered. - How it fails: None of the new tests run in CI, so a future regression in the recurrent-state transfer path (extra_bytes sizing overrun, peer-validation false-reject) lands undetected.
- Suggested fix: Register the applicable unit tests in an L0 list and wire (or explicitly document the gating job for) the integration harness before un-drafting.
- What is wrong: Substantial production code is added (bounce byte-gate in
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/_util.py:163- self-flagged unresolved correctness items- What is wrong: The PR's own Dev/QA review states the generation-side replay-cache seeding function is not invoked, and that peer validation may reject valid hybrid models when ranks hold different pipeline-stage layer subsets. The replay-cache seeding code is not in this diff.
- How it fails: A gen request whose recurrent state depends on the seeded replay cache would decode from uninitialized state after the ctx->gen handoff, producing wrong tokens with no error.
- Suggested fix: Confirm (with a test) that replay-cache seeding is invoked on the gen path, and resolve the PP peer-validation follow-up before approval.
Minor notes (non-blocking)
tests/unittest/disaggregated/test_kda_mamba_transfer.py:575-test_kda_transferruns a real NIXL loopback but has no GPU/NIXL marker; CPU-only collection will fail instead of skipping.tensorrt_llm/_torch/pyexecutor/_util.py:170- theand not is_disaggchange makes the kimi_linear branch fall through to unseen manager-selection code; add a test asserting the selected manager class for kimi_linear+disagg (theinfo_oncelog is not a check).
QA view
- Test coverage: partial - unit tests exist for the byte-gate, hybrid bounce reserve, aux zero-draft round-trip, and KDA peer validation, but none are registered in a test list so CI does not run them; the gen-side replay-cache path and the
_util.pydisagg fall-through have no test here. - SM coverage: code touches Blackwell (GB200/GB300, sm100) fabric-VMM bounce and the NIXL transfer path; tests run only on CPU (logic) and single-node loopback (unmarked). No CI run on the target arch within this PR - a real coverage gap.
- Test code:
test_kda_transfermissing GPU/NIXL marker; nothing registered intest_lists//test-db//qa/; parity harness is a standalone multi-node script (self-test only). - Test time: unknown - tests are not wired into CI, so no measurable impact now; wiring the loopback/Blackwell runs later would add non-trivial time.
- Needs
/qa-verify: yes - Blackwell-specific transfer path with no CI test on that arch, tests not list-wired, and unresolved replay-cache/PP follow-ups; a human should re-run token-parity/GSM8K on GB200/GB300 disagg before trusting this.
Possible new issues
extra_bytesregion sizing intransfer.pydispatch_task assumes sendercollect_fragssizes are slot-independent; a future slot-dependent sender would overrun the bounce slot into the neighbor. Guarded only by the K3-geometry unit test._check_peer_compatiblenow raisesValueErrorinstead of returningbool; on the sender listener thread exceptions are only logged, so an incompatible peer fails silently sender-side and relies on the receiver-side gate - a fragile mixed contract.
What I could not verify
- The manager-selection fall-through in
get_kv_cache_manager_clsbeyond the shown lines (which manager is actually chosen for kimi_linear+disagg). - The generation-side replay-cache seeding code (not in this diff) - cannot confirm it is invoked.
- Because the PR is stacked on #17269, the shown diff may include unrelated content that will shrink after rebase; runtime behavior on hardware is unverified.
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: Arch-specific (Blackwell) transfer path with no CI test on that arch, tests not registered in any list, and author-acknowledged unresolved follow-ups (gen-side replay-cache seeding not invoked, PP peer validation). A human QA should re-run token-parity/GSM8K on GB200/GB300 disagg before this is trusted.
Automated review by NVCortex Lite, run by @fredricz-20070104.
…rness error handling, launcher quote stripping - test_kda_mamba_transfer.py: shut down every GPU-backed manager created by _create_kda_managers (not just the ranks under test) and wrap run_kda_transfer_test in try/finally so managers and transceivers are released on early assertion failures. - kimi_k3_disagg_parity.py: _served_model now returns None on connection-level failures (URLError/OSError) and malformed responses, matching its documented contract; fix the mismatched-model NOTE to state the actual behavior (one model name sent to both endpoints). - start_server.sh / start_worker.sh: strip literal single quotes from TRTLLM_PATH_PREPEND / TRTLLM_PYTHONPATH_PREPEND before use, matching run_benchmark.sh. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
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 `@tests/integration/defs/kimi_k3_disagg_parity.py`:
- Around line 129-131: The _served_model() response handling must validate
malformed /v1/models payloads before indexing and catch json.JSONDecodeError
alongside existing connection failures, returning None consistently for invalid
JSON, missing or null data, and null entries. Add regression tests covering
these malformed payloads and connection failures, and register the tests in the
appropriate test-db/ or qa/ test list.
🪄 Autofix
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: 89bcf8c7-5597-4d03-b61f-cef603462277
📒 Files selected for processing (4)
examples/disaggregated/slurm/benchmark/start_server.shexamples/disaggregated/slurm/benchmark/start_worker.shtests/integration/defs/kimi_k3_disagg_parity.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/disaggregated/slurm/benchmark/start_server.sh
- examples/disaggregated/slurm/benchmark/start_worker.sh
- tests/unittest/disaggregated/test_kda_mamba_transfer.py
… synthetic rank info Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #64435 [ run ] completed with state
|
…or recurrent-state payloads Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…tent MambaPolicy.validate_peer_compatible now checks the global (per-rank bytes x mamba_tp) recurrent-state size, so the fixed-size synthetic mamba group in make_page_table() reads as a replicated state under heterogeneous TP and fails registration in the tp2-vs-tp1 registrar tests. Shard the fixture's mamba pools from a fixed global size by a mamba_tp parameter (default 2, matching make_rankinfo's default tp_size and preserving the previous byte values) and pass mamba_tp=1 for the tp=1 peers. Also restores the intended failure mode of test_peer_registrar_rejects_misaligned_subbyte_head_mismatch, which had been passing on the mamba mismatch instead of the alignment check. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Bounce-gate finding addressed: plain-KV payloads keep the original 96-block gate (min_blocks default restored to 96) and the 2 MiB byte gate now applies only when extra_bytes > 0, so an existing bounce deployment is unchanged. Also checked the narrower case Codex raised -- a non-Kimi hybrid model with bounce explicitly enabled -- and it cannot regress either: on main impl.py:230 bails any transfer whose layer group has no known slot size ("e.g. mamba"), so hybrids never reach the arena today; this PR is what first makes them eligible.
Approving over the 3 open CodeRabbit threads: the start_worker.sh copyright nit (whole directory is header-less), the kimi_k3_disagg_parity.py import of a module from #17327, and the _http_json error-handling nit -- the latter two are in a manually-invoked harness that pytest does not collect. The #17392 file collision still needs merge-order care.
test_on_update_kv_lens_rebuilds_stale_map builds its metadata with object.__new__, bypassing __init__ where in_mtp_draft_loop is initialized. Since NVIDIA#16925 made on_update_kv_lens() read that flag, the test fails with AttributeError on every pre-merge run. Set the __init__ default on the stub. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
… configs K3 ships a custom tokenizer, so a server started from the README commands without trust_remote_code comes up tokenizer-less and rejects every string prompt with 'tokenizer is required to tokenize string prompt' (found running the OpenAI-completions path end to end; the benchmark harness yaml in the same directory already sets it). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #64589 [ run ] triggered by Bot. Commit: |
|
PR_Github #64589 [ run ] completed with state
|
|
/bot run |
Validation summary (out-of-CI, Blackwell nodes)
|
|
PR_Github #64611 [ run ] triggered by Bot. Commit: |
…t malformed responses Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
@fredricz-20070104 thanks for the detailed review. Point-by-point on the concerns:
|
|
/bot run |
6 participants
Description
Adds disaggregated-serving support for Kimi K3 (KimiLinear), in four commits:
MambaPolicy state-region mapping for the KDA mixer, peer registration,
bounce-buffer config/impl updates sized for the KDA state payload, and
transceiver handling for hybrid (attention + recurrent-state) models.
pyexecutor/_util.pypreviouslyraised NotImplementedError for disaggregated serving (referencing this
ticket); it now routes through the shared hybrid transceiver validation.
updates, plus a Kimi K3 disagg logits-parity integration harness.
Stacked on #17269 — draft until that merges; the diff then shrinks to
these four commits after a rebase onto main.
Notes
the SA disagg test depend on [TRTLLM-14814][feat] Kimi K3 serving parsers, chat template, and speculative decoding (suffix automaton + DFlash scaffold) #17327 and follow once both PRs are in.
wired into any test list here; unit-suite results on Blackwell hardware
will be posted before un-drafting. Prior validation of this code on the
feature bring-up branch: token-level parity between disaggregated and
aggregated serving on the target model.
Test Coverage
PR Checklist
[TRTLLM-14815][feat]conventionDev Engineer Review
trtllm-servefallback handling.#17327.test-db/, orqa/changes are included.QA Engineer Review
kimi_k3_disagg_parity.pyparity harness with endpoint, token, logprob, GSM8K, reporting, CLI, and self-test coverage.tests/integration/test_lists/,test-db/, orqa/.