[TRTLLM-15264][test] Kimi K3 disagg review fixups: KDA test geometry, gate docs, example cleanup

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

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

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
 

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(

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