[Review only][NIXL] Simplify draft head slicing from #58471

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -79,6 +79,7 @@
 from vllm.utils.network_utils import make_zmq_path
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     CircularBufferSpec,
     FullAttentionSpec,
     KpoolTailSpec,
@@ -419,6 +420,44 @@ def _is_region_replicated(self, region_idx: int) -> bool:
         """
         return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]
 
+    def _resolve_head_sharded_draft_kv_heads(self) -> int | None:
+        """Total KV heads of a non-MLA draft under an MLA target, else None.
+
+        Such a draft's KV is head-sharded while the target's is replicated."""
+        if not self.use_mla or self._has_mamba or self._is_csa_linear:
+            return None
+        spec_config = self.vllm_config.speculative_config
+        draft_config = spec_config.draft_model_config if spec_config else None
+        if draft_config is None or draft_config.use_mla:
+            return None
+        return draft_config.get_total_num_kv_heads()
+
+    def _get_head_sharded_draft_regions(self) -> set[int]:
+        heads = self._head_sharded_draft_kv_heads
+        if heads is None:
+            return set()
+        spec_config = self.vllm_config.speculative_config
+        assert spec_config is not None and spec_config.draft_model_config is not None
+        target_layers = self.model_config.get_total_num_hidden_layers()
+        draft_layers = spec_config.draft_model_config.get_total_num_hidden_layers()
+        regions = set[int]()
+        for i, name in enumerate(self.region_names):
+            if self._is_region_replicated(i):
+                continue
+            # DFlash numbers draft layers after the target instead of prefixing them.
+            if not name.startswith("draft_model."):
+                match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", name)
+                if match is None or not (
+                    target_layers <= int(match.group(1)) < target_layers + draft_layers
+                ):
+                    continue
+            spec = self._layer_specs[name]
+            assert isinstance(spec, AttentionSpec) and spec.num_kv_heads <= heads, (
+                f"Draft region {name!r} is inconsistent with {heads} total KV heads"
+            )
+            regions.add(i)
+        return regions
+
     def _set_region_layers(self, region_layers: list[list[str]]) -> None:
         layer_names = [layer_name for region in region_layers for layer_name in region]
         assert len(layer_names) == len(set(layer_names)), (
@@ -926,6 +965,7 @@ def __init__(
         # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
         # combining both (e.g. GQA main + MLA Eagle-3 draft).
         self._region_is_mla = list[bool]()
+        self._head_sharded_draft_regions = set[int]()
         self._ssm_region_indices = list[int]()
         # Regions holding a scratch cache (the CSA compressor circular buffer).
         # Tracked explicitly because a scratch page shares its address with the
@@ -954,6 +994,9 @@ def __init__(
         self._transfer_layer_group_ids = tuple[int, ...]()
         # Per-engine TP mappings. Generated during handshake.
         self.tp_mappings: dict[EngineId, TPMapping] = {}
+        # GQA draft under an MLA target: its SPLIT regions use a draft TP mapping.
+        self._head_sharded_draft_kv_heads = self._resolve_head_sharded_draft_kv_heads()
+        self.draft_tp_mappings: dict[EngineId, TPMapping] = {}
 
         self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
             "enforce_handshake_compat", True
@@ -1789,6 +1832,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         )
         # Descriptor ids must be region-ordered, matching the remote side.
         self._scratch_region_indices.sort()
+        self._head_sharded_draft_regions = self._get_head_sharded_draft_regions()
 
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
         self.num_regions = len(seen_base_addresses)
@@ -2090,9 +2134,6 @@ def _build_fa_remote(
         fa_group_idx = next(
             i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
         )
-        # SPLIT regions read their head slice from this many remote ranks at a
-        # per-rank offset; REPLICATE regions read the whole block once.
-        split_reads = len(plan.source_ranks_per_group[fa_group_idx])
         region_num_blocks = nixl_agent_meta.region_num_blocks or [
             nixl_agent_meta.num_blocks
         ] * len(nixl_agent_meta.kv_caches_base_addr)
@@ -2113,11 +2154,17 @@ def _build_fa_remote(
                 local_block_len = remote_kv_block_len
 
             # REPLICATE reads the whole block once at offset 0; SPLIT gathers
-            # its head slice from `split_reads` remote ranks at a per-rank offset.
-            num_reads = 1 if replicated else split_reads
-            rank_offset = (
-                0 if replicated else plan.rank_offset_factor * remote_kv_block_len
-            )
+            # its head slice from each source rank at a per-rank offset.
+            if replicated:
+                num_reads, rank_offset = 1, 0
+            else:
+                region_plan = (
+                    self.draft_tp_mappings[nixl_agent_meta.engine_id]
+                    if local_region in self._head_sharded_draft_regions
+                    else plan
+                )
+                num_reads = len(region_plan.source_ranks_per_group[fa_group_idx])
+                rank_offset = region_plan.rank_offset_factor * remote_kv_block_len
             local_block_len = local_block_len // num_reads
 
             block_arange = np.arange(region_num_blocks[i], dtype=np.uint64)
@@ -2331,6 +2378,14 @@ def add_remote_agent(
             group_spec_types=self._group_spec_types,
             remote_dcp_size=remote_dcp_size,
         )
+        if self._head_sharded_draft_regions:
+            self.draft_tp_mappings[engine_id] = compute_tp_mapping(
+                transfer_topology=transfer_topo,
+                remote_tp_size=remote_tp_size,
+                group_spec_types=self._group_spec_types,
+                remote_dcp_size=remote_dcp_size,
+                head_sharded_kv_heads=self._head_sharded_draft_kv_heads,
+            )
 
         remote_agent_name = self.nixl_wrapper.add_remote_agent(
             nixl_agent_meta.agent_metadata
@@ -2517,6 +2572,25 @@ def _validate_remote_agent_handshake(
             assert not (
                 tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
             )
+        has_draft_regions = bool(self._head_sharded_draft_regions)
+        draft_local_heads = draft_remote_heads = 1
+        if has_draft_regions:
+            assert self._head_sharded_draft_kv_heads is not None
+            draft_local_heads = max(
+                1, self._head_sharded_draft_kv_heads // self.transfer_topo.tp_size
+            )
+            draft_remote_heads = max(
+                1, self._head_sharded_draft_kv_heads // remote_tp_size
+            )
+            if tp_ratio < 0 or self.dcp_size > 1 or remote_dcp_size > 1:
+                raise NotImplementedError(
+                    "NIXL head-sharded draft KV under an MLA target requires "
+                    "local TP >= remote TP and no DCP: "
+                    f"local TP={self.transfer_topo.tp_size}, "
+                    f"remote TP={remote_tp_size}, "
+                    f"local DCP={self.dcp_size}, remote DCP={remote_dcp_size}, "
+                    f"transfer_mode={self._TRANSFER_MODE}."
+                )
 
         remote_physical_per_logical = (
             nixl_agent_meta.physical_blocks_per_logical_kv_block
@@ -2609,6 +2683,22 @@ def _validate_remote_agent_handshake(
                 "Heterogeneous TP head-dimension splitting requires contiguous heads. "
                 "Use a block-contiguous layout (e.g. LBHNC) on the prefill side."
             )
+        # The MLA exemption above does not cover a head-sharded draft, which
+        # slices the prefill draft block by head.
+        if (
+            has_draft_regions
+            and draft_remote_heads != draft_local_heads
+            and (
+                nixl_agent_meta.kv_cache_layout != kv_cache_layout
+                or not KVCacheLayout[kv_cache_layout].is_block_contiguous
+            )
+        ):
+            raise RuntimeError(
+                "Head-sharded draft KV under an MLA target needs the same "
+                "block-contiguous kv_cache_layout (e.g. LBHNC) on prefill and "
+                f"decode, got local={kv_cache_layout}, "
+                f"remote={nixl_agent_meta.kv_cache_layout}."
+            )
 
         # Per-region block_len validation enforcing the P/D invariant.
         # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)
@@ -2642,12 +2732,19 @@ def _validate_remote_agent_handshake(
                 remote_engine_id
             )
             total_kv_heads = self.transfer_topo.total_num_kv_heads
-            local_heads = self.transfer_topo.local_physical_heads
-            remote_heads = max(1, total_kv_heads // remote_tp_size)
             for i, region_idx in enumerate(local_regions):
                 local_len = self.block_len_per_layer[region_idx]
-                replicated = model_replicated or self._is_region_replicated(region_idx)
                 remote_len = nixl_agent_meta.block_lens[i]
+                if region_idx in self._head_sharded_draft_regions:
+                    replicated = False
+                    local_heads = draft_local_heads
+                    remote_heads = draft_remote_heads
+                else:
+                    replicated = model_replicated or self._is_region_replicated(
+                        region_idx
+                    )
+                    local_heads = self.transfer_topo.local_physical_heads
+                    remote_heads = max(1, total_kv_heads // remote_tp_size)
                 if replicated:
                     assert local_len // block_size_ratio == remote_len, (
                         "KV cache sizes must match between P and D when "
@@ -3768,6 +3865,7 @@ def _cleanup_remote_engine(
         self.dst_uses_region_group_mapping.pop(engine_id, None)
         self.dst_region_mem_types.pop(engine_id, None)
         self.tp_mappings.pop(engine_id, None)
+        self.draft_tp_mappings.pop(engine_id, None)
         if self.transfer_topo is not None:
             self.transfer_topo.unregister_remote_engine(engine_id)
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

@@ -49,6 +49,7 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
     PUSH_REG_NOTIF_PREFIX,
+    NixlAgentMetadata,
     NixlConnectorMetadata,
     RemoteMeta,
     ReqId,
@@ -151,6 +152,25 @@ def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
             self._push_writer_thread.start()
             logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
 
+    def _validate_remote_agent_handshake(
+        self,
+        nixl_agent_meta: NixlAgentMetadata,
+        remote_tp_size: int,
+        remote_dcp_size: int = 1,
+    ) -> None:
+        super()._validate_remote_agent_handshake(
+            nixl_agent_meta, remote_tp_size, remote_dcp_size
+        )
+        assert self.transfer_topo is not None
+        if (
+            self._head_sharded_draft_regions
+            and self.transfer_topo.tp_size != remote_tp_size
+        ):
+            raise NotImplementedError(
+                "NIXL push mode does not support a head-sharded draft's KV "
+                "under an MLA target with heterogeneous TP. Use pull mode."
+            )
+
     def shutdown(self) -> None:
         self._push_writer_stop.set()
         # Unblock the writer if it's waiting in the no-active-state branch.

vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py

@@ -67,11 +67,26 @@ class TPMapping:
 # ======================================================================
 
 
+def compute_head_offset(
+    tp_rank: int, tp_size: int, remote_tp_size: int, total_num_kv_heads: int
+) -> int:
+    """Return the offset in local head slices within the remote KV block."""
+    if tp_size <= remote_tp_size:
+        return 0
+    if tp_size > total_num_kv_heads:
+        remote_rank = tp_rank * remote_tp_size // tp_size
+        local_head = tp_rank * total_num_kv_heads // tp_size
+        remote_head = remote_rank * total_num_kv_heads // remote_tp_size
+        return local_head - remote_head
+    return tp_rank % (tp_size // remote_tp_size)
+
+
 def compute_tp_mapping(
     transfer_topology: TransferTopology,
     remote_tp_size: int,
     group_spec_types: tuple[type[KVCacheSpec], ...],
     remote_dcp_size: int = 1,
+    head_sharded_kv_heads: int | None = None,
 ) -> TPMapping:
     """Build the complete local-to-remote TP mapping.
 
