diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index b7dd3ce70f..db7dbdab09 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -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 ] 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 a437a5ee97..b22d7c9485 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -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", 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 48cf893210..d92857c341 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -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 = {} 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 25f8195044..ed0e6eafe6 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 @@ -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