diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 1f7a62d2c9..0124f6516e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index e2fc41a022..6ac9ce030d 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/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) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index d670e5563b..adf0063701 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -43,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( 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 @@ class _StubWriterWorker(NixlPushConnectorWorker): 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]]] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 480ac8937d..12b9f0d937 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1093,7 +1093,7 @@ class NixlBaseConnectorWorker: 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 @@ class NixlBaseConnectorWorker: 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 @@ class NixlBaseConnectorWorker: ) 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 @@ class NixlBaseConnectorWorker: # 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" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 6bc48963c5..88977e3403 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -52,7 +52,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( 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 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): 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 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): 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: