[KV Connector] Support NIXL P/D for hybrid MLA+SSM models

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1380,3 +1380,199 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
     assert list(result) == expected_kernel_block_ids, (
         f"Expected {expected_kernel_block_ids}, got {result}"
     )
+
+
+# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
+
+
+def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
+    """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed
+    MambaSpec) groups whose layers share the same HMA tensors, with a
+    mamba-aligned unified page and an MLA kernel block smaller than the
+    logical block."""
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MambaSpec,
+        MLAAttentionSpec,
+    )
+
+    # 12-token logical blocks over a 4-token MLA kernel block.
+    mla_spec = MLAAttentionSpec(
+        block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16
+    )
+    unified_page = mla_spec.page_size_bytes
+    kda_spec = MambaSpec(
+        block_size=12,
+        # GDN-decomposable conv (Q|K|V = 2|2|4 cols x 3 rows) + fp32 temporal.
+        shapes=((8, 3), (1, 4, 4)),
+        dtypes=(torch.float16, torch.float32),
+        page_size_padded=unified_page,
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    assert kda_spec.page_size_bytes == unified_page
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_blocks * unified_page,
+                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+            )
+            for i in range(2)
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
+            KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
+            KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
+        ],
+    )
+
+
+@pytest.mark.cpu_test
+def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
+    """Hybrid MLA+KDA registration: HMA tensors shared by both layer types
+    must be flagged as MLA regions even when a KDA layer registers them
+    first, expose TP-independent kernel-granularity block lens, and build
+    FA + mamba descriptors for every region."""
+    from unittest.mock import MagicMock
+
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    kv_cache_config = _make_hybrid_mla_kv_cache_config()
+    unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
+    vllm_config = create_vllm_config(block_size=12)
+    # kv_buffer_device defaults to the *real* platform's device type, which on
+    # a CPU-only test host would make this a host-buffer worker: host xfer
+    # buffers are per-layer, so the HMA shared tensors would not be
+    # deduplicated. Pin it to the faked device type.
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [4]
+    fake_backend.get_name.return_value = "FLASHMLA"
+    fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    with (
+        patch.object(bw, "NixlWrapper"),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        patch(
+            "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
+            return_value="DS",
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
+        worker.use_mla = True  # opt-125m test config is not MLA; force the flag
+        worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
+
+        tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
+        # KDA layer first per tensor: exercises the dual-purpose flag merge.
+        worker.register_kv_caches(
+            {
+                "kda_a.0": tensors[0],
+                "mla.0": tensors[0],
+                "kda_b.0": tensors[0],
+                "kda_a.1": tensors[1],
+                "mla.1": tensors[1],
+                "kda_b.1": tensors[1],
+            }
+        )
+
+    # 12-token logical blocks over the 4-token MLA kernel block.
+    assert worker._physical_blocks_per_logical_kv_block == 3
+    assert worker.block_size == 4 and worker.num_blocks == 12
+    # Both shared tensors are dual-purpose: their FA view is MLA even though
+    # a KDA layer registered them first.
+    assert worker._region_is_mla == [True, True]
+    assert worker.num_regions == 2 and worker.num_descs == 24
+    # Kernel-granularity block lens; TP-independent for MLA hybrids.
+    assert worker.block_len_per_layer == [unified_page // 3] * 2
+    # Split handles must replicate every FA descriptor (MLA isn't head-sharded).
+    assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24
+    # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page.
+    # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
+    assert worker.src_blocks_data.shape == (24 + 32, 3)
+    fa_descs = worker.src_blocks_data[:24]
+    assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
+    assert all(size == unified_page // 3 for size in fa_descs[:, 1])
+
+
+@pytest.mark.cpu_test
+def test_push_write_hybrid_mla_replicates_attention():
+    """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be
+    written to every covered D rank (replicated MLA latent) while SSM state
+    is written per-rank through the split handles."""
+    import threading
+    from collections import defaultdict
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
+        NixlPushConnectorWorker,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
+        TPMapping,
+    )
+    from vllm.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec
+
+    worker = object.__new__(NixlPushConnectorWorker)
+    worker.shutdown = lambda: None  # skeleton worker: silence __del__
+    worker.use_mla = True
+    worker._has_mamba = True
+    worker._group_spec_types = (MLAAttentionSpec, MambaSpec)
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.tp_ratio.return_value = -2
+    remote_info = MagicMock()
+    remote_info.remote_physical_blocks_per_logical = 1
+    remote_info.remote_block_size = 4
+    worker.transfer_topo.get_engine_info.return_value = remote_info
+
+    engine_id = "remote-engine"
+    # Read-oriented mapping collapses the replicated attention group to one
+    # source rank; the SSM state is sharded across both covered D ranks.
+    worker.tp_mappings = {
+        engine_id: TPMapping(
+            source_ranks_per_group=((0,), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 0},
+            rank_offset_factor=0,
+        )
+    }
+    worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
+    worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
+    worker.src_xfer_handles_by_block_size = {4: 300}
+    worker._sending_transfers = defaultdict(list)
+    worker._sending_transfers_lock = threading.Lock()
+    worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
+    worker._xfer_blocks = MagicMock(return_value=1)
+
+    meta = MagicMock()
+    meta.remote.engine_id = engine_id
+    meta.remote.block_ids = [[7, 8], [3]]
+    meta.local_physical_block_ids = [[1, 2], [5]]
+
+    worker._xfer_blocks_for_req("req-1", meta)
+
+    calls = worker._xfer_blocks.call_args_list
+    assert len(calls) == 2
+    for call, rank, local_handle, remote_handle in zip(
+        calls, (0, 1), (200, 201), (100, 101)
+    ):
+        spec = call.kwargs["read_spec"]
+        assert spec.remote_rank == rank
+        # Attention group replicated to every rank, SSM by membership.
+        assert spec.local_block_ids == [[1, 2], [5]]
+        assert spec.remote_block_ids == [[7, 8], [3]]
+        assert call.kwargs["local_xfer_side_handle"] == local_handle
+        assert call.kwargs["remote_xfer_side_handle"] == remote_handle

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -617,3 +617,24 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens):
         num_tokens=num_tokens,
         tp_size=8,
     )
+
+
+@pytest.mark.cpu_test
+def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid():
+    """The MLA per-token page is TP-independent, so kernel block lengths
+    differing by anything other than the block-size ratio must fail the
+    handshake loudly rather than transfer at mismatched geometry."""
+    worker = _make_mla_hybrid_worker(
+        local_block_size=12, kernel_block_size=4, num_logical_blocks=8
+    )
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=8,
+        remote_kernel_block_size=4,
+        remote_num_logical=12,
+        remote_ssm_sizes=(24, 32),
+    )
+    # Equal kernel block sizes (ratio 1), but a half-sized per-token page.
+    meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer]
+    with pytest.raises((AssertionError, RuntimeError)):
+        worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2)

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -43,6 +43,7 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
     get_base_request_id,
 )
