diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index a6debd7dba..621b0f7c9d 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -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 diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index a41d5868b1..46cb0bc482 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -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 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 f2999bdb13..029ffafd3c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -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" 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 4d8ead28f3..7ce6666ae8 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -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 = () 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 f4a643d2b7..0cbdd786d8 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 @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index ca6f30d0cf..8d0aad7633 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -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