@@ -80,13 +95,17 @@ def compute_tp_mapping(
 
     DCP support is scoped to MLA only, with a side is either fully replicated or fully
     sharded. DCP-branch reuses the same rank set used at handshake selection.
+
+    ``head_sharded_kv_heads`` maps head-sharded KV with that many total KV heads
+    instead of the topology's, e.g. a GQA draft under an MLA target.
     """
     tp_rank = transfer_topology.tp_rank
     tp_size = transfer_topology.tp_size
-    total_num_kv_heads = transfer_topology.total_num_kv_heads
+    is_mla = transfer_topology.is_mla and head_sharded_kv_heads is None
+    total_num_kv_heads = head_sharded_kv_heads or transfer_topology.total_num_kv_heads
     # --- Attention source ranks ---
-    if transfer_topology.is_mla or tp_size >= remote_tp_size:
-        if transfer_topology.is_mla and remote_dcp_size > 1:
+    if is_mla or tp_size >= remote_tp_size:
+        if is_mla and remote_dcp_size > 1:
             attn_ranks = transfer_topology.dcp_source_ranks(
                 remote_tp_size, remote_dcp_size
             )
@@ -137,16 +156,11 @@ def compute_tp_mapping(
     }
 
     # --- Rank offset factor ---
-    if transfer_topology.is_mla or tp_size <= remote_tp_size:
-        # We don't index into remote for reading, no offset needed.
-        rank_offset_factor = 0
-    elif tp_size > total_num_kv_heads:
-        local_head = tp_rank * total_num_kv_heads // tp_size
-        p_start = attn_ranks[0] * total_num_kv_heads // remote_tp_size
-        rank_offset_factor = local_head - p_start
-    else:
-        # D TP > P TP: we index into remote to read different heads depending on rank.
-        rank_offset_factor = tp_rank % (tp_size // remote_tp_size)
+    rank_offset_factor = (
+        0
+        if is_mla
+        else compute_head_offset(tp_rank, tp_size, remote_tp_size, total_num_kv_heads)
+    )
 
     local_consumers = transfer_topology.dcp_consumer_count(
         remote_tp_size, remote_dcp_size