+from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput
 
 from .utils import make_nixl_push_scheduler
@@ -337,6 +338,9 @@ def fresh(cls) -> _StubWriterWorker:
         w.engine_id = "test-decode-engine"
         w._remote_agents = {}
         w._physical_blocks_per_logical_kv_block = 1
+        # Single non-hybrid attention group, matching the stub block id lists.
+        w._has_mamba = False
+        w._group_spec_types = (FullAttentionSpec,)
 
         # Track _do_start_push_kv invocations.
         calls: list[tuple[str, Any, dict[str, Any]]] = []

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

@@ -1093,7 +1093,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
 
         caches_data = []
         # With hybrid allocator, layers can share a kv cache tensor
-        seen_base_addresses = []
+        seen_base_addresses: list[int] = []
 
         # K and V are packed into the content dim, so each attention layer is a
         # single NIXL region whose block transfers as one unit. Mamba layers instead
@@ -1139,11 +1139,19 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             curr_tensor_size_bytes = num_blocks * physical_page_size
 
             base_addr = cache.data_ptr()
+            is_mla_region = isinstance(
+                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
+            )
             if base_addr in seen_base_addresses:
                 # NOTE (NickLucche) HMA employs memory pooling to share tensors
                 # across groups. This results in skipping all tensors but the ones
                 # pointed to by group0. Also, generally we will have more blocks
                 # per tensor but fewer regions.
+                # A shared tensor may back both SSM and attention layers (e.g.
+                # KDA+MLA in KimiLinear); the region's FA view is MLA whichever
+                # layer registered it first.
+                idx = seen_base_addresses.index(base_addr)
+                self._region_is_mla[idx] |= is_mla_region
                 logger.debug("Skipping %s because it's already seen", layer_name)
                 continue
             logger.debug(
@@ -1157,9 +1165,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 )
             else:
                 self.block_len_per_layer.append(physical_page_size)
