diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 9fcf33b739b9..fd551c7e3b74 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 @@ -306,6 +311,8 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: int, + peer_kv_heads: int, ): self._ri = self_ri self._peer_ri = peer_ri @@ -341,28 +348,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, ) @@ -492,6 +492,8 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: int, + peer_kv_heads: int, ) -> None: # Deliberately do not call HNDHeadMismatchMapper.__init__: its offsets # assume HND-contiguous heads. Initialize the three attributes consumed @@ -504,8 +506,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 + peer_heads = peer_kv_heads if self_buffers_per_layer != peer_buffers_per_layer: raise ValueError( "NHD buffer count per layer mismatch: " @@ -611,16 +613,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 @@ -742,30 +756,72 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool: ) ) - def _head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + @staticmethod + def view_kv_heads(pool_view: PoolView) -> int: + """Require the per-rank head count published by an attention view.""" + heads = pool_view.num_kv_heads + if heads is None or heads <= 0: + raise ValueError( + "Attention pool views must publish a positive num_kv_heads; " + f"pool={pool_view.pool_idx}, role={sorted(pool_view.pool_role)}, " + f"num_kv_heads={heads}" + ) + return heads + + def _head_counts( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView], + peer_pv: Optional[PoolView], + ) -> tuple[int, int]: + if self_pv is None and peer_pv is None: + # Topology checks without a view pair use whole-rank metadata. + return self._ri.attention.kv_heads_per_rank, peer_ri.attention.kv_heads_per_rank + if self_pv is None or peer_pv is None: + raise ValueError("Both pool views are required for per-view head geometry") + return self.view_kv_heads(self_pv), self.view_kv_heads(peer_pv) + + 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 + self_heads, peer_heads = self._head_counts(peer_ri, self_pv, peer_pv) + return self_heads * self_tp, peer_heads * peer_tp - def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]: - factor_self, factor_peer = self._head_factors(peer_ri) + 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. + + 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 - ) + self_heads, peer_heads = self._head_counts(peer_ri, self_pv, peer_pv) + head_match = self._ri.attention.is_mla or self_heads == peer_heads 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: + self_heads, peer_heads = self._head_counts(peer_ri, self_pv, peer_pv) + if self_heads == 0 or peer_heads == 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 @@ -832,7 +888,9 @@ def build_mapper( peer_bytes_per_layer, ) - head_match, _ = self.head_match(peer_ri) + if self_pv is None or peer_pv is None: + raise ValueError("Both pool views are required for per-view head geometry") + head_match, _ = self.head_match(peer_ri, self_pv, peer_pv) if head_match: return IntactMapper( self_layer_offsets, @@ -852,6 +910,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_pv), + peer_kv_heads=self.view_kv_heads(peer_pv), ) return HNDHeadMismatchMapper( @@ -863,4 +923,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_pv), + peer_kv_heads=self.view_kv_heads(peer_pv), ) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..2d8300e6341e 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: @@ -65,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] @@ -84,32 +94,35 @@ 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_pv), policy.view_kv_heads(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 +193,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 +209,39 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: return True + def validate_view_head_geometry(self, peer_ri: RankInfo) -> None: + """Require matching per-view head geometry before registering the peer.""" + 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] + bytes_per_head = [] + for side, lg_idx, pi, pool_view in ( + ("local", self_lg, self_pi, self_pv), + ("peer", peer_lg, peer_pi, peer_pv), + ): + heads = policy.view_kv_heads(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]: + 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]})" + ) + 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 +343,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 +357,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 @@ -487,7 +575,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] @@ -536,6 +625,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 +634,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 +704,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 10ae23a48dce..f3469dc581d9 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2988,6 +2988,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/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 279916d53001..fa3c9d2bce8b 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], int]] = None, ) -> List[PoolView]: """Bucket one slot-desc variant's coalesced buffers into pool views. @@ -243,6 +244,10 @@ 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). ``kv_heads_for_layer`` gives a layer's positive per-rank KV-head + count; head-sharded views also partition by that count so equal-sized + buffers with different head geometry never share a view. """ role_layouts = role_layouts or {} bucket_entries: Dict[tuple, list] = defaultdict(list) @@ -261,7 +266,13 @@ def _build_pool_views_for_variant( if view_partition is not None else None ) - bucket_key = (pool_idx, kind, partition) + num_kv_heads = ( + kv_heads_for_layer(int(buffer_id.layer_id)) + if kv_heads_for_layer is not None + and kind in (MapperKind.INDEXED, MapperKind.NHD) + else None + ) + bucket_key = (pool_idx, kind, partition, num_kv_heads) bucket_entries[bucket_key].append( (int(buffer_id.layer_id), offset, single_buffer_size) ) @@ -279,7 +290,7 @@ def _build_pool_views_for_variant( key=lambda key: (key[0], min(entry[1] for entry in bucket_entries[key])), ) for bucket_key in lg_bucket_keys: - pool_idx, mapper_kind, _ = bucket_key + pool_idx, mapper_kind, _, num_kv_heads = bucket_key roles = frozenset(bucket_roles[bucket_key]) context = ( f"View(layer_group={layer_group_id}, pool={pool_idx}, " @@ -330,6 +341,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, ring_slot_bytes=layout.ring_slot_bytes, ) ) @@ -404,6 +416,42 @@ def _window_size_for_layer(internal_layer_id: int): raise ValueError(f"Cannot resolve layer config for internal layer {internal_layer_id}") + virtual_layer_mapping = getattr(manager, "_layer_attn_to_layer_id", None) + model_layer_by_internal_id: dict[int, int] = {} + if virtual_layer_mapping is not None: + # Shared caches are declared by their owner before dependent layers. + for (model_layer, _attn_type), internal_layer_id in virtual_layer_mapping.items(): + model_layer_by_internal_id.setdefault(int(internal_layer_id), int(model_layer)) + + def _kv_heads_for_layer(internal_layer_id: int) -> int: + local_layer_id = internal_layer_id + if virtual_layer_mapping is not None: + model_layer = model_layer_by_internal_id.get(internal_layer_id) + if model_layer is None: + raise ValueError( + f"Cannot resolve KV-head count for internal layer {internal_layer_id}: " + "no owning model layer" + ) + local_layer_id = manager.layer_offsets.get(model_layer) + if local_layer_id is None: + raise ValueError( + f"Cannot resolve KV-head count for internal layer {internal_layer_id}: " + f"model layer {model_layer} is not local to this rank" + ) + heads = manager.num_kv_heads_per_layer + if not 0 <= local_layer_id < len(heads): + raise ValueError( + f"Cannot resolve KV-head count for internal layer {internal_layer_id}: " + f"local layer {local_layer_id} is outside the per-layer head list" + ) + num_kv_heads = int(heads[local_layer_id]) + if num_kv_heads <= 0: + raise ValueError( + f"Internal layer {internal_layer_id} must have a positive KV-head count, " + f"got {num_kv_heads}" + ) + return num_kv_heads + 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) @@ -452,10 +500,8 @@ def _window_size_for_layer(internal_layer_id: int): ) cache_kind = CacheKind.STATE if is_recurrent else CacheKind.PAGED - # Bucket buffer entries by (pool, mapper kind). One PoolView is - # emitted per bucket and spans every layer of that role class, - # so the view count per layer group is bounded by the number of - # role classes — never by the layer count. ``pool_role`` stays + # Bucket buffer entries by pool, mapper kind, manager partition + # and resolved KV-head count. ``pool_role`` stays # the manager-supplied equivalence label used for peer matching # without enumerating role names. Buffer offsets within a slot # follow ``buffer_ids`` order: the i-th buffer of a coalesced @@ -468,6 +514,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]: @@ -486,15 +533,12 @@ def _window_size_for_layer(internal_layer_id: int): ) continue - # Determine layer group metadata. - # For managers with virtual layers, internal layer_ids - # may exceed the length of num_kv_heads_per_layer. Use index 0 as - # all layers within a pool group share the same kv_heads count. + # The group field describes its first head-sharded view; each + # view carries the actual geometry used for transfer. first_local_layer = all_internal_layer_ids[0] - if first_local_layer < len(manager.num_kv_heads_per_layer): - num_kv_heads = manager.num_kv_heads_per_layer[first_local_layer] - else: - num_kv_heads = manager.num_kv_heads_per_layer[0] + num_kv_heads = next( + (view.num_kv_heads for view in pool_views if view.num_kv_heads is not None), 0 + ) sliding_window_size = _window_size_for_layer(first_local_layer) layer_groups_by_id[layer_group_id] = AttentionLayerGroup( diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index e0b0f5e78f42..a89ca455e9ab 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -205,7 +205,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. Required + for head-sharded attention views; ``None`` for replicated and + recurrent-state views, which do not use attention head counts. ring_slot_bytes: REPLICATED views only. When set, each layer's region is a ring of ``bytes_per_layer // ring_slot_bytes`` token slots, and the token at global position ``p`` lives in slot @@ -223,6 +228,7 @@ class PoolView: section_bytes: Optional[List[int]] = None bytes_per_head: Optional[int] = None ring_slot_bytes: Optional[int] = None + num_kv_heads: Optional[int] = None def to_dict(self) -> dict: return { @@ -239,6 +245,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, "ring_slot_bytes": ( int(self.ring_slot_bytes) if self.ring_slot_bytes is not None else None ), @@ -268,6 +275,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 + ), ring_slot_bytes=( int(data["ring_slot_bytes"]) if data.get("ring_slot_bytes") 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 a4f91d5b1e8f..3015e907d34f 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 @@ -4265,6 +4265,26 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + 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. + + 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: + # Virtual storage IDs may not index the model-layer head list. + local_layer_id = int(layer_id) + if hasattr(self, "_layer_attn_to_layer_id") or not ( + 0 <= local_layer_id < len(self.num_kv_heads_per_layer) + ): + return None + return KVCacheLayerOwner.TARGET.value, self.num_kv_heads_per_layer[local_layer_id] + 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.""" if self.layer_specs is None: diff --git a/tests/unittest/disaggregated/region/test_block.py b/tests/unittest/disaggregated/region/test_block.py index 241a4edd362d..e313c0614f51 100644 --- a/tests/unittest/disaggregated/region/test_block.py +++ b/tests/unittest/disaggregated/region/test_block.py @@ -171,6 +171,8 @@ def test_head_mismatch_mapper(): dst_layer_offsets=[peer_layer_off], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=self_bytes_per_layer, peer_bytes_per_layer=peer_bytes_per_layer, self_buffers_per_layer=buffers_per_layer, diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec6ef42ad0f7..e37e89572440 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 + # Absent wire metadata remains absent so consumers can reject it. + 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,)): @@ -313,12 +322,14 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l get_disagg_ignored_roles=lambda: frozenset(), impl=impl, pp_layers=[0, 1], + layer_offsets={0: 0, 1: 1}, num_kv_heads_per_layer=[1, 1], get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, get_disagg_role_layouts=lambda: {}, ) +@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 +352,290 @@ 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( + "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"), + pytest.param(None, (2, 2), 1, id="no-layer-specs-uniform-heads"), + pytest.param(None, (2, 1), 2, id="no-layer-specs-mixed-heads"), + ], +) +def test_v2_builder_partitions_coalesced_views_by_owner_and_heads( + owners: tuple[str, str] | None, 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) * 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=(4 * buffer_size,) + ) + manager.num_kv_heads_per_layer = list(head_counts) + manager.layer_specs = None + if owners is not None: + manager.layer_specs = tuple( + KVCacheLayerSpec( + global_layer_id=layer, + owner=KVCacheLayerOwner(owners[layer]), + dtype=DataType.HALF, + num_kv_heads=heads, + head_dim=128 // heads, + buffer_roles=(Role.KEY, Role.VALUE), + ) + for layer, heads in enumerate(head_counts) + ) + 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 + + 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, buffer_size], + [2 * buffer_size, 3 * buffer_size], + ] + else: + assert views[0].num_kv_heads == 2 + assert get_unique_layers(views[0]) == {0, 1} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("mapper_kind", [MapperKind.INDEXED, MapperKind.NHD]) +@pytest.mark.parametrize("internal_ids", [(0, 1), (0, 2), (2, 3)]) +@pytest.mark.parametrize("head_counts", [(2, 2), (2, 1)]) +def test_v2_builder_resolves_virtual_layer_heads( + mapper_kind: MapperKind, internal_ids: tuple[int, int], head_counts: tuple[int, int] +) -> None: + from enum import IntEnum + from types import SimpleNamespace + + class AttentionKind(IntEnum): + PRIMARY = 0 + SHARED = 1 + + attrs = { + (layer, role): SimpleNamespace( + life_cycle_id=0, pool_index=0, offset=(2 * index + role_index) * 128, size=128 + ) + for index, layer in enumerate(internal_ids) + for role_index, role in enumerate((Role.KEY, Role.VALUE)) + } + manager = _make_fake_v2_manager(attrs, {Role.ALL: mapper_kind}, slot_bytes_list=(512,)) + manager.layer_specs = None + manager.impl.layer_grouping = (internal_ids,) + manager.pp_layers = [20, 22] + manager.layer_offsets = {20: 0, 22: 1} + manager.num_kv_heads_per_layer = list(head_counts) + # Storage IDs can be in-range yet name a different model layer. The + # later alias must retain the first model layer's cache ownership. + manager._layer_attn_to_layer_id = { + (22, AttentionKind.PRIMARY): internal_ids[0], + (20, AttentionKind.SHARED): internal_ids[0], + (20, AttentionKind.PRIMARY): internal_ids[1], + } + 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 head_counts[0] == head_counts[1]: + assert len(views) == 1 + assert get_unique_layers(views[0]) == set(internal_ids) + assert views[0].num_kv_heads == 2 + assert views[0].buffer_entries["offset"].tolist() == [0, 128, 256, 384] + else: + assert len(views) == 2 + assert [get_unique_layers(view) for view in views] == [{layer} for layer in internal_ids] + assert [view.num_kv_heads for view in views] == [head_counts[1], head_counts[0]] + assert [view.buffer_entries["offset"].tolist() for view in views] == [ + [0, 128], + [256, 384], + ] + + +@pytest.mark.cpu_only +def test_v2_builder_resolves_mixed_precision_twin_heads() -> None: + from enum import IntEnum + from types import SimpleNamespace + + class AttentionKind(IntEnum): + PRIMARY = 0 + HIGH_PRECISION = 1 + + attrs = { + (3, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (3, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + } + manager = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(256,)) + manager.impl.layer_grouping = ((3,),) + manager.impl.init_config.layers.extend( + [SimpleNamespace(window_size=128), SimpleNamespace(window_size=128)] + ) + manager.pp_layers = [20, 22] + manager.layer_offsets = {20: 0, 22: 1} + manager.num_kv_heads_per_layer = [8, 2] + manager._layer_attn_to_layer_id = { + (20, AttentionKind.PRIMARY): 0, + (22, AttentionKind.PRIMARY): 1, + (20, AttentionKind.HIGH_PRECISION): 2, + (22, AttentionKind.HIGH_PRECISION): 3, + } + + layer_group = build_page_table_from_manager(manager).layer_groups[0] + + assert layer_group.kv_head_num_per_rank == 2 + assert layer_group.sliding_window_size == 128 + assert layer_group.local_layers[0].global_layer_id == 45 + assert layer_group.pool_views[0].num_kv_heads == 2 + + +@pytest.mark.cpu_only +def test_v2_builder_accepts_replicated_virtual_layer_without_head_count() -> None: + from types import SimpleNamespace + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + (2, Role.INDEX_KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=256, size=128), + } + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED}, + slot_bytes_list=(384,), + ) + manager.layer_specs = None + manager.impl.layer_grouping = ((0, 2),) + manager.get_disagg_global_layer_ids = lambda _: [0, 2] + 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 + + assert len(views) == 2 + assert get_unique_layers(views[0]) == {0} + assert views[0].num_kv_heads == 1 + assert get_unique_layers(views[1]) == {2} + assert views[1].mapper_kind == MapperKind.REPLICATED + assert views[1].num_kv_heads is None + assert views[1].buffer_entries["offset"].tolist() == [256] + + +@pytest.mark.cpu_only +def test_v2_builder_accepts_replicated_only_layer_group_without_head_count() -> None: + from types import SimpleNamespace + + attrs = { + (2, Role.INDEX_KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + } + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED}, + slot_bytes_list=(128,), + ) + manager.impl.layer_grouping = ((2,),) + manager.impl.init_config.layers.append(SimpleNamespace(window_size=None)) + manager.get_disagg_global_layer_ids = lambda _: [2] + manager.num_kv_heads_per_layer = [] + + layer_group = build_page_table_from_manager(manager).layer_groups[0] + + assert layer_group.kv_head_num_per_rank == 0 + assert len(layer_group.pool_views) == 1 + assert layer_group.pool_views[0].mapper_kind == MapperKind.REPLICATED + assert layer_group.pool_views[0].num_kv_heads is None + + +@pytest.mark.cpu_only +def test_v2_builder_splits_mixed_heads_without_view_partition() -> None: + 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] + views = build_page_table_from_manager(manager).layer_groups[0].pool_views + assert [view.num_kv_heads for view in views] == [2, 1] + assert [get_unique_layers(view) for view in views] == [{0}, {1}] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("head_counts", [(0, 1), (-1, 1), (1,)]) +def test_v2_builder_rejects_unresolved_or_nonpositive_heads(head_counts: tuple[int, ...]) -> None: + 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 = list(head_counts) + + with pytest.raises(ValueError, match="KV-head count"): + build_page_table_from_manager(manager) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("missing_owner", [False, True]) +def test_v2_builder_rejects_unresolved_virtual_layer_heads(missing_owner: bool) -> None: + from enum import IntEnum + from types import SimpleNamespace + + class AttentionKind(IntEnum): + PRIMARY = 0 + + attrs = { + (0, Role.KEY): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=0, size=128), + (0, Role.VALUE): SimpleNamespace(life_cycle_id=0, pool_index=0, offset=128, size=128), + } + manager = _make_fake_v2_manager(attrs, {Role.ALL: MapperKind.INDEXED}, slot_bytes_list=(256,)) + manager.impl.layer_grouping = ((0,),) + manager.get_disagg_global_layer_ids = lambda _: [20] + manager._layer_attn_to_layer_id = {} if missing_owner else {(20, AttentionKind.PRIMARY): 0} + manager.layer_offsets = {} + + with pytest.raises(ValueError, match="Cannot resolve KV-head count for internal layer 0"): + 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. @@ -628,12 +917,15 @@ def test_v2_builder_describes_recurrent_life_cycle_from_descriptors(): ) manager = _make_fake_v2_mamba_manager() + manager.num_kv_heads_per_layer = [0, 0, 1] page_table = build_page_table_from_manager(manager) assert len(page_table.layer_groups) == 2 recurrent, attention = page_table.layer_groups assert isinstance(recurrent, MambaLayerGroup) assert isinstance(attention, AttentionLayerGroup) + assert all(view.num_kv_heads is None for view in recurrent.pool_views) + assert attention.pool_views[0].num_kv_heads == 1 # The recurrent group uses the shared descriptor pool group, not a # dedicated tensor-derived one, and keeps V2 layer numbering. assert recurrent.pool_group_idx == 0 diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 15670e3de0c6..bf88fb1c970e 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 @@ -57,7 +60,9 @@ pytestmark = pytest.mark.cpu_only -def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mamba_tp=2): +def make_page_table( + pool_ptrs=None, block_bytes=None, global_layer_ids=None, mamba_tp=2, kv_heads=2 +): """Create a KVCachePageTable for testing. ``mamba_tp`` is the TP degree the synthetic recurrent-state group is @@ -88,7 +93,11 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mam entries.append((i, base_offset, buffer_size)) entries.append((i, base_offset + buffer_size, buffer_size)) pool_views.append( - PoolView(pool_idx=pi, buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE)) + PoolView( + pool_idx=pi, + buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), + num_kv_heads=kv_heads, + ) ) physical_pools = [ PhysicalPool(base_address=ptr, slot_bytes=bs, num_slots=128) @@ -97,7 +106,7 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mam attn_lg = AttentionLayerGroup( pool_group_idx=0, - kv_head_num_per_rank=2, + kv_head_num_per_rank=kv_heads, sliding_window_size=None, local_layers=local_layers, pool_views=pool_views, @@ -145,7 +154,7 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mam return KVCachePageTable( tokens_per_block=16, - layer_groups=[attn_lg, mamba_lg], + layer_groups=([attn_lg] if kv_heads else []) + [mamba_lg], pool_groups=pool_groups, ) @@ -202,7 +211,7 @@ def make_rankinfo( def _make_peer_registrar_and_peer_ri(self_ri, peer_ri): - pt = make_page_table() + pt = make_page_table(kv_heads=self_ri.attention.kv_heads_per_rank) real_kv_extractor = KVRegionExtractorV1(pt) reg = PeerRegistrar(self_ri, real_kv_extractor) if peer_ri.page_table is None: @@ -224,7 +233,7 @@ def _make_peer_registrar_and_peer_ri(self_ri, peer_ri): is_mla=peer_ri.attention.is_mla, enable_attention_dp=peer_ri.attention.enable_attention_dp, layer_num_per_pp=peer_ri.layer_num_per_pp, - page_table=make_page_table(), + page_table=make_page_table(kv_heads=peer_ri.attention.kv_heads_per_rank), ) return reg, peer_ri @@ -403,12 +412,370 @@ def test_multiple_overlap(): def _make_peer_registrar(self_rankinfo): - pt = self_rankinfo.page_table if self_rankinfo.page_table else make_page_table() + pt = self_rankinfo.page_table or make_page_table( + kv_heads=self_rankinfo.attention.kv_heads_per_rank + ) real_kv_extractor = KVRegionExtractorV1(pt) reg = PeerRegistrar(self_rankinfo, real_kv_extractor) 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 + + +@pytest.mark.parametrize("mapper_kind", [MapperKind.INDEXED, MapperKind.NHD]) +@pytest.mark.parametrize( + "view_idx,peer_tp_rank,src_head_idx", + [ + (0, 0, 0), + (0, 1, 2), + (0, 2, 4), + (0, 3, 6), + (1, 0, 0), + (1, 1, 0), + (1, 2, 1), + (1, 3, 1), + ], +) +def test_per_view_heads_select_correct_transfer_bytes_with_larger_gen_tp( + mapper_kind: MapperKind, view_idx: int, peer_tp_rank: int, src_head_idx: int +) -> None: + # TP1 -> TP4: the eight-head view sends two heads per receiver; the + # two-head view sends head 0 to ranks 0/1 and head 1 to ranks 2/3. + self_ri = make_rankinfo( + 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), + ) + peer_ri = make_rankinfo( + instance_name="peer", + tp_size=4, + tp_rank=peer_tp_rank, + kv_heads_per_rank=2, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], 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, view_idx), (0, view_idx)) + src_heads = [8, 2][view_idx] + dst_heads = [2, 1][view_idx] + pair = mapper.map( + SpecRegion(memory=MemRegionGroup(ptrs=np.array([1000]), bytes_per_region=16 * src_heads)), + SpecRegion(memory=MemRegionGroup(ptrs=np.array([2000]), bytes_per_region=16 * dst_heads)), + ) + + if mapper_kind == MapperKind.NHD: + assert isinstance(mapper, NHDHeadMismatchMapper) + # Four bytes per token/head, with a fragment for each K/V token. + expected_src_ptrs = [1000 + 4 * src_head_idx + 4 * src_heads * i for i in range(4)] + expected_dst_ptrs = [2000 + 4 * dst_heads * i for i in range(4)] + expected_fragment_bytes = 4 * dst_heads + else: + assert isinstance(mapper, HNDHeadMismatchMapper) + # Eight bytes per head across both tokens, with one fragment per K/V. + expected_src_ptrs = [1000 + 8 * src_head_idx, 1000 + 8 * src_heads + 8 * src_head_idx] + expected_dst_ptrs = [2000, 2000 + 8 * dst_heads] + expected_fragment_bytes = 8 * dst_heads + assert pair.src.memory.ptrs.tolist() == expected_src_ptrs + assert pair.dst.memory.ptrs.tolist() == expected_dst_ptrs + assert pair.src.memory.bytes_per_region == expected_fragment_bytes + assert pair.dst.memory.bytes_per_region == expected_fragment_bytes + + +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]] + + +@pytest.mark.parametrize("mapper_kind", [MapperKind.INDEXED, MapperKind.NHD]) +@pytest.mark.parametrize("invalid_heads", [None, 0, -1]) +@pytest.mark.parametrize("invalid_sides", [("local",), ("peer",), ("local", "peer")]) +def test_per_view_rejects_missing_or_invalid_head_counts( + mapper_kind: MapperKind, invalid_heads: int | None, invalid_sides: tuple[str, ...] +) -> None: + self_ri = make_rankinfo( + tp_size=1, + 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=2, + tokens_per_block=2, + dims_per_head=2, + page_table=_make_per_view_page_table([2, 1], mapper_kind), + ) + # This view matches the rank-wide count: a fallback would silently accept it. + for side, ri in (("local", self_ri), ("peer", peer_ri)): + if side in invalid_sides: + ri.page_table.layer_groups[0].pool_views[0].num_kv_heads = invalid_heads + reg = _make_peer_registrar(self_ri) + with pytest.raises(ValueError, match="must publish a positive num_kv_heads"): + 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(), + 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"), + [ + ("missing_heads", "must publish a positive num_kv_heads"), + ("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 == "missing_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() + 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)], @@ -623,7 +990,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], mamba_tp=1), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1, kv_heads=4), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) @@ -758,6 +1125,8 @@ def test_nhd_head_mismatch_mapper_slices_each_token(): dst_layer_offsets=[0], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=32, peer_bytes_per_layer=16, self_buffers_per_layer=2, @@ -798,6 +1167,8 @@ def test_nhd_head_mismatch_mapper_uses_scale_pool_geometry(): dst_layer_offsets=[0], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=8, peer_bytes_per_layer=4, self_buffers_per_layer=2, @@ -1073,6 +1444,8 @@ def test_nhd_mapper_rejects_non_divisible_region_geometry(): dst_layer_offsets=[0], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=17, peer_bytes_per_layer=32, self_buffers_per_layer=2, @@ -1090,6 +1463,8 @@ def test_nhd_mapper_rejects_tokens_per_block_mismatch(): dst_layer_offsets=[0], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=32, peer_bytes_per_layer=64, self_buffers_per_layer=2, @@ -1133,10 +1508,8 @@ def test_non_uniform_view_geometry_rejected_for_all_kinds(): self_ri = make_rankinfo(page_table=self_pt) peer_ri = make_rankinfo(instance_name="peer", page_table=peer_pt) reg = _make_peer_registrar(self_ri) - reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) - - with pytest.raises(ValueError, match="not uniform"): - reg.get_kv_map(peer_ri, (0, 0), (0, 0)) + with pytest.raises(ValueError, match="not evenly distributed"): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): @@ -1155,7 +1528,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], mamba_tp=1), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1, kv_heads=4), ) reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) @@ -1185,7 +1558,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], mamba_tp=1), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1, kv_heads=4), ) reg = _make_peer_registrar(self_ri) @@ -1195,7 +1568,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], mamba_tp=1) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1, kv_heads=4) 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( @@ -1219,7 +1592,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], mamba_tp=1) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1, kv_heads=4) 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) @@ -1311,6 +1684,8 @@ def test_indexed_head_mismatch_subbyte_geometry_is_entries_derived(): dst_layer_offsets=[0, 128], # 2 layers x 128B/layer (kv_factor 2 x 64B) self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=256, peer_bytes_per_layer=128, self_buffers_per_layer=2, @@ -1339,6 +1714,8 @@ def test_indexed_head_mismatch_inconsistent_slot_geometry_raises(): dst_layer_offsets=[0, 256], self_ri=self_ri, peer_ri=peer_ri, + self_kv_heads=self_ri.attention.kv_heads_per_rank, + peer_kv_heads=peer_ri.attention.kv_heads_per_rank, self_bytes_per_layer=256, peer_bytes_per_layer=256, # should be 128 for half the heads self_buffers_per_layer=2,