[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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user