iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Jared Wen <w13431838023@gmail.com>
This commit is contained in:
Nick Hill
2026-07-29 11:50:30 +08:00
committed by GitHub
co-authored by Jared Wen
parent 32a423ac0a
commit f37f03db4a
5 changed files with 281 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+196
View File
@@ -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
+21
View File
@@ -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
+4
View File
@@ -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]]] = []
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+25 -5
View File
@@ -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"
)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+35 -33
View File
@@ -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: