[KVConnector][NIXL] Support packed MLA KV layouts in pipeline-parallel push prefill (#50499)

Signed-off-by: zixi-qi <zixi@inferact.ai>
This commit is contained in:
qizixi
2026-09-28 18:50:46 +02:00
committed by GitHub
parent 94d1462924
commit 6d3ea3c2c9
6 changed files with 357 additions and 28 deletions
+29 -11
View File
@@ -195,35 +195,53 @@ combined with a hybrid (HMA) KV layout:
that represents a pooled region can differ between P and D.
`NixlPushConnector` handles this for PP-sharded producers by routing
**by layer-name (member) identity** instead of by region index. Each
**by layer name** instead of by region index. Each
worker advertises which layer names back each of its NIXL regions in the
handshake metadata (`NixlAgentMetadata.region_members`). A producer that
needs member routing derives its member-major layout once, when it
needs layer routing derives its transfer layer order once, when it
registers its KV caches, and every transfer it issues uses that order.
`add_remote_agent` then selects exactly the remote regions this stage
owns and reorders them to match, so both sides stay paired regardless of
how each remote rank happens to order its metadata.
Addresses, block lengths, strides, and per-region capacities follow the same
member order. Descriptor offsets use each member's region capacity, so P and
layer order. Descriptor offsets use each layer's region capacity, so P and
D need not allocate the same number of blocks. Physical allocations are still
registered once, even when multiple members share them.
registered once, even when multiple layers share them.
Packed MLA caches (block-outermost layouts such as DeepSeek-V4's) store the
pages of every layer for one block in a single row. A PP stage holds only its
own layers, so its rows are shorter and each layer sits at a different offset:
```text
D (PP=1), block b: | L0 | L1 | L2 | L3 | L4 | L5 |
P (stage 2), block b: | L3 | L4 | L5 |
```
`PP=1` peers transfer whole rows. They also advertise where each layer's page
sits in the row (`NixlAgentMetadata.packed_member_layouts`), so a PP producer
writes each of its pages to the matching slot in D's row. Page sizes must
match; offsets, row strides, and block counts may differ. Because PP stages can
list the same attention backends in a different order, packed push compares the
sorted set of backend names during the handshake. Pull mode is unchanged.
Invariants enforced when the remote regions are aligned:
* every locally owned member must be advertised exactly once by the
remote; a missing member fails the handshake rather than silently
leaving that layer's KV stale, and remote-only members (owned by other
* every locally owned layer must be advertised exactly once by the
remote; a missing layer fails the handshake rather than silently
leaving that layer's KV stale, and remote-only layers (owned by other
PP stages) are ignored;
* a remote that omits member metadata while the local layout requires
member routing fails the handshake instead of falling back to
* a remote that omits layer metadata while the local layout requires
layer routing fails the handshake instead of falling back to
region-index routing;
* member order is a property of the local layout alone, so the same local
* layer order is a property of the local layout alone, so the same local
source descriptors serve every remote engine and TP rank; only the
remote descriptor list is rebuilt per rank.
Decode-side PP is unsupported because completions are counted per
consumer rank. Mamba/SSM hybrids are unsupported under PP.
consumer rank. Mamba/SSM hybrids are unsupported under PP. Packed PP push
requires MLA caches and equal P/D block sizes. MLA pages are replicated
across TP ranks, so producer TP1 to consumer TP2 is supported.
## Scheduler-side responsibilities
@@ -120,7 +120,8 @@ Current push PP + HMA limitations:
- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported.
- Hybrid SSM/Mamba layouts are not supported under PP.
- HMA requires the same block size on P and D.
- Attention-HMA member routing requires decode TP to be no greater than prefill TP.
- Non-MLA attention-HMA layer routing requires decode TP to be no greater than prefill TP.
- Packed layouts under PP require MLA caches. Local and remote packed block strides may differ; layer pages must have equal sizes. MLA supports decode TP greater than prefill TP because its KV is replicated.
### Quantized KV cache
@@ -16,10 +16,27 @@ from threading import Event, Lock
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import msgspec
import numpy as np
import pytest
import torch
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.push_worker import (
NixlPushConnectorWorker,
)
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
KVCacheLayout,
KVCacheTensor,
MLAAttentionSpec,
create_kv_cache_views,
)
from .utils import create_vllm_config
@@ -102,6 +119,189 @@ class _RecordingNixl:
pass
def _make_packed_mla_view_worker(
layouts,
block_stride,
*,
num_blocks=4,
pp_size=1,
backend_names=("FLASHMLA",),
):
"""Register real strided views; only NIXL and distributed runtime are fake."""
block_size = 16
raw = torch.zeros(num_blocks * block_stride, dtype=torch.int8)
tensors, groups, caches = [], [], {}
for name, (offset, page_size) in layouts.items():
spec = MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=page_size // block_size,
dtype=torch.uint8,
)
tensor = KVCacheTensor(
size=raw.nbytes,
layers=[name],
layer_stride=page_size,
block_stride=block_stride,
offset=offset,
)
tensors.append(tensor)
groups.append(KVCacheGroupSpec([name], spec))
(caches[name],) = create_kv_cache_views(
raw, spec, num_blocks, KVCacheLayout.BLHNC, tensor
)
config = create_vllm_config(block_size=block_size)
config.parallel_config.pipeline_parallel_size = pp_size
config.kv_transfer_config.kv_role = "kv_producer"
config.cache_config.kv_cache_layout = "BLHNC"
backends = []
for name in backend_names:
backend = MagicMock()
backend.get_supported_kernel_block_sizes.return_value = [block_size]
backend.get_name.return_value = name
backend.full_cls_name.return_value = f"fake.{name}"
backends.append(backend)
with (
patch.object(bw, "NixlWrapper", _RecordingNixl),
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=backends),
patch("threading.Thread"),
):
worker = NixlPushConnectorWorker(
config,
"producer" if pp_size > 1 else "consumer",
KVCacheConfig(num_blocks, tensors, groups),
)
worker.use_mla = True
worker.register_kv_caches(caches)
return worker, raw
@pytest.mark.cpu_test
@pytest.mark.parametrize("num_layers", [1, 2])
def test_packed_mla_pp1_push_peer_transfers_whole_rows(num_layers):
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)
assert worker._registered_descs == [[(raw.data_ptr(), raw.nbytes, 0, "")]]
assert worker.src_blocks_data.tolist() == [
[raw.data_ptr() + block * stride, stride, 0] for block in range(4)
]
assert worker._transfer_layer_group_ids == ()
metadata = msgspec.msgpack.decode(
worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
)
assert metadata.packed_member_layouts == layouts
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"remote_backends, compatible",
[
(("INDEXER", "FLASHMLA"), True),
(("FLASHMLA", "OTHER"), False),
(("FLASHMLA",), False),
],
)
def test_packed_push_compatibility_hash_uses_all_backends_in_stable_order(
remote_backends, compatible
):
producer, _ = _make_packed_mla_view_worker(
{"L0": (0, 128)},
128,
pp_size=2,
backend_names=("FLASHMLA", "INDEXER"),
)
consumer, _ = _make_packed_mla_view_worker(
{"L0": (0, 128)}, 128, backend_names=remote_backends
)
assert (producer.compat_hash == consumer.compat_hash) is compatible
@pytest.mark.cpu_test
@pytest.mark.parametrize("p_num_blocks, d_num_blocks", [(4, 4), (4, 6), (6, 4)])
def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(
p_num_blocks, d_num_blocks
):
# Different cache groups overlay pages of different sizes at the same address.
# PP also changes both the placement and the stride of the matching D pages.
producer, p_raw = _make_packed_mla_view_worker(
{"L2": (0, 128), "L2.swa": (0, 64), "L3": (128, 64)},
192,
num_blocks=p_num_blocks,
pp_size=2,
)
consumer, d_raw = _make_packed_mla_view_worker(
{"L0": (0, 128), "L2": (128, 128), "L2.swa": (0, 64), "L3": (64, 64)},
256,
num_blocks=d_num_blocks,
)
assert producer.block_len_per_layer == [128, 64, 64]
assert producer._transfer_layer_group_ids == (0, 1, 2)
assert producer.num_descs == 3 * p_num_blocks
local_ids = producer._compute_desc_ids(
[[1], [2], [3]],
producer.num_blocks,
None,
1,
region_num_blocks=producer.dst_region_num_blocks[producer.engine_id],
)
assert producer.src_blocks_data[local_ids].tolist() == [
[p_raw.data_ptr() + 192, 128, 0],
[p_raw.data_ptr() + 2 * 192, 64, 0],
[p_raw.data_ptr() + 128 + 3 * 192, 64, 0],
]
# Each decoder TP rank receives the complete MLA page, without head slicing.
for rank in (0, 1):
metadata = msgspec.msgpack.decode(
consumer.xfer_handshake_metadata.agent_metadata_bytes,
type=NixlAgentMetadata,
)
producer.add_remote_agent(metadata, remote_tp_rank=rank, remote_tp_size=2)
assert metadata.region_num_blocks == [d_num_blocks] * 3
assert metadata.region_names == ["L2", "L2.swa", "L3"]
remote_ids = producer._compute_desc_ids(
[[3], [1], [2]],
metadata.num_blocks,
None,
1,
region_num_blocks=producer.dst_region_num_blocks[consumer.engine_id],
)
handle = producer.dst_xfer_side_handles[consumer.engine_id][rank]
remote_descs = producer.nixl_wrapper.dlists[handle]
assert len(remote_descs) == 3 * d_num_blocks
assert remote_descs[remote_ids].tolist() == [
[d_raw.data_ptr() + 128 + 3 * 256, 128, 0],
[d_raw.data_ptr() + 256, 64, 0],
[d_raw.data_ptr() + 64 + 2 * 256, 64, 0],
]
aligned = msgspec.msgpack.encode(metadata)
producer._align_remote_regions_by_layer(metadata)
assert msgspec.msgpack.encode(metadata) == aligned
@pytest.mark.cpu_test
@pytest.mark.parametrize("remote_block_size", [8, 32])
def test_packed_mla_rejects_unequal_block_sizes_before_peer_registration(
remote_block_size,
):
producer, _ = _make_packed_mla_view_worker({"L0": (0, 128)}, 128, pp_size=2)
metadata = msgspec.msgpack.decode(
producer.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
)
metadata.engine_id = "remote"
metadata.block_size = remote_block_size
with pytest.raises(NotImplementedError, match="identical P/D block sizes"):
producer.add_remote_agent(metadata)
assert len(producer.nixl_wrapper.dlists) == 1 # Only the local list exists.
assert producer.tp_mappings == {}
assert "remote" not in producer.dst_num_blocks
with pytest.raises(KeyError):
producer.transfer_topo.get_engine_info("remote")
@pytest.mark.cpu_test
def test_local_descriptors_follow_each_region_pool_capacity():
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
@@ -200,8 +400,8 @@ def _register_overlaid_mla_worker(
worker.src_xfer_handles_by_block_size = {}
worker.kv_caches_base_addr = defaultdict(dict)
worker._mamba_ssm_size = (0, 0)
worker.kv_cache_layout = "NHD"
worker.host_buffer_kv_cache_layout = "NHD"
worker.kv_cache_layout = "LBNHC"
worker.host_buffer_kv_cache_layout = "LBNHC"
worker._physical_blocks_per_logical_kv_block = 1
worker._logical_num_blocks = num_blocks
worker.region_mem_types = []
@@ -221,6 +421,7 @@ def _register_overlaid_mla_worker(
worker.device_kv_caches = {}
worker.pp_size = 2 if push_pp else 1
worker._is_hma_required = True
worker._has_packed_cache = False
worker.dcp_size = 1
worker.pcp_size = 1
worker.kv_buffer_device = "cuda"
@@ -368,6 +368,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
w._recving_transfers = defaultdict(list)
w._is_hma_required = False
w.xfer_stats = NixlKVConnectorStats()
w._has_packed_cache = False
w._reqs_to_process = set()
w._reqs_to_send = {}
w.consumer_notification_counts_by_req = defaultdict(int)
@@ -1843,6 +1844,40 @@ def test_layer_metadata_round_trip():
assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata
@pytest.mark.parametrize(
"layouts, error",
[
({"L0": (-1, 64)}, "escapes its block"),
({"L0": (0, 0)}, "escapes its block"),
({"L0": (200, 64)}, "escapes its block"),
({"L1": (0, 64)}, "has no layout"),
({"L0": (0, 32)}, "page sizes must match"),
],
)
def test_packed_layer_alignment_rejects_invalid_page_layout(layouts, error):
worker = _layer_routing_worker([["L0"]], {"L0": 0})
worker.block_len_per_layer = [64]
metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256])
metadata.packed_member_layouts = layouts
with pytest.raises(AssertionError, match=error):
worker._align_remote_regions_by_layer(metadata)
def test_layer_alignment_preserves_nonpacked_regions_alongside_packed_layers():
worker = _layer_routing_worker([["a"], ["b"]], {"a": 0, "b": 1})
worker.block_len_per_layer = [64, 128]
metadata = _agent_metadata(
[["a", "other"], ["b"]], [0x1000, 0x2000], [192, 128], [256, 128]
)
metadata.packed_member_layouts = {"a": (64, 64), "other": (128, 64)}
worker._align_remote_regions_by_layer(metadata)
assert metadata.kv_caches_base_addr == [0x1040, 0x2000]
assert metadata.block_lens == [64, 128]
assert metadata.block_strides == [256, 128]
def test_layer_identity_gate_preserves_the_non_hma_path():
assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == (
0,
@@ -1850,6 +1885,7 @@ def test_layer_identity_gate_preserves_the_non_hma_path():
# A base worker (pull) never routes by layer name.
pull = object.__new__(NixlBaseConnectorWorker)
pull.pp_size = 2
pull._transfer_layer_names = ()
pull._transfer_layer_region_indices = ()
pull._transfer_layer_group_ids = ()
@@ -165,8 +165,8 @@ class NixlBaseConnectorWorker:
# Overridden by NixlPushConnectorWorker.
_TRANSFER_MODE: str = "pull"
# Layer-name routing is supported only by NixlPushConnector for HMA KV
# caches under PP.
# Layer-name routing is supported only by NixlPushConnector for HMA or
# packed KV caches under PP.
_supports_pp_hma = False
def _compute_desc_ids(
@@ -432,15 +432,22 @@ class NixlBaseConnectorWorker:
group_by_layer[name] for name in layer_names
)
def _requires_layer_name_routing(self) -> bool:
"""Whether PP push must match HMA layers by name, not region index."""
def _tracks_region_layers(self) -> bool:
"""Whether regions advertise their member layer names.
Push workers do so for HMA and packed layouts, so a PP producer can
match layers by name instead of by region index.
"""
return (
self._supports_pp_hma
and self.pp_size > 1
and self._is_hma_required
and (self._is_hma_required or self._has_packed_cache)
and not self._has_mamba
)
def _requires_layer_name_routing(self) -> bool:
"""Whether PP push must match HMA/packed layers by name, not region index."""
return self.pp_size > 1 and self._tracks_region_layers()
def _align_remote_regions_by_layer(
self, nixl_agent_meta: NixlAgentMetadata
) -> None:
@@ -479,6 +486,10 @@ class NixlBaseConnectorWorker:
]
assert not missing, f"Remote is missing locally owned layers: {missing}"
packed_regions = {
remote_region_by_layer[name]
for name in nixl_agent_meta.packed_member_layouts
}
remote_regions = [
remote_region_by_layer[name] for name in self._transfer_layer_names
]
@@ -505,6 +516,24 @@ class NixlBaseConnectorWorker:
]
if nixl_agent_meta.region_names is not None:
nixl_agent_meta.region_names = list(self._transfer_layer_names)
if nixl_agent_meta.packed_member_layouts:
for i, layer_name in enumerate(self._transfer_layer_names):
if remote_regions[i] not in packed_regions:
continue
assert layer_name in nixl_agent_meta.packed_member_layouts, (
f"Remote packed layer {layer_name!r} has no layout"
)
offset, page_size = nixl_agent_meta.packed_member_layouts[layer_name]
assert (
0 <= offset < offset + page_size <= nixl_agent_meta.block_lens[i]
), f"Remote packed layer {layer_name!r} escapes its block"
local_region = self._transfer_layer_region_indices[i]
assert page_size == self.block_len_per_layer[local_region], (
f"Packed MLA page sizes must match for layer {layer_name!r}"
)
nixl_agent_meta.kv_caches_base_addr[i] += offset
nixl_agent_meta.block_lens[i] = page_size
nixl_agent_meta.packed_member_layouts = {}
# 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]
@@ -899,7 +928,10 @@ class NixlBaseConnectorWorker:
# within a block (BLHNC/BHLNC), where stride > block_len.
self.block_stride_per_layer = list[int]()
# Region layer membership is populated for HMA layouts.
# Block-outermost layouts (BLHNC/BHLNC) pack every layer's page into one
# block row of a shared allocation.
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 = []
# Local layer order, set by ``_set_region_layers`` only when
# routing by layer name. Empty means region-index routing.
@@ -1352,6 +1384,8 @@ class NixlBaseConnectorWorker:
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
"""Register the KV Cache data in nixl."""
use_layer_name_routing = self._requires_layer_name_routing()
route_packed_layers = self._has_packed_cache and use_layer_name_routing
self.transfer_topo = TransferTopology(
tp_rank=self.transfer_tp_rank,
tp_size=self.transfer_tp_size,
@@ -1369,9 +1403,15 @@ class NixlBaseConnectorWorker:
else None,
is_mamba=self._has_mamba,
)
backend_name = self.backend_name
if self._supports_pp_hma and self._has_packed_cache:
# PP can reorder the same backends across stages of a packed model.
backend_name = ",".join(
sorted(backend.get_name() for backend in self.attn_backends)
)
self.compat_hash = compute_nixl_compatibility_hash(
self.vllm_config,
self.backend_name,
backend_name,
transfer_mode=self._TRANSFER_MODE,
)
@@ -1413,6 +1453,11 @@ class NixlBaseConnectorWorker:
# 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
if route_packed_layers and any(
not isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
for spec in self._layer_specs.values()
):
raise NotImplementedError("PP push with packed KV caches requires MLA")
layer_specs: dict[str, KVCacheSpec] = {}
compressed_region_owners: dict[int, torch.Tensor] = {}
@@ -1439,10 +1484,9 @@ class NixlBaseConnectorWorker:
):
compressed_region_owners.setdefault(cache.data_ptr(), cache)
track_region_layers = (
self._supports_pp_hma and self._is_hma_required and not self._has_mamba
)
track_region_layers = self._tracks_region_layers()
region_layers: list[list[str]] = []
packed_member_layouts: dict[str, tuple[int, 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
@@ -1587,11 +1631,34 @@ class NixlBaseConnectorWorker:
virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
layer_spec, cache, physical_page_size, num_blocks
)
# Push workers address the pages of packed MLA rows by layer.
packed_mla_push = (
track_region_layers
and self._has_packed_cache
and packed_storage
and storage_is_block_major
and is_mla_region
)
if virtual_transfer_pages:
# A compressed kernel row can contain multiple NIXL transfer pages.
region_specs = [
(cache.data_ptr(), physical_page_size, physical_page_size)
]
elif packed_mla_push 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.
storage_block_len = registration_len // num_blocks
region_specs = [
(registration_base, storage_block_len, storage_block_len)
]
offset = cache.data_ptr() - storage_addr
assert offset >= 0 and offset + physical_page_size <= block_stride
packed_member_layouts[layer_name] = (offset, physical_page_size)
elif storage_is_block_major and (
not page_contiguous
and ((packed_storage and is_mla_region) or not self._is_csa_linear)
@@ -1633,7 +1700,9 @@ class NixlBaseConnectorWorker:
]
for base_addr, block_len, block_stride in region_specs:
if base_addr in seen_base_addresses:
# A packed PP producer keeps one region per layer, so aliases at
# one address (a layer and its SWA view) stay distinct.
if base_addr in seen_base_addresses and not route_packed_layers:
region_index = seen_base_addresses.index(base_addr)
assert region_mem_types[region_index] == mem_type
self._region_is_mla[region_index] |= is_mla_region
@@ -1717,7 +1786,7 @@ class NixlBaseConnectorWorker:
self._set_region_layers(region_layers)
if self.pp_size > 1 and not self._is_hma_required:
if self.pp_size > 1 and not use_layer_name_routing:
start_layer, end_layer = self.model_config.get_layers_start_end_indices(
self.vllm_config.parallel_config
)
@@ -1807,6 +1876,7 @@ class NixlBaseConnectorWorker:
dcp_size=self.dcp_size,
pcp_size=self.pcp_size,
region_members=self.region_members,
packed_member_layouts=packed_member_layouts,
)
# Wrap metadata in payload with hash for defensive decoding
assert self.compat_hash is not None
@@ -48,8 +48,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
# 10: Add dense virtual transfer pages for compressed MLA caches
# 11: Add per-region transfer geometry and memory types to NixlAgentMetadata
# 12: Add per-region member names for PP push
# 13: Add packed-member layouts and order-independent packed-push backend hashes
#
NIXL_CONNECTOR_VERSION: int = 12
NIXL_CONNECTOR_VERSION: int = 13
@dataclass
@@ -74,6 +75,8 @@ class NixlAgentMetadata:
pcp_size: int = 1
# Layer names sharing each advertised region, in region order.
region_members: list[list[str]] = field(default_factory=list)
# Packed member -> (byte offset in its region's block, bytes per page).
packed_member_layouts: dict[str, tuple[int, int]] = field(default_factory=dict)
@dataclass