[NIXL] Per-region KV transfer classification for mixed full-attn + MLA groups

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -1062,6 +1062,87 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
             # whole block is moved.
             worker.add_remote_agent(meta, remote_tp_size=1)
 
+    @patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
+        FakeNixlWrapper,
+    )
+    def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):
+        """Mixed full-attn (SPLIT) + MLA (REPLICATE) single KV group under
+        heterogeneous TP must NOT raise (previously a NotImplementedError),
+        and the per-region gate must still reject a wrong block_len.
+        """
+        vllm_config = create_vllm_config()
+        with patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size",  # noqa: E501
+            return_value=2,
+        ):
+            connector = NixlConnector(
+                vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+            )
+            connector.connector_worker = FakeNixlConnectorWorker(
+                vllm_config, connector.engine_id, hand_shake_latency=0
+            )
+            worker = connector.connector_worker
+
+            # Region 0: full-attn (SPLIT). Region 1: MLA (REPLICATE).
+            fa_len = 4096 * worker.block_size
+            idx_len = 512 * worker.block_size
+            worker.slot_size_per_layer = [4096, 512]
+            worker.block_len_per_layer = [fa_len, idx_len]
+            worker._region_is_mla = [False, True]
+            worker.num_blocks = 1
+            worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
+            worker.src_blocks_data = [
+                (0, fa_len, worker.tp_rank),
+                (0, idx_len, worker.tp_rank),
+            ]
+            worker.num_descs = len(worker.src_blocks_data)
+
+            # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio;
+            # REPLICATE region is unchanged.
+            tp_ratio = 2
+            meta = NixlAgentMetadata(
+                engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+                agent_metadata=FakeNixlWrapper.AGENT_METADATA,
+                kv_caches_base_addr=[0, 0],
+                device_id=0,
+                num_blocks=1,
+                block_lens=[fa_len * tp_ratio, idx_len],
+                kv_cache_layout=worker.kv_cache_layout,
+                block_size=worker.block_size,
+                ssm_sizes=(0, 0),
+                attn_backend_name=worker.backend_name,
+                physical_blocks_per_logical_kv_block=1,
+            )
+            worker.add_remote_agent(meta, remote_tp_size=1)
+            assert (
+                FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles
+            )
+            # Gate rejects an MLA region wrongly scaled by tp_ratio.
+            worker2 = FakeNixlConnectorWorker(
+                vllm_config, connector.engine_id, hand_shake_latency=0
+            )
+            worker2.block_len_per_layer = [fa_len, idx_len]
+            worker2._region_is_mla = [False, True]
+            worker2.num_blocks = 1
+            worker2.dst_num_blocks[worker2.engine_id] = worker2.num_blocks
+            bad_meta = NixlAgentMetadata(
+                engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+                agent_metadata=FakeNixlWrapper.AGENT_METADATA,
+                kv_caches_base_addr=[0, 0],
+                device_id=0,
+                num_blocks=1,
+                # WRONG: MLA region scaled by tp_ratio (it should be replicated).
+                block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],
+                kv_cache_layout=worker2.kv_cache_layout,
+                block_size=worker2.block_size,
+                ssm_sizes=(0, 0),
+                attn_backend_name=worker2.backend_name,
+                physical_blocks_per_logical_kv_block=1,
+            )
+            with pytest.raises(AssertionError):
+                worker2.add_remote_agent(bad_meta, remote_tp_size=1)
+
 
 # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which
 # we put here is important. First run ray, it will clean up the resources, then

tests/v1/kv_connector/unit/test_tp_mapping.py

@@ -73,9 +73,19 @@ def test_source_ranks_p_gt_d(self):
 
 
 def _make_mock_worker_for_splits(group_spec_types):
-    """Build a mock NixlConnectorWorker with _group_spec_types for split tests."""
+    """Build a mock NixlConnectorWorker with _group_spec_types for split tests.
+
+    No per-region transfer classes are configured (``block_len_per_layer`` empty
+    and ``num_regions == 0``), so ``_fa_desc_classes`` takes its early-return path
+    and treats every FA descriptor as SPLIT, matching the legacy behavior these
+    tests assert.
+    """
     worker = object.__new__(NixlConnectorWorker)
     worker._group_spec_types = group_spec_types
