[None][feat] Per View KV head transfer

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),
         )