-            is_mla_region = isinstance(
-                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
-            )
             self._region_is_mla.append(is_mla_region)
 
             if not is_mla_region:
@@ -1801,7 +1806,22 @@ def _validate_remote_agent_handshake(
         # the per-rank KV head ratio rather than the raw tp_ratio, because GQA
         # replication caps per-rank heads at 1 when tp > total_kv_heads
         # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.
-        if not self._has_mamba:
+        if self._has_mamba and self.use_mla:
+            # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are
+            # kernel-granularity views of the mamba-unified page. The MLA
+            # per-token page is TP-independent, so the block lengths must
+            # match up to the kernel block size ratio even under
+            # heterogeneous TP (remote kernel blocks may be smaller).
+            # SSM geometry is validated via ssm_sizes/conv offsets instead.
+            assert self.block_len_per_layer == [
+                block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens
+            ], (
+                "Hybrid MLA kernel-granularity block lengths must match "
+                f"between P and D (block_size_ratio={block_size_ratio}): "
+                f"local={self.block_len_per_layer}, "
+                f"remote={nixl_agent_meta.block_lens}."
+            )
+        elif not self._has_mamba:
             assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
                 "Number of KV layers must match between prefill and decode"
             )

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

@@ -52,7 +52,10 @@
     ReqMeta,
     TransferHandle,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
+    ReadSpec,
+    _is_attention_spec,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id
 from vllm.logger import init_logger
 
@@ -505,41 +508,37 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
         local_block_ids = meta.local_physical_block_ids
         num_groups = len(local_block_ids)
 
-        if self.use_mla and tp_ratio < 0:
-            # MLA latent is replicated across D's TP ranks: the tp-mapping
-            # collapses to one rank (fine for reads), but push must WRITE every
-            # D rank or the rest decode stale KV; only the dst differs per rank.
+        # MLA latent is replicated across D's TP ranks: the tp-mapping
+        # collapses it to one rank (fine for reads), but push must WRITE every
+        # D rank or the rest decode stale KV. For hybrid MLA+SSM the sharded
+        # SSM state already targets every covered D rank, so only the
+        # attention groups need widening; pure MLA writes to all handshaked
+        # ranks (only the dst differs per rank).
+        replicate_attn = self.use_mla and tp_ratio < 0
+        if replicate_attn and not self._has_mamba:
             assert len(plan.all_source_ranks) == 1
-            mla_local_ids = [list(ids) for ids in local_block_ids]
-            mla_remote_ids = [list(ids) for ids in remote_block_ids]
-            read_specs = [
-                ReadSpec(
-                    remote_rank=rank,
-                    local_block_ids=mla_local_ids,
-                    remote_block_ids=mla_remote_ids,
-                )
-                for rank in self.dst_xfer_side_handles[engine_id]
-            ]
+            write_ranks = sorted(self.dst_xfer_side_handles[engine_id])
         else:
-            read_specs = [
-                ReadSpec(
-                    remote_rank=rank,
-                    local_block_ids=[
-                        list(local_block_ids[g])
-                        if rank in plan.source_ranks_per_group[g]
-                        else []
-                        for g in range(num_groups)
-                    ],
-                    remote_block_ids=[
-                        list(remote_block_ids[g])
-                        if rank in plan.source_ranks_per_group[g]
-                        else []
-                        for g in range(num_groups)
-                    ],
-                )
-                for rank in plan.all_source_ranks
+            write_ranks = list(plan.all_source_ranks)
+
+        def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
+            return [
+                list(block_ids[g])
+                if (replicate_attn and _is_attention_spec(self._group_spec_types[g]))
+                or rank in plan.source_ranks_per_group[g]
+                else []
+                for g in range(num_groups)
             ]
 
+        read_specs = [
+            ReadSpec(
+                remote_rank=rank,
+                local_block_ids=group_ids(local_block_ids, rank),
+                remote_block_ids=group_ids(remote_block_ids, rank),
+            )
+            for rank in write_ranks
+        ]
+
         handles: list[int] = []
         for i, spec in enumerate(read_specs):
             remote_block_size = remote_info.remote_block_size
@@ -551,7 +550,10 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
                 remote_block_size,
                 req_id,
             )
-            if tp_ratio < 0 and not self.use_mla:
+            if tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1):
+                # Multiple targets: write each rank its chunk of local memory.
+                # Hybrid MLA+SSM also lands here: its split handles replicate
+                # the attention descriptors and chunk only the SSM state.
                 split_key = (tp_ratio, remote_block_size)
                 local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
             else: