[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3 - #17334
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. |
e0ba3cc to
89bc6e4
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #64767 [ run ] triggered by Bot. Commit: |
|
PR_Github #64677 [ run ] completed with state |
|
PR_Github #64767 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64781 [ run ] triggered by Bot. Commit: |
|
PR_Github #64781 [ run ] completed with state
|
|
/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." |
|
PR_Github #64798 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64798 [ skip ] completed with state |
| # the shared hybrid transceiver validation below: the Python NIXL | ||
| # transceiver selects the Mixed manager, whose KDA recurrent/conv | ||
| # states transfer through the bounce buffer. | ||
| if is_kimi_linear(config) and not use_v2 and not is_disagg: |
There was a problem hiding this comment.
Now that the fail-fast is gone, a K3 disagg deployment that leaves the transceiver runtime at its default resolves to the C++ path and reaches this block without any error.
There was a problem hiding this comment.
Confirmed, this is a real gap. KimiLinearForCausalLM does not override get_preferred_transceiver_runtime, so transceiver_runtime: auto resolves to the C++ transceiver, and this route then returns CppMambaHybridCacheManager with no error even though only the Python NIXL transceiver knows how to move the KDA state. Qwen3-Next already returns "PYTHON" from get_preferred_transceiver_runtime; K3 should do the same, plus a loud guard on this route for paths that skip model defaults (e.g. AutoDeploy resolves auto to the C++ fallback). Queued at the top of TRTLLM-15264: the failure mode is silent.
There was a problem hiding this comment.
Fix is up: #17480 (KimiLinearForCausalLM.get_preferred_transceiver_runtime -> 'PYTHON' plus a loud guard on the K3 disagg route, with routing unit tests).
| _MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" | ||
| # Test/advanced overrides for the size gates below (users only tune the bounce size). Read on the | ||
| # generation side, so set them there; unset uses the defaults. | ||
| _MIN_BYTES_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES" # byte gate for recurrent-state payloads |
There was a problem hiding this comment.
Since the two gates are selected by payload kind rather than by transfer cost, could a short note record what each one actually guards?
There was a problem hiding this comment.
Agreed, will add a note at the env-var definitions: min_bytes gates payloads carrying recurrent state (the fallback cost scales with bytes, and block counts are meaningless for the non-paged state), while min_blocks is the pre-existing gate kept for plain-KV payloads so already-deployed bounce users see no behavior change. Tracked under TRTLLM-15264, together with folding in the observation from your comment below (for K3 the byte gate never rejects; the arena capacity is the effective admission control).
There was a problem hiding this comment.
Fix is up: #17483 (added notes at the min-bytes / min-blocks env-var gates documenting what each admission control covers — byte gate for recurrent-state payloads, block gate kept for plain-KV behavior compatibility).
| # inter-node — 2 MiB in-place at that rate is ~5 ms vs well under 1 ms bounced. Below 2 MiB the | ||
| # fixed overhead dominates and arena slots are better kept for large transfers. Heuristic, tunable | ||
| # via TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES. | ||
| DEFAULT_MIN_BYTES = 2 * _MIB |
There was a problem hiding this comment.
For K3 every request carries the fixed recurrent payload, so this threshold can never reject anything and the arena size is the effective admission control.
There was a problem hiding this comment.
Right. K3's fixed KDA payload clears the 2 MiB gate on every request, so for this model admission control is effectively the arena capacity plus reservation backpressure (reserve falls back to the per-fragment path on oversize or timeout). The byte gate only bites for hybrid models with small recurrent payloads or very short transfers. Will record this in the gate documentation note (previous comment, TRTLLM-15264) so the next reader does not go looking for a rejection path that never fires for K3.
There was a problem hiding this comment.
Fix is up: #17483 (the gate note records that for K3 the fixed KDA payload clears the byte gate on every request, so arena capacity plus reservation backpressure is the effective admission control and there is no per-request rejection path to look for).
| warn_key="kv-bounce-unknown-slot-size", | ||
| ) | ||
| total += int(block_ids.size) * self._block_bytes_per_group[g] | ||
| if extra_bytes > 0 and num_writers > 1: |
There was a problem hiding this comment.
Separating the classification from the size would remove that ambiguity and make the gate selection explicit.
There was a problem hiding this comment.
Agreed. extra_bytes > 0 currently doubles as the payload classifier, so a hybrid transfer whose recurrent payload happens to compute to zero bytes (e.g. an empty layer intersection under PP) would be gated as plain KV. The fix is to pass an explicit carries-recurrent-state flag from dispatch_task alongside extra_bytes, so classification is by payload kind and the byte value is only a size. Queued with the TRTLLM-15194 gate hardening, which already covers whether the two gates can be unified.
| @@ -182,9 +187,10 @@ def _reserve_and_gather(self, write_meta, *, timeout): | |||
| total = int(write_meta.sizes.sum()) | |||
There was a problem hiding this comment.
The receiver's reservation size never reaches the sender, so a coalesced write larger than the region has nothing to stop it. Would some form of bound check be worth adding here?
There was a problem hiding this comment.
Worth adding, agreed. The receiver knows the reserved size but only ships bounce_dst_base to the sender, so nothing bounds the coalesced write. Plan: include the reserved byte count (per-writer share for fan-in) in the recv request so the sender can check total <= reservation before issuing the write, failing the transfer through the normal error path instead of overrunning the neighboring slot. This complements the receiver-side check you suggested in core.py (comparing reported bytes against the reservation at result-record time); both are tracked under TRTLLM-15194 with the rest of the bounce hardening.
| if peer_infos.page_table is None: | ||
| allow_bounce = False # cannot size the sender's recurrent-state payload | ||
| else: | ||
| extra_bytes = MambaPolicy.payload_bytes( |
There was a problem hiding this comment.
The rank info cached per info endpoint comes from context rank 0, whose page table is pipeline-stage-local, so under PP the layer sets can fail to intersect and this returns zero while the actual sending stage still appends its full recurrent payload.
There was a problem hiding this comment.
Confirmed; this is the PP sizing gap tracked under TRTLLM-15194 (also raised by @chienchunhung on the PR thread). Under PP the cached RankInfo is stage-local to context rank 0, so payload_bytes can size from the wrong stage, including a zero from an empty layer intersection, while the real sending stage appends its full payload. It cannot trigger today (K3 requires pp_size 1 and every shipped config uses PP1), but the plan under TRTLLM-15194 is to size from per-rank page tables or disable recurrent-state bounce when the peer has pp_size > 1, plus a PP regression test.
| # reject here — before REGISTER_RANK_INFO is even sent — to fail | ||
| # the first gen request loudly instead of hanging on a transfer | ||
| # the sender will never serve. | ||
| MambaPolicy.validate_peer_compatible( |
There was a problem hiding this comment.
This raises on the executor thread with nothing catching it upstream, so a single mismatched context instance would take down the whole generation worker, including traffic from healthy peers. Failing just the affected request through the existing transfer-error path would keep the blast radius local while still surfacing the diagnostic.
There was a problem hiding this comment.
Fair point. The raise is deliberately at first-contact registration (fail the first request loudly rather than hang on a transfer the sender will never serve), but you are right that it propagates on the executor thread and takes down the whole generation worker even when other context peers are healthy. Will reroute it through the existing transfer-error path so only requests targeting the incompatible peer fail, keeping the diagnostic message. Tracked under TRTLLM-15264.
There was a problem hiding this comment.
Fix is up: #17482 (validate_peer_compatible failures now raise a typed PeerIncompatibleError that Receiver.dispatch_task turns into a per-request DISAGG_TRANS_ERROR, so only requests targeting the incompatible peer fail; the peer is cached to fast-fail later requests without re-validating, healthy peers and in-flight requests are unaffected, and the diagnostic is preserved).
| if slice.mamba_state_index is not None: | ||
| total += len(lg.mamba_layer_offsets) * ( | ||
| lg.conv_states.slot_bytes + lg.ssm_states.slot_bytes | ||
| ) |
There was a problem hiding this comment.
A cross-reference would help, since the two byte accountings added in this PR already differ on the pool-view axis.
There was a problem hiding this comment.
Agreed, will cross-reference the two accountings. Note the pool-view difference is also a real sizing issue on the bounce side (block_bytes_per_group reads pool 0 only, per your comment there); once that sums pool views the two accountings agree on the axis, and the comments will point at each other. Tracked under TRTLLM-15264 (the pool-view sizing fix itself is under TRTLLM-15194).
There was a problem hiding this comment.
Fix is up: #17483 (added cross-referencing comments between _slice_num_bytes (sums all pool views) and bounce block_bytes_per_group (pool 0 only); the pool-0-only sizing fix itself is tracked under TRTLLM-15194).
| are pre-scaled by tp_size when attention-DP is off so the per-rank | ||
| state stays full-size (EP-only parallelism, replicated KDA state). | ||
| """ | ||
| state_tp = tp if not enable_attention_dp else 1 |
There was a problem hiding this comment.
This file isn't in any test list, and the CPU sweep ignores modules without the cpu-only marker at collection time. Besides, the fixture pre-scales the head count so the per-rank state stays full size, which cancels the manager's own division
There was a problem hiding this comment.
Both parts are right. On wiring: test-list registration and the GPU/NIXL markers for these suites are handled in #17332 so the gating lands together; the CPU-runnable subset gets its marker there too. On the fixture: the pre-scaling does cancel the manager's division, and the docstring's claim that this mirrors production is stale. The merged mapping (extract_mamba_kv_cache_params) passes lin["num_heads"] through unscaled, so with attention-DP off the manager shards heads across TP; the fixture instead models a replicated full-size per-rank state. The follow-up will build the fixture through the production mapping (shared helper rather than hand-built kwargs) and cover both the ADP-on (replicated) and ADP-off (sharded) geometries that production actually constructs. Tracked under TRTLLM-15264.
There was a problem hiding this comment.
Fix is up: #17483 (_create_kda_managers now builds its Mamba cache params via the production extract_mamba_kv_cache_params with unscaled global head counts, so the manager's own attention-DP gate does the sharding; the stale "mirrors production" docstring is corrected).
|
|
||
|
|
||
| def test_kda_hetero_tp_rejected(): | ||
| """Replicated KDA state + heterogeneous TP (ADP off) must be rejected. |
There was a problem hiding this comment.
Following from the fixture above, this asserts a rejection for a geometry production doesn't build.
There was a problem hiding this comment.
Following from the fixture discussion above: correct. With the production (unscaled) mapping, ADP-off heterogeneous TP produces sharded per-rank state whose TP-aggregated size matches, so validate_peer_compatible would accept it; the rejection this test asserts only arises from the fixture's pre-scaled (replicated) geometry, which production does not build for K3. The follow-up reworks this to assert what production actually produces: positive transfer coverage through the TP-mismatch mappers for the sharded case, and a config-level rejection for any geometry K3 genuinely does not support. Tracked under TRTLLM-15264.
There was a problem hiding this comment.
Fix is up: #17483 (test_kda_hetero_tp_rejected became test_kda_hetero_tp_sharded_accepted — under the production head-sharded mapping, ADP-off hetero TP is accepted, so the test now asserts concrete per-rank slot sizes before acceptance rather than a rejection production never builds).
| # NIXL/PYTHON dropped for now: V2 _exchange_rank_info mpi_allgather can | ||
| # hang under this harness (srun --mpi=pmix). | ||
| - {backend: "NIXL", runtime: "CPP"} | ||
| - {backend: "UCX", runtime: "CPP"} |
There was a problem hiding this comment.
This cache-transceiver test is intended to sweep UCX environment variables and identify a suitable UCX configuration. It is not a model benchmark framework. adding a model specific yaml config is not appropriate. Also, K3 requires a python runtime cache transceiver ,but this config only involve cpp runtime
There was a problem hiding this comment.
Agreed on both points. The config existed to reproduce the KDA-sized payload measurement, but it is model-specific and only sweeps the CPP runtime, which K3 cannot use. Will remove it from cache_transceiver_test under TRTLLM-15264; the payload-size derivation (the header math) moves into the K3 example docs where the same numbers are already used.
There was a problem hiding this comment.
Fix is up: #17483 (removed the config from cache_transceiver_test; the payload-size derivation is preserved in the K3 disagg README where the same numbers are used).
| prerequisites for every worker shell (see caveats below for why): | ||
|
|
||
| ```bash | ||
| export UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc # on clusters where verbs cannot |
There was a problem hiding this comment.
UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc is a workaround for a specific cluster where the verbs transports fail to initialize. It should not be hard-coded into the default K3 benchmark configuration.
This transport list excludes RDMA/verbs transports. On B200/B300 clusters whose inter-node communication depends on RDMA, it may cause initialization or registration failures,or force communication onto an unintended TCP path.
Could this be moved to an optional troubleshooting section, with the default recommendation being to leave UCX_TLS unset unless the cluster's validated UCX configuration
requires an override?
There was a problem hiding this comment.
Agreed. Will move the UCX_TLS override out of the default launch sequence into a troubleshooting subsection, with the default guidance being to leave UCX_TLS unset so verbs/RDMA transports stay available on clusters that depend on them, plus a note on the symptom that motivates the override (NIXL/UCX initialization failure where verbs transports cannot come up). Tracked under TRTLLM-15264.
There was a problem hiding this comment.
Fix is up: #17479 (default is now no UCX_TLS pin in the README launch sequence and the benchmark config; the override moved to the troubleshooting caveats, scoped to clusters where verbs transports cannot initialize).
|
|
||
|
|
||
| @pytest.mark.timeout(180) | ||
| @pytest.mark.parametrize( |
There was a problem hiding this comment.
we need test cases that ctx_tp != gen_tp and adp =false/true?
There was a problem hiding this comment.
Agreed, and this ties into the fixture rework from the thread above (the current fixture pre-scales heads, which does not match the merged production mapping). Planned parametrization for the follow-up: matched TP with ADP on and off, and heterogeneous ctx/gen TP with ADP on (replicated per-group state) and ADP off (sharded state through the TP-mismatch mappers), asserting transfer correctness or a loud rejection according to what production supports for each geometry. Tracked under TRTLLM-15264.
There was a problem hiding this comment.
Fix is up: #17483 (new transfer parametrizations: matched TP2 with attention-DP on and off, and heterogeneous 2->4 / 4->2 with bitwise output comparison).
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
14 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.
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/.