[KV Connector][NIXL] Register packed MLA KV rows as a single region

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -68,6 +68,7 @@
 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

@@ -25,6 +25,9 @@
 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

@@ -387,6 +387,7 @@ def fresh(cls) -> _StubWriterWorker:
         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 = {}