+    worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=False)
+    worker.block_len_per_layer = []
+    worker.num_regions = 0
+    worker._region_is_mla = []
     return worker
 
 

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

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-region transfer classes for the NIXL connector.
+
+Each region in a KV-cache group is tagged with one ``RegionTransferClass``,
+which owns the descriptor rules for that region kind:
+
+- ``SPLIT`` — full-attention (GQA), head-sharded across TP. Remote block holds
+  ``tp_ratio`` x local heads; a rank reads its head slice at a per-rank offset,
+  and under the FlashInfer "virtually split" layout K/V are two streams.
+- ``REPLICATE`` — MLA (``num_kv_heads==1``), identical on every rank: whole
+  block read from one rank at offset 0, key-only (no V stream).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class RegionTransferClass:
+    """How one KV region is split (or not) when transferred over NIXL."""
+
+    is_replicate: bool
+
+    def num_streams(self, virtually_split: bool) -> int:
+        """Descriptor streams emitted: 2 (K, V) for a virtually-split SPLIT
+        region, else 1 (REPLICATE is key-only)."""
+        if self.is_replicate or not virtually_split:
+            return 1
+        return 2
+
+    def remote_num_reads(self, split_reads: int) -> int:
+        """Remote ranks read from: 1 for REPLICATE; ``split_reads`` for SPLIT
+        (gathered head slice when P_TP > D_TP)."""
+        return 1 if self.is_replicate else split_reads
+
+    def remote_rank_offset(self, offset_factor: int, remote_kv_block_len: int) -> int:
+        """Byte offset into the remote block: 0 for REPLICATE (whole block),
+        else this rank's head slice."""
+        return 0 if self.is_replicate else offset_factor * remote_kv_block_len
+
+    def local_split_desc(
+        self,
+        addr: int,
+        local_len: int,
+        device: int,
+        head_slot: int,
+        num_splits: int,
+    ) -> tuple[int, int, int]:
+        """Local dest descriptor for one source-rank read when gathering from
+        multiple remote ranks (P_TP > D_TP): REPLICATE writes the whole block,
+        SPLIT writes only this rank's head slice."""
+        if self.is_replicate:
+            return (addr, local_len, device)
+        chunk = local_len // num_splits
+        return (addr + head_slot * chunk, chunk, device)
+
+    def validate_block_len(
+        self,
+        region_idx: int,
+        local_len: int,
+        remote_len: int,
+        tp_ratio: int,
+        block_size_ratio: int,
+    ) -> None:
+        """Assert the P/D block-length invariant for a region of this class."""
+        if self.is_replicate:
+            # Whole block copied; only the number of blocks may differ.
+            assert local_len // block_size_ratio == remote_len, (
+                "KV cache sizes must match between P and D when replicated "
+                f"(region {region_idx}: local={local_len}, remote={remote_len}, "
+                f"bsr={block_size_ratio})."
+            )
+        elif tp_ratio > 0:
+            # D_TP >= P_TP: remote holds tp_ratio x local heads.
+            assert remote_len == (local_len * tp_ratio) // block_size_ratio, (
+                f"SPLIT region {region_idx}: remote P KV block_len {remote_len} "
+                f"must equal local {local_len} * tp_ratio {tp_ratio} "
+                f"// block_size_ratio {block_size_ratio}."
+            )
+        else:
+            # P_TP > D_TP: local holds |tp_ratio| x remote heads.
+            assert block_size_ratio == 1, (
+                "Different local/remote block sizes are not supported when P TP > D TP."
+            )
+            assert remote_len == local_len // (-tp_ratio), (
+                f"SPLIT region {region_idx}: remote P KV block_len {remote_len} "
+                f"must equal local {local_len} // |tp_ratio| {-tp_ratio}."
+            )
+
+
+SPLIT_CLASS = RegionTransferClass(is_replicate=False)
+REPLICATE_CLASS = RegionTransferClass(is_replicate=True)