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

[KV Connector][NIXL] Register packed MLA KV rows as a single region (#61074)

Browse Source 
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Nicolò Lucchesi
2026-10-15 09:18:15 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 Yifan Qiao Codex
parent 80dd705c1d
commit 7020600b1d
4 changed files with 206 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector.py
+77 -25
View File
@@ -68,6 +68,7 @@ from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.output_processor import OutputProcessor
from vllm.v1.kv_cache_interface import (
AttentionSpec,
CircularBufferSpec,
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
@@ -2523,8 +2524,46 @@ def test_register_kv_caches(
assert connector.connector_worker.block_size == block_size
def _dsv4_mla_spec(record_bytes: int, alignment: int, tokens_per_state: int):
return MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
cache_dtype_str="fp8_ds_mla",
tokens_per_state=tokens_per_state,
alignment=alignment,
model_version="deepseek_v4",
state_content_bytes=record_bytes,
)
@pytest.mark.parametrize(
"groups",
[
# V4 record: padded pages, so a layer view is not a contiguous page.
[(_dsv4_mla_spec(584, 576, 4), 4)],
# V4.1 MXFP8 record: unpadded pages in rows that the compressed caches
# and a compressor-state ring alias.
[
(_dsv4_mla_spec(528, 512, 2), 2),
(_dsv4_mla_spec(528, 512, 2), 2),
(
CircularBufferSpec(
block_size=8,
num_kv_heads=1,
head_size=8448,
head_size_v=0,
dtype=torch.uint8,
),
2,
),
],
],
ids=["v4-padded-pages", "v41-aliased-groups"],
)
def test_register_packed_dsv4_mla_cache_as_single_region(
default_vllm_config, dist_init
default_vllm_config, dist_init, groups
):
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
@@ -2537,27 +2576,21 @@ def test_register_packed_dsv4_mla_cache_as_single_region(
):
mock_backends.return_value = [TritonAttentionBackend]
num_blocks = 2
num_layers = 4
block_size = 256
spec = MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
cache_dtype_str="fp8_ds_mla",
tokens_per_state=4,
alignment=576,
model_version="deepseek_v4",
state_content_bytes=584,
)
layer_names = [f"layer{idx}" for idx in range(num_layers)]
specs = [spec for spec, _ in groups]
group_layers = [
[f"group{group_idx}.layer{idx}" for idx in range(num_layers)]
for group_idx, (_, num_layers) in enumerate(groups)
]
kv_cache_config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
kv_cache_groups=[
KVCacheGroupSpec(layer_names, spec)
for layer_names, spec in zip(group_layers, specs)
],
)
vllm_config = create_vllm_config(attention_backend="TRITON_ATTN")
vllm_config.cache_config.block_size = block_size
vllm_config.cache_config.block_size = 256
vllm_config.cache_config.kv_cache_layout = "BLHNC"
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
connector.connector_worker = FakeNixlConnectorWorker(
@@ -2571,25 +2604,44 @@ def test_register_packed_dsv4_mla_cache_as_single_region(
connector.connector_worker.nixl_wrapper = wrapper
wrapper.get_agent_metadata.return_value = b"fake_agent_metadata"
mock_thread.return_value.is_alive.return_value = False
connector.connector_worker.use_mla = True
# Every group packs its layers' pages into the same block rows.
(packed_block_len,) = {
len(layer_names) * spec.page_size_bytes
for layer_names, spec in zip(group_layers, specs)
}
raw = torch.zeros(
spec.page_size_bytes * num_blocks * num_layers,
packed_block_len * num_blocks,
dtype=torch.int8,
device=current_platform.device_type,
)
views = dense_kv_cache_views(
raw, spec, num_blocks, num_layers, KVCacheLayout.BLHNC
)
connector.register_kv_caches(dict(zip(layer_names, views)))
kv_caches: dict[str, torch.Tensor] = {}
for layer_names, spec in zip(group_layers, specs):
views = dense_kv_cache_views(
raw, spec, num_blocks, len(layer_names), KVCacheLayout.BLHNC
)
kv_caches.update(zip(layer_names, views))
connector.register_kv_caches(kv_caches)
# MLA pages transfer as whole rows; a ring page past the row start keeps
# its own region.
ring_pages = [
(cache.data_ptr(), spec.page_size_bytes)
for layer_names, spec in zip(group_layers, specs)
if isinstance(spec, CircularBufferSpec)
for cache in (kv_caches[name] for name in layer_names)
if cache.data_ptr() != raw.data_ptr()
]
_agent, block_runs = wrapper.prep_xfer_dlist.call_args[0]
packed_block_len = num_layers * spec.page_size_bytes
assert connector.connector_worker.kv_caches_base_addr[
connector.connector_worker.engine_id
][0] == [raw.data_ptr()]
# One packed region -> one run of num_blocks packed rows.
][0] == [raw.data_ptr()] + [addr for addr, _ in ring_pages]
assert block_runs.tolist() == [
[raw.data_ptr(), packed_block_len, 0, packed_block_len, num_blocks]
] + [
[addr, page_size, 0, packed_block_len, num_blocks]
for addr, page_size in ring_pages
]
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+43 -3
View File
@@ -25,6 +25,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
NixlAgentMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
NixlPullConnectorWorker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
NixlPushConnectorWorker,
)
@@ -127,6 +130,7 @@ def _make_packed_mla_view_worker(
num_blocks=4,
pp_size=1,
backend_names=("FLASHMLA",),
worker_cls=NixlPushConnectorWorker,
):
"""Register real strided views; only NIXL and distributed runtime are fake."""
block_size = 16
@@ -170,7 +174,7 @@ def _make_packed_mla_view_worker(
patch.object(bw, "get_current_attn_backends", return_value=backends),
patch("threading.Thread"),
):
worker = NixlPushConnectorWorker(
worker = worker_cls(
config,
"producer" if pp_size > 1 else "consumer",
KVCacheConfig(num_blocks, tensors, groups),
@@ -181,11 +185,18 @@ def _make_packed_mla_view_worker(
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"worker_cls", [NixlPullConnectorWorker, NixlPushConnectorWorker]
)
@pytest.mark.parametrize("num_layers", [1, 2])
def test_packed_mla_pp1_push_peer_transfers_whole_rows(num_layers):
def test_packed_mla_pp1_transfers_whole_rows(num_layers, worker_cls):
# BLHNC packs every layer's page of a block in one row: register the row
# once so a block is one descriptor, not one per layer.
layouts = {f"L{i}": (i * 128, 128) for i in range(num_layers)}
stride = num_layers * 128
worker, raw = _make_packed_mla_view_worker(layouts, stride)
worker, raw = _make_packed_mla_view_worker(layouts, stride, worker_cls=worker_cls)
assert worker.num_regions == 1
assert worker.num_descs == 4
assert worker._registered_descs == [[(raw.data_ptr(), raw.nbytes, 0, "")]]
# One packed region -> one run of 4 whole rows.
assert worker.src_blocks_data.tolist() == [[raw.data_ptr(), stride, 0, stride, 4]]
@@ -196,6 +207,35 @@ def test_packed_mla_pp1_push_peer_transfers_whole_rows(num_layers):
assert metadata.packed_member_layouts == layouts
@pytest.mark.cpu_test
def test_per_layer_peer_addresses_pages_inside_packed_rows():
# E.g. a HiSparse decoder reading a BLHNC prefiller's whole-row region.
remote, raw = _make_packed_mla_view_worker({"L0": (64, 128)}, 192)
metadata = msgspec.msgpack.decode(
remote.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
)
assert metadata.block_lens == [192]
local = SimpleNamespace(region_names=["L0"])
bw.NixlBaseConnectorWorker._split_remote_packed_rows(local, metadata)
assert metadata.kv_caches_base_addr == [raw.data_ptr() + 64]
assert metadata.block_lens == [128]
assert metadata.block_strides == [192]
assert metadata.region_names == ["L0"]
assert metadata.packed_member_layouts == {}
@pytest.mark.cpu_test
def test_packed_rows_reject_a_peer_packing_rows_differently():
producer, _ = _make_packed_mla_view_worker({"L0": (0, 128), "L1": (128, 64)}, 192)
consumer, _ = _make_packed_mla_view_worker({"L0": (64, 128), "L1": (0, 64)}, 192)
metadata = msgspec.msgpack.decode(
producer.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
)
metadata.engine_id = "remote"
with pytest.raises(RuntimeError, match="Packed KV block rows differ"):
consumer.add_remote_agent(metadata)
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"remote_backends, compatible",
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+1
View File
@@ -387,6 +387,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
w.region_group_ids = [0]
w._transfer_layer_names = ()
w._transfer_layer_region_indices = ()
w._packed_member_layouts = {}
w._transfer_layer_group_ids = ()
w.dst_region_num_blocks = {}
w.dst_region_group_ids = {}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+85 -13
View File
@@ -510,6 +510,55 @@ class NixlBaseConnectorWorker:
# One layer per region keeps a second alignment pass a no-op.
nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names]
def _split_remote_packed_rows(self, nixl_agent_meta: NixlAgentMetadata) -> None:
"""Map the remote's whole-row region onto our per-layer regions.
E.g. a BLHNC prefiller registers one region whose block is a full row
holding every layer's page, while a HiSparse decoder registers one
region per layer. Using the advertised `packed_member_layouts`, the
remote row is rewritten in place as one virtual region per local layer:
remote row b: [ L0 page | L1 page | ... ] base=B, stride=S
^off0 ^off1
remote metadata, before: [(B, S, S)] # (base, len, stride)
remote metadata, after: [(B + off0, page0, S), # L0
(B + off1, page1, S)] # L1
Remote regions outside the row (e.g. a separate indexer) pass through.
"""
layouts = nixl_agent_meta.packed_member_layouts
remote_names = nixl_agent_meta.region_names
assert remote_names is not None, "Remote advertised no region_names"
(row,) = [i for i, name in enumerate(remote_names) if name in layouts]
if nixl_agent_meta.region_group_ids is not None and (
nixl_agent_meta.region_group_ids[row] == _SHARED_REGION_GROUP_ID
):
raise NotImplementedError(
"Reading multi-group packed KV rows into per-layer regions is not "
"supported; use the same KV cache layout on P and D."
)
remote_index = {name: i for i, name in enumerate(remote_names)}
regions = [
(row, *layouts[name])
if name in layouts
else (remote_index[name], 0, nixl_agent_meta.block_lens[remote_index[name]])
for name in self.region_names
]
nixl_agent_meta.kv_caches_base_addr = [
nixl_agent_meta.kv_caches_base_addr[i] + offset for i, offset, _ in regions
]
nixl_agent_meta.block_lens = [page_size for _, _, page_size in regions]
nixl_agent_meta.block_strides = [
nixl_agent_meta.block_strides[i] for i, _, _ in regions
]
for attr in ("region_num_blocks", "region_group_ids", "region_mem_types"):
values = getattr(nixl_agent_meta, attr)
if values is not None:
setattr(nixl_agent_meta, attr, [values[i] for i, _, _ in regions])
nixl_agent_meta.region_names = list(self.region_names)
nixl_agent_meta.packed_member_layouts = {}
def __init__(
self,
vllm_config: "VllmConfig",
@@ -920,6 +969,7 @@ class NixlBaseConnectorWorker:
self._has_packed_cache = KVCacheLayout[self.kv_cache_layout].is_block_outermost
# Region layer membership is populated for HMA and packed layouts.
self.region_members = []
self._packed_member_layouts: dict[str, tuple[int, int]] = {}
# Local layer order, set by ``_set_region_layers`` only when
# routing by layer name. Empty means region-index routing.
self._transfer_layer_names = tuple[str, ...]()
@@ -1439,10 +1489,14 @@ class NixlBaseConnectorWorker:
self._ple_block_len = None
self._ple_region_index = None
packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values()))
# CSA-linear needs separate logical regions for attention and state
# aliases even though every view shares one block-major allocation.
packed_storage = packed_storage and not self._is_csa_linear
# A shared packed row may transfer whole only if every page in it is
# replicated across TP ranks: MLA without SSM state. Head-sharded
# attention or SSM state (e.g. Qwen3.8-Flash-Next) must stay split.
packed_storage = (
_share_storage_and_block_stride(list(xfer_buffers.values()))
and self.use_mla
and not self._has_mamba
)
if route_packed_layers and any(
not isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
for spec in self._layer_specs.values()
@@ -1566,22 +1620,23 @@ class NixlBaseConnectorWorker:
and cache[0].is_contiguous()
and cache[0].nbytes == physical_page_size
)
# Push workers address the pages of packed MLA rows by layer.
packed_mla_push = (
track_region_layers
and self._has_packed_cache
# Block-outermost MLA: one block row holds every layer's page.
packed_mla = (
self._has_packed_cache
and not self._has_mamba
and packed_storage
and storage_is_block_major
and is_mla_region
)
if packed_mla_push and use_layer_name_routing:
if packed_mla and use_layer_name_routing:
# PP>1 prefill stage: register only this layer's page of each row.
region_specs = [
(cache.data_ptr(), physical_page_size, block_stride)
]
elif packed_mla_push:
# PP=1 decode or prefill: register whole rows, and advertise where
# this layer's page sits so a PP>1 prefill can write into it.
elif packed_mla and (self.pp_size == 1 or not page_contiguous):
# Register whole rows so a block is a single descriptor. Also,
# advertise where this layer's page sits so a remote with
# LB* can still address it (HiSparse).
storage_block_len = block_stride
region_specs = [
(registration_base, storage_block_len, storage_block_len)
@@ -1716,6 +1771,7 @@ class NixlBaseConnectorWorker:
self.region_num_blocks = [self.num_blocks] * self.num_regions
self._set_region_layers(region_layers)
self._packed_member_layouts = packed_member_layouts
if self.pp_size > 1 and not use_layer_name_routing:
start_layer, end_layer = self.model_config.get_layers_start_end_indices(
@@ -1811,7 +1867,7 @@ class NixlBaseConnectorWorker:
dcp_size=self.dcp_size,
pcp_size=self.pcp_size,
region_members=self.region_members,
packed_member_layouts=packed_member_layouts,
packed_member_layouts=self._packed_member_layouts,
)
# Wrap metadata in payload with hash for defensive decoding
assert self.compat_hash is not None
@@ -2264,6 +2320,11 @@ class NixlBaseConnectorWorker:
"prefill TP yet"
)
self._align_remote_regions_by_layer(nixl_agent_meta)
elif nixl_agent_meta.packed_member_layouts and not self._packed_member_layouts:
# With HiSparse, P may be on BL* while D is LB*, as only indexer is on HBM.
# TODO (NickLucche) we may consider not allowing this case at all (P BL*)
# and just clear this function.
self._split_remote_packed_rows(nixl_agent_meta)
elif (
self.pp_size > 1
and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions
@@ -2294,6 +2355,17 @@ class NixlBaseConnectorWorker:
start:end
]
if (
self._packed_member_layouts
and not self._transfer_layer_region_indices
and self._packed_member_layouts != nixl_agent_meta.packed_member_layouts
):
# Whole rows transfer as one unit, so both sides must pack them alike.
raise RuntimeError(
"Packed KV block rows differ between P and D. Use the same "
"VLLM_KV_CACHE_LAYOUT and block size on both sides."
)
### Register remote engine in TransferTopology (idempotent).
physical_blocks_per_logical = (
nixl_agent_meta.physical_blocks_per_logical_kv_block