From ca6978ea078631177ca0c7238029c43cf3d12e96 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 27 Oct 2026 11:09:44 -0700 Subject: [PATCH 1/4] [None][feat] Reshard KV transfers using per-view attention heads Signed-off-by: Allison Lim --- .../native/mixers/attention/peer.py | 113 +++++++--- .../_torch/disaggregation/native/peer.py | 211 +++++++++++++++--- .../_torch/disaggregation/native/transfer.py | 12 +- .../disaggregation/resource/kv_extractor.py | 28 +++ .../_torch/disaggregation/resource/page.py | 12 +- .../kv_cache/kv_cache_manager_v2.py | 15 ++ tests/unittest/disaggregated/test_bounce.py | 22 ++ .../unittest/disaggregated/test_extractor.py | 77 ++++++- tests/unittest/disaggregated/test_peer.py | 130 +++++++++++ 9 files changed, 549 insertions(+), 71 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index f141cce93ff4..86f3be81d92a 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -26,7 +26,12 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import ( + CacheKind, + KVCachePageTable, + MapperKind, + PoolView, +) from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch from tensorrt_llm._utils import nvtx_range @@ -179,9 +184,15 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: Optional[int] = None, + peer_kv_heads: Optional[int] = None, ): self._ri = self_ri self._peer_ri = peer_ri + # Head counts of the matched view pair; rank-wide values are the + # fallback for views from builders that do not record them. + self_kv_heads = self_kv_heads or self_ri.attention.kv_heads_per_rank + peer_kv_heads = peer_kv_heads or peer_ri.attention.kv_heads_per_rank self_tp_per_dp = self_ri.tp_size_per_dp_group peer_tp_per_dp = peer_ri.tp_size_per_dp_group @@ -214,28 +225,21 @@ def __init__( buffers_per_layer=peer_buffers_per_layer, side="peer", ) - bytes_per_head = self._bytes_per_head( - src_buffer_bytes, self._ri.attention.kv_heads_per_rank, side="local" - ) - peer_bytes_per_head = self._bytes_per_head( - dst_buffer_bytes, peer_ri.attention.kv_heads_per_rank, side="peer" - ) + bytes_per_head = self._bytes_per_head(src_buffer_bytes, self_kv_heads, side="local") + peer_bytes_per_head = self._bytes_per_head(dst_buffer_bytes, peer_kv_heads, side="peer") if bytes_per_head != peer_bytes_per_head: raise ValueError( f"HND bytes per head mismatch: local={bytes_per_head}, peer={peer_bytes_per_head}" ) - self._bytes_cont_heads = ( - min(self._ri.attention.kv_heads_per_rank, peer_ri.attention.kv_heads_per_rank) - * bytes_per_head - ) + self._bytes_cont_heads = min(self_kv_heads, peer_kv_heads) * bytes_per_head self._src_head_off, self._dst_head_off = self._compute_head_offsets( self_tp_per_dp, peer_tp_per_dp, self_tp_rank, peer_tp_rank, - self_kv_heads=self._ri.attention.kv_heads_per_rank, - peer_kv_heads=peer_ri.attention.kv_heads_per_rank, + self_kv_heads=self_kv_heads, + peer_kv_heads=peer_kv_heads, bytes_per_head=bytes_per_head, ) @@ -365,6 +369,8 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: Optional[int] = None, + peer_kv_heads: Optional[int] = None, ) -> None: # Deliberately do not call HNDHeadMismatchMapper.__init__: its offsets # assume HND-contiguous heads. Initialize the three attributes consumed @@ -377,8 +383,8 @@ def __init__( f"local={self_tpb}, peer={peer_tpb}" ) - self_heads = self_ri.attention.kv_heads_per_rank - peer_heads = peer_ri.attention.kv_heads_per_rank + self_heads = self_kv_heads or self_ri.attention.kv_heads_per_rank + peer_heads = peer_kv_heads or peer_ri.attention.kv_heads_per_rank if self_buffers_per_layer != peer_buffers_per_layer: raise ValueError( "NHD buffer count per layer mismatch: " @@ -484,16 +490,28 @@ def __init__(self, self_rank_info: RankInfo): self._ri = self_rank_info def should_send( - self, peer_overlap, peer_rank_info, *, mapper_kind: Optional[MapperKind] = None + self, + peer_overlap, + peer_rank_info, + *, + mapper_kind: Optional[MapperKind] = None, + duplicate_head_factor: Optional[int] = None, ) -> "bool | None": """Attention uses head-duplication routing. - REPLICATED views (index-key side caches) hold identical bytes on every - TP rank, so they return None to defer to fan-in election instead. + ``duplicate_head_factor`` is the sending-side factor of the matched + view pair; the overlap's rank-wide factor is the default for callers + without a view pair. REPLICATED views (index-key side caches) hold + identical bytes on every TP rank, so they return None to defer to + fan-in election instead. """ if mapper_kind == MapperKind.REPLICATED: return None - dup = peer_overlap.duplicate_head_factor + dup = ( + peer_overlap.duplicate_head_factor + if duplicate_head_factor is None + else duplicate_head_factor + ) if dup <= 1: return True tp_rank = self._ri.tp_rank % self._ri.tp_size_per_dp_group @@ -615,30 +633,55 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool: ) ) - def _head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + @staticmethod + def view_kv_heads(ri: RankInfo, pool_view: Optional[PoolView] = None) -> int: + """Per-rank KV heads of a view; the rank-wide count for legacy views.""" + heads = getattr(pool_view, "num_kv_heads", None) + return heads if heads else ri.attention.kv_heads_per_rank + + def _head_factors( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[int, int]: self_tp = self._tp_per_dp(self._ri) peer_tp = self._tp_per_dp(peer_ri) - a = self._ri.attention - b = peer_ri.attention - return a.kv_heads_per_rank * self_tp, b.kv_heads_per_rank * peer_tp + return ( + self.view_kv_heads(self._ri, self_pv) * self_tp, + self.view_kv_heads(peer_ri, peer_pv) * peer_tp, + ) + + def head_match( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[bool, bool]: + """Whether the matched views hold the same heads per rank. - def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]: - factor_self, factor_peer = self._head_factors(peer_ri) + Without a view pair this compares the rank-wide head counts. + """ + factor_self, factor_peer = self._head_factors(peer_ri, self_pv, peer_pv) is_dup_head = factor_self != factor_peer # Slot layout is compatible if kv_heads_per_rank matches (same slot size). - head_match = ( - self._ri.attention.is_mla - or self._ri.attention.kv_heads_per_rank == peer_ri.attention.kv_heads_per_rank - ) + head_match = self._ri.attention.is_mla or self.view_kv_heads( + self._ri, self_pv + ) == self.view_kv_heads(peer_ri, peer_pv) return head_match, is_dup_head - def duplicate_head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + def duplicate_head_factors( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[int, int]: # Head duplication only applies to attention KV. In particular, do # not divide 0 / 0 for an attention-free hybrid PP stage; MambaPolicy # computes its TP mapping separately. - if self._ri.attention.kv_heads_per_rank == 0 or peer_ri.attention.kv_heads_per_rank == 0: + if self.view_kv_heads(self._ri, self_pv) == 0 or self.view_kv_heads(peer_ri, peer_pv) == 0: return 1, 1 - factor_self, factor_peer = self._head_factors(peer_ri) + factor_self, factor_peer = self._head_factors(peer_ri, self_pv, peer_pv) dup_head = max(1, factor_self // factor_peer) peer_dup_head = max(1, factor_peer // factor_self) return dup_head, peer_dup_head @@ -683,7 +726,7 @@ def build_mapper( peer_bytes_per_layer, ) - head_match, _ = self.head_match(peer_ri) + head_match, _ = self.head_match(peer_ri, self_pv, peer_pv) if head_match: return IntactMapper( self_layer_offsets, @@ -703,6 +746,8 @@ def build_mapper( peer_bytes_per_layer=peer_bytes_per_layer, self_buffers_per_layer=self_buffers_per_layer, peer_buffers_per_layer=peer_buffers_per_layer, + self_kv_heads=self.view_kv_heads(self._ri, self_pv), + peer_kv_heads=self.view_kv_heads(peer_ri, peer_pv), ) return HNDHeadMismatchMapper( @@ -714,4 +759,6 @@ def build_mapper( peer_bytes_per_layer=peer_bytes_per_layer, self_buffers_per_layer=self_buffers_per_layer, peer_buffers_per_layer=peer_buffers_per_layer, + self_kv_heads=self.view_kv_heads(self._ri, self_pv), + peer_kv_heads=self.view_kv_heads(peer_ri, peer_pv), ) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..23c07b5598b5 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -45,6 +45,16 @@ class PeerOverlap: duplicate_head_factor: int = 1 peer_duplicate_head_factor: int = 1 ranks: List[int] = field(default_factory=list) + # Duplication factors of the matched head-sharded view pairs whose factors + # differ from the rank-wide ones above, keyed by (self layer group, self + # pool view). Views without an entry use the rank-wide factors. + view_duplicate_head_factors: Dict["LGPoolKey", Tuple[int, int]] = field(default_factory=dict) + + def duplicate_head_factors_for(self, pool_key: "LGPoolKey") -> Tuple[int, int]: + """(self, peer) duplication factors of one view pair.""" + return self.view_duplicate_head_factors.get( + pool_key, (self.duplicate_head_factor, self.peer_duplicate_head_factor) + ) class PeerRegistrar: @@ -84,32 +94,37 @@ def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): extractor = KVRegionExtractorV1(peer_ri.page_table) self._peer_ext_cache[key] = extractor - head_match, _ = self._get_policy(CacheKind.PAGED).head_match(peer_ri) - if not head_match: - self_page_table = self._self_ext_cache.page_table - nhd_fragments_per_token = sum( - len( - self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].buffer_entries - ) - for layer_group_id, pool_idx in self.get_pool_mapping(peer_ri) - if self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].mapper_kind - == MapperKind.NHD + policy = self._get_policy(CacheKind.PAGED) + self_page_table = self._self_ext_cache.page_table + nhd_fragments_per_token = 0 + nhd_heads: set[tuple[int, int]] = set() + for (layer_group_id, pool_idx), (peer_lg, peer_pi) in self.get_pool_mapping( + peer_ri + ).items(): + self_pv = self_page_table.layer_groups[layer_group_id].pool_views[pool_idx] + if self_pv.mapper_kind != MapperKind.NHD: + continue + peer_pv = peer_ri.page_table.layer_groups[peer_lg].pool_views[peer_pi] + if policy.head_match(peer_ri, self_pv, peer_pv)[0]: + continue + nhd_fragments_per_token += len(self_pv.buffer_entries) + nhd_heads.add( + (policy.view_kv_heads(self._ri, self_pv), policy.view_kv_heads(peer_ri, peer_pv)) + ) + if nhd_fragments_per_token: + heads_text = ", ".join( + f"local_kv_heads={local}, peer_kv_heads={peer}" for local, peer in sorted(nhd_heads) + ) + heads_key = "_".join(f"{local}-{peer}" for local, peer in sorted(nhd_heads)) + logger.warning_once( + "NHD head-mismatched disaggregated KV transfer has no " + "contiguous staging path and will emit approximately " + f"{nhd_fragments_per_token} NIXL descriptors per transferred " + "token per peer, excluding block-level replicated pools " + f"({heads_text}). " + "Long-context TEP/DEP transfers may have high latency.", + key=f"native-nhd-head-mismatch-{heads_key}-{nhd_fragments_per_token}", ) - if nhd_fragments_per_token: - local_heads = self._ri.attention.kv_heads_per_rank - peer_heads = peer_ri.attention.kv_heads_per_rank - logger.warning_once( - "NHD head-mismatched disaggregated KV transfer has no " - "contiguous staging path and will emit approximately " - f"{nhd_fragments_per_token} NIXL descriptors per transferred " - "token per peer, excluding block-level replicated pools " - f"(local_kv_heads={local_heads}, peer_kv_heads={peer_heads}). " - "Long-context TEP/DEP transfers may have high latency.", - key=( - "native-nhd-head-mismatch-" - f"{local_heads}-{peer_heads}-{nhd_fragments_per_token}" - ), - ) def peer_extractor(self, peer_name: str, peer_rank: int) -> KVRegionExtractorV1: return self._peer_ext_cache[self._unique_key(peer_name, peer_rank)] @@ -180,6 +195,7 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + self._validate_view_head_geometry(peer_ri) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -195,6 +211,61 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: return True + def _validate_view_head_geometry(self, peer_ri: RankInfo) -> None: + """Check that every matched head-sharded view pair can be resharded. + + Views carry their own KV-head count. A view without one (tables from + older builders) falls back to the rank-wide count, which describes it + only while the rank's views share one head geometry. When either side + publishes counts, the bytes per head of both sides, taken with those + counts, must therefore agree; a legacy side whose views differ in head + geometry fails this check and has to publish per-view counts. + """ + self_page_table = self._self_ext_cache.page_table + peer_page_table = peer_ri.page_table + policy = self._get_policy(CacheKind.PAGED) + for (self_lg, self_pi), (peer_lg, peer_pi) in self.get_pool_mapping(peer_ri).items(): + layer_group = self_page_table.layer_groups[self_lg] + self_pv = layer_group.pool_views[self_pi] + if layer_group.kind != CacheKind.PAGED or self_pv.mapper_kind == MapperKind.REPLICATED: + continue + peer_pv = peer_page_table.layer_groups[peer_lg].pool_views[peer_pi] + if self_pv.num_kv_heads is None and peer_pv.num_kv_heads is None: + continue + bytes_per_head = [] + for side, ri, lg_idx, pi, pool_view in ( + ("local", self._ri, self_lg, self_pi, self_pv), + ("peer", peer_ri, peer_lg, peer_pi, peer_pv), + ): + heads = policy.view_kv_heads(ri, pool_view) + buffers = self._get_buffers_per_layer(pool_view, layer_group_id=lg_idx, pool_idx=pi) + layer_bytes = get_layer_byte_ranges(pool_view)[1] + if heads <= 0 or layer_bytes % (buffers * heads) != 0: + raise ValueError( + f"PeerRegistrar: the {side} view {sorted(pool_view.pool_role)} holds " + f"{layer_bytes} bytes per layer, not a multiple of {buffers} buffers x " + f"{heads} KV heads" + ) + bytes_per_head.append(layer_bytes // (buffers * heads)) + if bytes_per_head[0] != bytes_per_head[1]: + missing = [ + side + for side, pool_view in (("local", self_pv), ("peer", peer_pv)) + if pool_view.num_kv_heads is None + ] + hint = ( + f"; the {' and '.join(missing)} view has no per-view KV-head count and " + "falls back to the rank-wide one, so both peers must publish per-view " + "KV-head counts" + if missing + else "" + ) + raise ValueError( + "PeerRegistrar: bytes per KV head differ for view " + f"{sorted(self_pv.pool_role)} (local={bytes_per_head[0]}, " + f"peer={bytes_per_head[1]}){hint}" + ) + def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: """Get mapping from (self_lg_idx, self_pool_idx) -> (peer_lg_idx, peer_pool_idx). @@ -296,7 +367,7 @@ def _peer_l2g(kind: CacheKind) -> Dict[int, int]: # overlapping layer set; other same-role peer pools cover # disjoint layers and fall out via the overlap filter. self_layer_set = set(pv_global_ids) - matched_peer_pi = None + candidates = [] for peer_pi, peer_pv in enumerate(peer_lg.pool_views): if peer_pv.pool_role != self_pv.pool_role: continue @@ -310,11 +381,52 @@ def _peer_l2g(kind: CacheKind) -> Dict[int, int]: f"(local={self_pv.mapper_kind.name}, " f"peer={peer_pv.mapper_kind.name}, peer_pool={peer_pi})" ) - matched_peer_pi = peer_pi - break - - if matched_peer_pi is not None: - mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) + candidates.append((peer_pi, set(peer_global_ids))) + + if not candidates: + unmatched = [ + sorted(peer_pv.pool_role) + for peer_pv in peer_lg.pool_views + if self_layer_set & set(get_pool_view_global_layer_ids(peer_pv, peer_lg)) + ] + if unmatched and self_lg.kind == CacheKind.PAGED: + raise ValueError( + "PeerRegistrar.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}, role={sorted(self_pv.pool_role)}) " + f"has no peer view with the same roles; the peer holds those layers " + f"as {unmatched}" + ) + continue + # The pairing must be exact: one view per side carrying the + # same layers of the role. A self view split over several peer + # views (or the reverse) means the peers partition their + # layers differently, e.g. target and draft layers pooled + # together on one side only; matching the first candidate would + # silently drop the rest of the layers. + if len(candidates) > 1: + raise ValueError( + "PeerRegistrar.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}, role={sorted(self_pv.pool_role)}) " + f"overlaps several peer views {[pi for pi, _ in candidates]}; " + "the peers split their layers into views differently" + ) + matched_peer_pi, peer_layer_set = candidates[0] + self_other_layers = { + gid + for other_pi, other_pv in enumerate(self_lg.pool_views) + if other_pi != self_pi and other_pv.pool_role == self_pv.pool_role + for gid in get_pool_view_global_layer_ids(other_pv, self_lg) + } + misplaced = (peer_layer_set & self_other_layers) - self_layer_set + if misplaced: + raise ValueError( + "PeerRegistrar.get_pool_mapping: peer view " + f"(lg={peer_lg_idx}, pool={matched_peer_pi}) also holds layers " + f"{sorted(misplaced)} that this rank keeps in another view of role " + f"{sorted(self_pv.pool_role)}; the peers split their layers into " + "views differently" + ) + mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) self._lg_pool_mapping_cache[key] = mapping return mapping @@ -536,6 +648,7 @@ def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerO ranks.append(pp * peer_ri.tp_size * peer_ri.cp_size + tp * peer_ri.cp_size + cp) dup_head, peer_dup_head = self._get_policy(CacheKind.PAGED).duplicate_head_factors(peer_ri) + view_dup_factors = self._view_duplicate_head_factors(peer_ri, (dup_head, peer_dup_head)) targets = PeerOverlap( overlap_pp_size=overlap_pp_size, @@ -544,10 +657,35 @@ def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerO duplicate_head_factor=dup_head, peer_duplicate_head_factor=peer_dup_head, ranks=ranks, + view_duplicate_head_factors=view_dup_factors, ) self._overlap_cache[key] = targets return targets + def _view_duplicate_head_factors( + self, peer_ri: RankInfo, rank_wide: Tuple[int, int] + ) -> Dict[LGPoolKey, Tuple[int, int]]: + """Head duplication factors of the matched view pairs that differ from ``rank_wide``. + + Each pair uses the head counts of its own two views, so views of one + rank that shard or replicate heads differently (target and draft + layers) get their own factor. + """ + policy = self._get_policy(CacheKind.PAGED) + self_page_table = self._self_ext_cache.page_table + peer_page_table = peer_ri.page_table + factors: Dict[LGPoolKey, Tuple[int, int]] = {} + for (self_lg, self_pi), (peer_lg, peer_pi) in self.get_pool_mapping(peer_ri).items(): + layer_group = self_page_table.layer_groups[self_lg] + self_pv = layer_group.pool_views[self_pi] + if layer_group.kind != CacheKind.PAGED or self_pv.mapper_kind == MapperKind.REPLICATED: + continue + peer_pv = peer_page_table.layer_groups[peer_lg].pool_views[peer_pi] + pair = policy.duplicate_head_factors(peer_ri, self_pv, peer_pv) + if pair != rank_wide: + factors[(self_lg, self_pi)] = pair + return factors + def _owns_tp_fan_in(self, peer_rank_info: RankInfo) -> bool: """Elect one owner when replicated bytes fan in across TP ranks. @@ -589,7 +727,16 @@ def should_send_pool( pool_view = layer_group.pool_views[pool_idx] # Delegate to the policy's should_send. None means fan-in election. policy = self._get_policy(layer_group.kind) - result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) + kwargs = {} + if layer_group.kind == CacheKind.PAGED: + # The sending-side duplication factor of this view pair, not one + # factor for the whole rank. + kwargs["duplicate_head_factor"] = peer_overlap.duplicate_head_factors_for( + (layer_group_id, pool_idx) + )[0] + result = policy.should_send( + peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind, **kwargs + ) return self._owns_tp_fan_in(peer_rank_info) if result is None else result def should_send_aux(self, peer_rank_info: RankInfo) -> bool: diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 29243078fb1c..14d1acb56535 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2764,8 +2764,9 @@ def _fanin_bounce_safe( ) -> bool: """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in - expected_transfers, so the live writers overflow their slots; + * every head duplication factor is 1, rank-wide and in each matched view pair (target + and draft views may replicate heads differently) -- else some ranks don't send KV yet + still count in expected_transfers, so the live writers overflow their slots; * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; for PP fan-in, every overlapping stage must hold the same number of layers (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage @@ -2773,7 +2774,12 @@ def _fanin_bounce_safe( Otherwise fall back to the per-fragment path (correct, just not coalesced). Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects the mismatched case, so this only needs the count to split evenly.""" - if overlap.duplicate_head_factor != 1: + # Any duplicated head, rank-wide or in a single view pair, leaves some + # writers without a share (sender side) or doubles one (receiver side). + view_factors = getattr(overlap, "view_duplicate_head_factors", {}).values() + if overlap.duplicate_head_factor != 1 or any( + factor != 1 for pair in view_factors for factor in pair + ): return False if overlap.overlap_pp_size > 1: lpp = getattr(peer_ri, "layer_num_per_pp", None) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index b1132c10ddc4..48680acc1749 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -236,6 +236,7 @@ def _build_pool_views_for_variant( role_layouts: Optional[Dict] = None, ignored_roles: frozenset[DataRole] = frozenset(), view_partition: Optional[Callable] = None, + kv_heads_for_layer: Optional[Callable[[int], Optional[int]]] = None, ) -> List[PoolView]: """Bucket one slot-desc variant's coalesced buffers into pool views. @@ -243,11 +244,17 @@ def _build_pool_views_for_variant( kind share a view (KEY+VALUE), roles with different kinds in the same physical pool get separate views (M3 coalesced index-K). ``role_layouts`` supplies the per-role resharding geometry carried onto the view. + ``view_partition`` further splits roles of one kind (e.g. by owning + model) so layers that differ in head geometry never share a view even + when their buffers share a physical pool. ``kv_heads_for_layer`` gives a + layer's per-rank KV-head count, which head-sharded views record; all + layers of a view must agree on it. """ role_layouts = role_layouts or {} bucket_entries: Dict[tuple, list] = defaultdict(list) bucket_roles: Dict[tuple, set] = defaultdict(set) bucket_layouts: Dict[tuple, set] = defaultdict(set) + bucket_layers: Dict[tuple, set] = defaultdict(set) for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): single_buffer_size = int(coalesced_buffer.single_buffer_size) offset = 0 @@ -267,6 +274,7 @@ def _build_pool_views_for_variant( ) bucket_roles[bucket_key].add(str(buffer_id.role)) bucket_layouts[bucket_key].add(role_layouts.get(buffer_id.role)) + bucket_layers[bucket_key].add(int(buffer_id.layer_id)) offset += single_buffer_size # All ordering below is canonicalization — the page table is @@ -310,6 +318,20 @@ def _build_pool_views_for_variant( f"{context} bytes_per_layer={bytes_per_layer} is not a multiple of " f"bytes_per_head={layout.bytes_per_head}" ) + num_kv_heads = None + if kv_heads_for_layer is not None and mapper_kind != MapperKind.REPLICATED: + layer_heads = { + heads + for heads in (kv_heads_for_layer(layer) for layer in bucket_layers[bucket_key]) + if heads + } + if len(layer_heads) > 1: + raise ValueError( + f"{context} mixes layers with different per-rank KV-head counts " + f"{sorted(layer_heads)}; layers with different head geometry must be " + "in separate views" + ) + num_kv_heads = layer_heads.pop() if layer_heads else None pool_views.append( PoolView( pool_idx=pool_idx, @@ -321,6 +343,7 @@ def _build_pool_views_for_variant( list(layout.section_bytes) if layout.section_bytes is not None else None ), bytes_per_head=layout.bytes_per_head, + num_kv_heads=num_kv_heads, ) ) @@ -394,6 +417,10 @@ def _window_size_for_layer(internal_layer_id: int): raise ValueError(f"Cannot resolve layer config for internal layer {internal_layer_id}") + def _kv_heads_for_layer(internal_layer_id: int) -> Optional[int]: + heads = manager.num_kv_heads_per_layer + return int(heads[internal_layer_id]) if internal_layer_id < len(heads) else None + pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) @@ -458,6 +485,7 @@ def _window_size_for_layer(internal_layer_id: int): role_layouts=role_layouts, ignored_roles=ignored_roles, view_partition=view_partition, + kv_heads_for_layer=None if is_recurrent else _kv_heads_for_layer, ) for pool_view in pool_views: if pool_view.mapper_kind not in _MAPPER_KINDS_BY_CACHE_KIND[cache_kind]: diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 1254c06001f6..46a6206c6640 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -201,7 +201,12 @@ class PoolView: bytes_per_head: INDEXED views whose head count is not derivable from ``RankInfo.attention`` (recurrent SSM state). ``bytes_per_layer // bytes_per_head`` is the per-rank head count. Attention K/V leaves - it ``None`` and uses ``kv_heads_per_rank`` instead. + it ``None``; its head count is ``num_kv_heads``. + num_kv_heads: Attention K/V views only. Per-rank KV-head count shared + by every layer of the view, so views of one rank (e.g. target and + draft layers) may shard or replicate heads differently. ``None`` + in tables from older builders and for views without KV heads; + consumers then fall back to ``AttentionInfo.kv_heads_per_rank``. """ pool_idx: int @@ -211,6 +216,7 @@ class PoolView: bytes_per_layer: Optional[int] = None section_bytes: Optional[List[int]] = None bytes_per_head: Optional[int] = None + num_kv_heads: Optional[int] = None def to_dict(self) -> dict: return { @@ -227,6 +233,7 @@ def to_dict(self) -> dict: "bytes_per_head": ( int(self.bytes_per_head) if self.bytes_per_head is not None else None ), + "num_kv_heads": int(self.num_kv_heads) if self.num_kv_heads is not None else None, } @staticmethod @@ -253,6 +260,9 @@ def from_dict(data: dict) -> "PoolView": bytes_per_head=( int(data["bytes_per_head"]) if data.get("bytes_per_head") is not None else None ), + num_kv_heads=( + int(data["num_kv_heads"]) if data.get("num_kv_heads") is not None else None + ), ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 3fb440b09505..a04a37a92f9c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3858,6 +3858,21 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + def get_disagg_pool_view_partition(self, layer_id: LayerId, role: DataRole): + """Keep target and draft layers in separate transfer views. + + Target and draft layers may share a physical pool, but each model has + its own head geometry, so the transfer layout must keep them apart + whatever the storage coalescing does. Managers without draft-owned + layers need no partition. + """ + layer_specs = getattr(self, "layer_specs", None) + if layer_specs is None or not any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs + ): + return None + return layer_specs[int(layer_id)].owner.value + def get_layer_cache_spec(self, global_layer_id: int) -> KVCacheLayerSpec: """Return an explicitly supplied rank-local layer specification.""" if self.layer_specs is None: diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 5e7af40dd760..d12fcc5fa45f 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -285,6 +285,28 @@ def pt(mapper_kind: MapperKind) -> SimpleNamespace: ) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "view_factors,expected", [((1, 1), True), ((2, 1), False), ((1, 2), False)] +) +def test_fanin_bounce_checks_per_view_head_duplication(view_factors, expected): + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap + + overlap = PeerOverlap( + duplicate_head_factor=1, + overlap_pp_size=1, + ranks=[0, 1], + view_duplicate_head_factors={(0, 1): view_factors}, + ) + assert ( + tfr.Receiver._fanin_bounce_safe( + overlap, SimpleNamespace(layer_num_per_pp=[2], page_table=None), None + ) + is expected + ) + + # --------------------------------------------------------------------------- # # NoBounceTransport — disabled no-op fallback # --------------------------------------------------------------------------- # diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec6ef42ad0f7..ddef6b09323b 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -38,7 +38,12 @@ PLE_NGRAM_CONTEXT, Qwen4ExpHybridCacheManager, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaRole from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, @@ -226,7 +231,7 @@ def test_layer_group_meta_serialization(): dtype=BUFFER_ENTRY_DTYPE, ) kv_pool = PhysicalPool(base_address=1000, slot_bytes=512, num_slots=10) - pv = PoolView(pool_idx=0, buffer_entries=entries) + pv = PoolView(pool_idx=0, buffer_entries=entries, num_kv_heads=4) local_layers = [ LocalLayer(local_layer_id=0, global_layer_id=0), LocalLayer(local_layer_id=1, global_layer_id=1), @@ -251,6 +256,10 @@ def test_layer_group_meta_serialization(): assert restored_lg.kv_head_num_per_rank == 4 assert len(restored_lg.local_layers) == 2 assert len(restored_lg.pool_views[0].buffer_entries) == 2 + assert restored_lg.pool_views[0].num_kv_heads == 4 + # Older page tables omit per-view metadata and retain the rank-wide fallback. + del d["layer_groups"][0]["pool_views"][0]["num_kv_heads"] + assert KVCachePageTable.from_dict(d).layer_groups[0].pool_views[0].num_kv_heads is None def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_list=(640,)): @@ -319,6 +328,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l ) +@pytest.mark.cpu_only def test_v2_builder_stamps_homogeneous_pool_kinds(): """Split KV / INDEX_KEY pools get NHD and REPLICATED views respectively.""" from types import SimpleNamespace @@ -341,12 +351,75 @@ def test_v2_builder_stamps_homogeneous_pool_kinds(): assert len(views) == 2 assert views[0].pool_role == frozenset({"key", "value"}) assert views[0].mapper_kind == MapperKind.NHD + assert views[0].num_kv_heads == 1 assert get_unique_layers(views[0]) == {0, 1} assert views[1].pool_role == frozenset({"index_key"}) assert views[1].mapper_kind == MapperKind.REPLICATED + assert views[1].num_kv_heads is None assert get_unique_layers(views[1]) == {1} +@pytest.mark.cpu_only +@pytest.mark.parametrize("mixed_owners", [False, True]) +def test_v2_builder_partitions_coalesced_views_by_owner(mixed_owners): + from types import SimpleNamespace + + attrs = { + (layer, role): SimpleNamespace( + life_cycle_id=0, pool_index=0, offset=(2 * layer + role_index) * 128, size=128 + ) + for layer in range(2) + for role_index, role in enumerate((Role.KEY, Role.VALUE)) + } + manager = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(512,)) + manager.num_kv_heads_per_layer = [2, 1] if mixed_owners else [2, 2] + manager.layer_specs = tuple( + KVCacheLayerSpec( + global_layer_id=layer, + owner=KVCacheLayerOwner.DRAFT + if mixed_owners and layer == 1 + else KVCacheLayerOwner.TARGET, + dtype=DataType.HALF, + num_kv_heads=heads, + head_dim=4 // heads, + buffer_roles=(Role.KEY, Role.VALUE), + ) + for layer, heads in enumerate(manager.num_kv_heads_per_layer) + ) + manager.get_disagg_pool_view_partition = ( + KVCacheManagerV2.get_disagg_pool_view_partition.__get__(manager) + ) + + views = build_page_table_from_manager(manager).layer_groups[0].pool_views + + if mixed_owners: + assert [view.num_kv_heads for view in views] == [2, 1] + assert [get_unique_layers(view) for view in views] == [{0}, {1}] + assert [view.pool_idx for view in views] == [0, 0] + assert [view.buffer_entries["offset"].tolist() for view in views] == [[0, 128], [256, 384]] + else: + assert len(views) == 1 + assert views[0].num_kv_heads == 2 + assert get_unique_layers(views[0]) == {0, 1} + + +@pytest.mark.cpu_only +def test_v2_builder_rejects_mixed_heads_without_view_partition(): + from types import SimpleNamespace + + attrs = { + (layer, role): SimpleNamespace( + life_cycle_id=0, pool_index=0, offset=(2 * layer + role_index) * 128, size=128 + ) + for layer in range(2) + for role_index, role in enumerate((Role.KEY, Role.VALUE)) + } + manager = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(512,)) + manager.num_kv_heads_per_layer = [2, 1] + with pytest.raises(ValueError, match="different per-rank KV-head counts"): + build_page_table_from_manager(manager) + + def test_v2_builder_splits_mixed_kind_pool_into_per_class_views(): """A pool coalescing several role classes yields one view per class. diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 7ff5f5e72e7c..5146e3381d77 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -408,6 +408,136 @@ def _make_peer_registrar(self_rankinfo): return reg +def _make_per_view_page_table(head_counts, mapper_kind): + """Two-token K/V buffers with four bytes per token/head.""" + views = [] + pools = [] + for layer, heads in enumerate(head_counts): + buffer_bytes = 8 * heads + views.append( + PoolView( + pool_idx=layer, + buffer_entries=np.array( + [(layer, 0, buffer_bytes), (layer, buffer_bytes, buffer_bytes)], + dtype=BUFFER_ENTRY_DTYPE, + ), + mapper_kind=mapper_kind, + num_kv_heads=heads, + ) + ) + pools.append( + PhysicalPool( + base_address=0x1000 * (layer + 1), slot_bytes=2 * buffer_bytes, num_slots=8 + ) + ) + return KVCachePageTable( + tokens_per_block=2, + layer_groups=[ + AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=head_counts[0], + sliding_window_size=None, + local_layers=[LocalLayer(layer, layer) for layer in range(len(head_counts))], + pool_views=views, + ) + ], + pool_groups=[PhysicalPoolGroup(pools=pools)], + ) + + +@pytest.mark.parametrize("mapper_kind", [MapperKind.INDEXED, MapperKind.NHD]) +def test_per_view_heads_select_correct_transfer_bytes(mapper_kind): + # TP4 -> TP1: the eight-head view shards all ranks, while the two-head + # view duplicates each head on two ranks. Rank 2 sends its second head. + self_ri = make_rankinfo( + tp_size=4, + tp_rank=2, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], mapper_kind), + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + kv_heads_per_rank=8, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([8, 2], mapper_kind), + ) + reg = _make_peer_registrar(self_ri) + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + mapper = reg.get_kv_map(peer_ri, (0, 1), (0, 1)) + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=16)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=32)), + ) + + if mapper_kind == MapperKind.NHD: + assert isinstance(mapper, NHDHeadMismatchMapper) + assert pair.src.memory.ptrs.tolist() == [1000, 1004, 1008, 1012] + assert pair.dst.memory.ptrs.tolist() == [2004, 2012, 2020, 2028] + assert pair.src.memory.bytes_per_region == 4 + else: + assert isinstance(mapper, HNDHeadMismatchMapper) + assert pair.src.memory.ptrs.tolist() == [1000, 1008] + assert pair.dst.memory.ptrs.tolist() == [2008, 2024] + assert pair.src.memory.bytes_per_region == 8 + assert pair.dst.memory.bytes_per_region == pair.src.memory.bytes_per_region + + +def test_per_view_head_duplication_elects_senders_independently(): + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + kv_heads_per_rank=8, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([8, 2], MapperKind.INDEXED), + ) + ownership = [[], []] + for rank in range(4): + self_ri = make_rankinfo( + tp_size=4, + tp_rank=rank, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], MapperKind.INDEXED), + ) + reg = _make_peer_registrar(self_ri) + overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + assert overlap.duplicate_head_factor == 1 + assert overlap.duplicate_head_factors_for((0, 0)) == (1, 1) + assert overlap.duplicate_head_factors_for((0, 1)) == (2, 1) + for view in range(2): + ownership[view].append(reg.should_send_pool(overlap, peer_ri, 0, view)) + + assert ownership == [[True, True, True, True], [True, False, True, False]] + + +def test_per_view_missing_heads_rejects_incompatible_legacy_geometry(): + self_ri = make_rankinfo( + tp_size=4, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], MapperKind.INDEXED), + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + kv_heads_per_rank=8, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([8, 2], MapperKind.INDEXED), + ) + peer_ri.page_table.layer_groups[0].pool_views[1].num_kv_heads = None + reg = _make_peer_registrar(self_ri) + with pytest.raises(ValueError, match="both peers must publish per-view KV-head counts"): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + + @pytest.mark.parametrize( ("local_replay", "peer_replay"), [(False, False), (True, True), (False, True), (True, False)], From 716257dd35b46e6c46b2e267cd49eb6ddea7cc02 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 30 Oct 2026 14:06:59 -0700 Subject: [PATCH 2/4] Fix per-view KV transfer compatibility and bounce safety Signed-off-by: Allison Lim --- .../_torch/disaggregation/native/peer.py | 6 +- .../_torch/disaggregation/native/transfer.py | 7 +- .../kv_cache/kv_cache_manager_v2.py | 19 +-- tests/unittest/disaggregated/test_bounce.py | 16 +- .../unittest/disaggregated/test_extractor.py | 48 ++++-- tests/unittest/disaggregated/test_peer.py | 160 ++++++++++++++++++ .../test_transfer_ownership_regressions.py | 1 + 7 files changed, 225 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 23c07b5598b5..85fad6ea1ed1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -195,7 +195,7 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) - self._validate_view_head_geometry(peer_ri) + self.validate_view_head_geometry(peer_ri) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -211,8 +211,8 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: return True - def _validate_view_head_geometry(self, peer_ri: RankInfo) -> None: - """Check that every matched head-sharded view pair can be resharded. + def validate_view_head_geometry(self, peer_ri: RankInfo) -> None: + """Validate pool pairing and head geometry without registering the peer. Views carry their own KV-head count. A view without one (tables from older builders) falls back to the rank-wide count, which describes it diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index deec96d52c7e..e2ae3d03d6de 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2861,8 +2861,10 @@ def _fanin_bounce_safe( # Any duplicated head, rank-wide or in a single view pair, leaves some # writers without a share (sender side) or doubles one (receiver side). view_factors = getattr(overlap, "view_duplicate_head_factors", {}).values() - if overlap.duplicate_head_factor != 1 or any( - factor != 1 for pair in view_factors for factor in pair + if ( + overlap.duplicate_head_factor != 1 + or overlap.peer_duplicate_head_factor != 1 + or any(factor != 1 for pair in view_factors for factor in pair) ): return False if overlap.overlap_pp_size > 1: @@ -3114,6 +3116,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: self._registrar.self_extractor.page_table, sender_info.page_table, ) + self._registrar.validate_view_head_geometry(sender_info) except ValueError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index dcc3061398be..9e7db3c8e3a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -4220,20 +4220,19 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() - def get_disagg_pool_view_partition(self, layer_id: LayerId, role: DataRole): - """Keep target and draft layers in separate transfer views. + def get_disagg_pool_view_partition( + self, layer_id: LayerId, role: DataRole + ) -> tuple[str, int] | None: + """Keep different owners and KV-head counts in separate transfer views. - Target and draft layers may share a physical pool, but each model has - its own head geometry, so the transfer layout must keep them apart - whatever the storage coalescing does. Managers without draft-owned - layers need no partition. + Layers may share a physical pool when their buffers have equal byte + sizes, even if their head counts differ within the same model. """ layer_specs = getattr(self, "layer_specs", None) - if layer_specs is None or not any( - layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs - ): + if layer_specs is None: return None - return layer_specs[int(layer_id)].owner.value + layer = layer_specs[int(layer_id)] + return layer.owner.value, layer.num_kv_heads def get_layer_cache_spec(self, global_layer_id: int) -> KVCacheLayerSpec: """Return an explicitly supplied rank-local layer specification.""" diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 5c741fb0fbc5..0670b07bd996 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -238,8 +238,15 @@ def test_fanin_bounce_safe_gate() -> None: safe = tfr.Receiver._fanin_bounce_safe - def ov(dup: int, pp: int, ranks: tuple[int, ...] = (0,)) -> SimpleNamespace: - return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp, ranks=list(ranks)) + def ov( + dup: int, pp: int, ranks: tuple[int, ...] = (0,), *, peer_dup: int = 1 + ) -> SimpleNamespace: + return SimpleNamespace( + duplicate_head_factor=dup, + peer_duplicate_head_factor=peer_dup, + overlap_pp_size=pp, + ranks=list(ranks), + ) def ri(lpp: list[int], page_table: SimpleNamespace | None = None) -> SimpleNamespace: return SimpleNamespace(layer_num_per_pp=lpp, page_table=page_table) @@ -248,10 +255,13 @@ def pt(mapper_kind: MapperKind) -> SimpleNamespace: view = SimpleNamespace(mapper_kind=mapper_kind) return SimpleNamespace(layer_groups=[SimpleNamespace(pool_views=[view])]) - # single PP stage (overlap_pp_size <= 1): only duplicate_head_factor matters + # A single PP stage is safe only when neither side duplicates heads. assert safe(ov(1, 1), ri([24]), None) is True assert safe(ov(1, 0), ri([24]), None) is True assert safe(ov(2, 1), ri([24]), None) is False # duplicate heads / MLA -> some don't send + # The sender can duplicate heads even when the receiver does not. With + # no per-view overrides, this must still disable the equal writer split. + assert safe(ov(1, 1, ranks=(0, 1), peer_dup=2), ri([24]), None) is False # EVEN PP fan-in (equal layers per overlapping stage) -> allowed assert safe(ov(1, 4), ri([20, 20, 20, 20]), None) is True # UNEVEN PP fan-in -> per-writer sizes differ -> fall back diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ddef6b09323b..d2f5b4f65864 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -360,31 +360,48 @@ def test_v2_builder_stamps_homogeneous_pool_kinds(): @pytest.mark.cpu_only -@pytest.mark.parametrize("mixed_owners", [False, True]) -def test_v2_builder_partitions_coalesced_views_by_owner(mixed_owners): +@pytest.mark.parametrize( + "owners, head_counts, expected_num_views", + [ + pytest.param(("target", "target"), (2, 2), 1, id="target-uniform-heads"), + pytest.param(("target", "target"), (2, 1), 2, id="target-mixed-heads"), + pytest.param(("target", "draft"), (2, 2), 2, id="mixed-owners-uniform-heads"), + pytest.param(("target", "draft"), (2, 1), 2, id="mixed-owners-mixed-heads"), + pytest.param(("draft", "draft"), (2, 1), 2, id="draft-mixed-heads"), + ], +) +def test_v2_builder_partitions_coalesced_views_by_owner_and_heads( + owners: tuple[str, str], head_counts: tuple[int, int], expected_num_views: int +) -> None: from types import SimpleNamespace + # FP16 (2 heads, dim 64) and (1 head, dim 128) share a 4096-byte K/V + # buffer size at 16 tokens per block. + buffer_size = 4096 attrs = { (layer, role): SimpleNamespace( - life_cycle_id=0, pool_index=0, offset=(2 * layer + role_index) * 128, size=128 + life_cycle_id=0, + pool_index=0, + offset=(2 * layer + role_index) * buffer_size, + size=buffer_size, ) for layer in range(2) for role_index, role in enumerate((Role.KEY, Role.VALUE)) } - manager = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(512,)) - manager.num_kv_heads_per_layer = [2, 1] if mixed_owners else [2, 2] + manager = _make_fake_v2_manager( + attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(4 * buffer_size,) + ) + manager.num_kv_heads_per_layer = list(head_counts) manager.layer_specs = tuple( KVCacheLayerSpec( global_layer_id=layer, - owner=KVCacheLayerOwner.DRAFT - if mixed_owners and layer == 1 - else KVCacheLayerOwner.TARGET, + owner=KVCacheLayerOwner(owners[layer]), dtype=DataType.HALF, num_kv_heads=heads, - head_dim=4 // heads, + head_dim=128 // heads, buffer_roles=(Role.KEY, Role.VALUE), ) - for layer, heads in enumerate(manager.num_kv_heads_per_layer) + for layer, heads in enumerate(head_counts) ) manager.get_disagg_pool_view_partition = ( KVCacheManagerV2.get_disagg_pool_view_partition.__get__(manager) @@ -392,13 +409,16 @@ def test_v2_builder_partitions_coalesced_views_by_owner(mixed_owners): views = build_page_table_from_manager(manager).layer_groups[0].pool_views - if mixed_owners: - assert [view.num_kv_heads for view in views] == [2, 1] + assert len(views) == expected_num_views + if expected_num_views == 2: + assert [view.num_kv_heads for view in views] == list(head_counts) assert [get_unique_layers(view) for view in views] == [{0}, {1}] assert [view.pool_idx for view in views] == [0, 0] - assert [view.buffer_entries["offset"].tolist() for view in views] == [[0, 128], [256, 384]] + assert [view.buffer_entries["offset"].tolist() for view in views] == [ + [0, buffer_size], + [2 * buffer_size, 3 * buffer_size], + ] else: - assert len(views) == 1 assert views[0].num_kv_heads == 2 assert get_unique_layers(views[0]) == {0, 1} diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index f2855653e269..16090f937a87 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace +from unittest.mock import Mock + import numpy as np import pytest @@ -539,6 +542,163 @@ def test_per_view_missing_heads_rejects_incompatible_legacy_geometry(): reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) +def _make_receiver_for_view_validation( + monkeypatch: pytest.MonkeyPatch, + registrar: PeerRegistrar, + sender_infos: dict[str, RankInfo], + enforce_physical_ownership: bool, +) -> tuple: + import tensorrt_llm._torch.disaggregation.native.transfer as transfer_module + + messengers = {} + + def make_messenger(*, mode: str, endpoint: str | None = None) -> Mock: + messenger = Mock() + messenger.endpoint = endpoint or "tcp://receiver" + if endpoint in sender_infos: + messenger.receive.return_value = [sender_infos[endpoint].to_bytes()] + messengers[endpoint] = messenger + return messenger + + messenger_factory = Mock(side_effect=make_messenger) + monkeypatch.setattr(transfer_module, "ZMQMessenger", messenger_factory) + receiver = transfer_module.Receiver( + registrar, + Mock(), + bounce=Mock(), + enforce_physical_ownership=enforce_physical_ownership, + ) + receiver._build_recv_req_info = Mock() + receiver._request_sender_data = Mock() + return receiver, messengers, messenger_factory + + +@pytest.mark.parametrize("enforce_physical_ownership", [False, True]) +@pytest.mark.parametrize( + ("mismatch", "diagnostic"), + [ + ("legacy_heads", "both peers must publish per-view KV-head counts"), + ("pool_roles", "has no peer view with the same roles"), + ], +) +def test_receiver_rejects_view_mismatch_before_registration( + monkeypatch: pytest.MonkeyPatch, + enforce_physical_ownership: bool, + mismatch: str, + diagnostic: str, +) -> None: + from tensorrt_llm._torch.disaggregation.base import TokenRange + from tensorrt_llm._torch.disaggregation.native.extent import Chunk + from tensorrt_llm._torch.disaggregation.native.transfer import ( + KVRecvTask, + MessageType, + PeerIncompatibleError, + TaskStatus, + ) + + local_ri, bad_ri, good_ri = [ + make_rankinfo( + instance_name=name, + tp_size=1, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], MapperKind.INDEXED), + ) + for name in ("receiver", "bad-sender", "good-sender") + ] + bad_view = bad_ri.page_table.layer_groups[0].pool_views[1] + if mismatch == "legacy_heads": + bad_view.num_kv_heads = None + else: + bad_view.pool_role = frozenset({"other_kv"}) + bad_ri.sender_endpoints = ["tcp://bad-sender"] + good_ri.sender_endpoints = ["tcp://good-sender"] + bad_endpoint = "tcp://bad-info" + good_endpoint = "tcp://good-info" + registrar = _make_peer_registrar(local_ri) + receiver, messengers, messenger_factory = _make_receiver_for_view_validation( + monkeypatch, + registrar, + {bad_endpoint: bad_ri, good_endpoint: good_ri}, + enforce_physical_ownership, + ) + chunk = Chunk( + block_ids_per_layer_groups=[np.array([], dtype=np.int64)], + kind_per_layer_group=[CacheKind.PAGED], + token_range=TokenRange(start=0, end=0), + is_last=True, + ) + + try: + for rid in (1, 2): + params = SimpleNamespace(ctx_info_endpoint=bad_endpoint, ctx_dp_rank=0) + task = KVRecvTask(rid, chunk, slice_id=0, params=params, aux_slot=None) + if enforce_physical_ownership: + task.begin_publication() + + receiver.dispatch_task(task) + + assert task.status == TaskStatus.ERROR + assert task.is_done + assert isinstance(task._exception, PeerIncompatibleError) + assert diagnostic in str(task._exception) + if enforce_physical_ownership: + assert task.resources_drained + + assert bad_endpoint in receiver._incompatible_peers + assert bad_endpoint not in receiver._sender_ep_instance_map + assert receiver._dealers == {} + assert messenger_factory.call_count == 2 # Receiver listener and one info lookup. + receiver._request_sender_data.assert_not_called() + receiver._bounce.reserve.assert_not_called() + messengers[bad_endpoint].send.assert_called_once_with([MessageType.REQUEST_INSTANCE_INFO]) + messengers[bad_endpoint].receive.assert_called_once_with() + messengers[bad_endpoint].stop.assert_called_once_with() + with pytest.raises(KeyError): + registrar.get_peer_rank_info(bad_ri.instance_name, bad_ri.instance_rank) + + # The same receiver can still register a different, compatible sender. + good_params = SimpleNamespace(ctx_info_endpoint=good_endpoint, ctx_dp_rank=0) + sender_info = receiver._get_sender_info(good_params) + assert sender_info.instance_name == good_ri.instance_name + assert registrar.get_pool_mapping(sender_info) == {(0, 0): (0, 0), (0, 1): (0, 1)} + messengers["tcp://good-sender"].send.assert_called_once_with( + [MessageType.REGISTER_RANK_INFO, local_ri.to_bytes()] + ) + assert receiver._get_sender_info(good_params) is sender_info + assert messenger_factory.call_count == 4 # One more info lookup and one sender dealer. + messengers[good_endpoint].receive.assert_called_once_with() + with pytest.raises(KeyError): + registrar.get_peer_rank_info(good_ri.instance_name, good_ri.instance_rank) + finally: + receiver.shutdown() + + +def test_view_validation_allows_nonoverlapping_pp_stage() -> None: + local_page_table = _make_per_view_page_table([2, 1], MapperKind.INDEXED) + local_page_table.layer_groups[0].local_layers = [LocalLayer(0, 2), LocalLayer(1, 3)] + local_ri = make_rankinfo( + tp_size=1, + pp_size=2, + pp_rank=1, + instance_rank=1, + page_table=local_page_table, + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=1, + pp_size=2, + page_table=_make_per_view_page_table([2, 1], MapperKind.INDEXED), + ) + registrar = _make_peer_registrar(local_ri) + + registrar.validate_view_head_geometry(peer_ri) + + assert registrar.get_pool_mapping(peer_ri) == {} + with pytest.raises(KeyError): + registrar.get_peer_rank_info(peer_ri.instance_name, peer_ri.instance_rank) + + @pytest.mark.parametrize( ("local_replay", "peer_replay"), [(False, False), (True, True), (False, True), (True, False)], diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index f415e387b063..64ca684945c7 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -694,6 +694,7 @@ def test_partial_bounced_publication_waits_for_queued_writer_success( overlap = SimpleNamespace( ranks=[0, 1], duplicate_head_factor=1, + peer_duplicate_head_factor=1, overlap_pp_size=1, ) receiver._registrar = SimpleNamespace( From 2774e0455d2c28f9e623c010147e80b2d38f8fc5 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 30 Oct 2026 14:26:49 -0700 Subject: [PATCH 3/4] Fix per-view overlap caching across peer ranks Signed-off-by: Allison Lim --- tensorrt_llm/_torch/disaggregation/native/peer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 85fad6ea1ed1..52092c626123 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -75,7 +75,7 @@ def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1 ] = {} # key: (peer_key, self_lg_pool_key, peer_lg_pool_key) self._self_ext_cache = self_extractor self._peer_ext_cache: Dict[str, KVRegionExtractorV1] = {} - self._overlap_cache: Dict[str, PeerOverlap] = {} + self._overlap_cache: dict[tuple[str, int, int], PeerOverlap] = {} self._aux_transfer_layout_cache: Dict[str, AuxTransferLayout] = {} self._lg_pool_mapping_cache: Dict[ str, Dict[LGPoolKey, LGPoolKey] @@ -599,7 +599,8 @@ def _find_overlap(self_val, peer_val, self_rank, peer_rank=None): def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerOverlap: peer_ri = peer_rank_info - key = self._unique_key(peer_ri.instance_name, peer_dp_rank) + # Per-view duplication factors depend on the specific peer rank's page table. + key = (peer_ri.instance_name, peer_ri.instance_rank, peer_dp_rank) if key in self._overlap_cache: return self._overlap_cache[key] From f8b996c6618127db17fd40bdf586793114fe5e93 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 30 Oct 2026 15:14:33 -0700 Subject: [PATCH 4/4] Disable multi-writer bounce for PP senders Signed-off-by: Allison Lim --- tensorrt_llm/_torch/disaggregation/native/transfer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index e2ae3d03d6de..56e69ca2aa7e 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2961,16 +2961,19 @@ def dispatch_task(self, task: KVRecvTask) -> None: else: task.expected_transfers = len(dp0_overlap.ranks) # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: - # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank + # _fanin_bounce_safe() (TP-by-head), and never under ADP broadcast (sender_dp_rank # None), where the real writer count exceeds expected_transfers and would overflow the slot. topo_overlap = peer_overlap if sender_dp_rank is not None else dp0_overlap # Helix writers contribute unequal shares (one KV owner, one state # sender, the rest empty), so the fan-in equal split would overflow # into neighboring reservations; use the per-fragment path. cp_involved = self._registrar.self_rank_info.cp_size > 1 or peer_infos.cp_size > 1 + # Rank 0's page table cannot establish equal writer shares for other + # PP stages, so multi-writer PP transfers use the per-fragment path. allow_bounce = task.expected_transfers == 1 or ( sender_dp_rank is not None and not cp_involved + and peer_infos.pp_size <= 1 and self._fanin_bounce_safe( topo_overlap, peer_infos,