diff --git a/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml b/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml deleted file mode 100644 index 674c275890a8..000000000000 --- a/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml +++ /dev/null @@ -1,83 +0,0 @@ -# Kimi K3 KDA-payload cache-transceiver micro-benchmark. -# Run from the repository root: -# python3 examples/disaggregated/slurm/cache_transceiver_test/submit.py \ -# -c examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml -# -# Synthetic KV geometry sized so ONE request's per-rank transfer equals the -# exact per-request KDA state payload of Kimi K3: 454,459,392 bytes -# (69 layers x (conv [3*96*128, 4] bf16 + delta [96,128,128] fp32)). -# Per-rank bytes/token = 69 layers * kvFactor(2) * (24/4 kv heads) * 128 * 2B -# = 211,968 B -> 2144 tokens == 454,459,392 B exactly. -# request_lengths 512 / 1024 also match MLA-latent payloads at ~4k / ~8k ISL -# (bf16) for the KDA-vs-MLA comparison. -# -# Fill in the <...> placeholders for your cluster before submitting. - -slurm: - partition: "" - account: "" - job_time: "00:30:00" - job_name: "ctt_kda_payload" - extra_args: "--gpus-per-node=4" - -hardware: - gpus_per_node: 4 - -environment: - container_image: "" - # Mount the TRT-LLM checkout into the container. - container_mount: "" # Format: path1:path1,path2:path2 - # TRT-LLM comes from the checkout's venv (PATH below) + checkout code - # (PYTHONPATH below); skip pip install. - trtllm_repo: "" - # Output directory (results.json, logs/, csv/); point at any writable - # scratch location. - work_dir: "" - trtllm_wheel_path: "" - build_wheel: false - cuda_architectures: "" - -test_matrix: - combinations: - # 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"} - cache_manager_versions: ["V1"] - # 512 tok = 108.5 MB/rank (~ MLA bf16 @4k ISL), 1024 = 217.1 MB (~ MLA @8k), - # 2144 = 454.46 MB = exact K3 per-request KDA payload. - request_lengths: [512, 1024, 2144] - num_requests_per_length: 8 - warmup_requests: 2 - -kv_cache: - num_layers: 69 - num_kv_heads: 24 # /tp4 -> 6 heads per rank - head_dim: 128 - tokens_per_block: 32 - dtype: "HALF" - max_tokens_in_buffer: 4096 - -parallel: - ctx_tp: 4 - ctx_pp: 1 - gen_tp: 4 - gen_pp: 1 - -ucx_env_sweep: - # Single sweep. PATH/PYTHONPATH are exported inside the container prelude so - # python3 resolves to the checkout's venv (system-site-packages=true, - # so container torch/mpi4py are visible) and the checkout on PYTHONPATH - # wins. Keep your cluster's SLURM bin dir on PATH so host-side srun works. - # The UCX_TLS pin (no verbs) is for clusters where verbs transports cannot - # initialize on the compute nodes; drop the pin where verbs works. - - name: "venv_no_verbs" - env: - UCX_TLS: "tcp,self,sm,cuda_copy,cuda_ipc" - PATH: ":/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:" - PYTHONPATH: "" - -run: - timeout_per_cell_s: 120 - max_sweep_s: 900 - capture_proto_info: true diff --git a/examples/kimi_k3/disagg/README.md b/examples/kimi_k3/disagg/README.md index f07a6d0d870c..b7a3901b1b64 100644 --- a/examples/kimi_k3/disagg/README.md +++ b/examples/kimi_k3/disagg/README.md @@ -19,11 +19,12 @@ for constraints. - **EP-only parallelism on BOTH sides**: `ep_size == tp_size`, no PP, no TP on linears. Deployed as DEP-N (`enable_attention_dp: true`). -- **Matched ctx/gen parallelism (DEP16 = DEP16)** for now. Heterogeneous - ctx/gen TP with attention-DP *off* would silently corrupt memory for - K3's replicated KDA state and is rejected at peer registration — do - not deviate. Hetero DEP with attention-DP on both sides is believed - correct but unvalidated. +- **Matched ctx/gen parallelism (DEP16 = DEP16)** for now: only this + geometry is validated end-to-end on hardware. Heterogeneous ctx/gen + parallelism passes peer validation and single-node loopback transfer + tests (with attention-DP off the KDA state is head-sharded across TP, + so the TP-mismatch mappers re-tile it; with attention-DP on it is + replicated), but no hetero geometry has been validated at scale. - **Ctx sizing = DEP16**: DEP16 is the smallest verified fit (~193 GiB weights/rank on GB300). A DEP8 ctx is estimated at ~273 GiB/rank for weights alone (extrapolating the 1.5 TB checkpoint: @@ -51,6 +52,15 @@ for constraints. tcp path (`[kv-bounce] in-place: transfer 649MiB exceeds the 512MiB bounce region`). +## KDA state payload size + +Per-request recurrent-state payload (fixed, token-count independent): +69 KDA layers x (conv `[3*96*128, 4]` bf16 + delta `[96, 128, 128]` fp32) += 454,459,392 bytes (~433 MiB). For synthetic-KV harnesses that size +transfers per token: at K3's per-rank 211,968 B/token +(69 layers x kvFactor 2 x 6 kv heads/rank x head_dim 128 x 2 B), 2144 +tokens reproduce this payload exactly. + ## Launch sequence (manual, single ctx + single gen) Each K3 worker spans 16 GPUs (4 NVL72 nodes at 4 GPUs/node). Leave @@ -129,8 +139,8 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ change together with K3 SA support (an SA `gen_config.yaml` variant ships with it). 3. **Matched-DP only.** Keep ctx and gen at identical DEP16 with - attention-DP on both sides; heterogeneous parallelism with - attention-DP off is rejected (see constraints above). + attention-DP on both sides; heterogeneous parallelism passes peer + validation but is not validated end-to-end (see constraints above). 4. **Cluster environment** (NVL72 nodes): the default is to leave `UCX_TLS` unset and let UCX pick transports (RDMA/verbs included). Two exceptions: never run V2 NIXL with a container-default diff --git a/examples/kimi_k3/disagg/ctx_config.yaml b/examples/kimi_k3/disagg/ctx_config.yaml index 4be9ff24473b..4fdb8dfd8a58 100644 --- a/examples/kimi_k3/disagg/ctx_config.yaml +++ b/examples/kimi_k3/disagg/ctx_config.yaml @@ -8,8 +8,10 @@ # GB300); a DEP8 ctx would need ~273 GiB/rank for weights alone # (extrapolated from the 1.5 TB checkpoint), which does not leave # activation headroom on GB300 (288 GiB) and cannot fit GB200 (186 GiB). -# Ctx/gen parallelism must match: heterogeneous ctx/gen TP is rejected -# for K3's replicated KDA recurrent state. +# Keep ctx/gen parallelism matched: heterogeneous ctx/gen TP passes +# peer validation (the KDA recurrent state is head-sharded with +# attention-DP off) but only matched DEP16=DEP16 is validated +# end-to-end. # # Hard K3 constraints baked in: no chunked prefill, no KV block reuse # (KDA state), tokens_per_block=64 (MLA latent layout), and diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py index ae6691ce6da6..6a6976ab6ae8 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py @@ -25,6 +25,13 @@ # 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 gates payloads that carry recurrent (mamba/KDA) state: the fallback cost scales +# with bytes, not block count, so the gate is byte-denominated. +# - min_blocks is the legacy plain-KV gate, kept so existing bounce deployments see no behavior +# change. +# For Kimi K3 the byte gate never rejects: the fixed ~433 MiB per-request KDA payload always +# clears the 2 MiB default, so arena capacity plus reservation backpressure is the effective +# admission control. _MIN_BYTES_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES" # byte gate for recurrent-state payloads _MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" # block-count gate for plain-KV payloads diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index f7952c29ad1a..c2b9a13e0b60 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -363,12 +363,15 @@ def validate_peer_compatible( The core invariant checked is *global* (TP-aggregated) state size: for a TP-sharded state, ``per_rank_bytes * mamba_tp`` is TP-invariant, so it must match between peers even when their TP - sizes differ. Models with *replicated* recurrent state (e.g. Kimi K3 - KDA, whose per-rank state is pre-scaled to full size when - attention-DP is off) violate this invariant under heterogeneous TP — - exactly the configuration where the TP-mismatch mappers would - compute shard offsets past the end of the slot, silently corrupting - the replicated state — and are therefore rejected here. + sizes differ. Kimi K3 KDA satisfies it: with attention-DP off the + cache manager head-shards the state across tp_size (matching the + model's head-sharded KDA compute), so heterogeneous ctx/gen TP + passes; under attention-DP the state is replicated and ``_mamba_tp`` + is 1 on that side. A model that kept a replicated full-size per-rank + state while reporting ``mamba_tp > 1`` would violate the invariant + under heterogeneous TP — exactly the configuration where the + TP-mismatch mappers would compute shard offsets past the end of the + slot, silently corrupting the state — and is rejected here. """ self_mlg = MambaPolicy._find_mamba_layer_group(self_page_table) peer_mlg = MambaPolicy._find_mamba_layer_group(peer_page_table) @@ -411,9 +414,9 @@ def _check_global(field: str, self_bytes: int, peer_bytes: int) -> None: f"peer {peer_bytes} bytes/rank x mamba_tp={peer_tp}. Per-rank state " "sizes are inconsistent with a TP-sharded layout across the two " "sides; either the state shape/dtype differs, or the model keeps a " - "replicated (non-TP-sharded) recurrent state (e.g. Kimi K3 KDA), " - "which supports heterogeneous ctx/gen TP only with attention-DP " - "enabled on both sides." + "replicated (non-TP-sharded) per-rank recurrent state while " + "reporting mamba_tp > 1, which supports heterogeneous ctx/gen TP " + "only with attention-DP enabled on both sides." ) _check_global( diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index fd8db9d0e9cc..5d3990627c56 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -339,7 +339,13 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: def _slice_num_bytes(self, slice: KVSlice) -> int: """Local-rank KV bytes covered by a slice (sum of num_valid_blocks * pool.slot_bytes), enough to populate - kv_cache_size and unblock the perf-metric timestamps that gate on it.""" + kv_cache_size and unblock the perf-metric timestamps that gate on it. + + Counterpart accounting: the bounce reserve sizing (bounce/impl.py block_bytes_per_group) + computes per-block bytes for the same layer groups but reads pool 0 only, while this sums + every pool view of a group. The pool-0-only sizing gap for multi-pool attention groups is + tracked under TRTLLM-15194; keep the two accountings in mind together when changing either. + """ pt = self._page_table if pt is None: return 0 diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 51e5a7ded365..8b38c36e2b45 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -17,27 +17,34 @@ Synthetic loopback coverage of the KDA recurrent-state transfer path: * KDA slot layout, mapped onto the Mamba cache-manager parametrization the - way ``_util.py`` does for ``kimi_linear``: + way ``_util.py`` does for ``kimi_linear`` (params come from + ``extract_mamba_kv_cache_params`` UNSCALED — global head counts): - short-conv slot ``[3*H*hd, W]`` **bf16** (qwen3_next ``[Q|K|V]`` 3-section layout, all sections equal width), - delta-rule slot ``[H, hd, hd]`` **fp32** (``state_size == head_dim``), - - EP pre-scale: ``num_heads``/``n_groups`` multiplied by ``tp_size`` when - attention-DP is off, so the per-rank state is full-size (replicated). + - TP semantics: the manager itself gates on the mapping + (``tp_size = 1`` under attention-DP, else ``mapping.tp_size``) and + divides ``num_heads``/``n_groups``/``conv_dim`` by it. So the per-rank + state is a head shard when attention-DP is off (matching the model's + head-sharded KDA compute in ``KimiKDARuntime``) and a full-size + replica under attention-DP. * ``test_kda_layer_group_descriptors`` checks that the V2 page table describes BOTH slots with exact byte sizes and that the matched-TP ``MambaPolicy`` descriptors tile each layer's slot bytes exactly. * ``test_kda_transfer`` performs a real single-node NIXL loopback transfer - and bitwise-compares both dtypes on the gen side. + and bitwise-compares both dtypes on the gen side: matched and + heterogeneous ctx/gen TP, attention-DP on and off. -* ``test_kda_hetero_tp_rejected`` asserts the peer-registration - guard: with replicated (pre-scaled) state, heterogeneous - ctx/gen TP (attention-DP off) is rejected instead of producing fragment - pointers outside the slot; supported layouts still validate. +* ``test_kda_hetero_tp_sharded_accepted`` asserts that peer registration + accepts heterogeneous ctx/gen TP with attention-DP off for what + production actually builds (head-sharded per-rank state, so the global + TP-aggregated size invariant holds). """ import uuid +from types import SimpleNamespace from typing import Dict, List import pytest @@ -56,6 +63,7 @@ build_page_table_from_manager, ) from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup +from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -81,23 +89,54 @@ _KDA_MASK = [False] + [True] * NUM_KDA_LAYERS _ATTN_MASK = [True] + [False] * NUM_KDA_LAYERS -# Full per-layer slot byte sizes (replicated on every rank for K3). +# Global (TP-aggregated) per-layer slot byte sizes. Per-rank slot bytes are +# these divided by tp_size when attention-DP is off (head-sharded state) and +# equal to these under attention-DP (replicated state). CONV_SLOT_BYTES = 3 * KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_W * CONV_DTYPE.itemsize SSM_SLOT_BYTES = KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize +def _kimi_linear_config(): + """Minimal ``kimi_linear`` HF-config stand-in for the production mapping. + + Just enough attributes for ``extract_mamba_kv_cache_params`` (the + ``is_kimi_linear`` route): KDA layers are 1..NUM_KDA_LAYERS+1 (1-indexed), + layer 1 is the dummy full-attention layer. + """ + kda_layers = [i + 1 for i, m in enumerate(_KDA_MASK) if m] + full_attn_layers = [i + 1 for i, m in enumerate(_KDA_MASK) if not m] + return SimpleNamespace( + model_type="kimi_linear", + linear_attn_config={ + "head_dim": KDA_HEAD_DIM, + "short_conv_kernel_size": KDA_W, + "num_heads": KDA_NUM_HEADS, + "kda_layers": kda_layers, + "full_attn_layers": full_attn_layers, + }, + num_hidden_layers=_NUM_TOTAL_LAYERS, + torch_dtype=CONV_DTYPE, + ) + + def _create_kda_managers( tp: int, enable_attention_dp: bool = False, max_batch_size: int = MAX_BATCH_SIZE ): """Create MixedMambaHybridCacheManagers with K3-style KDA slots. - Mirrors the ``is_kimi_linear`` route in ``_util.py:1855-1915``: KDA is - mapped onto (state_size=hd, conv_kernel=W+1, num_heads=H, n_groups=H, - head_dim=hd, model_type='qwen3_next'), and ``num_heads``/``n_groups`` - 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). + Mirrors the ``is_kimi_linear`` route in ``_util.py``: the mamba params + come from ``extract_mamba_kv_cache_params`` UNSCALED (global + ``num_heads``/``n_groups``); the manager applies its own TP gate + (``tp_size = 1`` under attention-DP, else ``mapping.tp_size``) and + divides ``num_heads``/``n_groups``/``conv_dim`` by it — the same + head-shard semantics the model runtime uses (``KimiKDARuntime`` + constructs the mixer with ``num_heads // tp_size``). """ - state_tp = tp if not enable_attention_dp else 1 + params = extract_mamba_kv_cache_params(_kimi_linear_config()) + assert params.mamba_layer_mask == _KDA_MASK + assert params.target_full_attention_layer_mask == _ATTN_MASK + assert params.dtype == CONV_DTYPE + assert params.mamba_ssm_cache_dtype == SSM_DTYPE # kimi_linear forces fp32 managers = [] for rank in range(tp): mapping = Mapping( @@ -108,15 +147,15 @@ def _create_kda_managers( enable_attention_dp=enable_attention_dp, ) mgr = MixedMambaHybridCacheManager( - mamba_d_state=KDA_HEAD_DIM, - mamba_d_conv=KDA_W + 1, - mamba_num_heads=KDA_NUM_HEADS * state_tp, - mamba_n_groups=KDA_NUM_HEADS * state_tp, - mamba_head_dim=KDA_HEAD_DIM, - mamba_num_layers=NUM_KDA_LAYERS, - mamba_layer_mask=_KDA_MASK, - mamba_cache_dtype=CONV_DTYPE, - mamba_ssm_cache_dtype=SSM_DTYPE, + mamba_d_state=params.state_size, + mamba_d_conv=params.conv_kernel, + mamba_num_heads=params.num_heads, + mamba_n_groups=params.n_groups, + mamba_head_dim=params.head_dim, + mamba_num_layers=params.num_mamba_layers, + mamba_layer_mask=params.mamba_layer_mask, + mamba_cache_dtype=params.dtype, + mamba_ssm_cache_dtype=params.mamba_ssm_cache_dtype, # dummy attention layer (page table scaffolding) kv_cache_config=KvCacheConfig( max_tokens=256 * max_batch_size, @@ -245,13 +284,14 @@ def test_kda_layer_group_descriptors(enable_attention_dp): mgr.shutdown() -def test_kda_hetero_tp_rejected(): - """Replicated KDA state + heterogeneous TP (ADP off) must be rejected. +def test_kda_hetero_tp_sharded_accepted(): + """Heterogeneous ctx/gen TP with ADP off passes peer validation. - Hetero ctx/gen TP would silently corrupt replicated state: the TP-mismatch - mappers assume sharded state and would compute shard offsets past the - end of K3's replicated (pre-scaled) slots. The peer-registration gate - (``MambaPolicy.validate_peer_compatible``) must reject this loudly. + With attention-DP off, production builds a head-sharded per-rank KDA + state (the manager divides num_heads/n_groups/conv_dim by tp_size, + matching ``KimiKDARuntime``'s head-sharded compute), so + ``per_rank_slot_bytes * tp`` is TP-invariant and the global-size gate in + ``MambaPolicy.validate_peer_compatible`` accepts hetero ctx/gen TP. """ ctx_mgrs = _create_kda_managers(2, enable_attention_dp=False) gen_mgrs = _create_kda_managers(4, enable_attention_dp=False) @@ -259,16 +299,22 @@ def test_kda_hetero_tp_rejected(): try: ctx_pt = build_page_table_from_manager(ctx_mgr) gen_pt = build_page_table_from_manager(gen_mgr) + + # Production per-rank slots are the global sizes divided by tp. + ctx_mlg = _get_mamba_layer_group(ctx_pt) + gen_mlg = _get_mamba_layer_group(gen_pt) + assert ctx_mlg.conv_states.slot_bytes == CONV_SLOT_BYTES // 2 + assert ctx_mlg.ssm_states.slot_bytes == SSM_SLOT_BYTES // 2 + assert gen_mlg.conv_states.slot_bytes == CONV_SLOT_BYTES // 4 + assert gen_mlg.ssm_states.slot_bytes == SSM_SLOT_BYTES // 4 + ctx_ri = RankInfo.from_kv_cache_manager("kda_ctx", ctx_mgr, device_id=0) gen_ri = RankInfo.from_kv_cache_manager("kda_gen", gen_mgr, device_id=0) - - with pytest.raises(ValueError, match="TP-aggregated"): - MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) # And via the registrar entry point used at runtime. registrar = PeerRegistrar(ctx_ri, KVRegionExtractorV1(ctx_pt)) - with pytest.raises(ValueError, match="TP-aggregated"): - registrar.register("kda_gen", 1, gen_ri) + registrar.register("kda_gen", 1, gen_ri) finally: for mgr in ctx_mgrs + gen_mgrs: mgr.shutdown() @@ -277,13 +323,17 @@ def test_kda_hetero_tp_rejected(): @pytest.mark.parametrize( "ctx_cfg,gen_cfg", [ - ((2, False), (2, False)), # matched TP, ADP off (EP pre-scaled) + ((2, False), (2, False)), # matched TP, ADP off (head-sharded) ((2, True), (4, True)), # heterogeneous DEP with ADP on both sides ], ids=["matched_tp2_adp_off", "hetero_dep_adp_on"], ) def test_kda_peer_validation_accepts_supported_shapes(ctx_cfg, gen_cfg): - """Matched-TP and ADP-on-both-sides layouts must pass peer validation.""" + """Matched-TP and ADP-on-both-sides layouts must pass peer validation. + + Hetero TP with ADP off (also supported for K3's sharded state) is + covered by ``test_kda_hetero_tp_sharded_accepted``. + """ ctx_tp, ctx_adp = ctx_cfg gen_tp, gen_adp = gen_cfg ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=ctx_adp) @@ -346,25 +396,30 @@ def _synthetic_rank_info(tp: int, adp: bool): @pytest.mark.parametrize( "ctx,gen,ok", [ - # (tp, adp, slot_scale_denominator): replicated K3 state = full slots. - ((2, False, 1), (4, False, 1), False), # hetero TP, ADP off: would corrupt state + # (tp, adp, slot_scale_denominator). denom == tp models K3's actual + # head-sharded per-rank slots (slot_bytes = full // tp when ADP off); + # denom == 1 with tp > 1 and ADP off models a hypothetical model that + # keeps a replicated full-size per-rank state while reporting + # mamba_tp > 1 — the layout the global-size gate must reject under + # hetero TP (shard offsets would land past the end of the slot). + ((2, False, 1), (4, False, 1), False), # replicated slots, hetero TP: reject ((4, False, 1), (2, False, 1), False), # ...both directions - ((2, True, 1), (2, False, 1), False), # mixed ADP, replicated - ((2, False, 1), (2, False, 1), True), # matched TP - ((2, True, 1), (4, True, 1), True), # hetero DEP, ADP on both - ((2, False, 2), (4, False, 4), True), # sharded state, hetero TP + ((2, True, 1), (2, False, 1), False), # ADP on vs sharded-claim, full slots both + ((2, False, 1), (2, False, 1), True), # matched TP: sizes agree either way + ((2, True, 1), (4, True, 1), True), # hetero DEP, ADP on both (replicated) + ((2, False, 2), (4, False, 4), True), # K3 layout: sharded slots, hetero TP ], ids=[ - "reject_hetero_tp_adp_off", - "reject_hetero_tp_adp_off_rev", - "reject_mixed_adp_replicated", + "reject_hetero_tp_replicated_slots", + "reject_hetero_tp_replicated_slots_rev", + "reject_mixed_adp_full_slots", "accept_matched_tp", "accept_hetero_dep_adp_on", "accept_sharded_hetero_tp", ], ) def test_kda_peer_validation_synthetic_cpu(ctx, gen, ok): - """CPU-only reject/accept matrix for the gap-F1 guard (no CUDA manager).""" + """CPU-only reject/accept matrix for the global-size guard (no CUDA manager).""" full_ssm = KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize full_conv = 3 * KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_W * CONV_DTYPE.itemsize @@ -426,8 +481,27 @@ def test_kda_peer_validation_allows_pipeline_parallel_layer_split(): # --------------------------------------------------------------------------- # Loopback transfer over a real NIXL agent (single node) # --------------------------------------------------------------------------- +def _shard_kda_ssm(full_ssm: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Head shard of the [H, hd, hd] delta-rule state (identity for tp=1).""" + n = KDA_NUM_HEADS // tp + return full_ssm[rank * n : (rank + 1) * n].clone() + + +def _shard_kda_conv(full_conv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Per-section shard of the [3*H*hd, W] conv state (identity for tp=1). + + qwen3_next [Q | K | V] sectioning with three equal H*hd sections; each + section is sharded independently across tp, like the manager's + conv_section_dims. + """ + sec = KDA_NUM_HEADS * KDA_HEAD_DIM + n = sec // tp + parts = [full_conv[s * sec + rank * n : s * sec + (rank + 1) * n] for s in range(3)] + return torch.cat(parts, dim=0).clone() + + def _generate_ground_truth(num_requests: int, seed: int = 20260722): - """Full replicated KDA states per (request, kda_layer), two dtypes.""" + """Full (global, unsharded) KDA states per (request, kda_layer), two dtypes.""" gen = torch.Generator(device="cpu").manual_seed(seed) results = [] for _ in range(num_requests): @@ -452,7 +526,24 @@ def _generate_ground_truth(num_requests: int, seed: int = 20260722): def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = False): - """Loopback: matched ctx/gen TP, replicated full-size KDA state per rank.""" + """Loopback KDA transfer, matched or heterogeneous ctx/gen TP. + + ADP off: per-rank state is a head shard (state_tp = tp) and every TP rank + participates in each request; the TP-mismatch mappers re-tile shards + between the two TP layouts. ADP on: per-rank state is a full-size replica + (state_tp = 1) and each request belongs to ONE dp rank per side (the + production scheduler routes it), so only that rank moves and verifies it + (same routing as ``test_kv_transfer.add_and_verify_request``). + """ + ctx_state_tp = 1 if enable_attention_dp else ctx_tp + gen_state_tp = 1 if enable_attention_dp else gen_tp + + def _ctx_ranks(req_idx: int) -> List[int]: + return [req_idx % ctx_tp] if enable_attention_dp else list(range(ctx_tp)) + + def _gen_ranks(req_idx: int) -> List[int]: + return [req_idx % gen_tp] if enable_attention_dp else list(range(gen_tp)) + ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=enable_attention_dp) gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=enable_attention_dp) ctx_tcs, gen_tcs = [], [] @@ -466,8 +557,12 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = transceiver_runtime="PYTHON", max_tokens_in_buffer=512, ) - ctx_tcs = _create_transceivers(ctx_tp, ctx_mgrs, config) - gen_tcs = _create_transceivers(gen_tp, gen_mgrs, config) + ctx_tcs = _create_transceivers( + ctx_tp, ctx_mgrs, config, enable_attention_dp=enable_attention_dp + ) + gen_tcs = _create_transceivers( + gen_tp, gen_mgrs, config, enable_attention_dp=enable_attention_dp + ) ctx_endpoint = ctx_tcs[0]._context_info_endpoint sampling_params = SamplingParams() @@ -497,7 +592,7 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = ) gen_req.py_disaggregated_params = DisaggregatedParams( ctx_request_id=ctx_rid, - ctx_dp_rank=0, + ctx_dp_rank=_ctx_ranks(req_idx)[0] if enable_attention_dp else 0, ctx_info_endpoint=ctx_endpoint, disagg_request_id=unique_rid, ) @@ -520,21 +615,29 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = for mgr in gen_mgrs: mgr.update_resources(gen_batch) - # Ground truth: identical full state on every ctx rank (replicated). + # Ground truth: each ctx rank gets its shard of the global state + # (the full state on every rank when replicated under ADP). ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) - for mgr in ctx_mgrs: + for rank, mgr in enumerate(ctx_mgrs): + # Under ADP state_tp is 1 and every rank holds shard 0 (the full + # state); with ADP off rank r holds shard r. + shard_rank = rank % ctx_state_tp for req_idx, rid in enumerate(ctx_rids): slot = mgr.mamba_cache_index[rid] for layer_idx in mgr._impl.mamba_layer_offsets: full = ground_truth[req_idx][layer_idx] - mgr.get_conv_states(layer_idx)[slot] = full["conv"] - mgr.get_ssm_states(layer_idx)[slot] = full["ssm"] - - for rank in range(gen_tp): - for req in gen_reqs: + mgr.get_conv_states(layer_idx)[slot] = _shard_kda_conv( + full["conv"], ctx_state_tp, shard_rank + ) + mgr.get_ssm_states(layer_idx)[slot] = _shard_kda_ssm( + full["ssm"], ctx_state_tp, shard_rank + ) + + for req_idx, req in enumerate(gen_reqs): + for rank in _gen_ranks(req_idx): gen_tcs[rank].request_and_receive_async(req) - for rank in range(ctx_tp): - for req in ctx_reqs: + for req_idx, req in enumerate(ctx_reqs): + for rank in _ctx_ranks(req_idx): ctx_tcs[rank].respond_and_send_async(req) _run_concurrent( ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) @@ -543,28 +646,39 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = # Transfer-size metric must include the fixed-size KDA state — the actual # transferred bytes must cover the computed payload size: per rank, - # num_layers * (conv + ssm) slot bytes on top of any KV bytes. - kda_bytes_per_rank = NUM_KDA_LAYERS * (CONV_SLOT_BYTES + SSM_SLOT_BYTES) + # num_layers * (conv + ssm) per-rank slot bytes on top of any KV bytes. + # The transceiver's _kv_size_rank_factor is 1 under attention-DP (the + # local slot already holds the full replicated state) and gen_tp + # otherwise (the request total is the sum of the per-rank shards). rank_factor = 1 if enable_attention_dp else gen_tp + kda_bytes_per_rank = NUM_KDA_LAYERS * (CONV_SLOT_BYTES + SSM_SLOT_BYTES) // gen_state_tp for req in gen_reqs: assert req.py_kv_cache_xfer_bytes >= kda_bytes_per_rank * rank_factor, ( f"kv_cache_xfer_bytes={req.py_kv_cache_xfer_bytes} misses the KDA " f"state payload ({kda_bytes_per_rank} bytes/rank x {rank_factor})" ) - # Bitwise comparison on every gen rank (state is replicated). + # Bitwise comparison of every participating gen rank's shard (the + # single owning dp rank under ADP; all TP ranks otherwise). for gen_rank, mgr in enumerate(gen_mgrs): + shard_rank = gen_rank % gen_state_tp for req_idx, rid in enumerate(gen_rids): + if gen_rank not in _gen_ranks(req_idx): + continue slot = mgr.mamba_cache_index[rid] for layer_idx in mgr._impl.mamba_layer_offsets: full = ground_truth[req_idx][layer_idx] + expected = { + "conv": _shard_kda_conv(full["conv"], gen_state_tp, shard_rank), + "ssm": _shard_kda_ssm(full["ssm"], gen_state_tp, shard_rank), + } for name, getter in ( ("conv", mgr.get_conv_states), ("ssm", mgr.get_ssm_states), ): torch.testing.assert_close( getter(layer_idx)[slot].cpu(), - full[name], + expected[name], rtol=0, atol=0, msg=lambda m, n=name, r=gen_rank, ri=req_idx, li=layer_idx: ( @@ -582,10 +696,29 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = @pytest.mark.timeout(180) @pytest.mark.parametrize( - "ctx_tp,gen_tp", - [(1, 1), (2, 2)], - ids=["tp1_tp1", "tp2_tp2_ep_prescaled"], + "ctx_tp,gen_tp,enable_attention_dp", + [ + (1, 1, False), + (2, 2, False), + (2, 2, True), + (2, 4, True), + (2, 4, False), + (4, 2, False), + ], + ids=[ + "tp1_tp1", + "matched_tp2_adp_off_sharded", + "matched_tp2_adp_on_replicated", + "hetero_dep_adp_on", + "hetero_ctx_tp2_gen_tp4_adp_off_sharded", + "hetero_ctx_tp4_gen_tp2_adp_off_sharded", + ], ) -def test_kda_transfer(ctx_tp, gen_tp): - """KDA two-dtype state transfer, matched parallelism, real NIXL loopback.""" - run_kda_transfer_test(ctx_tp, gen_tp) +def test_kda_transfer(ctx_tp, gen_tp, enable_attention_dp): + """KDA two-dtype state transfer over real NIXL loopback. + + Covers matched TP with ADP on/off, hetero DEP with ADP on, and hetero + ctx/gen TP with ADP off (head-sharded state re-tiled by the TP-mismatch + mappers, both expand and contract directions). + """ + run_kda_transfer_test(ctx_tp, gen_tp, enable_attention_dp=enable_attention_dp) diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index dd701a0bc5ba..dfdfa0addeea 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -25,6 +25,9 @@ # injects; see test_kv_transfer.py for the full rationale. os.environ["UCX_TLS"] = "^ib,gdr_copy" os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" +# Deliberately do NOT clear UCX_NET_DEVICES: cluster images pin it to a +# working interface (e.g. eth0), and without the pin UCX tries every +# interface and dies on unbindable ones (rdma_vf_rail0 IPv6 bind failure). import tensorrt_llm import tensorrt_llm.bindings @@ -166,14 +169,20 @@ def _worker(idx, item): raise err -def _create_transceivers(tp, managers, config): +def _create_transceivers(tp, managers, config, enable_attention_dp=False): shared = {"barrier": threading.Barrier(tp), "lock": threading.Lock()} results = [None] * tp errors = [None] * tp def _init(rank): try: - mapping = Mapping(world_size=tp, rank=rank, tp_size=tp, pp_size=1) + mapping = Mapping( + world_size=tp, + rank=rank, + tp_size=tp, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) dist = ThreadSafeDistributed(rank, tp, shared) results[rank] = KvCacheTransceiverV2( mapping=mapping,