@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f37f03db4a

[KV Connector] Support NIXL P/D for hybrid MLA+SSM models (#49762)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Jared Wen <w13431838023@gmail.com>

Nick Hill committed Jul 28, 2026 at 20:50 UTC f37f03db4af69e4969242767734db3d9175055f7
5 files changed +281 -38
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+196
@@ -1380,3 +1380,199 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
1380 assert list(result) == expected_kernel_block_ids, (
1381 f"Expected {expected_kernel_block_ids}, got {result}"
1382 )
1383 +
1384 +
1385 +# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
1386 +
1387 +
1388 +def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
1389 + """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed
1390 + MambaSpec) groups whose layers share the same HMA tensors, with a
1391 + mamba-aligned unified page and an MLA kernel block smaller than the
1392 + logical block."""
1393 + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
1394 + from vllm.v1.kv_cache_interface import (
1395 + KVCacheConfig,
1396 + KVCacheGroupSpec,
1397 + KVCacheTensor,
1398 + MambaSpec,
1399 + MLAAttentionSpec,
1400 + )
1401 +
1402 + # 12-token logical blocks over a 4-token MLA kernel block.
1403 + mla_spec = MLAAttentionSpec(
1404 + block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16
1405 + )
1406 + unified_page = mla_spec.page_size_bytes
1407 + kda_spec = MambaSpec(
1408 + block_size=12,
1409 + # GDN-decomposable conv (Q|K|V = 2|2|4 cols x 3 rows) + fp32 temporal.
1410 + shapes=((8, 3), (1, 4, 4)),
1411 + dtypes=(torch.float16, torch.float32),
1412 + page_size_padded=unified_page,
1413 + mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
1414 + )
1415 + assert kda_spec.page_size_bytes == unified_page
1416 + return KVCacheConfig(
1417 + num_blocks=num_blocks,
1418 + kv_cache_tensors=[
1419 + KVCacheTensor(
1420 + size=num_blocks * unified_page,
1421 + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
1422 + )
1423 + for i in range(2)
1424 + ],
1425 + kv_cache_groups=[
1426 + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
1427 + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
1428 + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
1429 + ],
1430 + )
1431 +
1432 +
1433 +@pytest.mark.cpu_test
1434 +def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
1435 + """Hybrid MLA+KDA registration: HMA tensors shared by both layer types
1436 + must be flagged as MLA regions even when a KDA layer registers them
1437 + first, expose TP-independent kernel-granularity block lens, and build
1438 + FA + mamba descriptors for every region."""
1439 + from unittest.mock import MagicMock
1440 +
1441 + from vllm.config import set_current_vllm_config
1442 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
1443 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
1444 + NixlConnectorWorker,
1445 + )
1446 +
1447 + kv_cache_config = _make_hybrid_mla_kv_cache_config()
1448 + unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
1449 + vllm_config = create_vllm_config(block_size=12)
1450 + # kv_buffer_device defaults to the *real* platform's device type, which on
1451 + # a CPU-only test host would make this a host-buffer worker: host xfer
1452 + # buffers are per-layer, so the HMA shared tensors would not be
1453 + # deduplicated. Pin it to the faked device type.
1454 + vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
1455 +
1456 + fake_backend = MagicMock()
1457 + fake_backend.get_supported_kernel_block_sizes.return_value = [4]
1458 + fake_backend.get_name.return_value = "FLASHMLA"
1459 + fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
1460 + fake_platform = MagicMock()
1461 + fake_platform.device_type = "cuda"
1462 + fake_platform.get_nixl_memory_type.return_value = "VRAM"
1463 +
1464 + with (
1465 + patch.object(bw, "NixlWrapper"),
1466 + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
1467 + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
1468 + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
1469 + patch.object(bw, "current_platform", fake_platform),
1470 + patch(
1471 + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
1472 + return_value="DS",
1473 + ),
1474 + set_current_vllm_config(vllm_config),
1475 + ):
1476 + worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
1477 + worker.use_mla = True # opt-125m test config is not MLA; force the flag
1478 + worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
1479 +
1480 + tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
1481 + # KDA layer first per tensor: exercises the dual-purpose flag merge.
1482 + worker.register_kv_caches(
1483 + {
1484 + "kda_a.0": tensors[0],
1485 + "mla.0": tensors[0],
1486 + "kda_b.0": tensors[0],
1487 + "kda_a.1": tensors[1],
1488 + "mla.1": tensors[1],
1489 + "kda_b.1": tensors[1],
1490 + }
1491 + )
1492 +
1493 + # 12-token logical blocks over the 4-token MLA kernel block.
1494 + assert worker._physical_blocks_per_logical_kv_block == 3
1495 + assert worker.block_size == 4 and worker.num_blocks == 12
1496 + # Both shared tensors are dual-purpose: their FA view is MLA even though
1497 + # a KDA layer registered them first.
1498 + assert worker._region_is_mla == [True, True]
1499 + assert worker.num_regions == 2 and worker.num_descs == 24
1500 + # Kernel-granularity block lens; TP-independent for MLA hybrids.
1501 + assert worker.block_len_per_layer == [unified_page // 3] * 2
1502 + # Split handles must replicate every FA descriptor (MLA isn't head-sharded).
1503 + assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24
1504 + # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page.
1505 + # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
1506 + assert worker.src_blocks_data.shape == (24 + 32, 3)
1507 + fa_descs = worker.src_blocks_data[:24]
1508 + assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
1509 + assert all(size == unified_page // 3 for size in fa_descs[:, 1])
1510 +
1511 +
1512 +@pytest.mark.cpu_test
1513 +def test_push_write_hybrid_mla_replicates_attention():
1514 + """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be
1515 + written to every covered D rank (replicated MLA latent) while SSM state
1516 + is written per-rank through the split handles."""
1517 + import threading
1518 + from collections import defaultdict
1519 + from unittest.mock import MagicMock
1520 +
1521 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
1522 + NixlPushConnectorWorker,
1523 + )
1524 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
1525 + TPMapping,
1526 + )
1527 + from vllm.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec
1528 +
1529 + worker = object.__new__(NixlPushConnectorWorker)
1530 + worker.shutdown = lambda: None # skeleton worker: silence __del__
1531 + worker.use_mla = True
1532 + worker._has_mamba = True
1533 + worker._group_spec_types = (MLAAttentionSpec, MambaSpec)
1534 + worker.transfer_topo = MagicMock()
1535 + worker.transfer_topo.tp_ratio.return_value = -2
1536 + remote_info = MagicMock()
1537 + remote_info.remote_physical_blocks_per_logical = 1
1538 + remote_info.remote_block_size = 4
1539 + worker.transfer_topo.get_engine_info.return_value = remote_info
1540 +
1541 + engine_id = "remote-engine"
1542 + # Read-oriented mapping collapses the replicated attention group to one
1543 + # source rank; the SSM state is sharded across both covered D ranks.
1544 + worker.tp_mappings = {
1545 + engine_id: TPMapping(
1546 + source_ranks_per_group=((0,), (0, 1)),
1547 + all_source_ranks=(0, 1),
1548 + rank_to_attention_slot={0: 0, 1: 0},
1549 + rank_offset_factor=0,
1550 + )
1551 + }
1552 + worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
1553 + worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
1554 + worker.src_xfer_handles_by_block_size = {4: 300}
1555 + worker._sending_transfers = defaultdict(list)
1556 + worker._sending_transfers_lock = threading.Lock()
1557 + worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
1558 + worker._xfer_blocks = MagicMock(return_value=1)
1559 +
1560 + meta = MagicMock()
1561 + meta.remote.engine_id = engine_id
1562 + meta.remote.block_ids = [[7, 8], [3]]
1563 + meta.local_physical_block_ids = [[1, 2], [5]]
1564 +
1565 + worker._xfer_blocks_for_req("req-1", meta)
1566 +
1567 + calls = worker._xfer_blocks.call_args_list
1568 + assert len(calls) == 2
1569 + for call, rank, local_handle, remote_handle in zip(
1570 + calls, (0, 1), (200, 201), (100, 101)
1571 + ):
1572 + spec = call.kwargs["read_spec"]
1573 + assert spec.remote_rank == rank
1574 + # Attention group replicated to every rank, SSM by membership.
1575 + assert spec.local_block_ids == [[1, 2], [5]]
1576 + assert spec.remote_block_ids == [[7, 8], [3]]
1577 + assert call.kwargs["local_xfer_side_handle"] == local_handle
1578 + assert call.kwargs["remote_xfer_side_handle"] == remote_handle
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+21
@@ -617,3 +617,24 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens):
617 num_tokens=num_tokens,
618 tp_size=8,
619 )
620 +
621 +
622 +@pytest.mark.cpu_test
623 +def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid():
624 + """The MLA per-token page is TP-independent, so kernel block lengths
625 + differing by anything other than the block-size ratio must fail the
626 + handshake loudly rather than transfer at mismatched geometry."""
627 + worker = _make_mla_hybrid_worker(
628 + local_block_size=12, kernel_block_size=4, num_logical_blocks=8
629 + )
630 + meta_r = _make_remote_meta(
631 + worker,
632 + remote_block_size=8,
633 + remote_kernel_block_size=4,
634 + remote_num_logical=12,
635 + remote_ssm_sizes=(24, 32),
636 + )
637 + # Equal kernel block sizes (ratio 1), but a half-sized per-token page.
638 + meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer]
639 + with pytest.raises((AssertionError, RuntimeError)):
640 + worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2)
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+4
@@ -43,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
43 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
44 get_base_request_id,
45 )
46 +from vllm.v1.kv_cache_interface import FullAttentionSpec
47 from vllm.v1.outputs import KVConnectorOutput
48
49 from .utils import make_nixl_push_scheduler
@@ -337,6 +338,9 @@ class _StubWriterWorker(NixlPushConnectorWorker):
338 w.engine_id = "test-decode-engine"
339 w._remote_agents = {}
340 w._physical_blocks_per_logical_kv_block = 1
341 + # Single non-hybrid attention group, matching the stub block id lists.
342 + w._has_mamba = False
343 + w._group_spec_types = (FullAttentionSpec,)
344
345 # Track _do_start_push_kv invocations.
346 calls: list[tuple[str, Any, dict[str, Any]]] = []
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+25 -5
@@ -1093,7 +1093,7 @@ class NixlBaseConnectorWorker:
1093
1094 caches_data = []
1095 # With hybrid allocator, layers can share a kv cache tensor
1096 - seen_base_addresses = []
1096 + seen_base_addresses: list[int] = []
1097
1098 # K and V are packed into the content dim, so each attention layer is a
1099 # single NIXL region whose block transfers as one unit. Mamba layers instead
@@ -1139,11 +1139,19 @@ class NixlBaseConnectorWorker:
1139 curr_tensor_size_bytes = num_blocks * physical_page_size
1140
1141 base_addr = cache.data_ptr()
1142 + is_mla_region = isinstance(
1143 + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
1144 + )
1145 if base_addr in seen_base_addresses:
1146 # NOTE (NickLucche) HMA employs memory pooling to share tensors
1147 # across groups. This results in skipping all tensors but the ones
1148 # pointed to by group0. Also, generally we will have more blocks
1149 # per tensor but fewer regions.
1150 + # A shared tensor may back both SSM and attention layers (e.g.
1151 + # KDA+MLA in KimiLinear); the region's FA view is MLA whichever
1152 + # layer registered it first.
1153 + idx = seen_base_addresses.index(base_addr)
1154 + self._region_is_mla[idx] |= is_mla_region
1155 logger.debug("Skipping %s because it's already seen", layer_name)
1156 continue
1157 logger.debug(
@@ -1157,9 +1165,6 @@ class NixlBaseConnectorWorker:
1165 )
1166 else:
1167 self.block_len_per_layer.append(physical_page_size)
1160 - is_mla_region = isinstance(
1161 - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
1162 - )
1168 self._region_is_mla.append(is_mla_region)
1169
1170 if not is_mla_region:
@@ -1801,7 +1806,22 @@ class NixlBaseConnectorWorker:
1806 # the per-rank KV head ratio rather than the raw tp_ratio, because GQA
1807 # replication caps per-rank heads at 1 when tp > total_kv_heads
1808 # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.
1804 - if not self._has_mamba:
1809 + if self._has_mamba and self.use_mla:
1810 + # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are
1811 + # kernel-granularity views of the mamba-unified page. The MLA
1812 + # per-token page is TP-independent, so the block lengths must
1813 + # match up to the kernel block size ratio even under
1814 + # heterogeneous TP (remote kernel blocks may be smaller).
1815 + # SSM geometry is validated via ssm_sizes/conv offsets instead.
1816 + assert self.block_len_per_layer == [
1817 + block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens
1818 + ], (
1819 + "Hybrid MLA kernel-granularity block lengths must match "
1820 + f"between P and D (block_size_ratio={block_size_ratio}): "
1821 + f"local={self.block_len_per_layer}, "
1822 + f"remote={nixl_agent_meta.block_lens}."
1823 + )
1824 + elif not self._has_mamba:
1825 assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
1826 "Number of KV layers must match between prefill and decode"
1827 )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+35 -33
@@ -52,7 +52,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
52 ReqMeta,
53 TransferHandle,
54 )
55 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec
55 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
56 + ReadSpec,
57 + _is_attention_spec,
58 +)
59 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id
60 from vllm.logger import init_logger
61
@@ -505,41 +508,37 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
508 local_block_ids = meta.local_physical_block_ids
509 num_groups = len(local_block_ids)
510
508 - if self.use_mla and tp_ratio < 0:
509 - # MLA latent is replicated across D's TP ranks: the tp-mapping
510 - # collapses to one rank (fine for reads), but push must WRITE every
511 - # D rank or the rest decode stale KV; only the dst differs per rank.
511 + # MLA latent is replicated across D's TP ranks: the tp-mapping
512 + # collapses it to one rank (fine for reads), but push must WRITE every
513 + # D rank or the rest decode stale KV. For hybrid MLA+SSM the sharded
514 + # SSM state already targets every covered D rank, so only the
515 + # attention groups need widening; pure MLA writes to all handshaked
516 + # ranks (only the dst differs per rank).
517 + replicate_attn = self.use_mla and tp_ratio < 0
518 + if replicate_attn and not self._has_mamba:
519 assert len(plan.all_source_ranks) == 1
513 - mla_local_ids = [list(ids) for ids in local_block_ids]
514 - mla_remote_ids = [list(ids) for ids in remote_block_ids]
515 - read_specs = [
516 - ReadSpec(
517 - remote_rank=rank,
518 - local_block_ids=mla_local_ids,
519 - remote_block_ids=mla_remote_ids,
520 - )
521 - for rank in self.dst_xfer_side_handles[engine_id]
522 - ]
520 + write_ranks = sorted(self.dst_xfer_side_handles[engine_id])
521 else:
524 - read_specs = [
525 - ReadSpec(
526 - remote_rank=rank,
527 - local_block_ids=[
528 - list(local_block_ids[g])
529 - if rank in plan.source_ranks_per_group[g]
530 - else []
531 - for g in range(num_groups)
532 - ],
533 - remote_block_ids=[
534 - list(remote_block_ids[g])
535 - if rank in plan.source_ranks_per_group[g]
536 - else []
537 - for g in range(num_groups)
538 - ],
539 - )
540 - for rank in plan.all_source_ranks
522 + write_ranks = list(plan.all_source_ranks)
523 +
524 + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
525 + return [
526 + list(block_ids[g])
527 + if (replicate_attn and _is_attention_spec(self._group_spec_types[g]))
528 + or rank in plan.source_ranks_per_group[g]
529 + else []
530 + for g in range(num_groups)
531 ]
532
533 + read_specs = [
534 + ReadSpec(
535 + remote_rank=rank,
536 + local_block_ids=group_ids(local_block_ids, rank),
537 + remote_block_ids=group_ids(remote_block_ids, rank),
538 + )
539 + for rank in write_ranks
540 + ]
541 +
542 handles: list[int] = []
543 for i, spec in enumerate(read_specs):
544 remote_block_size = remote_info.remote_block_size
@@ -551,7 +550,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
550 remote_block_size,
551 req_id,
552 )
554 - if tp_ratio < 0 and not self.use_mla:
553 + if tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1):
554 + # Multiple targets: write each rank its chunk of local memory.
555 + # Hybrid MLA+SSM also lands here: its split handles replicate
556 + # the attention descriptors and chunk only the SSM state.
557 split_key = (tp_ratio, remote_block_size)
558 local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
559 else: