From 5f81eed469ce8caa8db4c0e47e2559b1ad358640 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 5 Aug 2026 19:56:27 -0700 Subject: [PATCH 01/14] [TRTLLM-14815][feat] Support KDA/hybrid recurrent-state transfer in native 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 --- .../disaggregation/native/bounce/config.py | 54 +++++--- .../disaggregation/native/bounce/core.py | 13 +- .../disaggregation/native/bounce/impl.py | 104 ++++++++++---- .../disaggregation/native/mixers/ssm/peer.py | 130 ++++++++++++++++++ .../_torch/disaggregation/native/peer.py | 11 ++ .../_torch/disaggregation/native/transfer.py | 32 ++++- .../_torch/disaggregation/transceiver.py | 37 ++++- 7 files changed, 334 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py index cfdcd4e45557..64c218385acc 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py @@ -23,26 +23,25 @@ _MIB = 1024 * 1024 -# Test/advanced override for the block-count gate below (users only tune the bounce size). Read on -# the generation side, so set it there; unset uses the default. -_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" # the operative gate, in bytes +_MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" # legacy block-count gate -def _env_min_blocks(default: int) -> int: - """Read the gate from the env, defensively: unset or malformed falls back to the default (never +def _env_int_gate(name: str, default: int) -> int: + """Read a gate from the env, defensively: unset or malformed falls back to the default (never crashing), and the value is clamped to at least 1.""" - raw = os.environ.get(_MIN_BLOCKS_ENV) + raw = os.environ.get(name) if raw is None: return default try: value = int(raw) except ValueError: - logger.warning(f"{_MIN_BLOCKS_ENV}={raw!r} is not an integer; using default {default}") + logger.warning(f"{name}={raw!r} is not an integer; using default {default}") return default if value < 1: - logger.warning( - f"{_MIN_BLOCKS_ENV}={value} < 1; clamping to 1 (bounce always clears the gate)" - ) + logger.warning(f"{name}={value} < 1; clamping to 1 (bounce always clears the gate)") return 1 return value @@ -103,20 +102,41 @@ def fit_within_free( return capacity_bytes +# Skip bounce below this many bytes. The cost this gate guards scales with BYTES, not blocks: an +# earlier block-count gate (96 blocks, calibrated for 128-token blocks) silently skipped a 433 MiB +# Kimi-K3 transfer (67 blocks of 32 tokens) and dropped it onto the ~0.4 GB/s host-staged fallback, +# a ~1000x cliff. Break-even is small: bounce adds one gather plus one scatter copy (device-local, +# ~hundreds of GB/s) and ~0.1 ms of fixed launch/reservation overhead, while the in-place path can +# be as slow as ~0.4 GB/s 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 + + @dataclass class Config: sizing: Sizing = field(default_factory=FixedSizing) # how much memory to reserve (pluggable) chunk_mb: int = 32 # physical chunk size; a large chunk keeps the write to a single descriptor - # skip bounce below this many blocks (roughly 12k tokens at 128 per block); heuristic, tunable - min_blocks: int = 96 + # skip bounce below this many bytes (the operative gate; see DEFAULT_MIN_BYTES for the rationale) + min_bytes: int = DEFAULT_MIN_BYTES + # legacy block-count gate, kept for back-compat (TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS): both gates + # must pass, and the default of 1 makes this one vacuous so the byte gate decides + min_blocks: int = 1 -def config_from_size(size_mb: int, min_blocks: Optional[int] = None) -> Optional[Config]: +def config_from_size( + size_mb: int, min_blocks: Optional[int] = None, min_bytes: Optional[int] = None +) -> Optional[Config]: """Build a bounce config from a per-region size in MiB, or None to leave bounce off (size <= 0). - Size is both the capacity and the on/off switch. min_blocks is the gate below which a transfer - stays on the per-block path; when unset it comes from the env, else the default.""" + Size is both the capacity and the on/off switch. min_bytes (and the legacy min_blocks) is the + gate below which a transfer stays on the per-block path; when unset it comes from the env, else + the default.""" if size_mb is None or size_mb <= 0: return None if min_blocks is None: - min_blocks = _env_min_blocks(Config.min_blocks) - return Config(sizing=FixedSizing(capacity_mb=size_mb), min_blocks=min_blocks) + min_blocks = _env_int_gate(_MIN_BLOCKS_ENV, Config.min_blocks) + if min_bytes is None: + min_bytes = _env_int_gate(_MIN_BYTES_ENV, Config.min_bytes) + return Config( + sizing=FixedSizing(capacity_mb=size_mb), min_bytes=min_bytes, min_blocks=min_blocks + ) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index 1878fa90fd3a..0eb621d5d189 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -192,8 +192,17 @@ def release_send(self, slot_id) -> None: """Release a send region after its write completes.""" @abstractmethod - def reserve(self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = None) -> bool: - """Reserve a region and record its address for the senders. False falls back to per-fragment.""" + def reserve( + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = None, + extra_bytes: int = 0, + ) -> bool: + """Reserve a region and record its address for the senders. False falls back to per-fragment. + ``extra_bytes`` sizes the sender's non-paged payload (mamba/KDA recurrent state) that rides + the same coalesced write as the KV blocks.""" @abstractmethod def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index a9597b7cbb2a..0d2bc5ca0323 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -37,7 +37,7 @@ from tensorrt_llm._utils import CUASSERT from .buffer import SlotAllocator -from .config import SizingContext, fit_within_free +from .config import DEFAULT_MIN_BYTES, SizingContext, fit_within_free from .core import BounceTransport, Disposition, Settlement, TransferContext from .gather_scatter import Plan, gather_contiguous, scatter_contiguous @@ -83,6 +83,7 @@ def from_config( capacity_bytes=capacity_bytes, phys_chunk_size=chunk, block_bytes_per_group=block_bytes_per_group, + min_bytes=cfg.min_bytes, min_blocks=cfg.min_blocks, ) @@ -94,7 +95,8 @@ def __init__( capacity_bytes: int, phys_chunk_size: int, block_bytes_per_group: List[int], - min_blocks: int = 96, + min_bytes: int = DEFAULT_MIN_BYTES, + min_blocks: int = 1, quarantine_grace_s: float = _QUARANTINE_GRACE_S, name: str = "kv_bounce", ): @@ -102,8 +104,10 @@ def __init__( self._device_id = device_id # The byte size of one cache block, listed for each attention layer group. self._block_bytes_per_group = list(block_bytes_per_group) - # Below this many blocks, skip bounce: coalescing only pays off for long context (the default - # is roughly twelve thousand tokens; a heuristic, and tunable). + # Below this many bytes, skip bounce: coalescing only pays off once the transfer is large + # enough to beat the gather+scatter overhead (see config.DEFAULT_MIN_BYTES for the + # rationale). min_blocks is the legacy block-count gate, vacuous at its default of 1. + self._min_bytes = min_bytes self._min_blocks = min_blocks # how long an orphaned region is held out of reuse; must outlast the worst in-flight write self._quarantine_grace_s = quarantine_grace_s @@ -182,9 +186,10 @@ def _reserve_and_gather(self, write_meta, *, timeout): total = int(write_meta.sizes.sum()) res = self._send_alloc.reserve(total, timeout=timeout) if res is None: - logger.debug( + logger.warning_once( f"[kv-bounce] in-place: no send region space for {total // _MIB}MiB within {timeout}s " - f"(sender backpressure); falling back" + f"(sender backpressure); falling back", + key="kv-bounce-send-backpressure", ) return None slot_id, src_addr = res @@ -209,29 +214,71 @@ def release_send(self, slot_id) -> None: self._send_alloc.release(slot_id) @staticmethod - def _skip_bounce(reason: str, *, warn_key: Optional[str] = None) -> bool: + def _skip_bounce(reason: str, *, warn_key: str) -> bool: """Log why a transfer falls back to the per-fragment path and return False, so the guards - above stay one line each.""" - msg = f"[kv-bounce] in-place: {reason}" - logger.warning_once(msg, key=warn_key) if warn_key else logger.debug(msg) + above stay one line each. Every reason logs at warning once per key: silently skipping + bounce can be a ~1000x bandwidth cliff (host-staged tcp vs cuda_ipc), so the first skip per + distinct reason must be visible at the default log level.""" + logger.warning_once(f"[kv-bounce] in-place: {reason}", key=warn_key) return False def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + extra_bytes: int = 0, ) -> bool: """Reserve a region and create its state, recording the address for the senders. Returns False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total - must divide across the writers.""" - nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) - if nblocks < self._min_blocks: - return self._skip_bounce(f"{nblocks} blocks < min {self._min_blocks} (too small)") + must divide across the writers. ``extra_bytes`` is the non-paged payload the sender appends + to the same coalesced write (mamba/KDA recurrent state, sized by the receiver via + ``MambaPolicy.payload_bytes``); the region must cover it or the write would overrun into the + neighboring slot.""" total = 0 for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups): - if g >= len(self._block_bytes_per_group): - return self._skip_bounce(f"layer group {g} has no known slot size (e.g. mamba)") + if int(block_ids.size) == 0: + # Nothing to transfer for this group. Hybrid models (Kimi K3: KDA + MLA) always + # carry a trailing empty entry for the mamba layer group — its state is not paged + # and rides extra_bytes instead — so an empty group must not disable bounce. + continue + known = g < len(self._block_bytes_per_group) and self._block_bytes_per_group[g] + if not known: + return self._skip_bounce( + f"layer group {g} has blocks but no known slot size", + 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: + # Each fan-in writer appends its own recurrent-state fragments, whose sizes may differ + # per writer (PP stages hold different mamba layers), breaking the equal region split. + return self._skip_bounce( + f"fan-in across {num_writers} senders with {extra_bytes}B of recurrent state; the " + f"equal split cannot account for per-writer state fragments", + warn_key="kv-bounce-mamba-fanin", + ) + total += int(extra_bytes) if total <= 0: - return self._skip_bounce(f"computed transfer size {total} <= 0") + return self._skip_bounce( + f"computed transfer size {total} <= 0", warn_key="kv-bounce-nonpositive-size" + ) + # The size gate is expressed in BYTES: the cost it guards (falling back to the slow + # per-fragment path) scales with bytes, not blocks, and blocks vary ~20x in size across + # models. min_blocks is the legacy gate, vacuous by default. + nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) + if total < self._min_bytes: + return self._skip_bounce( + f"{total}B ({nblocks} blocks) < min {self._min_bytes}B (too small; tune " + f"TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES)", + warn_key="kv-bounce-below-min-bytes", + ) + if nblocks < self._min_blocks: + return self._skip_bounce( + f"{nblocks} blocks < min {self._min_blocks} (too small; legacy " + f"TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS gate)", + warn_key="kv-bounce-below-min-blocks", + ) if num_writers > 1 and total % num_writers != 0: return self._skip_bounce( f"fan-in {total}B across {num_writers} senders is not an even split " @@ -262,7 +309,8 @@ def reserve( res = self._recv_alloc.reserve(total, timeout=timeout) if res is None: return self._skip_bounce( - f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)" + f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)", + warn_key="kv-bounce-recv-backpressure", ) slot_id, addr = res recv_req.bounce_dst_base = addr @@ -436,7 +484,12 @@ def release_send(self, slot_id) -> None: pass def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + extra_bytes: int = 0, ) -> bool: return False @@ -526,8 +579,12 @@ def decode_result_tail(message): def block_bytes_per_group(page_table) -> list: - """Byte size of one cache block for each leading attention layer group, stopping at the first - non-attention group.""" + """Byte size of one cache block for each layer group, aligned with the layer-group indices a + recv request uses. Non-attention groups (mamba/KDA recurrent state) hold ``None``: they carry + no paged blocks (their KVSlice entry is always empty — see ``_create_kv_slice``) and their + payload is sized separately via ``MambaPolicy.payload_bytes``. Keeping them as placeholders + instead of truncating means a trailing (or hypothetically interleaved) mamba group can never + shift an attention group off the end of this list and poison the bounce gate.""" from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool @@ -535,6 +592,7 @@ def block_bytes_per_group(page_table) -> list: out: list = [] for lg_idx, lg in enumerate(page_table.layer_groups): if not isinstance(lg, AttentionLayerGroup): - break + out.append(None) + continue out.append(int(get_physical_pool(page_table, lg_idx, 0).slot_bytes)) return out diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index a96fc3e8f496..1e55deae36e9 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -335,6 +335,107 @@ class MambaPolicy: Dispatch mappers and build frags for Mamba state transfer. """ + @staticmethod + def _find_mamba_layer_group( + page_table: Optional[KVCachePageTable], + ) -> Optional[MambaLayerGroup]: + if page_table is None: + return None + return next( + (lg for lg in page_table.layer_groups if isinstance(lg, MambaLayerGroup)), + None, + ) + + @staticmethod + def validate_peer_compatible( + self_ri: RankInfo, + peer_ri: RankInfo, + self_page_table: Optional[KVCachePageTable], + peer_page_table: Optional[KVCachePageTable], + ) -> None: + """Validate recurrent-state (Mamba/KDA) layout compatibility with a peer. + + Analogue of the C++ ``rnnCacheFormatter inquireSupport`` gate for the + V2 path: reject at peer-registration time instead of corrupting + memory at transfer time. Raises ``ValueError`` naming the mismatched + field. + + 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. + """ + self_mlg = MambaPolicy._find_mamba_layer_group(self_page_table) + peer_mlg = MambaPolicy._find_mamba_layer_group(peer_page_table) + if self_mlg is None and peer_mlg is None: + return + if (self_mlg is None) != (peer_mlg is None): + raise ValueError( + "MambaPolicy.validate_peer_compatible: one side has " + f"recurrent-state pools and the other does not (local={self_mlg is not None}, " + f"peer={peer_mlg is not None})" + ) + + if set(self_mlg.mamba_layer_offsets.keys()) != set(peer_mlg.mamba_layer_offsets.keys()): + raise ValueError( + "MambaPolicy.validate_peer_compatible: mamba layer sets differ " + f"(local={sorted(self_mlg.mamba_layer_offsets)}, " + f"peer={sorted(peer_mlg.mamba_layer_offsets)})" + ) + + if ( + self_mlg.ssm_bytes_per_head is not None + and peer_mlg.ssm_bytes_per_head is not None + and self_mlg.ssm_bytes_per_head != peer_mlg.ssm_bytes_per_head + ): + # TP-invariant: head_dim * d_state * element_size. A mismatch + # means different state shape or SSM cache dtype. + raise ValueError( + "MambaPolicy.validate_peer_compatible: ssm_bytes_per_head differs " + f"(local={self_mlg.ssm_bytes_per_head}, peer={peer_mlg.ssm_bytes_per_head}); " + "check head_dim / d_state / mamba_ssm_cache_dtype" + ) + + self_tp, _ = MambaPolicy._mamba_tp(self_ri) + peer_tp, _ = MambaPolicy._mamba_tp(peer_ri) + + def _check_global(field: str, self_bytes: int, peer_bytes: int) -> None: + if self_bytes * self_tp != peer_bytes * peer_tp: + raise ValueError( + f"MambaPolicy.validate_peer_compatible: global (TP-aggregated) {field} " + f"differs: local {self_bytes} bytes/rank x mamba_tp={self_tp} vs " + 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." + ) + + _check_global( + "ssm slot_bytes", self_mlg.ssm_states.slot_bytes, peer_mlg.ssm_states.slot_bytes + ) + _check_global( + "conv slot_bytes", self_mlg.conv_states.slot_bytes, peer_mlg.conv_states.slot_bytes + ) + + if self_mlg.conv_section_bytes is not None and peer_mlg.conv_section_bytes is not None: + if len(self_mlg.conv_section_bytes) != len(peer_mlg.conv_section_bytes): + raise ValueError( + "MambaPolicy.validate_peer_compatible: conv section count differs " + f"(local={len(self_mlg.conv_section_bytes)}, " + f"peer={len(peer_mlg.conv_section_bytes)})" + ) + for i, (s, p) in enumerate( + zip(self_mlg.conv_section_bytes, peer_mlg.conv_section_bytes) + ): + _check_global(f"conv_section_bytes[{i}]", s, p) + @staticmethod def _mamba_tp(ri: RankInfo) -> Tuple[int, int]: """Return (mamba_effective_tp_size, mamba_effective_tp_rank). @@ -492,6 +593,35 @@ def build_mamba_frags( return src_frags, dst_frags, kv_sizes + @staticmethod + def payload_bytes( + sender_page_table: KVCachePageTable, + receiver_page_table: KVCachePageTable, + dst_slot: Optional[int], + sender_ri: RankInfo, + receiver_ri: RankInfo, + ) -> int: + """Total recurrent-state bytes one sender appends to the KV write for a request. + + Mirrors the sender's ``collect_frags`` call in ``_build_kv_write_meta`` with the same + argument roles (sender as self/src), so the receiver can size a bounce region for the + exact bytes the coalesced write will carry. Slot indices only shift pointers, never + fragment sizes, so a dummy ``src_slot`` stands in for the sender's state index (which + the receiver does not know at reserve time). Returns 0 when either side has no mamba + layer group or the receiver has no state slot for the request. + """ + if dst_slot is None: + return 0 + _, _, sizes = MambaPolicy.collect_frags( + self_page_table=sender_page_table, + peer_page_table=receiver_page_table, + src_slot=0, + dst_slot=dst_slot, + self_ri=sender_ri, + peer_ri=receiver_ri, + ) + return int(sum(sizes)) + @staticmethod def collect_frags( self_page_table: KVCachePageTable, diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 6b8f40232f1d..485432787d0d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( @@ -152,6 +153,16 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: if not self._attention_policy.check_peer_compatible(peer_ri): return False + # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a + # field-level diagnostic instead of returning False, so the precise + # mismatch reaches the caller of register(). + MambaPolicy.validate_peer_compatible( + self._ri, + peer_ri, + self._self_ext_cache.page_table if self._self_ext_cache is not None else None, + peer_ri.page_table, + ) + self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) if self_layers != peer_layers: diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 4f7fc7edf19e..673d088e8443 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1659,7 +1659,24 @@ def dispatch_task(self, task: KVRecvTask): allow_bounce = task.expected_transfers == 1 or ( sender_dp_rank is not None and self._fanin_bounce_safe(topo_overlap, peer_infos) ) - bounced = allow_bounce and self._bounce.reserve(receiver_req, task.expected_transfers) + # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks (the sender + # appends its MambaPolicy fragments in _build_kv_write_meta), so the bounce region must be + # sized for those bytes too or the write would overrun into the neighboring slot. + extra_bytes = 0 + if receiver_req.mamba_state_index is not None: + if peer_infos.page_table is None: + allow_bounce = False # cannot size the sender's recurrent-state payload + else: + extra_bytes = MambaPolicy.payload_bytes( + sender_page_table=peer_infos.page_table, + receiver_page_table=self._registrar.self_extractor.page_table, + dst_slot=receiver_req.mamba_state_index, + sender_ri=peer_infos, + receiver_ri=self._registrar.self_rank_info, + ) + bounced = allow_bounce and self._bounce.reserve( + receiver_req, task.expected_transfers, extra_bytes=extra_bytes + ) session = self._get_session(task._unique_rid) if session is None: raise RuntimeError( @@ -1716,6 +1733,19 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() + # Recurrent-state (Mamba/KDA) layout gate on the receiver side. + # The sender-side check (PeerRegistrar.register) runs in the + # sender's listener thread, where exceptions are only logged, so + # 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( + self._registrar.self_rank_info, + sender_info, + self._registrar.self_extractor.page_table, + sender_info.page_table, + ) + for endpoint in sender_info.sender_endpoints: dealer = self._get_or_connect_dealer(endpoint) rank_info = self._registrar.self_rank_info diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 6968159a90d7..d065534396a0 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -94,7 +94,18 @@ def __init__( self._device_id = torch.cuda.current_device() logger.info(f"device_id: {self._device_id} in KvCacheTransceiverV2") + # Setup interleaves MPI collectives (broadcast/allgather below) with + # per-rank native NIXL/UCX initialization inside TransferWorker. A rank + # that blocks or dies in the native phase leaves its peers stuck in the + # next collective, so log each phase per rank to make the blocking + # rank/phase identifiable from the logs (see the 'setup:' lines). + rank = self._dist.rank + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} broadcast instance name (collective)") self._instance_name = self._broadcast_instance_name() + logger.info( + f"KvCacheTransceiverV2 setup: rank={rank} creating TransferWorker " + "(native NIXL agent init + KV memory registration)" + ) self._transfer_worker = TransferWorker( TransferWorkerConfig( kv_cache_manager=kv_cache_manager, @@ -106,14 +117,21 @@ def __init__( max_concurrent_sessions=max(1, int(kv_cache_manager.max_batch_size)) * 20000, tx_timeout_s=self._sender_future_timeout_ms / 1000.0, rx_timeout_s=self.kv_transfer_timeout_ms / 1000.0, - # Size 0 turns bounce off; the block-count gate is internal (tuned via env). + # Size 0 turns bounce off; the byte-size gate is internal (tuned via env: + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES, plus the legacy ..._MIN_BLOCKS). bounce=bounce_config_from_size(cache_transceiver_config.kv_cache_bounce_size_mb), ) ) + logger.info( + f"KvCacheTransceiverV2 setup: rank={rank} TransferWorker ready; " + "broadcast context endpoint (collective)" + ) self._dp_rank = mapping.tp_rank if mapping.enable_attention_dp else 0 self._context_info_endpoint = self._broadcast_context_endpoint() self._init_sync_policy() + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} exchange rank info (collective)") self._exchange_rank_info() + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} complete") self._send_sessions: Dict[int, TxSessionBase] = {} self._recv_sessions: Dict[int, RxSessionBase] = {} @@ -317,10 +335,21 @@ def _slice_num_bytes(self, slice: KVSlice) -> int: return 0 total = 0 for lg_id, block_ids in enumerate(slice.block_ids_per_layer_groups): - if block_ids is None or block_ids.size == 0: - continue lg = pt.layer_groups[lg_id] if isinstance(lg, MambaLayerGroup): + # Fixed-size recurrent state (mamba/KDA): one slot per layer in + # each of the conv and ssm pools, independent of token count. + # For hybrid models (e.g. Kimi K3) this blob can dominate + # short-prompt transfers, so it must be counted. The caller's + # tp_size scaling then yields total bytes moved across ranks + # (exact for sharded state; for replicated state every rank + # pair moves a full copy, so it matches bytes on the wire). + if slice.mamba_state_index is not None: + total += len(lg.mamba_layer_offsets) * ( + lg.conv_states.slot_bytes + lg.ssm_states.slot_bytes + ) + continue + if block_ids is None or block_ids.size == 0: continue n = int((block_ids >= 0).sum()) if n == 0: From d681f2eef4f14ee903a3de39b5d56fc44d447934 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 5 Aug 2026 19:56:55 -0700 Subject: [PATCH 02/14] [TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3 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 --- tensorrt_llm/_torch/pyexecutor/_util.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dc7b753a89ef..4c04d7a9e507 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -163,14 +163,11 @@ def get_kv_cache_manager_cls( # Mixed manager (separate KV / recurrent-state pools) stays the # default. SA speculative decoding is validated on the Mixed # manager's SpeculativeState scratch path only; reuse + SA is - # unvalidated. - if is_kimi_linear(config) and not use_v2: - if is_disagg: - # Fail fast instead of bypassing the disagg transceiver - # validation below with an unvalidated route. - raise NotImplementedError( - "Disaggregated serving is not supported for Kimi K3 yet " - "(TRTLLM-14815).") + # unvalidated. Disaggregated serving (TRTLLM-14815) routes through + # 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: if kv_cache_config.enable_block_reuse: logger.info( "Using CppMambaHybridCacheManager for Kimi K3 hybrid " @@ -591,6 +588,10 @@ def _get_model_kv_cache_manager_cls( cache_transceiver_config=self._cache_transceiver_config) cls = self._fallback_if_unsupported_kv_cache_manager_v2( cls, model_config, kv_cache_config) + if is_hybrid_linear(model_config.pretrained_config): + logger.info_once( + f"Selected hybrid KV cache manager: {cls.__name__}", + key=f"hybrid_kv_cache_manager_{cls.__name__}") # Compatibility managers do not support MTP block reuse. Warn at the # routing site so users see the concrete manager selected for the # incompatible combination. From b00d33bd5caad4d561d9909ee86b48176f06b694 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 5 Aug 2026 20:01:10 -0700 Subject: [PATCH 03/14] [TRTLLM-14815][test] Add KDA/hybrid disaggregation unit tests and Kimi 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 --- ...lap_transceiver_runtime_python_bounce.yaml | 7 +- .../integration/defs/kimi_k3_disagg_parity.py | 750 ++++++++++++++++++ .../unittest/disaggregated/region/test_aux.py | 59 +- tests/unittest/disaggregated/test_bounce.py | 230 +++++- .../disaggregated/test_kda_mamba_transfer.py | 543 +++++++++++++ 5 files changed, 1570 insertions(+), 19 deletions(-) create mode 100644 tests/integration/defs/kimi_k3_disagg_parity.py create mode 100644 tests/unittest/disaggregated/test_kda_mamba_transfer.py diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml index c6573753cbea..4c0bfee7e5d9 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml @@ -3,9 +3,10 @@ # Same as disagg_config_overlap_transceiver_runtime_python.yaml, plus the bounce switch on both the # context (sender) and generation (receiver) cache_transceiver_config: # kv_cache_bounce_size_mb: >0 turns bounce on and sizes the per-region fabric-VMM arena. -# The block-count gate below which a transfer keeps the per-block path is lowered to 1 via the -# TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS env (set by the test) so the ordinary short test prompts still -# take the coalesced-bounce WRITE path (the production default of 96 would need a ~2k-token prompt). +# The byte gate below which a transfer keeps the per-block path is lowered to 1 via the +# TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES env (set by the test) so the ordinary short test prompts still +# take the coalesced-bounce WRITE path (the production default of 2 MiB may exceed a short prompt's +# KV footprint on a tiny model). # GB200/GB300 only, since the bounce arena is fabric (MNNVL) VMM memory. model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost diff --git a/tests/integration/defs/kimi_k3_disagg_parity.py b/tests/integration/defs/kimi_k3_disagg_parity.py new file mode 100644 index 000000000000..ae5759e473d9 --- /dev/null +++ b/tests/integration/defs/kimi_k3_disagg_parity.py @@ -0,0 +1,750 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 two-endpoint parity harness (aggregated vs disaggregated). + +Compares a REFERENCE deployment against a CANDIDATE deployment through +their OpenAI-compatible ``/v1/completions`` endpoints and produces a +parity report: + +* per-prompt first-token agreement (in disagg the first token is produced + by the ctx server and handed off — a mismatch points at the + KV/KDA-state transfer or the first-token handoff), +* longest-common-prefix (LCP) length of the greedy token streams, +* logprob drift along the shared prefix + near-tie classification at the + first divergence (reuses ``_compare_logits_parity`` from + ``kimi_k3_sa_harness.py`` — same tolerances, same aggregate + "non-ties must not dominate" rule), +* optional GSM8K accuracy diff (runs ``lm_eval --model local-completions`` + against both endpoints, the same flow as + ``examples/disaggregated/slurm/benchmark/submit.py``). + +The harness only needs HTTP access to the two endpoints — no GPUs, no +tensorrt_llm runtime — so it runs from a login node against servers +launched elsewhere. + +DRY-RUN RECIPE (two aggregated servers; debugs the harness before disagg +transfer works — expect near-perfect parity): + + # server A (reference) and server B (candidate): identical aggregated + # deployments of the same checkpoint on two ports, e.g. + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8000 \ + --config examples/kimi_k3/eval_extra_llm_options.yaml & + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8001 \ + --config examples/kimi_k3/eval_extra_llm_options.yaml & + + python tests/integration/defs/kimi_k3_disagg_parity.py \ + --reference http://localhost:8000 --candidate http://localhost:8001 \ + --extra-prompts 16 --report-json parity_dryrun.json + +REAL RECIPE (aggregated vs disagg proxy): + + # reference: aggregated DEP16 deployment (as in + # examples/kimi_k3/run_gsm8k_kimi_k3.sbatch, but served) + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8000 ... + + # candidate: ctx + gen workers behind the disagg proxy + trtllm-serve disaggregated -c disagg_config.yaml # port 9000 + + python tests/integration/defs/kimi_k3_disagg_parity.py \ + --reference http://ref-host:8000 --candidate http://proxy-host:9000 \ + --extra-prompts 48 --gsm8k --gsm8k-limit 200 \ + --report-json parity_disagg.json + +Notes: +* Requests are sent sequentially (batch of 1) to both endpoints so batch + composition cannot perturb the comparison. +* Per-token logprobs are requested via the completions ``logprobs`` field + (PyTorch backend supports top-k dicts). Endpoints that reject + ``detokenize=false`` or ``logprobs`` degrade gracefully: the harness + falls back to decoded-token-string comparison, then to text-only, and + reports which capability level each endpoint provided. One-engine + spec-dec samplers (SA) do not emit per-token logprobs — run the parity + probe with spec dec off first, or accept token-level-only + parity for SA-on runs. +* GSM8K mode shells out to ``lm_eval`` (must be installed, e.g. + ``pip install lm-eval[api]``) and diffs the exact_match metrics between + the two endpoints (default tolerance 0.02). + +Self-test (no servers, canned responses — exercises the comparison and +classification logic): + + python tests/integration/defs/kimi_k3_disagg_parity.py --self-test + +Exit code 0 = PASS, 1 = FAIL (any parity failure, endpoint error, or +GSM8K delta above tolerance). +""" + +import argparse +import glob +import json +import os +import subprocess +import sys +import tempfile +import urllib.error +import urllib.request +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from kimi_k3_sa_harness import ( # noqa: E402 + PROMPTS_AND_CHECKS, + _compare_logits_parity, + _parity_prompts, +) + +# Request-payload variants, tried in order per endpoint until one is +# accepted (capability is cached per endpoint afterwards): +# ids+logprobs : token_ids (detokenize=false) + per-token logprobs +# str+logprobs : decoded token strings + per-token logprobs +# text-only : plain text completion (weakest comparison) +_VARIANTS = ("ids+logprobs", "str+logprobs", "text-only") + + +def _http_json(url, payload=None, timeout=300): + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode() + headers["Content-Type"] = "application/json" + req = urllib.request.Request(url, data=data, headers=headers) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read()) + + +def _served_model(base_url, timeout): + """Model id from /v1/models, or None if the endpoint does not expose it. + + The `trtllm-serve disaggregated` proxy only implements the completions + routes (no /v1/models -> HTTP 404), so a missing model listing must not + be fatal: the caller falls back to the other endpoint's model name or + to --model. (Dry runs against two aggregated servers do not hit this: + both expose /v1/models.) + """ + try: + models = _http_json(f"{base_url}/v1/models", timeout=timeout) + return models["data"][0]["id"] + except urllib.error.HTTPError as e: + print(f"[parity] NOTE: {base_url}/v1/models unavailable (HTTP {e.code})") + return None + + +class Endpoint: + """One OpenAI-compatible completions endpoint with capability fallback.""" + + def __init__(self, base_url, model, max_tokens, logprobs_k, timeout): + self.base_url = base_url.rstrip("/") + self.model = model + self.max_tokens = max_tokens + self.logprobs_k = logprobs_k + self.timeout = timeout + self.variant = None # discovered on first successful request + + def _payload(self, prompt, variant): + payload = { + "model": self.model, + "prompt": prompt, + "max_tokens": self.max_tokens, + "temperature": 0.0, + "stream": False, + } + if variant != "text-only" and self.logprobs_k > 0: + payload["logprobs"] = self.logprobs_k + if variant == "ids+logprobs": + payload["detokenize"] = False + return payload + + def complete(self, prompt): + """Return a normalized completion record for one prompt.""" + variants = [self.variant] if self.variant else list(_VARIANTS) + if self.logprobs_k <= 0: + variants = ["text-only"] + last_error = None + for variant in variants: + try: + rsp = _http_json( + f"{self.base_url}/v1/completions", + self._payload(prompt, variant), + timeout=self.timeout, + ) + except urllib.error.HTTPError as e: + # 4xx = capability rejection -> try the next variant; + # anything else is a real endpoint failure. + if 400 <= e.code < 500 and variant != variants[-1]: + last_error = e + continue + raise + if self.variant is None: + self.variant = variant + if variant != _VARIANTS[0]: + print( + f"[parity] NOTE: {self.base_url} degraded to " + f"'{variant}' requests ({last_error})" + ) + return _normalize_choice(rsp["choices"][0]) + raise last_error + + +def _normalize_choice(choice): + """Normalize a CompletionResponseChoice dict for comparison. + + Returns a namespace with: + text str ('' when detokenize=false) + ids per-position token identities: ints (token_ids), + else decoded token strings, else None + token_logprobs chosen-token logprob per position, or None + top_logprobs per-position {token_str: logprob} dict list, or None + """ + lp = choice.get("logprobs") or {} + tokens = lp.get("tokens") or None + token_logprobs = lp.get("token_logprobs") or None + top_logprobs = lp.get("top_logprobs") or None + ids = choice.get("token_ids") or tokens + return SimpleNamespace( + text=choice.get("text", ""), + ids=ids, + tokens=tokens, + token_logprobs=token_logprobs, + top_logprobs=top_logprobs, + ) + + +def _position_logprob_dict(chosen_key, chosen_lp, top_map): + """Build a per-position {token: logprob} dict for _compare_logits_parity. + + ``top_map`` (decoded-token-string keyed, from CompletionLogProbs) + includes the chosen token itself; when re-keying the chosen entry + under ``chosen_key`` (a token id), drop exactly one entry carrying + the chosen logprob so the top-2 near-tie gap is not computed against + a duplicate of the winner. Alternate entries get synthetic keys — + only the chosen key is ever looked up. + """ + d = {} + if top_map: + remaining = list(top_map.items()) + for i, (_, v) in enumerate(remaining): + if v == chosen_lp: + del remaining[i] + break + for j, (_, v) in enumerate(remaining): + d[("alt", j)] = v + d[chosen_key] = chosen_lp + return d + + +def _to_parity_namespace(comp): + """Adapt a normalized completion to _compare_logits_parity's shape. + + Returns a namespace with token_ids + per-position logprob dicts, or + None when the endpoint returned no aligned per-token logprobs. + """ + if comp.ids is None or comp.token_logprobs is None or len(comp.token_logprobs) != len(comp.ids): + return None + logprobs = [] + for i, (tid, lp) in enumerate(zip(comp.ids, comp.token_logprobs)): + top = comp.top_logprobs[i] if comp.top_logprobs and i < len(comp.top_logprobs) else None + logprobs.append(_position_logprob_dict(tid, lp, top)) + return SimpleNamespace(token_ids=list(comp.ids), logprobs=logprobs) + + +def _lcp(a, b): + n = 0 + while n < min(len(a), len(b)) and a[n] == b[n]: + n += 1 + return n + + +def compare_pair(ref, cand, prompt, failures, lp_tol, tie_tol): + """Compare one prompt's completions; returns the per-prompt record. + + Classification (reference logprobs available): + identical / benign (near-tie flip) / non_tie / drift — exactly the + kimi_k3_sa_harness semantics. Without logprobs, only token-level + facts are recorded (outcome 'divergence_unclassified'). + """ + ref_ids = list(ref.ids) if ref.ids is not None else None + cand_ids = list(cand.ids) if cand.ids is not None else None + record = {"prompt": prompt} + if ref_ids is None or cand_ids is None: + # Text-only endpoints: character-level LCP is the best available. + match = ref.text == cand.text + lcp = _lcp(ref.text, cand.text) + record.update( + comparison="text", + text_match=match, + char_lcp=lcp, + outcome="identical" if match else "divergence_unclassified", + ) + if not match: + print( + f"[parity] {prompt!r}: text mismatch at char {lcp} " + "(token/logprob detail unavailable)" + ) + return record + + lcp = _lcp(ref_ids, cand_ids) + full = lcp == len(ref_ids) == len(cand_ids) + record.update( + comparison="tokens", + first_token_match=lcp > 0, + lcp=lcp, + ref_len=len(ref_ids), + cand_len=len(cand_ids), + full_match=full, + ) + + ref_ns = _to_parity_namespace(ref) + cand_ns = _to_parity_namespace(cand) + if ref_ns is not None and cand_ns is not None: + outcome = _compare_logits_parity( + ref_ns, cand_ns, prompt, failures, tol=lp_tol, tie_tol=tie_tol + ) + # Mean |delta| of chosen-token logprobs over the shared prefix — + # a drift trend indicator below the hard tolerance. + deltas = [abs(ref.token_logprobs[i] - cand.token_logprobs[i]) for i in range(lcp)] + if deltas: + record["mean_abs_lp_delta"] = sum(deltas) / len(deltas) + record["max_abs_lp_delta"] = max(deltas) + elif full: + outcome = "identical" + else: + outcome = "divergence_unclassified" + print( + f"[parity] {prompt!r}: token divergence at position {lcp} " + "(logprobs unavailable — cannot classify near-tie vs real)" + ) + record["outcome"] = outcome + + if lcp == 0: + # In disagg the first token comes from the ctx server via the + # handoff; a confident first-token flip is the KV/KDA-transfer + # bug signature. Only a certified near-tie is excusable. + if outcome == "benign": + print( + f"[parity] WARNING: first-token near-tie flip for " + f"{prompt!r} (benign, but watch the aggregate)" + ) + else: + failures.append( + f"first-token mismatch for {prompt!r} " + f"(outcome={outcome}; disagg first-token handoff or " + "KV/state transfer suspect)" + ) + return record + + +def _aggregate(outcomes, failures): + """Apply the kimi_k3_sa_harness systemic check. + + Scattered non-tie flips are rounding; a real transfer/state bug makes + them dominate. + """ + divergences = [o for o in outcomes if o in ("benign", "non_tie")] + non_ties = outcomes.count("non_tie") + unclassified = outcomes.count("divergence_unclassified") + print( + f"[parity] summary: {len(outcomes)} prompts, " + f"{outcomes.count('identical')} identical, " + f"{len(divergences)} classified divergences ({non_ties} non-tie), " + f"{unclassified} unclassified, {outcomes.count('drift')} drift" + ) + if non_ties >= 2 and non_ties > 0.25 * max(len(divergences), 1): + failures.append( + f"non-tie divergences dominate: {non_ties}/{len(divergences)} " + "divergences exceeded the tie bound (systemic — suspect " + "transfer/state bug)" + ) + + +def run_gsm8k(base_url, model, out_dir, limit, concurrency, timeout, trust_remote_code=False): + """Run lm_eval GSM8K against one endpoint; returns {metric: value}. + + Mirrors the accuracy flow of examples/disaggregated/slurm/benchmark/ + submit.py (local-completions against /v1/completions). ``model`` must + be tokenizer-resolvable for lm_eval (a local checkpoint dir or HF repo + id) — the bare served-model name from /v1/models usually is not; pass + --tokenizer in that case. + """ + os.makedirs(out_dir, exist_ok=True) + model_args = ( + f"model={model},base_url={base_url}/v1/completions," + f"num_concurrent={concurrency},max_retries=3," + f"tokenized_requests=false,timeout={timeout}," + "max_gen_toks=256,max_length=4096" + ) + if trust_remote_code: + # e.g. Kimi K3's tiktoken-based tokenizer ships as checkpoint code + model_args += ",trust_remote_code=true" + cmd = [ + "lm_eval", + "--model", + "local-completions", + "--tasks", + "gsm8k", + "--model_args", + model_args, + "--log_samples", + "--output_path", + out_dir, + ] + if limit: + cmd += ["--limit", str(limit)] + print(f"[parity] running: {' '.join(cmd)}") + subprocess.run(cmd, check=True) + results = sorted( + glob.glob(os.path.join(out_dir, "**", "results*.json"), recursive=True), + key=os.path.getmtime, + ) + if not results: + raise RuntimeError(f"lm_eval produced no results json in {out_dir}") + with open(results[-1]) as f: + gsm8k = json.load(f)["results"]["gsm8k"] + return {k: v for k, v in gsm8k.items() if k.startswith("exact_match") and "stderr" not in k} + + +def _diff_gsm8k(ref_scores, cand_scores, tol, failures): + report = {} + for metric, ref_val in sorted(ref_scores.items()): + cand_val = cand_scores.get(metric) + delta = None if cand_val is None else cand_val - ref_val + report[metric] = {"reference": ref_val, "candidate": cand_val, "delta": delta} + print( + f"[parity] gsm8k {metric}: reference={ref_val:.4f} " + f"candidate={cand_val:.4f} delta={delta:+.4f}" + if cand_val is not None + else f"[parity] gsm8k {metric}: candidate missing" + ) + if delta is None or abs(delta) > tol: + failures.append( + f"gsm8k {metric} delta {delta} exceeds tolerance {tol} " + f"(reference {ref_val} vs candidate {cand_val})" + ) + return report + + +def _parse_args(argv): + parser = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--reference", + help="base URL of the reference deployment (e.g. aggregated: http://host:8000)", + ) + parser.add_argument( + "--candidate", + help="base URL of the candidate deployment (e.g. disagg proxy: http://host:9000)", + ) + parser.add_argument( + "--model", default=None, help="served model name (default: query /v1/models)" + ) + parser.add_argument( + "--tokenizer", + default=None, + help="tokenizer path/repo for lm_eval GSM8K (a local checkpoint " + "dir; default: the --model value, which must then be " + "tokenizer-resolvable)", + ) + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument( + "--extra-prompts", + type=int, + default=16, + help="extra deterministic parity prompts on top of " + "the builtin set (kimi_k3_sa_harness pool)", + ) + parser.add_argument( + "--prompts-file", + default=None, + help="newline-separated prompt file replacing the builtin prompt set", + ) + parser.add_argument( + "--logprobs", + type=int, + default=5, + help="top-k logprobs to request per token (0 disables logprob-level parity)", + ) + parser.add_argument( + "--lp-tol", + type=float, + default=float(os.environ.get("KIMI_K3_SPEC_LP_TOL", "1.0")), + help="shared-prefix logprob drift tolerance", + ) + parser.add_argument( + "--tie-tol", + type=float, + default=float(os.environ.get("KIMI_K3_SPEC_TIE_TOL", "0.3")), + help="near-tie top-2 logprob gap bound at divergence", + ) + parser.add_argument( + "--tie-strict", action="store_true", help="fail (not warn) on any single non-tie divergence" + ) + parser.add_argument("--timeout", type=int, default=600, help="per-request timeout in seconds") + parser.add_argument( + "--gsm8k", + action="store_true", + help="also run lm_eval GSM8K against both endpoints and diff scores", + ) + parser.add_argument( + "--gsm8k-only", action="store_true", help="skip the token/logprob probe; GSM8K diff only" + ) + parser.add_argument( + "--gsm8k-limit", type=int, default=None, help="lm_eval --limit (default: full GSM8K)" + ) + parser.add_argument("--gsm8k-concurrency", type=int, default=8) + parser.add_argument( + "--gsm8k-trust-remote-code", + action="store_true", + help="pass trust_remote_code=true to lm_eval (checkpoints whose " + "tokenizer ships as remote code, e.g. Kimi K3)", + ) + parser.add_argument( + "--gsm8k-tol", type=float, default=0.02, help="max tolerated |accuracy delta|" + ) + parser.add_argument( + "--work-dir", default=None, help="directory for lm_eval outputs (default: temp dir)" + ) + parser.add_argument( + "--report-json", default=None, help="write the full parity report to this path" + ) + parser.add_argument( + "--self-test", + action="store_true", + help="run the built-in comparison-logic self-test (no servers needed)", + ) + args = parser.parse_args(argv) + if not args.self_test and not (args.reference and args.candidate): + parser.error("--reference and --candidate are required (unless --self-test)") + return args + + +def main(argv=None) -> int: + args = _parse_args(argv) + if args.self_test: + return _self_test() + if args.tie_strict: + os.environ["KIMI_K3_SPEC_TIE_STRICT"] = "1" + + failures = [] + report = {"reference": args.reference, "candidate": args.candidate} + + if args.model: + model = args.model + else: + model = _served_model(args.reference, args.timeout) + cand_model = _served_model(args.candidate, args.timeout) + # Either side may lack /v1/models (e.g. the disagg proxy); fall + # back to the side that reports one. + model = model if model is not None else cand_model + if model is None: + print("[parity] ERROR: neither endpoint exposes /v1/models; pass --model explicitly") + return 1 + if cand_model is not None and cand_model != model: + print( + f"[parity] NOTE: endpoints serve different model names " + f"({model!r} vs {cand_model!r}); using each server's own" + ) + report["model"] = model + + if not args.gsm8k_only: + if args.prompts_file: + with open(args.prompts_file) as f: + prompts = [line.rstrip("\n") for line in f if line.strip()] + else: + prompts = [p for p, _ in PROMPTS_AND_CHECKS] + [ + p for p, _ in _parity_prompts(args.extra_prompts) + ] + print( + f"[parity] probing {len(prompts)} prompts x " + f"{args.max_tokens} max_tokens, logprobs={args.logprobs}" + ) + + endpoints = {} + for role, url in (("reference", args.reference), ("candidate", args.candidate)): + endpoints[role] = Endpoint(url, model, args.max_tokens, args.logprobs, args.timeout) + + records, outcomes = [], [] + for prompt in prompts: + try: + ref = endpoints["reference"].complete(prompt) + cand = endpoints["candidate"].complete(prompt) + except (urllib.error.URLError, OSError, KeyError) as e: + failures.append(f"endpoint error for {prompt!r}: {e}") + records.append({"prompt": prompt, "error": str(e)}) + continue + record = compare_pair(ref, cand, prompt, failures, args.lp_tol, args.tie_tol) + records.append(record) + outcomes.append(record["outcome"]) + _aggregate(outcomes, failures) + report["prompts"] = records + report["capability"] = {role: ep.variant for role, ep in endpoints.items()} + + if args.gsm8k or args.gsm8k_only: + work_dir = args.work_dir or tempfile.mkdtemp(prefix="k3-parity-") + scores = {} + for role, url in (("reference", args.reference), ("candidate", args.candidate)): + scores[role] = run_gsm8k( + url, + args.tokenizer or model, + os.path.join(work_dir, f"gsm8k_{role}"), + args.gsm8k_limit, + args.gsm8k_concurrency, + args.timeout, + args.gsm8k_trust_remote_code, + ) + report["gsm8k"] = _diff_gsm8k( + scores["reference"], scores["candidate"], args.gsm8k_tol, failures + ) + + if args.report_json: + with open(args.report_json, "w") as f: + json.dump(report, f, indent=2, default=str) + print(f"[parity] report written to {args.report_json}") + + if failures: + print("[parity] FAIL") + for f in failures: + print(f" - {f}") + return 1 + print("[parity] PASS") + return 0 + + +# --------------------------------------------------------------------------- +# Self-test: canned completions through the real comparison pipeline. +# --------------------------------------------------------------------------- + + +def _fake(ids, lps, top=None, text=""): + return SimpleNamespace(text=text, ids=ids, tokens=None, token_logprobs=lps, top_logprobs=top) + + +def _self_test() -> int: + lp_tol, tie_tol = 1.0, 0.3 + checks = [] + + def check(name, cond): + checks.append((name, cond)) + print(f"[self-test] {'PASS' if cond else 'FAIL'}: {name}") + + top_confident = [{"a": -0.05, "b": -3.5}] * 4 + top_tie = [{"a": -0.60, "b": -0.72}] * 4 + + # 1. identical streams + logprobs -> identical, no failures. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + "identical", + f, + lp_tol, + tie_tol, + ) + check( + "identical outcome", + r["outcome"] == "identical" and not f and r["full_match"] and r["lcp"] == 3, + ) + + # 2. near-tie divergence mid-stream -> benign, no failures. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.6], top_tie), + _fake([1, 2, 4], [-0.1, -0.2, -0.7], top_tie), + "near-tie", + f, + lp_tol, + tie_tol, + ) + check( + "near-tie -> benign", + r["outcome"] == "benign" and not f and r["lcp"] == 2 and r["first_token_match"], + ) + + # 3. logprob drift on the shared prefix -> drift + failure. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + _fake([1, 2, 3], [-0.1, -5.2, -0.3], top_confident), + "drift", + f, + lp_tol, + tie_tol, + ) + check("drift fails", r["outcome"] == "drift" and len(f) == 1) + + # 4. confident (non-tie) divergence mid-stream -> warning only. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.05], top_confident), + _fake([1, 2, 4], [-0.1, -0.2, -3.5], top_confident), + "non-tie", + f, + lp_tol, + tie_tol, + ) + check("single non-tie warns", r["outcome"] == "non_tie" and not f) + + # 5. confident first-token mismatch -> hard failure. + f = [] + r = compare_pair( + _fake([1, 2], [-0.05, -0.2], top_confident), + _fake([9, 2], [-3.5, -0.2], top_confident), + "first-token", + f, + lp_tol, + tie_tol, + ) + check("first-token mismatch fails", r["lcp"] == 0 and len(f) >= 1) + + # 6. no logprobs available -> unclassified divergence, first-token + # mismatch still fails (cannot be excused without logprobs). + f = [] + r = compare_pair(_fake([5, 6], None), _fake([7, 6], None), "no-logprobs", f, lp_tol, tie_tol) + check("logprob-less mismatch fails", r["outcome"] == "divergence_unclassified" and len(f) == 1) + + # 7. text-only endpoints -> character comparison. + f = [] + r = compare_pair( + _fake(None, None, text="hello world"), + _fake(None, None, text="hello there"), + "text-only", + f, + lp_tol, + tie_tol, + ) + check( + "text-only unclassified", + r["comparison"] == "text" and r["outcome"] == "divergence_unclassified", + ) + + # 8. aggregate rule: dominating non-ties -> systemic failure. + f = [] + _aggregate(["non_tie", "non_tie", "benign"], f) + check("dominating non-ties fail", len(f) == 1) + f = [] + _aggregate(["benign"] * 8 + ["non_tie", "identical"], f) + check("scattered non-tie passes", not f) + + # 9. chosen-token duplicate is dropped from the near-tie gap. + d = _position_logprob_dict(42, -0.05, {"a": -0.05, "b": -3.5}) + top = sorted(d.values(), reverse=True) + check( + "top-2 gap excludes chosen duplicate", len(d) == 2 and abs((top[0] - top[1]) - 3.45) < 1e-9 + ) + + # 10. gsm8k diff: within-tolerance passes, above-tolerance fails. + f = [] + _diff_gsm8k({"exact_match,strict-match": 0.90}, {"exact_match,strict-match": 0.89}, 0.02, f) + ok_within = not f + f = [] + _diff_gsm8k({"exact_match,strict-match": 0.90}, {"exact_match,strict-match": 0.80}, 0.02, f) + check("gsm8k tolerance", ok_within and len(f) == 1) + + failed = [name for name, cond in checks if not cond] + if failed: + print(f"[self-test] FAIL ({len(failed)}/{len(checks)}): {failed}") + return 1 + print(f"[self-test] PASS ({len(checks)} checks)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittest/disaggregated/region/test_aux.py b/tests/unittest/disaggregated/region/test_aux.py index 54d07252cbd6..7c2528fbf81c 100644 --- a/tests/unittest/disaggregated/region/test_aux.py +++ b/tests/unittest/disaggregated/region/test_aux.py @@ -3,7 +3,12 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBuffer, AuxBufferMeta, AuxSlot +from tensorrt_llm._torch.disaggregation.native.auxiliary import ( + AuxBuffer, + AuxBufferMeta, + AuxSlot, + build_aux_transfer_layout, +) pytestmark = pytest.mark.cpu_only @@ -128,3 +133,55 @@ def test_fill_slot_unallocated_raises(): with pytest.raises(ValueError, match="not currently allocated"): buf.fill_slot(0, mock_request) + + +def test_aux_buffer_zero_max_draft_len_round_trip(): + """AuxBuffer round-trips with max_draft_len=0. + + A ctx server without speculative_config builds an AuxBuffer with + max_draft_len=0 (ctx/gen spec-config split). fill/get must round-trip + with an empty draft-token list. + """ + buf = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=0, device="cpu") + slot = buf.alloc_slot() + + mock_request = MagicMock() + mock_request.get_last_tokens.return_value = [42] + mock_request.py_draft_tokens = [] + mock_request.prompt_len = 16 + mock_request.cached_tokens = 0 + mock_request.py_disaggregated_params = None + + buf.fill_slot(slot.id, mock_request) + first_tokens, draft_tokens = buf.get_slot_tokens(slot.id) + assert first_tokens == [42] + assert draft_tokens == [] + + # The draft buffer's per-slot item size is 0. + assert buf.meta.item_sizes[1] == 0 + + # Overfilling draft tokens must be rejected, not silently truncated. + mock_request.py_draft_tokens = [1] + with pytest.raises(ValueError, match="exceeds `max_draft_len`"): + buf.fill_slot(slot.id, mock_request) + + +def test_aux_transfer_layout_ctx_no_spec_gen_sa(): + """Ctx/gen spec split: ctx max_draft_len=0, gen max_draft_len=2. + + A ctx server without speculative_config publishes a zero-size draft-token + buffer while an SA/NGram gen server sizes it by max_draft_len. The + empty-source entry must be dropped from the transfer layout (nothing to + send), while the other buffers keep each side's own slot stride. + """ + ctx = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=0, device="cpu") + gen = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=2, device="cpu") + layout = build_aux_transfer_layout(ctx.meta, gen.meta) + + # Buffer order: [first_tokens, draft_tokens, token_counts, prompt_token_counts]; + # the draft_tokens entry (index 1) is dropped. + expected_keep = [0, 2, 3] + np.testing.assert_array_equal(layout.src_item_sizes, ctx.meta.item_sizes[expected_keep]) + np.testing.assert_array_equal(layout.dst_item_sizes, gen.meta.item_sizes[expected_keep]) + np.testing.assert_array_equal(layout.src_base_ptrs, ctx.meta.ptrs[expected_keep]) + np.testing.assert_array_equal(layout.dst_base_ptrs, gen.meta.ptrs[expected_keep]) diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index be94d873ed61..38d72164b4d3 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -36,11 +36,24 @@ try: from tensorrt_llm._torch.disaggregation.native.bounce import buffer as bbuf from tensorrt_llm._torch.disaggregation.native.bounce import impl as btr + from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy _HAVE_TRANSPORT = True except ImportError: # pragma: no cover - CPU-only env without CUDA bindings _HAVE_TRANSPORT = False +# page.py is pure numpy dataclasses, always importable on CPU. +from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + AttentionLayerGroup, + KVCachePageTable, + LocalLayer, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + PoolView, +) + _MIB = 1024 * 1024 @@ -90,7 +103,10 @@ def test_fit_within_free_none_when_too_small(self): def test_config_defaults(self): cfg = bcfg.Config() assert isinstance(cfg.sizing, bcfg.FixedSizing) - assert cfg.chunk_mb == 32 and cfg.min_blocks == 96 + assert cfg.chunk_mb == 32 + # The operative gate is the byte one; the legacy block gate defaults to 1 (vacuous). + assert cfg.min_bytes == bcfg.DEFAULT_MIN_BYTES == 2 * _MIB + assert cfg.min_blocks == 1 # --------------------------------------------------------------------------- # @@ -106,24 +122,41 @@ def test_positive_enables_with_capacity(self): assert isinstance(cfg, bcfg.Config) assert cfg.sizing.capacity_mb == 2048 - def test_min_blocks_defaults_and_overrides(self, monkeypatch): + def test_min_bytes_defaults_and_overrides(self, monkeypatch): + monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", raising=False) + assert bcfg.config_from_size(2048).min_bytes == 2 * _MIB # keeps the Config default + assert bcfg.config_from_size(2048, min_bytes=1).min_bytes == 1 # explicit arg overrides + assert bcfg.config_from_size(2048, min_bytes=64 * _MIB).min_bytes == 64 * _MIB + # The gate is tuned for tests via the env override (no user-facing config field). + monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "1") + assert bcfg.config_from_size(2048).min_bytes == 1 # env override + assert bcfg.config_from_size(2048, min_bytes=250).min_bytes == 250 # arg beats env + + def test_min_blocks_backcompat_defaults_and_overrides(self, monkeypatch): + # The legacy block-count gate is kept for back-compat; it defaults to 1 (vacuous, so the + # byte gate decides) and still honors the explicit arg and the env override. monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", raising=False) - assert bcfg.config_from_size(2048).min_blocks == 96 # keeps the Config default - assert bcfg.config_from_size(2048, 1).min_blocks == 1 # explicit arg still overrides - assert bcfg.config_from_size(2048, 250).min_blocks == 250 - # The gate is lowered for tests via the env override (no user-facing config field). - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "1") - assert bcfg.config_from_size(2048).min_blocks == 1 # env override + assert bcfg.config_from_size(2048).min_blocks == 1 # keeps the Config default + assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg still overrides + monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "96") + assert bcfg.config_from_size(2048).min_blocks == 96 # env override assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg beats env - def test_min_blocks_env_is_parsed_defensively(self, monkeypatch): + @pytest.mark.parametrize( + "env,attr,default", + [ + ("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "min_bytes", 2 * _MIB), + ("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "min_blocks", 1), + ], + ) + def test_gate_env_is_parsed_defensively(self, monkeypatch, env, attr, default): # A bad value must not crash setup (falls back to the default); a non-positive value clamps to 1. for bad in ("", "auto", "1.5"): - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", bad) - assert bcfg.config_from_size(2048).min_blocks == 96 + monkeypatch.setenv(env, bad) + assert getattr(bcfg.config_from_size(2048), attr) == default for nonpos in ("0", "-5"): - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", nonpos) - assert bcfg.config_from_size(2048).min_blocks == 1 + monkeypatch.setenv(env, nonpos) + assert getattr(bcfg.config_from_size(2048), attr) == 1 # --------------------------------------------------------------------------- # @@ -276,6 +309,7 @@ def __init__(self, capacity_bytes, phys_chunk_size, name="kv_bounce"): self.next_id = 0 self.released = [] self.quarantined = [] + self.reserved_sizes = [] @property def capacity(self): @@ -286,6 +320,7 @@ def reserve(self, size, timeout=None): return None sid = self.next_id self.next_id += 1 + self.reserved_sizes.append(size) return sid, self.base def release(self, slot_id): @@ -301,7 +336,9 @@ def reg_descs(self): return [] -def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_blocks=1): +def _make_transport( + monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bytes=1, min_blocks=1 +): monkeypatch.setattr(btr, "SlotAllocator", _FakeAlloc) monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) monkeypatch.setattr( @@ -316,6 +353,7 @@ def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bl capacity_bytes=capacity, phys_chunk_size=32 * _MIB, block_bytes_per_group=block_bytes_per_group, + min_bytes=min_bytes, min_blocks=min_blocks, ) @@ -369,9 +407,34 @@ def test_reserve_single_writer_ok(self, monkeypatch): req = _recv_req([1]) # total = 3, num_writers=1 -> no even-split requirement assert t.reserve(req, num_writers=1) is True - def test_reserve_too_small_falls_back(self, monkeypatch): + def test_reserve_below_min_bytes_falls_back(self, monkeypatch): + # The gate is in BYTES: 4 blocks x 100B = 400B < 1000B falls back, regardless of how many + # blocks that is (the old block-count gate silently skipped huge small-block transfers). + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + req = _recv_req([4]) + assert t.reserve(req, num_writers=1) is False + assert req.bounce_dst_base is None + + def test_reserve_at_min_bytes_bounces(self, monkeypatch): + # exactly at the threshold (10 x 100B = 1000B >= 1000B) the transfer bounces + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + req = _recv_req([10]) + assert t.reserve(req, num_writers=1) is True + assert req.bounce_dst_base == 0x100000 + + def test_reserve_few_large_blocks_bounce(self, monkeypatch): + # The K3 regression shape: FEW but LARGE blocks must clear a byte gate that a block-count + # gate of 96 would have silently failed (67 blocks x 6.5 MiB = 433 MiB). + t = _make_transport( + monkeypatch, block_bytes_per_group=[int(6.5 * _MIB)], min_bytes=2 * _MIB + ) + assert t.reserve(_recv_req([67]), num_writers=1) is True + + def test_reserve_legacy_min_blocks_backcompat(self, monkeypatch): + # the legacy block-count gate still applies when raised explicitly (back-compat) t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_blocks=96) assert t.reserve(_recv_req([4]), num_writers=1) is False # 4 < 96 blocks + assert t.reserve(_recv_req([96]), num_writers=1) is True def test_reserve_unknown_slot_size_falls_back(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) # only 1 group known @@ -548,6 +611,143 @@ def test_orphan_reservation_quarantines_and_is_idempotent(self, monkeypatch): assert t._recv_alloc.quarantined == [0] +# --------------------------------------------------------------------------- # +# Hybrid layouts (Kimi K3: KDA mamba + MLA attention) — page-table gate, +# trailing-empty-group handling, and recurrent-state (extra_bytes) sizing +# --------------------------------------------------------------------------- # + +# Kimi K3 geometry: 24 MLA layers in one attention layer group; 69 KDA layers +# whose per-layer recurrent state is a bf16 short-conv slot + fp32 delta slot, +# replicated per rank. +_K3_MLA_BLOCK_BYTES = 32 * 576 * 2 * 24 # tpb x (kv_lora_rank+rope) x bf16 x 24 layers = 884,736 +_K3_KDA_LAYERS = 69 +_K3_CONV_SLOT_BYTES = 294_912 # [3*H*hd, W] bf16 per layer +_K3_SSM_SLOT_BYTES = 6_291_456 # [H, hd, hd] fp32 per layer (95.5% of the state) +_K3_KDA_PAYLOAD_BYTES = ( + 454_459_392 # 69 x (conv + delta) per request per rank, from the geometry above +) + + +def _k3_page_table() -> KVCachePageTable: + """A K3-shaped page table exactly as the builders produce it. + + Attention layer group(s) first and the mamba layer group appended LAST + (kv_extractor.py, both ``build_page_table`` and ``_build_page_table_v2``, + append it after every attention group). + """ + attn = AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=1, + sliding_window_size=None, + local_layers=[LocalLayer(i, i) for i in range(24)], + pool_views=[PoolView(pool_idx=0, buffer_entries=np.array([], dtype=BUFFER_ENTRY_DTYPE))], + ) + mamba = MambaLayerGroup( + pool_group_idx=1, # dangling by construction: mamba pools live on the LG itself + mamba_layer_offsets={gl: i for i, gl in enumerate(range(24, 24 + _K3_KDA_LAYERS))}, + conv_states=PhysicalPool(0x200000, _K3_CONV_SLOT_BYTES, 8), + ssm_states=PhysicalPool(0x300000, _K3_SSM_SLOT_BYTES, 8), + conv_section_bytes=[_K3_CONV_SLOT_BYTES // 3] * 3, + ssm_bytes_per_head=_K3_SSM_SLOT_BYTES // 4, + ) + return KVCachePageTable( + tokens_per_block=32, + layer_groups=[attn, mamba], + pool_groups=[PhysicalPoolGroup(pools=[PhysicalPool(0x100000, _K3_MLA_BLOCK_BYTES, 128)])], + ) + + +def _k3_rank_info(): + # MambaPolicy only reads tp_size / tp_rank / attention.enable_attention_dp. + return SimpleNamespace(tp_size=1, tp_rank=0, attention=None) + + +@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") +class TestHybridK3Bounce: + def test_block_bytes_per_group_keeps_mamba_placeholder(self): + # The mamba group must stay in the list as a placeholder (None), aligned with the + # layer-group indices a recv request uses, instead of truncating the list. + assert btr.block_bytes_per_group(_k3_page_table()) == [_K3_MLA_BLOCK_BYTES, None] + + def test_reserve_engages_on_k3_mixed_layout(self, monkeypatch): + # Regression pin: a K3 recv request always carries a trailing EMPTY entry for + # the mamba layer group (transceiver._create_kv_slice), which used to trip the + # unknown-slot-size guard and silently push every K3 request onto the per-fragment + # (~0.4 GB/s host-staged) path. It must engage bounce, sized for MLA KV + KDA state. + t = _make_transport( + monkeypatch, + block_bytes_per_group=btr.block_bytes_per_group(_k3_page_table()), + min_bytes=2 * _MIB, + ) + req = _recv_req([67, 0]) # 67 MLA blocks (~2144 tokens) + the empty mamba entry + assert t.reserve(req, num_writers=1, extra_bytes=_K3_KDA_PAYLOAD_BYTES) is True + assert req.bounce_dst_base == 0x100000 + # the region covers the MLA KV blocks plus the appended KDA state + assert t._recv_alloc.reserved_sizes == [67 * _K3_MLA_BLOCK_BYTES + _K3_KDA_PAYLOAD_BYTES] + + def test_reserve_nonempty_group_with_unknown_size_still_falls_back(self, monkeypatch): + # Safety direction of the gate is preserved: a group that HAS blocks but no known slot + # size (None placeholder) still disables bounce for the request. + t = _make_transport(monkeypatch, block_bytes_per_group=[_K3_MLA_BLOCK_BYTES, None]) + assert t.reserve(_recv_req([2, 1]), num_writers=1) is False + + def test_reserve_extra_bytes_counts_toward_min_bytes(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=599) is False # 999 < 1000 + assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=600) is True # exactly 1000 + + def test_reserve_extra_bytes_counts_toward_capacity(self, monkeypatch): + # extra_bytes must be part of the reservation, or the sender's coalesced write + # (KV + recurrent state) would overrun the region into the neighboring slot. + t = _make_transport(monkeypatch, block_bytes_per_group=[100], capacity=500) + assert t.reserve(_recv_req([2]), num_writers=1, extra_bytes=400) is False # 600 > 500 + + def test_reserve_extra_bytes_fanin_falls_back(self, monkeypatch): + # Fan-in splits the region equally by writer; per-writer recurrent-state fragments can + # differ (PP stages hold different mamba layers), so state + fan-in falls back. + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + assert t.reserve(_recv_req([2, 0]), num_writers=2, extra_bytes=64) is False + assert t.reserve(_recv_req([2, 0]), num_writers=2) is True # no state -> fan-in fine + + def test_mamba_payload_bytes_matches_k3_geometry(self): + # Matched-TP replicated KDA state: payload_bytes must reproduce the per-request + # per-rank number derived from K3's KDA geometry (constants above) exactly. + pt = _k3_page_table() + ri = _k3_rank_info() + got = MambaPolicy.payload_bytes( + sender_page_table=pt, receiver_page_table=pt, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + assert got == _K3_KDA_PAYLOAD_BYTES + + def test_mamba_payload_bytes_matches_collect_frags(self): + # payload_bytes mirrors the sender's collect_frags call: same sizes, slot-independent. + pt = _k3_page_table() + ri = _k3_rank_info() + _, _, sizes = MambaPolicy.collect_frags( + self_page_table=pt, peer_page_table=pt, src_slot=5, dst_slot=3, self_ri=ri, peer_ri=ri + ) + got = MambaPolicy.payload_bytes( + sender_page_table=pt, receiver_page_table=pt, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + assert got == sum(sizes) > 0 + + def test_mamba_payload_bytes_zero_without_slot_or_group(self): + pt = _k3_page_table() + ri = _k3_rank_info() + assert MambaPolicy.payload_bytes(pt, pt, dst_slot=None, sender_ri=ri, receiver_ri=ri) == 0 + pure_attn = KVCachePageTable( + tokens_per_block=32, + layer_groups=[pt.layer_groups[0]], + pool_groups=pt.pool_groups, + ) + assert ( + MambaPolicy.payload_bytes( + pure_attn, pure_attn, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + == 0 + ) + + # --------------------------------------------------------------------------- # # SlotAllocator — first-fit reuses out-of-order freed holes (real Buffer mocked) # --------------------------------------------------------------------------- # diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py new file mode 100644 index 000000000000..110f7365df4e --- /dev/null +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -0,0 +1,543 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3 KDA-shaped recurrent-state transfer through the Python V2 transceiver. + +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``: + - 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). + +* ``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. + +* ``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. +""" + +import uuid +from typing import Dict, List + +import pytest +import torch +from test_mamba_transfer import _create_transceivers, _run_concurrent + +import tensorrt_llm +import tensorrt_llm.bindings +import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 +from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + KVRegionExtractorV1, + build_page_table_from_manager, +) +from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup +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 +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, KvCacheConfig + +# --------------------------------------------------------------------------- +# KDA parameters (small stand-ins for Kimi K3's per-layer KDA state) +# --------------------------------------------------------------------------- +NUM_KDA_LAYERS = 4 +KDA_NUM_HEADS = 4 # H +KDA_HEAD_DIM = 32 # hd; KDA state_size == head_dim +KDA_W = 3 # short_conv_kernel_size; manager d_conv = W + 1 +CONV_DTYPE = torch.bfloat16 +SSM_DTYPE = torch.float32 +MAX_BATCH_SIZE = 4 +REQUEST_LENGTHS = [16, 32] + +# Internal: layer 0 is a dummy attention layer required by page table infra; +# layers 1..NUM_KDA_LAYERS are KDA (under test). +_NUM_TOTAL_LAYERS = NUM_KDA_LAYERS + 1 +_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). +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 _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). + """ + state_tp = tp if not enable_attention_dp else 1 + managers = [] + for rank in range(tp): + mapping = Mapping( + world_size=tp, + rank=rank, + tp_size=tp, + pp_size=1, + 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, + # dummy attention layer (page table scaffolding) + kv_cache_config=KvCacheConfig( + max_tokens=256 * max_batch_size, + enable_block_reuse=False, + event_buffer_max_size=0, + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + layer_mask=_ATTN_MASK, + num_kv_heads=4, + head_dim=64, + tokens_per_block=8, + max_seq_len=256, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=DataType.FLOAT, + model_type="qwen3_next", + ) + managers.append(mgr) + return managers + + +def _get_mamba_layer_group(page_table) -> MambaLayerGroup: + mlgs = [lg for lg in page_table.layer_groups if isinstance(lg, MambaLayerGroup)] + assert len(mlgs) == 1, f"expected exactly one MambaLayerGroup, got {len(mlgs)}" + return mlgs[0] + + +def _layer_slot_start(pool, local_layer_idx: int, slot: int) -> int: + return pool.base_address + (local_layer_idx * pool.num_slots + slot) * pool.slot_bytes + + +def _assert_frags_tile_slots(frags, mlg: MambaLayerGroup, slot: int): + """Assert (ptr, size) frags exactly tile every layer's slot bytes.""" + ptrs, sizes = frags + covered: Dict[int, List[tuple]] = {} + for ptr, size in zip(ptrs, sizes): + pool = None + for cand in (mlg.conv_states, mlg.ssm_states): + lo = cand.base_address + hi = lo + len(mlg.mamba_layer_offsets) * cand.num_slots * cand.slot_bytes + if lo <= ptr < hi: + pool = cand + break + assert pool is not None, f"frag ptr {ptr} outside both mamba pools" + rel = ptr - pool.base_address + layer = rel // (pool.num_slots * pool.slot_bytes) + in_layer = rel - layer * pool.num_slots * pool.slot_bytes + frag_slot = in_layer // pool.slot_bytes + off = in_layer - frag_slot * pool.slot_bytes + assert frag_slot == slot, f"frag targets slot {frag_slot}, expected {slot}" + assert off + size <= pool.slot_bytes, ( + f"frag [{off}, {off + size}) exceeds slot_bytes {pool.slot_bytes}" + ) + covered.setdefault((id(pool), int(layer)), []).append((off, off + size)) + + # Exact tiling per (pool, layer): sorted intervals must be contiguous + # from 0 to slot_bytes with no overlap. + n_layers = len(mlg.mamba_layer_offsets) + for pool in (mlg.conv_states, mlg.ssm_states): + for layer in range(n_layers): + intervals = sorted(covered.get((id(pool), layer), [])) + assert intervals, f"layer {layer} of pool not covered" + pos = 0 + for lo, hi in intervals: + assert lo == pos, f"gap/overlap at byte {pos} (next frag at {lo})" + pos = hi + assert pos == pool.slot_bytes, ( + f"layer {layer}: covered {pos} of {pool.slot_bytes} slot bytes" + ) + + +# --------------------------------------------------------------------------- +# Descriptor-level tests (no NIXL transfer) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("enable_attention_dp", [True, False], ids=["adp_on", "adp_off_tp1"]) +def test_kda_layer_group_descriptors(enable_attention_dp): + """V2 page table must describe BOTH KDA slots with exact byte extents.""" + mgr = _create_kda_managers(1, enable_attention_dp=enable_attention_dp)[0] + try: + pt = build_page_table_from_manager(mgr) + mlg = _get_mamba_layer_group(pt) + + conv = mgr._impl.mamba_cache.conv + ssm = mgr._impl.mamba_cache.temporal + + # Shapes/dtypes of the backing tensors are KDA-shaped. + assert conv.dtype == CONV_DTYPE + assert ssm.dtype == SSM_DTYPE + assert tuple(conv.shape[2:]) == (3 * KDA_NUM_HEADS * KDA_HEAD_DIM, KDA_W) + assert tuple(ssm.shape[2:]) == (KDA_NUM_HEADS, KDA_HEAD_DIM, KDA_HEAD_DIM) + + # Both pools present, byte-exact (bf16 conv, fp32 delta). + assert mlg.conv_states is not None and mlg.ssm_states is not None + assert mlg.conv_states.base_address == conv.data_ptr() + assert mlg.ssm_states.base_address == ssm.data_ptr() + assert mlg.conv_states.slot_bytes == CONV_SLOT_BYTES + assert mlg.ssm_states.slot_bytes == SSM_SLOT_BYTES + assert mlg.conv_states.num_slots == conv.shape[1] + assert mlg.ssm_states.num_slots == ssm.shape[1] + + # qwen3_next 3-sectioning: equal sections summing to the conv slot. + assert mlg.conv_section_bytes == [CONV_SLOT_BYTES // 3] * 3 + assert sum(mlg.conv_section_bytes) == mlg.conv_states.slot_bytes + assert mlg.ssm_bytes_per_head == KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + assert mlg.ssm_states.slot_bytes // mlg.ssm_bytes_per_head == KDA_NUM_HEADS + + # Layer offsets cover exactly the KDA layers. + assert sorted(mlg.mamba_layer_offsets.keys()) == [i for i, m in enumerate(_KDA_MASK) if m] + + # Matched-parallelism frags: 2 pools x layers, full-slot, exact tiling. + ri = RankInfo.from_kv_cache_manager("kda_test", mgr, device_id=0) + src_slot, dst_slot = 0, 1 + src_frags, dst_frags, sizes = MambaPolicy.build_mamba_frags( + mlg, mlg, src_slot, dst_slot, ri, ri + ) + assert len(src_frags) == len(dst_frags) == len(sizes) == 2 * NUM_KDA_LAYERS + _assert_frags_tile_slots((src_frags, sizes), mlg, src_slot) + _assert_frags_tile_slots((dst_frags, sizes), mlg, dst_slot) + + # Full-slot copies at the expected per-layer addresses. + for glid, lid in sorted(mlg.mamba_layer_offsets.items()): + assert _layer_slot_start(mlg.conv_states, lid, src_slot) in src_frags + assert _layer_slot_start(mlg.ssm_states, lid, src_slot) in src_frags + finally: + mgr.shutdown() + + +def test_kda_hetero_tp_rejected(): + """Replicated KDA state + heterogeneous TP (ADP off) must be rejected. + + 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. + """ + ctx_mgr = _create_kda_managers(2, enable_attention_dp=False)[0] + gen_mgr = _create_kda_managers(4, enable_attention_dp=False)[1] + try: + ctx_pt = build_page_table_from_manager(ctx_mgr) + gen_pt = build_page_table_from_manager(gen_mgr) + 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) + + # 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) + finally: + ctx_mgr.shutdown() + gen_mgr.shutdown() + + +@pytest.mark.parametrize( + "ctx_cfg,gen_cfg", + [ + ((2, False), (2, False)), # matched TP, ADP off (EP pre-scaled) + ((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.""" + ctx_tp, ctx_adp = ctx_cfg + gen_tp, gen_adp = gen_cfg + ctx_mgr = _create_kda_managers(ctx_tp, enable_attention_dp=ctx_adp)[0] + gen_mgr = _create_kda_managers(gen_tp, enable_attention_dp=gen_adp)[-1] + try: + ctx_pt = build_page_table_from_manager(ctx_mgr) + gen_pt = build_page_table_from_manager(gen_mgr) + 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) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + finally: + ctx_mgr.shutdown() + gen_mgr.shutdown() + + +def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int): + """MambaLayerGroup-only page table with fake addresses (no CUDA needed).""" + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, PhysicalPool + + mlg = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={glid: i for i, glid in enumerate(range(1, NUM_KDA_LAYERS + 1))}, + conv_states=PhysicalPool(base_address=0x1000, slot_bytes=conv_slot_bytes, num_slots=8), + ssm_states=PhysicalPool(base_address=0x2000000, slot_bytes=ssm_slot_bytes, num_slots=8), + conv_section_bytes=[conv_slot_bytes // 3] * 3, + ssm_bytes_per_head=KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize, + ) + return KVCachePageTable(tokens_per_block=8, layer_groups=[mlg], pool_groups=[]) + + +def _synthetic_rank_info(tp: int, adp: bool): + from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo + + return RankInfo( + instance_name="syn", + instance_rank=0, + tp_size=tp, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[_NUM_TOTAL_LAYERS], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=b"", + attention=AttentionInfo( + kv_heads_per_rank=4, + tokens_per_block=8, + dims_per_head=64, + element_bytes=2, + enable_attention_dp=adp, + is_mla=False, + ), + ) + + +@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 + ((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 + ], + ids=[ + "reject_hetero_tp_adp_off", + "reject_hetero_tp_adp_off_rev", + "reject_mixed_adp_replicated", + "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).""" + 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 + + def build(cfg): + tp, adp, denom = cfg + return ( + _synthetic_rank_info(tp, adp), + _synthetic_kda_page_table(full_ssm // denom, full_conv // denom), + ) + + ctx_ri, ctx_pt = build(ctx) + gen_ri, gen_pt = build(gen) + if ok: + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + else: + with pytest.raises(ValueError, match="TP-aggregated"): + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + +# --------------------------------------------------------------------------- +# Loopback transfer over a real NIXL agent (single node) +# --------------------------------------------------------------------------- +def _generate_ground_truth(num_requests: int, seed: int = 20260722): + """Full replicated KDA states per (request, kda_layer), two dtypes.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + results = [] + for _ in range(num_requests): + layers = {} + for i, is_kda in enumerate(_KDA_MASK): + if not is_kda: + continue + layers[i] = { + "conv": torch.rand(3 * KDA_NUM_HEADS * KDA_HEAD_DIM, KDA_W, generator=gen).to( + CONV_DTYPE + ), + "ssm": torch.rand( + KDA_NUM_HEADS, + KDA_HEAD_DIM, + KDA_HEAD_DIM, + generator=gen, + dtype=SSM_DTYPE, + ), + } + results.append(layers) + return results + + +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.""" + 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) + for mgr in ctx_mgrs + gen_mgrs: + mgr._impl.mamba_cache.conv.zero_() + mgr._impl.mamba_cache.temporal.zero_() + + config = CacheTransceiverConfig( + backend="NIXL", + 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_endpoint = ctx_tcs[0]._context_info_endpoint + + sampling_params = SamplingParams() + ctx_rids, gen_rids, ctx_reqs, gen_reqs = [], [], [], [] + for req_idx, req_len in enumerate(REQUEST_LENGTHS): + unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + ctx_rid, gen_rid = req_idx * 2, req_idx * 2 + 1 + ctx_rids.append(ctx_rid) + gen_rids.append(gen_rid) + sc = tensorrt_llm.bindings.SamplingConfig(sampling_params._get_sampling_config()) + ctx_req = LlmRequest( + request_id=ctx_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + ) + ctx_req.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + gen_req = LlmRequest( + request_id=gen_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + gen_req.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=ctx_rid, + ctx_dp_rank=0, + ctx_info_endpoint=ctx_endpoint, + disagg_request_id=unique_rid, + ) + ctx_reqs.append(ctx_req) + gen_reqs.append(gen_req) + + ctx_batch = ScheduledRequests() + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) + gen_batch = ScheduledRequests() + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) + for req in ctx_reqs + gen_reqs: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + for mgr in ctx_mgrs: + mgr.update_resources(ctx_batch) + for mgr in gen_mgrs: + mgr.update_resources(gen_batch) + + # Ground truth: identical full state on every ctx rank (replicated). + ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) + for mgr in ctx_mgrs: + 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: + gen_tcs[rank].request_and_receive_async(req) + for rank in range(ctx_tp): + for req in ctx_reqs: + ctx_tcs[rank].respond_and_send_async(req) + _run_concurrent(ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True)) + _run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) + + # 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) + rank_factor = 1 if enable_attention_dp else gen_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). + for gen_rank, mgr in enumerate(gen_mgrs): + for req_idx, rid in enumerate(gen_rids): + slot = mgr.mamba_cache_index[rid] + for layer_idx in mgr._impl.mamba_layer_offsets: + full = ground_truth[req_idx][layer_idx] + 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], + rtol=0, + atol=0, + msg=lambda m, n=name, r=gen_rank, ri=req_idx, li=layer_idx: ( + f"{n} mismatch: gen_rank={r} req={ri} layer={li} " + f"ctx_tp={ctx_tp} gen_tp={gen_tp}: {m}" + ), + ) + + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() + for tc in ctx_tcs + gen_tcs: + tc.shutdown() + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp", + [(1, 1), (2, 2)], + ids=["tp1_tp1", "tp2_tp2_ep_prescaled"], +) +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) From dcfbe78fbf9c4fa944e4d27fd17cec6e0bbb0c50 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 5 Aug 2026 20:05:13 -0700 Subject: [PATCH 04/14] [TRTLLM-14815][doc] Add Kimi K3 disaggregated-serving example configs 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 --- .../slurm/benchmark/run_benchmark.sh | 14 ++ .../slurm/benchmark/start_server.sh | 19 +- .../slurm/benchmark/start_worker.sh | 30 +++- .../configs/kda_payload_kimi_k3.yaml | 83 +++++++++ .../run_cache_transceiver_test.py | 40 ++++- examples/kimi_k3/disagg/README.md | 165 ++++++++++++++++++ .../disagg/benchmark_kimi_k3_dep16.yaml | 136 +++++++++++++++ examples/kimi_k3/disagg/ctx_config.yaml | 63 +++++++ .../kimi_k3/disagg/disagg_proxy_config.yaml | 23 +++ examples/kimi_k3/disagg/gen_config_no_sa.yaml | 46 +++++ 10 files changed, 614 insertions(+), 5 deletions(-) create mode 100644 examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml create mode 100644 examples/kimi_k3/disagg/README.md create mode 100644 examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml create mode 100644 examples/kimi_k3/disagg/ctx_config.yaml create mode 100644 examples/kimi_k3/disagg/disagg_proxy_config.yaml create mode 100644 examples/kimi_k3/disagg/gen_config_no_sa.yaml diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark.sh b/examples/disaggregated/slurm/benchmark/run_benchmark.sh index 8fa89815704c..7ccae84a6e69 100644 --- a/examples/disaggregated/slurm/benchmark/run_benchmark.sh +++ b/examples/disaggregated/slurm/benchmark/run_benchmark.sh @@ -5,6 +5,20 @@ set -e set -u trap 'echo "Error occurred at line $LINENO"; exit 1' ERR +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +# The client_cmds srun passes these via --export with single-quoted values +# (submit.py convert_envs_to_str), and srun keeps the quotes literal — strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + # Add parameter validation if [ "$#" -lt 10 ]; then echo "Error: Missing required arguments, got $# arguments, args: $@" diff --git a/examples/disaggregated/slurm/benchmark/start_server.sh b/examples/disaggregated/slurm/benchmark/start_server.sh index ff8e5aa90277..570b9ceb30dc 100644 --- a/examples/disaggregated/slurm/benchmark/start_server.sh +++ b/examples/disaggregated/slurm/benchmark/start_server.sh @@ -5,4 +5,21 @@ set -x config_file=$1 -trtllm-serve disaggregated -c ${config_file} -t 7200 -r 7200 +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console +# script; fall back to the module entry point in that case. +trtllm_serve_cmd="trtllm-serve" +if ! command -v trtllm-serve >/dev/null 2>&1; then + trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve" +fi + +${trtllm_serve_cmd} disaggregated -c ${config_file} -t 7200 -r 7200 diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 0a5b5897b773..0c8653893958 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -38,8 +38,25 @@ else export CUDA_VISIBLE_DEVICES=${cuda_devices} fi -# Clear UCX_TLS for specific clusters -unset UCX_TLS +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +# Clear UCX_TLS for specific clusters. Some clusters instead need an +# explicit transport list (e.g. NVL72 nodes whose verbs transports cannot +# initialize): set TRTLLM_WORKER_UCX_TLS in worker_env_var to re-pin +# UCX_TLS here, after the container-provided value is cleared. +if [ -n "${TRTLLM_WORKER_UCX_TLS:-}" ]; then + export UCX_TLS="${TRTLLM_WORKER_UCX_TLS}" +else + unset UCX_TLS +fi echo "SLURM_PROCID: ${SLURM_PROCID}, hostname: $(hostname), instance_id: ${instance_id}" echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}" @@ -63,7 +80,14 @@ else nsys_prefix="nsys profile -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none" fi +# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console +# script; fall back to the module entry point in that case. +trtllm_serve_cmd="trtllm-serve" +if ! command -v trtllm-serve >/dev/null 2>&1; then + trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve" +fi + ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \ - trtllm-serve ${model_path} \ + ${trtllm_serve_cmd} ${model_path} \ --host $(hostname) --port ${port} \ --config ${config_file} 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 new file mode 100644 index 000000000000..674c275890a8 --- /dev/null +++ b/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml @@ -0,0 +1,83 @@ +# 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/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index c4ebd8c6460c..3bfb64a8dcbc 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -287,6 +287,35 @@ class _TransferError(Exception): pass +def _wait_ctx_complete(xcvr, rid, runtime): + """Block until this ctx request's send finishes (or errors). + + * C++ transceiver: check_context_transfer_status(None) is a true block-all; + a single call suffices (see the comment at the call site). + * PYTHON (V2) transceiver: even under block_all, each TxSession wait is + bounded by kv_transfer_sender_future_timeout_ms (default 1000 ms). A + slower peer handshake makes the call log "TxSession ... timed out" and + return with the request still DISAGG_CONTEXT_TRANS_IN_PROGRESS -- in the + real executor that is benign (the session stays open and is re-polled + every iteration), but returning here would free and refill the KV blocks + mid-flight, so the receiver reads the NEXT request's pattern (verify FAIL + on every request except the last). So poll until + this rid lands in the completed/failed lists. Collectively safe: every + ctx rank loops on the same rid and the per-call consensus makes all + ranks observe completion on the same iteration. The per-cell + signal.alarm and the hang detector bound the loop. + """ + if runtime != "PYTHON": + xcvr.check_context_transfer_status(None) + return + while True: + completed, failed = xcvr.check_context_transfer_status(None) + if rid in failed: + raise _TransferError(f"ctx transfer failed for rid={rid}") + if rid in completed: + return + + def _wait_gen_complete(xcvr, req, runtime): """Block until this gen request's receive finishes (or errors). @@ -372,7 +401,9 @@ def run_one_request( # transfer is still in progress. NIXL/UCX cold-start connection setup can # exceed that, so the harness would free the request mid-transfer, leaving # the gen side hung and the ctx sender thread asserting on a freed session. - xcvr.check_context_transfer_status(None) + # The PYTHON runtime additionally needs a poll loop on top of block_all; + # see _wait_ctx_complete. + _wait_ctx_complete(xcvr, rid, runtime) state = req.state tensorrt_llm.logger.info(f"[ctx r{rank}] rid={rid}: transfer DONE (send), state={state}") free_sequence(kvm, req, kv_handle, use_v2) @@ -648,6 +679,13 @@ def arm_watchdog(combination_idx, reqlen_idx, what): backend=backend, transceiver_runtime=(None if runtime == "CPP" else "PYTHON"), max_tokens_in_buffer=cfg["kv_cache"]["max_tokens_in_buffer"], + # PYTHON (V2) only; 0 keeps bounce off. With bounce on, the KV data + # rides a fabric-VMM staging buffer (CU_MEM_HANDLE_TYPE_FABRIC), which + # is what lets UCX pick cuda_ipc across NVL72 nodes -- direct + # pool-to-pool transfers from non-fabric allocations fall back to + # much slower host-staged tcp, so enable bounce for cross-node + # transfers inside an NVLink domain. + kv_cache_bounce_size_mb=int(cfg["kv_cache"].get("bounce_size_mb", 0)), ) # Build the cache manager + transceiver ONCE per case (the manager is diff --git a/examples/kimi_k3/disagg/README.md b/examples/kimi_k3/disagg/README.md new file mode 100644 index 000000000000..8597bd8fe643 --- /dev/null +++ b/examples/kimi_k3/disagg/README.md @@ -0,0 +1,165 @@ +# Kimi K3 disaggregated serving (ctx/gen split) + +Configuration pair + deployment wiring for running Kimi K3 with separate +context (prefill) and generation (decode) servers. Status: **validated +end-to-end on hardware** (GB300 NVL72, 1 ctx + 1 gen, DEP16 both sides, +GSM8K accuracy parity with aggregated serving) — see the caveats section +for constraints. + +## Files + +| File | Purpose | +|---|---| +| `ctx_config.yaml` | Context-server extra LLM-API options (DEP16, overlap scheduler off, no spec decode) | +| `gen_config_no_sa.yaml` | Generation-server options, no speculative decoding (CUDA graphs ON by default: GSM8K 96.89, 765/2138 tok/s @c64/c256 vs aggregated 643/1972; null `cuda_graph_config` for token-parity debugging). A suffix-automaton (SA) speculative-decoding variant lands together with K3 SA support. | +| `disagg_proxy_config.yaml` | `trtllm-serve disaggregated` proxy config (1 ctx + 1 gen) | +| `benchmark_kimi_k3_dep16.yaml` | Config for the SLURM benchmark harness (`examples/disaggregated/slurm/benchmark/submit.py`) | + +## K3 constraints baked into the configs + +- **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. +- **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: + replicated share ~113 GiB + experts/8), leaving no activation headroom + on GB300 (288 GiB) and not fitting GB200 (186 GiB). Treat DEP8-ctx as + ruled out on GB200 and an open (likely negative) question on GB300. +- **`transceiver_runtime: PYTHON` is mandatory**: `auto` resolves to the + C++ transceiver, which throws at construction for K3's + `MixedMambaHybridCacheManager`. +- `disable_overlap_scheduler: true` on the ctx server (disagg + requirement) and on the gen server (keeps the smoke runs maximally + comparable across configurations). +- `enable_block_reuse: false`, `tokens_per_block: 64`, no chunked + prefill, beam width 1 (model requirements). +- `max_tokens_in_buffer: 8448` covers the target max ISL of 8192; raise + it together with `max_num_tokens`/`max_seq_len` for longer ISL. +- **`kv_cache_bounce_size_mb: 1024` on both sides**: the V2 transceiver's default + pool-to-pool path cannot use inter-node cuda_ipc on MNNVL (the KV pool + is a plain, non-fabric allocation) and falls back to ~0.4 GB/s + host-staged tcp; the fabric-VMM bounce buffer restores cuda_ipc/MNNVL + eligibility (measured ~455 GB/s/GPU). Bounce engages automatically for payloads above `TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES` (default 2 MiB) — always true for K3's ~433 MiB per-request state. + The region must fit ONE request's full payload: fixed 433 MiB KDA + state + ~27 KB/token MLA latent (649 MiB at 8k ISL). A + 512 MiB value makes every 8k transfer fall back to the per-fragment + tcp path (`[kv-bounce] in-place: transfer 649MiB exceeds the 512MiB + bounce region`). + +## Launch sequence (manual, single ctx + single gen) + +Each K3 worker spans 16 GPUs (4 NVL72 nodes at 4 GPUs/node). Environment +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 + # initialize; a container-default + # UCX_TLS=tcp breaks V2 NIXL +``` + +1. Start the context server (16-rank MPI world across its 4 nodes): + + ```bash + trtllm-llmapi-launch trtllm-serve $MODEL_PATH \ + --host --port 8001 \ + --config examples/kimi_k3/disagg/ctx_config.yaml + ``` + +2. Start the generation server: + + ```bash + trtllm-llmapi-launch trtllm-serve $MODEL_PATH \ + --host --port 8002 \ + --config examples/kimi_k3/disagg/gen_config_no_sa.yaml + ``` + +3. Edit `disagg_proxy_config.yaml` (worker URLs = the head nodes above), + then start the proxy: + + ```bash + trtllm-serve disaggregated -c examples/kimi_k3/disagg/disagg_proxy_config.yaml + ``` + +4. Send OpenAI-compatible requests to the proxy (port 8000). + +## SLURM benchmark harness + +`benchmark_kimi_k3_dep16.yaml` drives the full orchestration (worker +config generation, node allocation, proxy, benchmark client): + +```bash +python3 examples/disaggregated/slurm/benchmark/submit.py \ + -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml --dry-run # inspect +python3 examples/disaggregated/slurm/benchmark/submit.py \ + -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml # submit +``` + +- Set `benchmark.dataset_file` before an e2e submission. +- **Gen-only baseline**: set `benchmark.mode: gen_only_no_context` + (submit.py exports `TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1` to the + workers) to measure the decode-side ceiling without KV transfer. +- The harness's `start_worker.sh` clears `UCX_TLS`; the config carries + the transport pin via `TRTLLM_WORKER_UCX_TLS`, which `start_worker.sh` + re-exports as `UCX_TLS` after the clear. +- pyxis/enroot resets image-defined variables (notably `PATH`) at + container start, so the config injects the in-place TRT-LLM venv via + `TRTLLM_PATH_PREPEND` / `TRTLLM_PYTHONPATH_PREPEND`, applied inside + the container by `start_worker.sh` / `start_server.sh` / + `run_benchmark.sh`. + +## Current caveats (read before running) + +1. **V2 transceiver "MPI hang" — root-caused, environmental (RESOLVED + with the env pins).** A reported multi-node hang in + `KvCacheTransceiverV2._exchange_rank_info` → `mpi_allgather` was a + downstream symptom of `UCX_TLS=all` on nodes where `ud_verbs` cannot + initialize: the broken transport wedges native NIXL/UCX agent init + asymmetrically per rank, and the healthy ranks park forever in the + setup MPI collectives. Not an MPI/pmix or V2 code bug; with + `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` V2 NIXL passes multi-node + with no code change. +2. **No speculative decoding yet.** These configs run the gen server + without spec decode. Suffix-automaton (SA) speculative decoding for + K3 disagg is validated on the feature branch and lands in a separate + 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). +4. **Cluster environment** (NVL72 nodes): on clusters where verbs + transports cannot initialize, pin + `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` (`UCX_TLS=all` hangs setup, + see caveat 1) and never run V2 NIXL with a container-default + `UCX_TLS=tcp` (breaks V2 NIXL VRAM registration) — unset/override it. + No bounce env override is needed: the byte gate + (`TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES`, default 2 MiB) is always cleared + by K3 payloads (constraints section above). +5. **Transfer payload**: each request moves a fixed 433.4 MiB (~454.5 MB) + KDA state blob ctx → gen in addition to the MLA latent KV + (~27 KB/token). + Within an NVL72 domain this is ~0.9 ms/request (measured; not a + bottleneck), but off-fabric paths would pay 11–23 ms — keep ctx and + gen inside one NVL72 domain. +6. **Bounce-buffer sizing cliff (silent).** Size `kv_cache_bounce_size_mb` + to the largest single request's full KV payload (fixed KDA state plus + the per-token MLA latent; ≥1024 MB for 8k ISL). An undersized region + does not error — every transfer silently falls back to a much slower + host-staged TCP path. +7. **Prefill capacity and TTFT under burst.** Without chunked prefill, + context-server throughput is limited and queued prefills grow TTFT + roughly linearly under closed-loop bursts. Rate-match the ctx:gen + instance ratio to the expected traffic instead of oversubscribing a + single context server. +8. **Startup time.** Weight loading takes tens of minutes per 16-GPU + instance before the first token; set health-check, idle-reaper, and + job time limits accordingly. The disaggregated proxy does not serve + `/v1/models` (404) — point readiness probes at a different endpoint. +9. **`max_num_tokens` coupling.** The generation side must cover + `max_batch_size × (1 + max_draft_len)` scheduled tokens (with spec + decode off, `max_draft_len` is 0); the context side needs + `max_tokens_in_buffer` ≥ max ISL (see constraints above). diff --git a/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml new file mode 100644 index 000000000000..bb0db632d279 --- /dev/null +++ b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml @@ -0,0 +1,136 @@ +# Kimi K3 disaggregated benchmark-harness config for +# examples/disaggregated/slurm/benchmark/submit.py, modeled on +# examples/wide_ep/slurm_scripts/kimi-k2-thinking.yaml with K3's +# constraints baked in (EP-only DEP16 on both sides, no chunked prefill, +# no KV block reuse, tokens_per_block 64, NIXL + Python V2 transceiver). +# Fill in the <...> placeholders for your cluster before submitting. +# +# Usage (from the repository root): +# python3 examples/disaggregated/slurm/benchmark/submit.py \ +# -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml [--dry-run] +# +# Gen-only baseline: set benchmark.mode to "gen_only_no_context" +# (submit.py then exports TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1). + +slurm: + script_file: "disaggr_torch.slurm" + partition: "" + account: "" + job_time: "02:00:00" + job_name: "k3-disagg-dep16" + extra_args: "--gpus-per-node=4" + numa_bind: true # GB200/GB300 NVL72 + +benchmark: + mode: "e2e" # e2e | gen_only | gen_only_no_context (gen-only baseline) + use_nv_sa_benchmark: false + multi_round: 4 + benchmark_ratio: 0.8 + streaming: true + concurrency_list: "16" + input_length: 8192 + output_length: 1024 + # Point at a benchmark dataset (e.g. one generated by + # benchmarks/cpp/prepare_dataset.py) before submitting an e2e run. + dataset_file: "" + # Benchmark client env: same in-container venv injection as the workers + # (see the environment section notes). + env_var: + TRTLLM_PATH_PREPEND: "" + TRTLLM_PYTHONPATH_PREPEND: "" + +hardware: + gpus_per_node: 4 + num_ctx_servers: 1 + num_gen_servers: 1 + +environment: + # Mount the TRT-LLM checkout and the K3 checkpoint into the container. + container_mount: "" # Format: path1:path1,path2:path2 + container_image: "" + model_path: "" + # TRT-LLM comes from the checkout's in-place install (venv on PATH, + # repo root on PYTHONPATH via worker/server env below); skip the + # harness's pip install. + trtllm_repo: "" + build_wheel: false + trtllm_wheel_path: "" + work_dir: "" + # Cluster notes: + # - On clusters where verbs UCX transports cannot initialize on the + # compute nodes, UCX_TLS must be pinned — UCX_TLS=all can wedge + # native NIXL init asymmetrically, leaving the surviving ranks hung + # in the V2 setup MPI collectives. start_worker.sh clears + # UCX_TLS (a container-provided UCX_TLS=tcp also breaks V2 NIXL VRAM + # registration), so the pin is carried via TRTLLM_WORKER_UCX_TLS and + # re-exported by start_worker.sh after the clear. + # - pairs with the workers' + # kv_cache_bounce_size_mb (fabric-VMM bounce; measured + # ~455 GB/s/GPU). The default gate of 96 blocks would silently skip + # bounce for requests under 6144 tokens at tokens_per_block=64, + # falling back to ~0.4 GB/s host-staged tcp. + # - pyxis/enroot reset image-defined PATH at container start, so the + # in-place venv is injected via TRTLLM_PATH_PREPEND / + # TRTLLM_PYTHONPATH_PREPEND, which the harness scripts apply inside + # the container (setting PATH= here would be silently dropped). + worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 ENROOT_ALLOW_DEV=yes TRTLLM_WORKER_UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc TRTLLM_PATH_PREPEND= TRTLLM_PYTHONPATH_PREPEND=" + server_env_var: "TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_PATH_PREPEND= TRTLLM_PYTHONPATH_PREPEND=" + +worker_config: + # GEN: DEP16, no speculative decoding, CUDA graphs on (matches + # examples/kimi_k3/disagg/gen_config_no_sa.yaml). + gen: + tensor_parallel_size: 16 + moe_expert_parallel_size: 16 + enable_attention_dp: true + pipeline_parallel_size: 1 + disable_overlap_scheduler: true + enable_chunked_prefill: false + cuda_graph_config: + max_batch_size: 32 + max_batch_size: 32 + max_num_tokens: 256 + max_seq_len: 9236 + print_iter_log: true + kv_cache_config: + enable_block_reuse: false # mandatory for K3 (KDA state) + tokens_per_block: 64 # mandatory for K3 MLA kernels + free_gpu_memory_fraction: 0.25 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON # mandatory: C++ path throws for K3 + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce: needed for cross-node V2 transfers on clusters + # where inter-node cuda_ipc requires fabric allocations (else + # ~0.4 GB/s host-staged tcp); measured ~455 GB/s/GPU with + # (worker_env_var above). + # 1024 fits one full request payload (433 MiB KDA state + ~27 KB/token + # MLA latent = 649 MiB at 8k ISL); at 512 every 8k transfer falls back + # to the per-fragment tcp path. + kv_cache_bounce_size_mb: 1024 + trust_remote_code: true + # CTX: matched DEP16 (see ../disagg/ctx_config.yaml for the sizing + # rationale; hetero ctx/gen parallelism is not supported yet). + ctx: + tensor_parallel_size: 16 + moe_expert_parallel_size: 16 + enable_attention_dp: true + pipeline_parallel_size: 1 + disable_overlap_scheduler: true # mandatory for disagg ctx servers + enable_chunked_prefill: false + cuda_graph_config: null + max_batch_size: 1 + max_num_tokens: 8448 + max_seq_len: 8212 + print_iter_log: true + kv_cache_config: + enable_block_reuse: false + tokens_per_block: 64 + free_gpu_memory_fraction: 0.25 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8448 + # Must match the gen side; see the gen block's comment. + kv_cache_bounce_size_mb: 1024 + trust_remote_code: true diff --git a/examples/kimi_k3/disagg/ctx_config.yaml b/examples/kimi_k3/disagg/ctx_config.yaml new file mode 100644 index 000000000000..f3b3a3b80482 --- /dev/null +++ b/examples/kimi_k3/disagg/ctx_config.yaml @@ -0,0 +1,63 @@ +# Kimi K3 disaggregated serving - CONTEXT (prefill) server extra LLM-API +# options (`trtllm-serve --config ctx_config.yaml`). +# +# Sizing: DEP16 (attention data-parallel, experts sharded across the +# 16-rank EP group). K3 is EP-only (ep_size == tp_size, no PP, no TP on +# linears), so the ctx server must also be a DEP-N deployment. DEP16 is +# the smallest deployment verified to fit (~193 GiB weights/rank on +# 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. +# +# Hard K3 constraints baked in: no chunked prefill, no KV block reuse +# (KDA state), tokens_per_block=64 (MLA latent layout), and +# disable_overlap_scheduler=true (mandatory for any disagg ctx server). +tensor_parallel_size: 16 +moe_expert_parallel_size: 16 +enable_attention_dp: true +pipeline_parallel_size: 1 +disable_overlap_scheduler: true +enable_chunked_prefill: false +cuda_graph_config: null +# Prefill-shaped limits: one 8k-ISL request per attention-DP rank. +# max_num_tokens must cover the full ISL (no chunked prefill for K3). +max_batch_size: 1 +max_num_tokens: 8448 +max_seq_len: 8212 +kv_cache_config: + # Mandatory for K3 (per-request KDA recurrent state). + enable_block_reuse: false + # Mandatory for K3 MLA (576, 512) trtllm-gen kernels. + tokens_per_block: 64 + # Conservative starting point proven by the aggregated DEP16 evals; the + # ctx server only holds in-flight prefills (MLA latent KV is small: + # ~27 KB/token/request), so keep this low to leave prefill activation + # headroom. Tune upward only if KV allocation fails. + free_gpu_memory_fraction: 0.25 +cache_transceiver_config: + backend: NIXL + # K3 requires the Python V2 transceiver: `auto` resolves to the C++ + # runtime, which throws at construction for K3's + # MixedMambaHybridCacheManager (LinearAttentionMetadata not found). + transceiver_runtime: PYTHON + # Must be >= the target max ISL (8192 here). + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce buffer, needed for cross-node transfers on NVLink + # domains: V2 sends pool-to-pool by default, but the KV pool is a + # plain (non-fabric) allocation, so inter-node cuda_ipc/MNNVL is + # impossible and UCX falls back to ~0.4 GB/s host-staged tcp. 512 MiB + # measured at ~455 GB/s/GPU. Engages automatically for payloads above + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES (default 2 MiB) — always true for + # K3's ~433 MiB per-request recurrent state; no env override needed. + # + # Sizing: the region must fit ONE request's full payload = fixed KDA + # state (433 MiB) + MLA latent for the ISL (~27 KB/token, ~216 MiB at + # 8k). At 512 every 8k request is 649 MiB and silently falls back to + # the per-fragment tcp path ("[kv-bounce] in-place: transfer 649MiB + # exceeds the 512MiB bounce region"). 1024 + # covers the 8448-token max ISL with headroom. + kv_cache_bounce_size_mb: 1024 +# No speculative_config on the ctx server: SA drafting runs on the +# generation side only (ctx produces the prompt KV + first token). diff --git a/examples/kimi_k3/disagg/disagg_proxy_config.yaml b/examples/kimi_k3/disagg/disagg_proxy_config.yaml new file mode 100644 index 000000000000..b8e83983e694 --- /dev/null +++ b/examples/kimi_k3/disagg/disagg_proxy_config.yaml @@ -0,0 +1,23 @@ +# Kimi K3 disaggregated serving - proxy configuration for +# `trtllm-serve disaggregated -c disagg_proxy_config.yaml`. +# +# Modeled on examples/disaggregated/disagg_config.yaml. The proxy routes +# each request to a context server for prefill, then to a generation +# server for decode; the KV + KDA state moves ctx -> gen through the +# cache transceiver configured in the worker configs (NIXL, Python V2 +# runtime - see ctx_config.yaml / gen_config_no_sa.yaml). +# +# Replace hostnames/ports with the actual worker endpoints (each K3 +# worker spans 16 GPUs = 4 NVL72 nodes; the URL is rank 0's node). +hostname: localhost +port: 8000 +model: +backend: "pytorch" +context_servers: + num_instances: 1 + urls: + - "localhost:8001" +generation_servers: + num_instances: 1 + urls: + - "localhost:8002" diff --git a/examples/kimi_k3/disagg/gen_config_no_sa.yaml b/examples/kimi_k3/disagg/gen_config_no_sa.yaml new file mode 100644 index 000000000000..e5f3f92b046a --- /dev/null +++ b/examples/kimi_k3/disagg/gen_config_no_sa.yaml @@ -0,0 +1,46 @@ +# Kimi K3 disaggregated serving - GENERATION (decode) server extra LLM-API +# options WITHOUT speculative decoding +# (`trtllm-serve --config gen_config_no_sa.yaml`). +# +# Default: CUDA graphs ON. Validated end-to-end: GSM8K 96.89 (in the +# aggregated band) and 765/2138 out tok/s at c64/c256 — ~3x the eager +# TPOT and above aggregated serving throughput despite the overlap +# scheduler being off (a disagg requirement on the ctx side). +# For strict token-parity debugging (attribution-clean SA-on/SA-off or +# disagg-vs-agg comparisons), null out the cuda graph config to match +# the SA config's eager regime, and place ctx and gen in SEPARATE NVL72 +# domains (co-located placement selects different comm algorithms and +# diverges token-level on a fixed prompt subset; accuracy-neutral). +tensor_parallel_size: 16 +moe_expert_parallel_size: 16 +enable_attention_dp: true +pipeline_parallel_size: 1 +disable_overlap_scheduler: true +enable_chunked_prefill: false +cuda_graph_config: + max_batch_size: 32 +# No SA SpeculativeState buffer cap here; 32 matches the aggregated +# non-SA DEP16 eval config. +max_batch_size: 32 +max_num_tokens: 256 +max_seq_len: 9236 +kv_cache_config: + enable_block_reuse: false # mandatory for K3 (KDA state) + tokens_per_block: 64 # mandatory for K3 MLA kernels + free_gpu_memory_fraction: 0.25 +cache_transceiver_config: + backend: NIXL + # Mandatory for K3: `auto` resolves to the C++ transceiver, which + # throws for K3's MixedMambaHybridCacheManager. + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce buffer (recv side); must match the ctx server. + # Without it, cross-node V2 pool-to-pool transfers fall back to + # ~0.4 GB/s host-staged tcp; 512 MiB measured at ~455 GB/s/GPU. + # Bounce engages automatically above TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES + # (default 2 MiB) — always true for K3 payloads. + # at tokens_per_block=64). See ctx_config.yaml. + # Sized to fit one full request payload (fixed 433 MiB KDA state + + # ~27 KB/token MLA latent; 649 MiB at 8k ISL) -- see ctx_config.yaml; + # at 512 every 8k transfer falls back to the per-fragment tcp path. + kv_cache_bounce_size_mb: 1024 From cd5ab68ca6eced236e4171a00ac6abc76f4fddfe Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 6 Aug 2026 17:55:52 -0700 Subject: [PATCH 05/14] [TRTLLM-14815][fix] Do not require identical recurrent layer sets across 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 --- .../disaggregation/native/mixers/ssm/peer.py | 27 ++++++------ .../disaggregated/test_kda_mamba_transfer.py | 44 ++++++++++++++++++- 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 1e55deae36e9..f7952c29ad1a 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -372,22 +372,21 @@ def validate_peer_compatible( """ self_mlg = MambaPolicy._find_mamba_layer_group(self_page_table) peer_mlg = MambaPolicy._find_mamba_layer_group(peer_page_table) - if self_mlg is None and peer_mlg is None: + if self_mlg is None or peer_mlg is None: + # Under pipeline parallelism each rank publishes only its own + # stage's layers, so a hybrid model can pair a rank holding + # recurrent layers with a peer stage holding none. The transfer + # path (build_mamba_frags) intersects the two layer sets and + # moves nothing when either side has no recurrent layers, so + # there is nothing to validate for this pair. return - if (self_mlg is None) != (peer_mlg is None): - raise ValueError( - "MambaPolicy.validate_peer_compatible: one side has " - f"recurrent-state pools and the other does not (local={self_mlg is not None}, " - f"peer={peer_mlg is not None})" - ) - - if set(self_mlg.mamba_layer_offsets.keys()) != set(peer_mlg.mamba_layer_offsets.keys()): - raise ValueError( - "MambaPolicy.validate_peer_compatible: mamba layer sets differ " - f"(local={sorted(self_mlg.mamba_layer_offsets)}, " - f"peer={sorted(peer_mlg.mamba_layer_offsets)})" - ) + # Layer sets are NOT required to match: with pipeline parallelism the + # two sides may partition layers differently, and the transfer covers + # exactly the intersection (empty intersection moves no state). The + # invariants below are per-layer-slot quantities, uniform across a + # model's recurrent layers, so they apply regardless of which layers + # overlap. if ( self_mlg.ssm_bytes_per_head is not None and peer_mlg.ssm_bytes_per_head is not None diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 110f7365df4e..f241eeb39e54 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -298,13 +298,15 @@ def test_kda_peer_validation_accepts_supported_shapes(ctx_cfg, gen_cfg): gen_mgr.shutdown() -def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int): +def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int, layer_ids=None): """MambaLayerGroup-only page table with fake addresses (no CUDA needed).""" from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, PhysicalPool + if layer_ids is None: + layer_ids = range(1, NUM_KDA_LAYERS + 1) mlg = MambaLayerGroup( pool_group_idx=0, - mamba_layer_offsets={glid: i for i, glid in enumerate(range(1, NUM_KDA_LAYERS + 1))}, + mamba_layer_offsets={glid: i for i, glid in enumerate(layer_ids)}, conv_states=PhysicalPool(base_address=0x1000, slot_bytes=conv_slot_bytes, num_slots=8), ssm_states=PhysicalPool(base_address=0x2000000, slot_bytes=ssm_slot_bytes, num_slots=8), conv_section_bytes=[conv_slot_bytes // 3] * 3, @@ -380,6 +382,44 @@ def build(cfg): MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) +def test_kda_peer_validation_allows_pipeline_parallel_layer_split(): + """Peer validation must not require identical layer sets. + + Under pipeline parallelism each rank publishes only its own stage's + layers, and the transfer path takes the intersection of the two layer + sets. Peers with partially overlapping, disjoint, or one-sided + recurrent-layer sets are all legitimate stage pairings and must pass + validation as long as the per-slot size invariants hold. + """ + 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 + ctx_ri = _synthetic_rank_info(2, False) + gen_ri = _synthetic_rank_info(2, False) + + # Partially overlapping stages (ctx PP split differs from gen's). + ctx_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[1, 2, 3]) + gen_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[3, 4, 5]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + # Disjoint stages: this pair simply moves no recurrent state. + gen_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[7, 8]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + # One side holds a recurrent-layer-free stage of the hybrid model + # (no MambaLayerGroup at all): nothing to validate. + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable + + empty_pt = KVCachePageTable(tokens_per_block=8, layer_groups=[], pool_groups=[]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, empty_pt) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, empty_pt, gen_pt) + + # Size invariants still apply on the overlap: mismatched global state + # sizes are rejected even when the layer sets differ. + bad_gen_pt = _synthetic_kda_page_table(full_ssm // 2, full_conv // 2, layer_ids=[3, 4, 5]) + with pytest.raises(ValueError, match="TP-aggregated"): + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, bad_gen_pt) + + # --------------------------------------------------------------------------- # Loopback transfer over a real NIXL agent (single node) # --------------------------------------------------------------------------- From ce960671c80a30a744585e02f3589cd3108d60e1 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 6 Aug 2026 20:04:34 -0500 Subject: [PATCH 06/14] Address trivial review comments Signed-off-by: Brian Nguyen --- examples/disaggregated/slurm/benchmark/run_benchmark.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark.sh b/examples/disaggregated/slurm/benchmark/run_benchmark.sh index 7ccae84a6e69..154bf7840d7c 100644 --- a/examples/disaggregated/slurm/benchmark/run_benchmark.sh +++ b/examples/disaggregated/slurm/benchmark/run_benchmark.sh @@ -10,7 +10,9 @@ trap 'echo "Error occurred at line $LINENO"; exit 1' ERR # Allow the launcher config to prepend entries from inside the container. # The client_cmds srun passes these via --export with single-quoted values # (submit.py convert_envs_to_str), and srun keeps the quotes literal — strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" From 8d566d72230004bcdc48b87f091b50e7bb12c131 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 6 Aug 2026 20:16:03 -0700 Subject: [PATCH 07/14] [TRTLLM-14815][fix] Address review comments: test teardown, parity-harness 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 --- .../slurm/benchmark/start_server.sh | 5 + .../slurm/benchmark/start_worker.sh | 5 + .../integration/defs/kimi_k3_disagg_parity.py | 6 +- .../disaggregated/test_kda_mamba_transfer.py | 251 +++++++++--------- 4 files changed, 144 insertions(+), 123 deletions(-) diff --git a/examples/disaggregated/slurm/benchmark/start_server.sh b/examples/disaggregated/slurm/benchmark/start_server.sh index 570b9ceb30dc..690748a7d816 100644 --- a/examples/disaggregated/slurm/benchmark/start_server.sh +++ b/examples/disaggregated/slurm/benchmark/start_server.sh @@ -8,6 +8,11 @@ config_file=$1 # Container runtimes (pyxis/enroot) reset image-defined variables like PATH # at container start, so values passed via srun --export are lost for them. # Allow the launcher config to prepend entries from inside the container. +# srun --export keeps any quotes in the exported values literal; strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" fi diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 0c8653893958..3f695bd48a32 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -41,6 +41,11 @@ fi # Container runtimes (pyxis/enroot) reset image-defined variables like PATH # at container start, so values passed via srun --export are lost for them. # Allow the launcher config to prepend entries from inside the container. +# srun --export keeps any quotes in the exported values literal; strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" fi diff --git a/tests/integration/defs/kimi_k3_disagg_parity.py b/tests/integration/defs/kimi_k3_disagg_parity.py index ae5759e473d9..b692f33cd1d7 100644 --- a/tests/integration/defs/kimi_k3_disagg_parity.py +++ b/tests/integration/defs/kimi_k3_disagg_parity.py @@ -126,6 +126,9 @@ def _served_model(base_url, timeout): except urllib.error.HTTPError as e: print(f"[parity] NOTE: {base_url}/v1/models unavailable (HTTP {e.code})") return None + except (urllib.error.URLError, OSError, KeyError, IndexError) as e: + print(f"[parity] NOTE: {base_url}/v1/models unavailable ({e})") + return None class Endpoint: @@ -540,7 +543,8 @@ def main(argv=None) -> int: if cand_model is not None and cand_model != model: print( f"[parity] NOTE: endpoints serve different model names " - f"({model!r} vs {cand_model!r}); using each server's own" + f"({model!r} vs {cand_model!r}); sending {model!r} to both, " + "pass --model to override" ) report["model"] = model diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index f241eeb39e54..778f0d878cd0 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -253,8 +253,9 @@ def test_kda_hetero_tp_rejected(): end of K3's replicated (pre-scaled) slots. The peer-registration gate (``MambaPolicy.validate_peer_compatible``) must reject this loudly. """ - ctx_mgr = _create_kda_managers(2, enable_attention_dp=False)[0] - gen_mgr = _create_kda_managers(4, enable_attention_dp=False)[1] + ctx_mgrs = _create_kda_managers(2, enable_attention_dp=False) + gen_mgrs = _create_kda_managers(4, enable_attention_dp=False) + ctx_mgr, gen_mgr = ctx_mgrs[0], gen_mgrs[1] try: ctx_pt = build_page_table_from_manager(ctx_mgr) gen_pt = build_page_table_from_manager(gen_mgr) @@ -269,8 +270,8 @@ def test_kda_hetero_tp_rejected(): with pytest.raises(ValueError, match="TP-aggregated"): registrar.register("kda_gen", 1, gen_ri) finally: - ctx_mgr.shutdown() - gen_mgr.shutdown() + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() @pytest.mark.parametrize( @@ -285,8 +286,9 @@ def test_kda_peer_validation_accepts_supported_shapes(ctx_cfg, gen_cfg): """Matched-TP and ADP-on-both-sides layouts must pass peer validation.""" ctx_tp, ctx_adp = ctx_cfg gen_tp, gen_adp = gen_cfg - ctx_mgr = _create_kda_managers(ctx_tp, enable_attention_dp=ctx_adp)[0] - gen_mgr = _create_kda_managers(gen_tp, enable_attention_dp=gen_adp)[-1] + ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=ctx_adp) + gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=gen_adp) + ctx_mgr, gen_mgr = ctx_mgrs[0], gen_mgrs[-1] try: ctx_pt = build_page_table_from_manager(ctx_mgr) gen_pt = build_page_table_from_manager(gen_mgr) @@ -294,8 +296,8 @@ def test_kda_peer_validation_accepts_supported_shapes(ctx_cfg, gen_cfg): gen_ri = RankInfo.from_kv_cache_manager("kda_gen", gen_mgr, device_id=0) MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) finally: - ctx_mgr.shutdown() - gen_mgr.shutdown() + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int, layer_ids=None): @@ -452,124 +454,129 @@ def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = """Loopback: matched ctx/gen TP, replicated full-size KDA state per rank.""" 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) - for mgr in ctx_mgrs + gen_mgrs: - mgr._impl.mamba_cache.conv.zero_() - mgr._impl.mamba_cache.temporal.zero_() - - config = CacheTransceiverConfig( - backend="NIXL", - 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_endpoint = ctx_tcs[0]._context_info_endpoint - - sampling_params = SamplingParams() - ctx_rids, gen_rids, ctx_reqs, gen_reqs = [], [], [], [] - for req_idx, req_len in enumerate(REQUEST_LENGTHS): - unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF - ctx_rid, gen_rid = req_idx * 2, req_idx * 2 + 1 - ctx_rids.append(ctx_rid) - gen_rids.append(gen_rid) - sc = tensorrt_llm.bindings.SamplingConfig(sampling_params._get_sampling_config()) - ctx_req = LlmRequest( - request_id=ctx_rid, - max_new_tokens=1, - input_tokens=list(range(req_len)), - sampling_config=sc, - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, - ) - ctx_req.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) - gen_req = LlmRequest( - request_id=gen_rid, - max_new_tokens=1, - input_tokens=list(range(req_len)), - sampling_config=sc, - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_tcs, gen_tcs = [], [] + try: + for mgr in ctx_mgrs + gen_mgrs: + mgr._impl.mamba_cache.conv.zero_() + mgr._impl.mamba_cache.temporal.zero_() + + config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + max_tokens_in_buffer=512, ) - gen_req.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, - ctx_dp_rank=0, - ctx_info_endpoint=ctx_endpoint, - disagg_request_id=unique_rid, + ctx_tcs = _create_transceivers(ctx_tp, ctx_mgrs, config) + gen_tcs = _create_transceivers(gen_tp, gen_mgrs, config) + ctx_endpoint = ctx_tcs[0]._context_info_endpoint + + sampling_params = SamplingParams() + ctx_rids, gen_rids, ctx_reqs, gen_reqs = [], [], [], [] + for req_idx, req_len in enumerate(REQUEST_LENGTHS): + unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + ctx_rid, gen_rid = req_idx * 2, req_idx * 2 + 1 + ctx_rids.append(ctx_rid) + gen_rids.append(gen_rid) + sc = tensorrt_llm.bindings.SamplingConfig(sampling_params._get_sampling_config()) + ctx_req = LlmRequest( + request_id=ctx_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + ) + ctx_req.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + gen_req = LlmRequest( + request_id=gen_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + gen_req.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=ctx_rid, + ctx_dp_rank=0, + ctx_info_endpoint=ctx_endpoint, + disagg_request_id=unique_rid, + ) + ctx_reqs.append(ctx_req) + gen_reqs.append(gen_req) + + ctx_batch = ScheduledRequests() + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) + gen_batch = ScheduledRequests() + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) + for req in ctx_reqs + gen_reqs: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + for mgr in ctx_mgrs: + mgr.update_resources(ctx_batch) + for mgr in gen_mgrs: + mgr.update_resources(gen_batch) + + # Ground truth: identical full state on every ctx rank (replicated). + ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) + for mgr in ctx_mgrs: + 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: + gen_tcs[rank].request_and_receive_async(req) + for rank in range(ctx_tp): + for req in ctx_reqs: + ctx_tcs[rank].respond_and_send_async(req) + _run_concurrent( + ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) ) - ctx_reqs.append(ctx_req) - gen_reqs.append(gen_req) - - ctx_batch = ScheduledRequests() - ctx_batch.reset_context_requests(ctx_reqs) - for mgr in ctx_mgrs: - mgr.prepare_resources(ctx_batch) - gen_batch = ScheduledRequests() - gen_batch.reset_context_requests(gen_reqs) - for mgr in gen_mgrs: - mgr.prepare_resources(gen_batch) - for req in ctx_reqs + gen_reqs: - req.context_current_position = req.prompt_len - req.add_new_token(req.prompt_len, 0) - for mgr in ctx_mgrs: - mgr.update_resources(ctx_batch) - for mgr in gen_mgrs: - mgr.update_resources(gen_batch) - - # Ground truth: identical full state on every ctx rank (replicated). - ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) - for mgr in ctx_mgrs: - 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): + _run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) + + # 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) + rank_factor = 1 if enable_attention_dp else gen_tp for req in gen_reqs: - gen_tcs[rank].request_and_receive_async(req) - for rank in range(ctx_tp): - for req in ctx_reqs: - ctx_tcs[rank].respond_and_send_async(req) - _run_concurrent(ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True)) - _run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) - - # 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) - rank_factor = 1 if enable_attention_dp else gen_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})" - ) + 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). - for gen_rank, mgr in enumerate(gen_mgrs): - for req_idx, rid in enumerate(gen_rids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: - full = ground_truth[req_idx][layer_idx] - 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], - rtol=0, - atol=0, - msg=lambda m, n=name, r=gen_rank, ri=req_idx, li=layer_idx: ( - f"{n} mismatch: gen_rank={r} req={ri} layer={li} " - f"ctx_tp={ctx_tp} gen_tp={gen_tp}: {m}" - ), - ) - - for mgr in ctx_mgrs + gen_mgrs: - mgr.shutdown() - for tc in ctx_tcs + gen_tcs: - tc.shutdown() + # Bitwise comparison on every gen rank (state is replicated). + for gen_rank, mgr in enumerate(gen_mgrs): + for req_idx, rid in enumerate(gen_rids): + slot = mgr.mamba_cache_index[rid] + for layer_idx in mgr._impl.mamba_layer_offsets: + full = ground_truth[req_idx][layer_idx] + 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], + rtol=0, + atol=0, + msg=lambda m, n=name, r=gen_rank, ri=req_idx, li=layer_idx: ( + f"{n} mismatch: gen_rank={r} req={ri} layer={li} " + f"ctx_tp={ctx_tp} gen_tp={gen_tp}: {m}" + ), + ) + + finally: + for tc in ctx_tcs + gen_tcs: + tc.shutdown() + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() @pytest.mark.timeout(180) From 084f8b7245b014391d0b739cbe0382d98b4fa6df Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 6 Aug 2026 21:19:04 -0700 Subject: [PATCH 08/14] [TRTLLM-14815][test] Drop removed RankInfo.server_endpoint field from synthetic rank info Signed-off-by: Brian Nguyen --- tests/unittest/disaggregated/test_kda_mamba_transfer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 778f0d878cd0..fcef4587c05b 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -329,7 +329,6 @@ def _synthetic_rank_info(tp: int, adp: bool): pp_rank=0, layer_num_per_pp=[_NUM_TOTAL_LAYERS], sender_endpoints=[], - server_endpoint="", self_endpoint="", transfer_engine_info=b"", attention=AttentionInfo( From 887fa864c0bffceb7a4dc5076d5f12362dfdd599 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 01:39:07 -0700 Subject: [PATCH 09/14] [TRTLLM-14815][fix] Restore plain-KV bounce gate; relaxed gate only for recurrent-state payloads Signed-off-by: Brian Nguyen --- .../disaggregation/native/bounce/config.py | 38 ++++++++------ .../disaggregation/native/bounce/impl.py | 41 +++++++++------ .../_torch/disaggregation/transceiver.py | 5 +- ...lap_transceiver_runtime_python_bounce.yaml | 7 ++- tests/unittest/disaggregated/test_bounce.py | 52 +++++++++++-------- 5 files changed, 81 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py index 64c218385acc..ae6691ce6da6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py @@ -25,8 +25,8 @@ # 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" # the operative gate, in bytes -_MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" # legacy block-count gate +_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 def _env_int_gate(name: str, default: int) -> int: @@ -102,14 +102,15 @@ def fit_within_free( return capacity_bytes -# Skip bounce below this many bytes. The cost this gate guards scales with BYTES, not blocks: an -# earlier block-count gate (96 blocks, calibrated for 128-token blocks) silently skipped a 433 MiB -# Kimi-K3 transfer (67 blocks of 32 tokens) and dropped it onto the ~0.4 GB/s host-staged fallback, -# a ~1000x cliff. Break-even is small: bounce adds one gather plus one scatter copy (device-local, -# ~hundreds of GB/s) and ~0.1 ms of fixed launch/reservation overhead, while the in-place path can -# be as slow as ~0.4 GB/s 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. +# Byte gate for transfers that carry recurrent (mamba/KDA) state. The cost this gate guards scales +# with BYTES, not blocks: the block-count gate (96 blocks, calibrated for 128-token blocks) +# silently skipped a 433 MiB Kimi-K3 transfer (67 blocks of 32 tokens plus the non-paged KDA state) +# and dropped it onto the ~0.4 GB/s host-staged fallback, a ~1000x cliff. Break-even is small: +# bounce adds one gather plus one scatter copy (device-local, ~hundreds of GB/s) and ~0.1 ms of +# fixed launch/reservation overhead, while the in-place path can be as slow as ~0.4 GB/s +# 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 @@ -117,20 +118,23 @@ def fit_within_free( class Config: sizing: Sizing = field(default_factory=FixedSizing) # how much memory to reserve (pluggable) chunk_mb: int = 32 # physical chunk size; a large chunk keeps the write to a single descriptor - # skip bounce below this many bytes (the operative gate; see DEFAULT_MIN_BYTES for the rationale) + # Which gate applies depends on the payload (see VmmBounceTransport.reserve): transfers that + # carry recurrent state use min_bytes; plain-KV transfers keep the original min_blocks gate so + # pre-existing bounce deployments see no behavior change. + # byte gate for recurrent-state payloads (see DEFAULT_MIN_BYTES for the rationale) min_bytes: int = DEFAULT_MIN_BYTES - # legacy block-count gate, kept for back-compat (TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS): both gates - # must pass, and the default of 1 makes this one vacuous so the byte gate decides - min_blocks: int = 1 + # block-count gate for plain-KV payloads (roughly 12k tokens at 128 per block); heuristic, + # tunable via TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS + min_blocks: int = 96 def config_from_size( size_mb: int, min_blocks: Optional[int] = None, min_bytes: Optional[int] = None ) -> Optional[Config]: """Build a bounce config from a per-region size in MiB, or None to leave bounce off (size <= 0). - Size is both the capacity and the on/off switch. min_bytes (and the legacy min_blocks) is the - gate below which a transfer stays on the per-block path; when unset it comes from the env, else - the default.""" + Size is both the capacity and the on/off switch. min_bytes (recurrent-state payloads) and + min_blocks (plain-KV payloads) are the gates below which a transfer stays on the per-block + path; when unset they come from the env, else the defaults.""" if size_mb is None or size_mb <= 0: return None if min_blocks is None: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 0d2bc5ca0323..9085a1214d53 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -96,7 +96,7 @@ def __init__( phys_chunk_size: int, block_bytes_per_group: List[int], min_bytes: int = DEFAULT_MIN_BYTES, - min_blocks: int = 1, + min_blocks: int = 96, quarantine_grace_s: float = _QUARANTINE_GRACE_S, name: str = "kv_bounce", ): @@ -104,9 +104,10 @@ def __init__( self._device_id = device_id # The byte size of one cache block, listed for each attention layer group. self._block_bytes_per_group = list(block_bytes_per_group) - # Below this many bytes, skip bounce: coalescing only pays off once the transfer is large - # enough to beat the gather+scatter overhead (see config.DEFAULT_MIN_BYTES for the - # rationale). min_blocks is the legacy block-count gate, vacuous at its default of 1. + # Size gates below which bounce is skipped: coalescing only pays off once the transfer is + # large enough to beat the gather+scatter overhead (see config.DEFAULT_MIN_BYTES for the + # rationale). min_bytes applies to payloads carrying recurrent (mamba/KDA) state; + # min_blocks applies to plain-KV payloads (see the gate in reserve()). self._min_bytes = min_bytes self._min_blocks = min_blocks # how long an orphaned region is held out of reuse; must outlast the worst in-flight write @@ -263,20 +264,28 @@ def reserve( return self._skip_bounce( f"computed transfer size {total} <= 0", warn_key="kv-bounce-nonpositive-size" ) - # The size gate is expressed in BYTES: the cost it guards (falling back to the slow - # per-fragment path) scales with bytes, not blocks, and blocks vary ~20x in size across - # models. min_blocks is the legacy gate, vacuous by default. + # Which size gate applies depends on the payload. Payloads carrying recurrent (mamba/KDA) + # state gate on BYTES: the cost the gate guards (falling back to the slow per-fragment + # path) scales with bytes, and a block count is meaningless for the non-paged state (the + # Kimi-K3 regression: a 433 MiB transfer of 67 small blocks plus KDA state failed a + # 96-block gate and fell onto the ~0.4 GB/s host-staged path). Plain-KV payloads keep the + # original block-count gate so pre-existing bounce deployments (opted in via + # kv_cache_bounce_size_mb) see no change in which transfers use the arena. + # TODO(TRTLLM followup, ticket to be filed): investigate whether the byte-only gate is + # safe (or better) for plain-KV payloads too, so this special case can be removed and + # both payload kinds share one gate. nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) - if total < self._min_bytes: - return self._skip_bounce( - f"{total}B ({nblocks} blocks) < min {self._min_bytes}B (too small; tune " - f"TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES)", - warn_key="kv-bounce-below-min-bytes", - ) - if nblocks < self._min_blocks: + if extra_bytes > 0: + if total < self._min_bytes: + return self._skip_bounce( + f"{total}B ({nblocks} blocks + recurrent state) < min {self._min_bytes}B " + f"(too small; tune TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES)", + warn_key="kv-bounce-below-min-bytes", + ) + elif nblocks < self._min_blocks: return self._skip_bounce( - f"{nblocks} blocks < min {self._min_blocks} (too small; legacy " - f"TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS gate)", + f"{nblocks} blocks < min {self._min_blocks} (too small; tune " + f"TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS)", warn_key="kv-bounce-below-min-blocks", ) if num_writers > 1 and total % num_writers != 0: diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index d065534396a0..0d7e3429424b 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -117,8 +117,9 @@ def __init__( max_concurrent_sessions=max(1, int(kv_cache_manager.max_batch_size)) * 20000, tx_timeout_s=self._sender_future_timeout_ms / 1000.0, rx_timeout_s=self.kv_transfer_timeout_ms / 1000.0, - # Size 0 turns bounce off; the byte-size gate is internal (tuned via env: - # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES, plus the legacy ..._MIN_BLOCKS). + # Size 0 turns bounce off; the per-transfer size gates are internal (tuned via + # env: TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads). bounce=bounce_config_from_size(cache_transceiver_config.kv_cache_bounce_size_mb), ) ) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml index 4c0bfee7e5d9..b4b3cd1234b9 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml @@ -3,10 +3,9 @@ # Same as disagg_config_overlap_transceiver_runtime_python.yaml, plus the bounce switch on both the # context (sender) and generation (receiver) cache_transceiver_config: # kv_cache_bounce_size_mb: >0 turns bounce on and sizes the per-region fabric-VMM arena. -# The byte gate below which a transfer keeps the per-block path is lowered to 1 via the -# TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES env (set by the test) so the ordinary short test prompts still -# take the coalesced-bounce WRITE path (the production default of 2 MiB may exceed a short prompt's -# KV footprint on a tiny model). +# Plain-KV transfers gate on block count; the gate is lowered to 1 via the +# TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS env (set by the test) so the ordinary short test prompts still +# take the coalesced-bounce WRITE path (the production default of 96 would need a ~2k-token prompt). # GB200/GB300 only, since the bounce arena is fabric (MNNVL) VMM memory. model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 38d72164b4d3..84a7eda4fe1a 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -104,9 +104,9 @@ def test_config_defaults(self): cfg = bcfg.Config() assert isinstance(cfg.sizing, bcfg.FixedSizing) assert cfg.chunk_mb == 32 - # The operative gate is the byte one; the legacy block gate defaults to 1 (vacuous). + # Byte gate for recurrent-state payloads; block gate for plain-KV payloads. assert cfg.min_bytes == bcfg.DEFAULT_MIN_BYTES == 2 * _MIB - assert cfg.min_blocks == 1 + assert cfg.min_blocks == 96 # --------------------------------------------------------------------------- # @@ -132,21 +132,21 @@ def test_min_bytes_defaults_and_overrides(self, monkeypatch): assert bcfg.config_from_size(2048).min_bytes == 1 # env override assert bcfg.config_from_size(2048, min_bytes=250).min_bytes == 250 # arg beats env - def test_min_blocks_backcompat_defaults_and_overrides(self, monkeypatch): - # The legacy block-count gate is kept for back-compat; it defaults to 1 (vacuous, so the - # byte gate decides) and still honors the explicit arg and the env override. + def test_min_blocks_defaults_and_overrides(self, monkeypatch): + # The block-count gate (plain-KV payloads) keeps its original default of 96 and honors + # the explicit arg and the env override. monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", raising=False) - assert bcfg.config_from_size(2048).min_blocks == 1 # keeps the Config default + assert bcfg.config_from_size(2048).min_blocks == 96 # keeps the Config default assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg still overrides - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "96") - assert bcfg.config_from_size(2048).min_blocks == 96 # env override + monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "48") + assert bcfg.config_from_size(2048).min_blocks == 48 # env override assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg beats env @pytest.mark.parametrize( "env,attr,default", [ ("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "min_bytes", 2 * _MIB), - ("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "min_blocks", 1), + ("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "min_blocks", 96), ], ) def test_gate_env_is_parsed_defensively(self, monkeypatch, env, attr, default): @@ -408,31 +408,37 @@ def test_reserve_single_writer_ok(self, monkeypatch): assert t.reserve(req, num_writers=1) is True def test_reserve_below_min_bytes_falls_back(self, monkeypatch): - # The gate is in BYTES: 4 blocks x 100B = 400B < 1000B falls back, regardless of how many - # blocks that is (the old block-count gate silently skipped huge small-block transfers). + # Recurrent-state payloads gate on BYTES: 4 blocks x 100B + 100B state = 500B < 1000B + # falls back, regardless of how many blocks that is. t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) req = _recv_req([4]) - assert t.reserve(req, num_writers=1) is False + assert t.reserve(req, num_writers=1, extra_bytes=100) is False assert req.bounce_dst_base is None def test_reserve_at_min_bytes_bounces(self, monkeypatch): - # exactly at the threshold (10 x 100B = 1000B >= 1000B) the transfer bounces + # exactly at the threshold (9 x 100B + 100B state = 1000B >= 1000B) the transfer bounces t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) - req = _recv_req([10]) - assert t.reserve(req, num_writers=1) is True + req = _recv_req([9]) + assert t.reserve(req, num_writers=1, extra_bytes=100) is True assert req.bounce_dst_base == 0x100000 - def test_reserve_few_large_blocks_bounce(self, monkeypatch): - # The K3 regression shape: FEW but LARGE blocks must clear a byte gate that a block-count - # gate of 96 would have silently failed (67 blocks x 6.5 MiB = 433 MiB). + def test_reserve_recurrent_payload_ignores_block_gate(self, monkeypatch): + # The K3 regression shape: FEW but LARGE blocks plus recurrent state must clear the byte + # gate even though a block-count gate of 96 fails (67 blocks x 6.5 MiB = 433 MiB). t = _make_transport( - monkeypatch, block_bytes_per_group=[int(6.5 * _MIB)], min_bytes=2 * _MIB + monkeypatch, + block_bytes_per_group=[int(6.5 * _MIB)], + min_bytes=2 * _MIB, + min_blocks=96, ) - assert t.reserve(_recv_req([67]), num_writers=1) is True + assert t.reserve(_recv_req([67]), num_writers=1, extra_bytes=1024) is True - def test_reserve_legacy_min_blocks_backcompat(self, monkeypatch): - # the legacy block-count gate still applies when raised explicitly (back-compat) - t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_blocks=96) + def test_reserve_plain_kv_uses_block_gate(self, monkeypatch): + # Plain-KV payloads (no recurrent state) keep the original block-count gate, and the byte + # gate does not apply to them (96 x 100B = 9600B passes despite min_bytes far above it). + t = _make_transport( + monkeypatch, block_bytes_per_group=[100], min_blocks=96, min_bytes=1 << 20 + ) assert t.reserve(_recv_req([4]), num_writers=1) is False # 4 < 96 blocks assert t.reserve(_recv_req([96]), num_writers=1) is True From 29da01d7d38a4df24a692ce92233e1261834f3da Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 01:41:50 -0700 Subject: [PATCH 10/14] [TRTLLM-14815][chore] Reference TRTLLM-15194 in the bounce-gate TODO Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/disaggregation/native/bounce/impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 9085a1214d53..5d88485ec8b1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -271,7 +271,7 @@ def reserve( # 96-block gate and fell onto the ~0.4 GB/s host-staged path). Plain-KV payloads keep the # original block-count gate so pre-existing bounce deployments (opted in via # kv_cache_bounce_size_mb) see no change in which transfers use the arena. - # TODO(TRTLLM followup, ticket to be filed): investigate whether the byte-only gate is + # TODO(TRTLLM-15194): investigate whether the byte-only gate is # safe (or better) for plain-KV payloads too, so this special case can be removed and # both payload kinds share one gate. nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) From 8f670a5e1648d16eae2b624359b8f4b3f8905cfe Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 02:14:14 -0700 Subject: [PATCH 11/14] [TRTLLM-14815][test] Make test_peer synthetic mamba fixture TP-consistent 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 --- tests/unittest/disaggregated/test_peer.py | 28 +++++++++++++++-------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 96a3a95186aa..604d4ab8d93e 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -47,8 +47,16 @@ pytestmark = pytest.mark.cpu_only -def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None): - """Create a KVCachePageTable for testing.""" +def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mamba_tp=2): + """Create a KVCachePageTable for testing. + + ``mamba_tp`` is the TP degree the synthetic recurrent-state group is + sharded for: per-rank sizes are the ``mamba_tp`` shard of a fixed global + state, so page tables built for different TP degrees stay + peer-compatible under MambaPolicy's global (TP-aggregated) size check. + Pass the rank's effective mamba TP (``make_rankinfo`` defaults to + ``tp_size=2``). + """ if pool_ptrs is None: pool_ptrs = [1234] if block_bytes is None: @@ -87,9 +95,9 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None): mamba_lg = MambaLayerGroup( pool_group_idx=1, mamba_layer_offsets={100: 0, 101: 1}, - conv_states=PhysicalPool(base_address=0xA000, slot_bytes=2048, num_slots=128), - ssm_states=PhysicalPool(base_address=0xB000, slot_bytes=4096, num_slots=128), - conv_section_bytes=[512, 256, 256], + conv_states=PhysicalPool(base_address=0xA000, slot_bytes=4096 // mamba_tp, num_slots=128), + ssm_states=PhysicalPool(base_address=0xB000, slot_bytes=8192 // mamba_tp, num_slots=128), + conv_section_bytes=[1024 // mamba_tp, 512 // mamba_tp, 512 // mamba_tp], ssm_bytes_per_head=64, ) pool_groups = [PhysicalPoolGroup(pools=physical_pools)] @@ -558,7 +566,7 @@ def test_peer_registrar_get_kv_map_head_mismatch(): tokens_per_block=16, dims_per_head=8, layer_num_per_pp=[2], - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) @@ -1026,7 +1034,7 @@ def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): element_bytes=0.5, kv_heads_per_rank=4, tp_size=1, - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) @@ -1056,7 +1064,7 @@ def test_peer_registrar_rejects_misaligned_subbyte_head_mismatch(): tokens_per_block=1, dims_per_head=1, tp_size=1, - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg = _make_peer_registrar(self_ri) @@ -1066,7 +1074,7 @@ def test_peer_registrar_rejects_misaligned_subbyte_head_mismatch(): def test_peer_registrar_dispatches_nhd_mapper(): self_pt = make_page_table() - peer_pt = make_page_table(block_bytes=[2048]) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1) self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD self_ri = make_rankinfo( @@ -1090,7 +1098,7 @@ def test_peer_registrar_dispatches_nhd_mapper(): def test_peer_registrar_warns_for_nhd_head_mismatch(monkeypatch): self_pt = make_page_table() - peer_pt = make_page_table(block_bytes=[2048]) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1) self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD self_ri = make_rankinfo(kv_heads_per_rank=2, tp_size=2, page_table=self_pt) From 9660f978cfa2c8138a7293cd25d11e616826d0a7 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 04:37:50 -0700 Subject: [PATCH 12/14] [TRTLLM-14815][doc] Set trust_remote_code in the disagg example serve 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 --- examples/kimi_k3/disagg/ctx_config.yaml | 4 ++++ examples/kimi_k3/disagg/gen_config_no_sa.yaml | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/examples/kimi_k3/disagg/ctx_config.yaml b/examples/kimi_k3/disagg/ctx_config.yaml index f3b3a3b80482..4be9ff24473b 100644 --- a/examples/kimi_k3/disagg/ctx_config.yaml +++ b/examples/kimi_k3/disagg/ctx_config.yaml @@ -14,6 +14,10 @@ # Hard K3 constraints baked in: no chunked prefill, no KV block reuse # (KDA state), tokens_per_block=64 (MLA latent layout), and # disable_overlap_scheduler=true (mandatory for any disagg ctx server). +# K3 ships a custom (trust-remote-code) tokenizer; without this the +# server starts tokenizer-less and 400s every string prompt with +# "tokenizer is required to tokenize string prompt". +trust_remote_code: true tensor_parallel_size: 16 moe_expert_parallel_size: 16 enable_attention_dp: true diff --git a/examples/kimi_k3/disagg/gen_config_no_sa.yaml b/examples/kimi_k3/disagg/gen_config_no_sa.yaml index e5f3f92b046a..ccb1f7812a2f 100644 --- a/examples/kimi_k3/disagg/gen_config_no_sa.yaml +++ b/examples/kimi_k3/disagg/gen_config_no_sa.yaml @@ -11,6 +11,10 @@ # the SA config's eager regime, and place ctx and gen in SEPARATE NVL72 # domains (co-located placement selects different comm algorithms and # diverges token-level on a fixed prompt subset; accuracy-neutral). +# K3 ships a custom (trust-remote-code) tokenizer; without this the +# server starts tokenizer-less and 400s every string prompt with +# "tokenizer is required to tokenize string prompt". +trust_remote_code: true tensor_parallel_size: 16 moe_expert_parallel_size: 16 enable_attention_dp: true From 4b80e74f7e607ddb9be17dd22126fae120b19493 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 08:55:19 -0700 Subject: [PATCH 13/14] [TRTLLM-14815][test] Harden parity-harness /v1/models handling against malformed responses Signed-off-by: Brian Nguyen --- .../integration/defs/kimi_k3_disagg_parity.py | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/kimi_k3_disagg_parity.py b/tests/integration/defs/kimi_k3_disagg_parity.py index b692f33cd1d7..f327899870a2 100644 --- a/tests/integration/defs/kimi_k3_disagg_parity.py +++ b/tests/integration/defs/kimi_k3_disagg_parity.py @@ -126,7 +126,11 @@ def _served_model(base_url, timeout): except urllib.error.HTTPError as e: print(f"[parity] NOTE: {base_url}/v1/models unavailable (HTTP {e.code})") return None - except (urllib.error.URLError, OSError, KeyError, IndexError) as e: + # URLError/OSError: connection-level failures. KeyError/IndexError/ + # TypeError: malformed payload shapes ([], {"data": None}, + # {"data": [None]}, ...). ValueError covers json.JSONDecodeError from + # non-JSON response bodies. + except (urllib.error.URLError, OSError, KeyError, IndexError, TypeError, ValueError) as e: print(f"[parity] NOTE: {base_url}/v1/models unavailable ({e})") return None @@ -742,6 +746,40 @@ def check(name, cond): _diff_gsm8k({"exact_match,strict-match": 0.90}, {"exact_match,strict-match": 0.80}, 0.02, f) check("gsm8k tolerance", ok_within and len(f) == 1) + # 11. _served_model returns None for malformed /v1/models payloads, + # invalid JSON, and connection failures; extracts the id otherwise. + global _http_json + orig_http_json = _http_json + bad_responses = [ + [], + {"data": None}, + {"data": [None]}, + {"data": [{}]}, + json.JSONDecodeError("Expecting value", "not-json", 0), + urllib.error.URLError("connection refused"), + ] + try: + results = [] + for rsp in bad_responses: + + def _canned(url, payload=None, timeout=0, _rsp=rsp): + if isinstance(_rsp, Exception): + raise _rsp + return _rsp + + _http_json = _canned + results.append(_served_model("http://stub", timeout=1)) + + def _good(url, payload=None, timeout=0): + return {"data": [{"id": "the-model"}]} + + _http_json = _good + served = _served_model("http://stub", timeout=1) + finally: + _http_json = orig_http_json + check("_served_model malformed payloads -> None", all(r is None for r in results)) + check("_served_model well-formed payload", served == "the-model") + failed = [name for name, cond in checks if not cond] if failed: print(f"[self-test] FAIL ({len(failed)}/{len(checks)}): {failed}") From 89bc6e4d9b5b10c4bac2e845833dfe3f6d83411d Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 23:06:10 -0700 Subject: [PATCH 14/14] [TRTLLM-14815][doc] Remove the disaggregated-serving limitation from 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 --- examples/kimi_k3/README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 00349c6d6d43..e58c618b5628 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -193,4 +193,3 @@ default cache manager. TRTLLM-14904. - FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported. - Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). Combining speculation with disaggregated serving is not yet supported. -- Disaggregated serving is not yet supported.