[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. |
… 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>
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.
| _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?
| # 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.
| 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.
| @@ -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?
| 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.
| # 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.
| 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.
| 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
|
|
||
|
|
||
| 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.
| # 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
| 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?
|
|
||
|
|
||
| @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?
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/.