[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3 - #17334
Merged
Conversation
Collaborator
|
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. |
brnguyen2
marked this pull request as ready for review
August 6, 2026 23:56
brnguyen2
requested review from
2ez4bz,
Barry-Delaney,
PerkzZheng,
Shixiaowei02,
Superjomn,
Tabrizian,
YihuiLu512,
chang-l,
eopXD,
fredricz-20070104,
kris1025,
ruodil,
xxi-nv and
yizhang-nv
August 6, 2026 23:56
yuanjingx87
approved these changes
Aug 7, 2026
…ative disaggregation Extend the Python-native disaggregation framework to transfer Kimi K3 KDA (Kimi Delta Attention) recurrent and conv states between context and generation instances: SSM mixer peer descriptors, per-request auxiliary state payloads, bounce-buffer staging for non-fabric pools, and transceiver routing for hybrid linear-attention models. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Replace the Kimi K3 disaggregated-serving fail-fast in the cache manager routing with the shared hybrid transceiver validation: the Python NIXL transceiver selects MixedMambaHybridCacheManager, whose KDA recurrent/conv states transfer through the bounce buffer. Also log the selected hybrid cache manager class once at routing time. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…i K3 parity harness - test_kda_mamba_transfer.py: KDA recurrent/conv state transfer through the native disaggregation path. - test_bounce.py / region/test_aux.py: cover bounce-buffer staging and auxiliary-state payloads for hybrid models. - kimi_k3_disagg_parity.py: two-endpoint aggregated-vs-disaggregated parity harness (multi-node; not wired into any test list here). - Update the overlap transceiver-runtime python bounce test config. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
… and benchmark wiring - examples/kimi_k3/disagg/: ctx/gen/proxy configs, SLURM benchmark harness config, and a README covering K3 disagg constraints (matched DEP16, Python NIXL transceiver, bounce-buffer sizing, UCX transport pins). Spec-decode (SA) variants land with K3 SA support. - slurm/benchmark harness: worker/server env plumbing (TRTLLM_WORKER_UCX_TLS, PATH/PYTHONPATH prepends) used by the configs. - cache_transceiver_test: K3-shaped KDA payload config and harness support. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…oss peers The peer-compatibility gate rejected any pair of ranks whose mamba/KDA layer sets differ. With pipeline parallelism each rank publishes only its own stage's layers, so the sets legitimately differ (or are disjoint, or one stage holds no recurrent layers at all) while the transfer path intersects the two sets on purpose. Drop the set-equality requirement and treat a missing recurrent layer group on either side as nothing to validate; keep the per-slot size invariants, which are layer-agnostic. Add a regression test covering partial overlap, disjoint stages, a recurrent-layer-free stage, and a size mismatch on the overlap. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…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>
… synthetic rank info Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…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>
… 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>
…t malformed responses Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…the K3 README Disaggregated serving for Kimi K3 lands in this change; drop the corresponding line from the current-limitations list. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
brnguyen2
force-pushed
the
k3/14815-disagg
branch
from
August 8, 2026 06:07
e0ba3cc to
89bc6e4
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
brnguyen2
enabled auto-merge (squash)
August 8, 2026 06:09
Collaborator
|
PR_Github #64767 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64677 [ run ] completed with state |
Collaborator
|
PR_Github #64767 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64781 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64781 [ run ] completed with state
|
Collaborator
Author
|
/bot skip --comment "Union-green across two full runs on head 89bc6e4: pipeline 52625 ran single-GPU fully green and multi-GPU 140/142 stages green; the sole red stage (DGX_B200-8_GPUs-PyTorch-1, one SLURM job) failed 3 accuracy tests unrelated to this disaggregation-gated diff (NemotronV3 Super/Ultra, DeepSeekV4ProDSpark) with init-time environment signatures (c10d store timeout, CUDA OOM during symmetric-memory allocation); the NemotronV3Ultra case is already waived on B300 (nvbugs/6474894). Pipeline 52611 on the same head was green except one unrelated pip-install network failure in a test setup fixture." |
Collaborator
|
PR_Github #64798 [ skip ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64798 [ skip ] completed with state |
Open
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
13 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
Base: #17269 (KimiLinear model) has merged; this PR is rebased onto
main and carries only the disaggregated-serving changes.
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/.