[NIXL] Prepare KV transfer descriptors as NIXL 1.5 strided lists (#61032)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Ilia Yastrebov <iyastrebov@nvidia.com>

docker/Dockerfile.rock

@@ -385,7 +385,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # NIXL/UCX build stages
 FROM base AS build_nixl
-ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
+ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
 ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
 ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
 ARG UCX_REPO="https://github.com/openucx/ucx.git"
@@ -461,6 +461,8 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # so configure the ROCm wheel variant directly through Meson.
 # Exclude libcore and libpull from auditwheel: transitive dependencies
 # that are not shipped in the wheel and vary across base images.
+# Tag manylinux_2_35 to match the Ubuntu 22.04 base: NIXL >= 1.5 links
+# GLIBCXX_3.4.30 symbols, which manylinux_2_34 does not allow.
 RUN cd /opt/nixl && \
     ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
     CC="ccache gcc" CXX="ccache g++" \
@@ -476,7 +478,7 @@ RUN cd /opt/nixl && \
         --exclude 'libcore*' \
         --exclude 'libpull*' \
         /tmp/nixl_wheels/nixl_rocm*.whl \
-        --plat manylinux_2_34_$(uname -m) \
+        --plat manylinux_2_35_$(uname -m) \
         --wheel-dir /tmp/nixl_wheels/repaired && \
     ./contrib/wheel_add_ucx_plugins.py \
         --ucx-plugins-dir ${UCX_HOME}/lib/ucx \

docker/Dockerfile.rocm

@@ -382,7 +382,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # NIXL/UCX build stages
 FROM base AS build_nixl
-ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
+ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
 ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
 ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
 ARG UCX_REPO="https://github.com/openucx/ucx.git"
@@ -459,6 +459,8 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 # so configure the ROCm wheel variant directly through Meson.
 # Exclude libcore and libpull from auditwheel: transitive dependencies
 # that are not shipped in the wheel and vary across base images.
+# Tag manylinux_2_35 to match the Ubuntu 22.04 base: NIXL >= 1.5 links
+# GLIBCXX_3.4.30 symbols, which manylinux_2_34 does not allow.
 RUN cd /opt/nixl && \
     ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
     CC="ccache gcc" CXX="ccache g++" \
@@ -474,7 +476,7 @@ RUN cd /opt/nixl && \
         --exclude 'libcore*' \
         --exclude 'libpull*' \
         /tmp/nixl_wheels/nixl_rocm*.whl \
-        --plat manylinux_2_34_$(uname -m) \
+        --plat manylinux_2_35_$(uname -m) \
         --wheel-dir /tmp/nixl_wheels/repaired && \
     ./contrib/wheel_add_ucx_plugins.py \
         --ucx-plugins-dir ${UCX_HOME}/lib/ucx \

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -83,6 +83,7 @@ from .utils import (
     create_request,
     create_scheduler,
     create_vllm_config,
+    expand_strided_descs,
     make_kv_cache_config,
     maybe_update_block_size,
 )
@@ -159,7 +160,9 @@ class FakeNixlWrapper:
     def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
         return [str(uuid.uuid4()) for _ in blocks_data]
 
-    def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
+    def prep_xfer_dlist(
+        self, agent_name: str, descs, mem_type: str | None = None
+    ) -> int:
         return uuid.uuid4().int
 
     def get_agent_metadata(self) -> bytes:
@@ -554,7 +557,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self.kv_cache_layout = kv_cache_layout
         # Mock register_kv_caches attributes needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
-        self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
+        self.src_blocks_data = np.empty((0, 5), dtype=np.uint64)
         rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
         test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1)
         self.transfer_topo = TransferTopology(
@@ -929,11 +932,12 @@ class TestNixlHandshake:
         worker.block_len_per_layer = [4096 * worker.block_size]
         worker.num_blocks = 1
         worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
+        block_len = worker.block_len_per_layer[0]
         worker.src_blocks_data = np.array(
-            [(0, worker.block_len_per_layer[0], worker.tp_rank)],
+            [(0, block_len, worker.tp_rank, block_len, worker.num_blocks)],
             dtype=np.uint64,
         )
-        worker.num_descs = len(worker.src_blocks_data)
+        worker.num_descs = int(worker.src_blocks_data[:, 4].sum())
 
         def check_handshake(remote_tp_size: int):
             tp_ratio = remote_tp_size // local_tp_size
@@ -1333,7 +1337,7 @@ class TestNixlHandshake:
         )
 
         assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
-            [0x1000 + local_block_len, local_block_len, 0]
+            [0x1000 + local_block_len, local_block_len, 0, remote_block_len, 1]
         ]
 
     @patch(
@@ -1371,12 +1375,12 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
             worker.src_blocks_data = np.array(
                 [
-                    (0, fa_len, worker.tp_rank),
-                    (0, idx_len, worker.tp_rank),
+                    (0, fa_len, worker.tp_rank, fa_len, worker.num_blocks),
+                    (0, idx_len, worker.tp_rank, idx_len, worker.num_blocks),
                 ],
                 dtype=np.uint64,
             )
-            worker.num_descs = len(worker.src_blocks_data)
+            worker.num_descs = int(worker.src_blocks_data[:, 4].sum())
 
             # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio;
             # REPLICATE region is unchanged.
@@ -2142,21 +2146,11 @@ def test_mixed_memory_local_descriptors_split_by_memory_type():
     worker._desc_pos_by_block_size = {}
     worker._dram_src_handles_by_block_size = {}
     worker.nixl_memory_type = "VRAM"
+    # One dense run of two 10-byte blocks per region: (addr, len, dev, stride, count).
     worker._build_fa_local = MagicMock(  # type: ignore[method-assign]
-        return_value=np.array(
-            [
-                [100, 10, 3],
-                [110, 10, 3],
-                [200, 10, 3],
-                [210, 10, 3],
-            ]
-        )
+        return_value=np.array([[100, 10, 3, 10, 2], [200, 10, 3, 10, 2]])
     )
     worker.nixl_wrapper = MagicMock()
-    worker.nixl_wrapper.get_xfer_descs.side_effect = lambda blocks, memory_type: (
-        memory_type,
-        blocks,
-    )
     worker.nixl_wrapper.prep_xfer_dlist.side_effect = [11, 22]
 
     handle, blocks = worker.register_local_xfer_handler(worker.block_size)
@@ -2165,11 +2159,17 @@ def test_mixed_memory_local_descriptors_split_by_memory_type():
     assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
     # Must stay an array: P_TP > D_TP splits call .tolist() on it.
     assert isinstance(blocks, np.ndarray)
-    assert blocks[:, 2].tolist() == [0, 0, 3, 3]
-    memory_types = [
-        call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
+    # DRAM runs are re-homed to device 0; the per-descriptor maps consumed by
+    # the mixed READ path still cover every block.
+    assert blocks[:, 2].tolist() == [0, 3]
+    assert worker._desc_is_dram_by_block_size[16].tolist() == [True, True, False, False]
+    assert worker._desc_pos_by_block_size[16].tolist() == [0, 1, 0, 1]
+    dlist_calls = worker.nixl_wrapper.prep_xfer_dlist.call_args_list
+    assert [call.kwargs["mem_type"] for call in dlist_calls] == ["DRAM", "VRAM"]
+    assert [call.args[1].tolist() for call in dlist_calls] == [
+        [[100, 10, 0, 10, 2]],
+        [[200, 10, 3, 10, 2]],
     ]
-    assert memory_types == ["DRAM", "VRAM"]
 
 
 @pytest.fixture
@@ -2362,8 +2362,8 @@ def test_register_kv_caches(
     This test verifies:
     1. nixl_wrapper.get_reg_descs() is called with caches_data containing
        tensor metadata
-    2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing
-       block layout info
+    2. nixl_wrapper.prep_xfer_dlist() is called with strided block runs whose
+       expansion covers exactly the registered blocks
     """
     block_size = maybe_update_block_size(16)
     vllm_config = create_vllm_config(
@@ -2479,9 +2479,12 @@ def test_register_kv_caches(
             assert size == raw.nbytes
             assert base_addr == raw.data_ptr()
 
-        # Verify get_xfer_descs was called with blocks_data
-        assert mock_wrapper_instance.get_xfer_descs.called
-        blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
+        # Verify the local dlist was prepared from one strided run per region;
+        # expand it to per-block windows for the layout-blind checks below.
+        assert mock_wrapper_instance.prep_xfer_dlist.called
+        _agent, block_runs = mock_wrapper_instance.prep_xfer_dlist.call_args[0]
+        assert block_runs.shape[1] == 5
+        blocks_data = expand_strided_descs(block_runs)
 
         # Layout-blind contract: whatever regions the worker carves out,
         # transferring "block b" must move exactly logical block b's bytes for
@@ -2515,6 +2518,7 @@ def test_register_kv_caches(
             start for start, _len, _tp in blocks_data if owner[start] == 0
         }
         assert len(blocks_data) == num_blocks * len(base_addrs)
+        assert len(block_runs) == len(base_addrs)
 
         assert connector.connector_worker.block_size == block_size
 
@@ -2578,14 +2582,14 @@ def test_register_packed_dsv4_mla_cache_as_single_region(
         )
         connector.register_kv_caches(dict(zip(layer_names, views)))
 
-        blocks_data, _ = wrapper.get_xfer_descs.call_args[0]
+        _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()]
-        assert blocks_data.tolist() == [
-            [raw.data_ptr() + block_idx * packed_block_len, packed_block_len, 0]
-            for block_idx in range(num_blocks)
+        # One packed region -> one run of num_blocks packed rows.
+        assert block_runs.tolist() == [
+            [raw.data_ptr(), packed_block_len, 0, packed_block_len, num_blocks]
         ]
 
 

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -37,6 +37,7 @@ from vllm.v1.kv_cache_interface import (
 from .utils import (
     create_request,
     create_vllm_config,
+    expand_strided_descs,
     make_kv_cache_config,
     make_nixl_scheduler,
     maybe_update_block_size,
@@ -2201,7 +2202,8 @@ def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size
         )
         for block in range(count)
     ]
-    assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
+    local_descs = expand_strided_descs(worker.src_blocks_data)
+    assert local_descs[:, 0].tolist() == expected_addrs
 
 
 @pytest.mark.cpu_test
@@ -2284,19 +2286,21 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     assert worker.num_regions == 2 and worker.num_descs == 24
     # Kernel-granularity block lens; TP-independent for MLA hybrids.
     assert worker.block_len_per_layer == [unified_page // 3] * 2
-    # Split handles must replicate every FA descriptor (MLA isn't head-sharded).
-    assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24
-    # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page.
-    # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
-    assert worker.src_blocks_data.shape == (24 + 32, 3)
-    fa_descs = worker.src_blocks_data[:24]
+    # FA runs: one per region x 12 kernel blocks, page stride = block_stride.
+    # Mamba runs: 2 regions x (3 conv sub-projections + 1 ssm), 4 blocks each.
+    assert worker.src_blocks_data.shape == (2 + 8, 5)
+    fa_runs = worker.src_blocks_data[:2]
+    assert fa_runs[:, 1].tolist() == [unified_page // 3] * 2
+    assert fa_runs[:, 3].tolist() == [block_stride] * 2
+    assert fa_runs[:, 4].tolist() == [12] * 2
+    assert worker.src_blocks_data[2:, 4].tolist() == [4] * 8
     expected_addrs = [
         tensor.data_ptr() + block * block_stride
         for tensor in tensors
         for block in range(12)
     ]
+    fa_descs = expand_strided_descs(fa_runs)
     assert fa_descs[:, 0].tolist() == expected_addrs
-    assert all(size == unified_page // 3 for size in fa_descs[:, 1])
     worker.nixl_wrapper.register_memory.assert_called_once()
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes,

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -37,7 +37,7 @@ from vllm.v1.kv_cache_interface import (
     create_kv_cache_views,
 )
 
-from .utils import create_vllm_config
+from .utils import create_vllm_config, expand_strided_descs
 
 
 class _RecordingNixl:
@@ -65,10 +65,11 @@ class _RecordingNixl:
     def get_xfer_descs(self, blocks_data, mem_type):
         return blocks_data
 
-    def prep_xfer_dlist(self, agent, descs):
+    def prep_xfer_dlist(self, agent, descs, mem_type=None):
         handle = self._next_handle
         self._next_handle += 1
-        self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3)
+        # Store per-block (addr, len, dev) rows so desc ids index directly.
+        self.dlists[handle] = expand_strided_descs(np.asarray(descs, dtype=np.uint64))
         return handle
 
     def add_remote_agent(self, metadata):
@@ -186,9 +187,8 @@ def test_packed_mla_pp1_push_peer_transfers_whole_rows(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)
-    ]
+    # One packed region -> one run of 4 whole rows.
+    assert worker.src_blocks_data.tolist() == [[raw.data_ptr(), stride, 0, stride, 4]]
     assert worker._transfer_layer_group_ids == ()
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata
@@ -248,7 +248,8 @@ def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(
         1,
         region_num_blocks=producer.dst_region_num_blocks[producer.engine_id],
     )
-    assert producer.src_blocks_data[local_ids].tolist() == [
+    local_descs = expand_strided_descs(producer.src_blocks_data)
+    assert local_descs[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],
@@ -318,7 +319,8 @@ def test_local_descriptors_follow_each_region_pool_capacity():
 
     descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1)
 
-    assert descriptors[:, 0].tolist() == [100, 116, 1000, 1016, 1032]
+    # One (addr, len, dev, stride, count) run per region, sized to its pool.
+    assert descriptors.tolist() == [[100, 16, 0, 16, 2], [1000, 16, 0, 16, 3]]
 
 
 def _register_overlaid_mla_worker(
@@ -479,7 +481,8 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
         allocation.data_ptr() + block * block_stride for block in range(num_blocks)
     ]
     num_desc_regions = 2 if push_pp else 1
-    assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions
+    local_descs = expand_strided_descs(worker.src_blocks_data)
+    assert local_descs[:, 0].tolist() == expected_addrs * num_desc_regions
     assert worker.num_descs == num_blocks * num_desc_regions
     assert (
         worker.dst_region_num_blocks[worker.engine_id]
@@ -500,11 +503,12 @@ def _descriptor_geometry(registered) -> dict:
     """The geometry handed to NIXL, addressed relative to the allocation."""
     worker = registered.worker
     base = registered.allocation.data_ptr()
+    descs = expand_strided_descs(worker.src_blocks_data)
     return {
         "block_len_per_layer": list(worker.block_len_per_layer),
         "block_stride_per_layer": list(worker.block_stride_per_layer),
-        "desc_offsets": [addr - base for addr in worker.src_blocks_data[:, 0].tolist()],
-        "desc_lens": worker.src_blocks_data[:, 1].tolist(),
+        "desc_offsets": [addr - base for addr in descs[:, 0].tolist()],
+        "desc_lens": descs[:, 1].tolist(),
     }
 
 
@@ -1365,10 +1369,11 @@ def test_csa_linear_ple_descriptor_is_not_split():
     bases = [0x10000, 0x20000, 0x30000, 0x40000]
     attention = worker._build_fa_local(bases, block_size_ratio=1)
     mamba = worker._build_mamba_local(bases)
-    assert attention.shape == (8, 3)
-    assert mamba.shape == (10, 3)
-    assert mamba[-2:, 0].tolist() == [0x10000, 0x10100]
-    assert mamba[-2:, 1].tolist() == [256, 256]
+    # One (addr, len, dev, stride, count) run per region / sub-region.
+    assert attention.shape == (4, 5)
+    assert mamba.shape == (5, 5)
+    # PLE run: 2 whole 256-byte pages at 0x10000, 0x10100.
+    assert mamba[-1].tolist() == [0x10000, 256, 0, 256, 2]
 
     mapping = TPMapping(
         source_ranks_per_group=((0,), (0,), (0, 1), (0, 1)),
@@ -1380,14 +1385,15 @@ def test_csa_linear_ple_descriptor_is_not_split():
         worker._build_local_splits_from_plan(
             mapping,
             np.concatenate([attention, mamba]),
-            num_fa_descs=len(attention),
+            num_fa_descs=int(attention[:, 4].sum()),
         )
     )
     assert len(splits) == 2
-    assert [row[1] for row in splits[0][-2:]] == [256, 256]
-    assert [row[1] for row in splits[1][-2:]] == [256, 256]
-    assert splits[0][-3][1] * 2 == mamba[-3, 1]
-    assert splits[1][-3][1] * 2 == mamba[-3, 1]
+    for split in splits:
+        # The PLE run is copied whole to every source rank, while the sharded
+        # SSM run right before it is halved per source.
+        assert split[-1].tolist() == mamba[-1].tolist()
+        assert split[-2, 1] * 2 == mamba[-2, 1]
 
 
 @pytest.mark.cpu_test
@@ -1424,9 +1430,9 @@ def test_csa_linear_remote_ple_is_copied_whole():
         tp_ratio=-2,
         transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
     )
-    assert descriptors.shape == (10, 3)
-    assert descriptors[-2:, 0].tolist() == [0x10000, 0x10100]
-    assert descriptors[-2:, 1].tolist() == [256, 256]
+    assert descriptors.shape == (5, 5)
+    # PLE run: 2 whole 256-byte pages at 0x10000, 0x10100.
+    assert descriptors[-1].tolist() == [0x10000, 256, 0, 256, 2]
 
     metadata.ple_block_len = 128
     with pytest.raises(ValueError, match="PLE pages require identical"):
@@ -1451,8 +1457,8 @@ def test_csa_linear_ple_page_wider_than_the_shared_region_page():
 
     bases = [0x10000, 0x20000, 0x30000, 0x40000]
     mamba = worker._build_mamba_local(bases)
-    assert mamba[-2:, 0].tolist() == [0x10000, 0x10000 + 384]
-    assert mamba[-2:, 1].tolist() == [384, 384]
+    # PLE run: 2 whole 384-byte pages at 0x10000, 0x10000 + 384.
+    assert mamba[-1].tolist() == [0x10000, 384, 0, 384, 2]
 
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -60,7 +60,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput
 
-from .utils import create_request, make_nixl_push_scheduler
+from .utils import create_request, expand_strided_descs, make_nixl_push_scheduler
 
 # ----------------------------------------------------------------- #
 #  Helpers / fakes                                                   #
@@ -2008,8 +2008,13 @@ def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(
     assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000]
 
     plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0)
-    local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1)
-    remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1)
+    # Builders emit one strided run per region; expand to per-block descriptors.
+    local_descs = expand_strided_descs(
+        worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1)
+    )
+    remote_descs = expand_strided_descs(
+        worker._build_fa_remote(plan, consumer, block_size_ratio=1)
+    )
     local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices]
     local_ids = worker._compute_desc_ids(
         [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts

tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py

@@ -75,7 +75,9 @@ class FakeNixlWrapper:
     def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
         return [str(uuid.uuid4()) for _ in blocks_data]
 
-    def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
+    def prep_xfer_dlist(
+        self, agent_name: str, descs, mem_type: str | None = None
+    ) -> int:
         return uuid.uuid4().int
 
     def get_agent_metadata(self) -> bytes:

tests/v1/kv_connector/unit/test_tp_mapping.py

@@ -145,9 +145,8 @@ def _make_mock_worker_for_splits(group_spec_types):
     """Build a mock NixlConnectorWorker with _group_spec_types for split tests.
 
     No per-region replicate flags are configured (``block_len_per_layer`` empty
-    and ``num_regions == 0``), so ``_fa_desc_replicated`` takes its early-return
-    path and treats every FA descriptor as SPLIT, matching the legacy behavior
-    these tests assert.
+    and ``num_regions == 0``), so ``_build_local_splits_from_plan`` treats every
+    FA descriptor as SPLIT, matching the legacy behavior these tests assert.
     """
     worker = object.__new__(NixlConnectorWorker)
     worker._group_spec_types = group_spec_types
@@ -174,11 +173,9 @@ class TestBuildSrcSplitHandles:
         )
 
         worker = _make_mock_worker_for_splits((FullAttentionSpec,))
-        src_blocks_data = np.array(
-            [(0x2000 + i * 1024, 1024, 0) for i in range(8)],
-            dtype=np.uint64,
-        )
-        num_descs = len(src_blocks_data)
+        # One region of 8 dense 1024-byte blocks as a single strided run.
+        src_blocks_data = np.array([(0x2000, 1024, 0, 1024, 8)], dtype=np.uint64)
+        num_descs = 8
         splits = list(
             worker._build_local_splits_from_plan(
                 plan,
@@ -189,9 +186,10 @@ class TestBuildSrcSplitHandles:
 
         assert len(splits) == remote_tp_size
         for handle in splits:
-            assert len(handle) == len(src_blocks_data)
-            for _, length, _ in handle:
-                assert length == 1024 // remote_tp_size
+            assert handle.shape == src_blocks_data.shape
+            # The head slice shrinks the length; block stride and count are kept.
+            assert handle[:, 1].tolist() == [1024 // remote_tp_size]
+            assert handle[:, 3:].tolist() == [[1024, 8]]
 
 
 class TestMambaPlanSplitHandles:
@@ -209,12 +207,12 @@ class TestMambaPlanSplitHandles:
         )
 
         worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
-        # 2 FA descs + 1 SSM desc
+        # 2 FA runs + 1 SSM run of one block each: (addr, len, dev, stride, count)
         src_blocks_data = np.array(
             [
-                (1000, 200, 0),  # FA desc 0
-                (2000, 200, 0),  # FA desc 1
-                (3000, 400, 0),  # SSM desc 0
+                (1000, 200, 0, 200, 1),  # FA desc 0
+                (2000, 200, 0, 200, 1),  # FA desc 1
+                (3000, 400, 0, 400, 1),  # SSM desc 0
             ],
             dtype=np.uint64,
         )
@@ -224,14 +222,22 @@ class TestMambaPlanSplitHandles:
         assert len(splits) == 2  # 2 source ranks
 
         # Rank 0 (FA source, p_idx=0):
-        # FA: chunk=200//1=200, slot=0 → (1000, 200, 0), (2000, 200, 0)
-        # SSM: chunk=400//2=200, idx=0 → (3000, 200, 0)
-        assert splits[0] == [(1000, 200, 0), (2000, 200, 0), (3000, 200, 0)]
+        # FA: chunk=200//1=200, slot=0 → (1000, 200), (2000, 200)
+        # SSM: chunk=400//2=200, idx=0 → (3000, 200); block stride stays 400
+        assert splits[0].tolist() == [
+            [1000, 200, 0, 200, 1],
+            [2000, 200, 0, 200, 1],
+            [3000, 200, 0, 400, 1],
+        ]
 
         # Rank 1 (not FA source, p_idx=1):
-        # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0)
-        # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0)
-        assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)]
+        # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200), (2000, 200)
+        # SSM: chunk=400//2=200, idx=1 → (3200, 200); block stride stays 400
+        assert splits[1].tolist() == [
+            [1000, 200, 0, 200, 1],
+            [2000, 200, 0, 200, 1],
+            [3200, 200, 0, 400, 1],
+        ]
 
     def test_hetero_block_size_splits(self):
         """With a block-size ratio, single-source FA sub-block descs pass
@@ -244,14 +250,13 @@ class TestMambaPlanSplitHandles:
         )
 
         worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
-        # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded).
+        # 2 FA blocks x ratio 2 sub-blocks (one run of 2 per block) + 1 SSM desc
+        # (never expanded).
         src_blocks_data = np.array(
             [
-                (1000, 100, 0),
-                (1100, 100, 0),
-                (2000, 100, 0),
-                (2100, 100, 0),
-                (3000, 400, 0),
+                (1000, 100, 0, 100, 2),
+                (2000, 100, 0, 100, 2),
+                (3000, 400, 0, 400, 1),
             ],
             dtype=np.uint64,
         )
@@ -259,14 +264,9 @@ class TestMambaPlanSplitHandles:
         splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2))
 
         assert len(splits) == 2
-        fa_passthrough = [
-            (1000, 100, 0),
-            (1100, 100, 0),
-            (2000, 100, 0),
-            (2100, 100, 0),
-        ]
-        assert splits[0] == fa_passthrough + [(3000, 200, 0)]
-        assert splits[1] == fa_passthrough + [(3200, 200, 0)]
+        fa_passthrough = [[1000, 100, 0, 100, 2], [2000, 100, 0, 100, 2]]
+        assert splits[0].tolist() == fa_passthrough + [[3000, 200, 0, 400, 1]]
+        assert splits[1].tolist() == fa_passthrough + [[3200, 200, 0, 400, 1]]
 
     def test_hetero_block_size_head_sharded_asserts(self):
         """Head-sharded FA reads (multiple FA sources) are incompatible with
@@ -280,7 +280,7 @@ class TestMambaPlanSplitHandles:
 
         worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
         src_blocks_data = np.array(
-            [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)],
+            [(1000, 100, 0, 100, 2), (3000, 400, 0, 400, 1)],
             dtype=np.uint64,
         )
 

tests/v1/kv_connector/unit/utils.py

@@ -8,6 +8,7 @@ from itertools import chain, count
 from typing import Any, Literal
 from unittest.mock import Mock
 
+import numpy as np
 import torch
 
 from vllm import SamplingParams
@@ -682,3 +683,18 @@ def make_moriio_writer(fake_worker: Any) -> Any:
     writer._defer_timeout = 60.0
     writer.ensure_worker_started = lambda: None
     return writer
+
+
+def expand_strided_descs(strided_descs: np.ndarray) -> np.ndarray:
+    """Expand Nx5 (addr, len, dev, stride, count) rows to one row per block.
+
+    Mirrors how NIXL indexes a strided descriptor list, so desc ids computed
+    by the connector can be resolved to (addr, len, dev) byte ranges.
+    """
+    strided_descs = np.asarray(strided_descs, dtype=np.uint64).reshape(-1, 5)
+    counts = strided_descs[:, 4].astype(np.intp)
+    starts = np.cumsum(strided_descs[:, 4]) - strided_descs[:, 4]
+    out = np.repeat(strided_descs[:, :3], counts, axis=0)
+    offsets = np.arange(len(out), dtype=np.uint64) - np.repeat(starts, counts)
+    out[:, 0] += offsets * np.repeat(strided_descs[:, 3], counts)
+    return out

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -65,7 +65,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import
     MambaConvSplitInfo,
     derive_mamba_conv_split,
 )
-from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
+from vllm.distributed.nixl_utils import (
+    NixlWrapper,
+    check_nixl_version,
+    nixl_agent_config,
+)
 from vllm.distributed.parallel_state import (
     get_pcp_group,
     get_tensor_model_parallel_rank,
@@ -277,20 +281,23 @@ class NixlBaseConnectorWorker:
         src_blocks_data: np.ndarray,
         num_fa_descs: int,
         block_size_ratio: int = 1,
-    ) -> Iterator[list[tuple[int, int, int]]]:
+    ) -> Iterator[np.ndarray]:
         """Build split handle data for P_TP > D_TP scenario.
 
         num_fa_descs is the boundary between FA and SSM descriptors.
         Split counts are derived from source_ranks_per_group lengths.
         FA uses rank_to_attention_slot for the slot offset;
         SSM uses the rank's positional index.
+        Each Nx5 strided desc covers one region, so slicing it slices every block.
         """
         fa_idx = next(
             i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
         )
         fa_num_splits = len(plan.source_ranks_per_group[fa_idx])
 
-        has_ssm_descs = num_fa_descs < len(src_blocks_data)
+        starts = self._strided_desc_starts(src_blocks_data)
+        num_descs = int(src_blocks_data[:, 4].sum())
+        has_ssm_descs = num_fa_descs < num_descs
         ssm_idx = next(
             (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)),
             None,
@@ -300,38 +307,34 @@ class NixlBaseConnectorWorker:
             if has_ssm_descs and ssm_idx is not None
             else 0
         )
-
-        # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
-        fa_desc_replicated = self._fa_desc_replicated(num_fa_descs, block_size_ratio)
-        sharded_desc_end = len(src_blocks_data) - (
+        sharded_desc_end = num_descs - (
             self._logical_num_blocks if self._ple_region_index is not None else 0
         )
         assert num_fa_descs <= sharded_desc_end
 
-        assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), (
+        # REPLICATE (MLA) descs are read whole on every rank; SPLIT (full-attn)
+        # descs take this rank's head slice. SSM descs before the PLE pages are
+        # sharded by the rank's positional index; PLE pages are copied whole.
+        is_fa = starts < num_fa_descs
+        fa_replicated = self._fa_desc_replicated(
+            starts[is_fa], num_fa_descs, block_size_ratio
+        )
+        assert block_size_ratio == 1 or fa_num_splits == 1 or fa_replicated.all(), (
             "Head-sharded attention reads with P_TP > D_TP and heterogeneous "
             "block sizes are not supported"
         )
-        src_blocks_list = src_blocks_data.tolist()
+        fa_split = is_fa.copy()
+        fa_split[is_fa] = ~fa_replicated
+        ssm_split = ~is_fa & (starts < sharded_desc_end)
 
         for p_idx, p_rank in enumerate(plan.all_source_ranks):
             fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)
-
-            handle: list[tuple[int, int, int]] = []
-            for j, (addr, local_len, dev) in enumerate(src_blocks_list):
-                if j < num_fa_descs:
-                    if fa_desc_replicated[j]:
-                        # REPLICATE (MLA): whole block written on every rank.
-                        handle.append((addr, local_len, dev))
-                    else:
-                        # SPLIT (full-attn): this rank's head slice.
-                        chunk = local_len // fa_num_splits
-                        handle.append((addr + fa_slot * chunk, chunk, dev))
-                elif j < sharded_desc_end:
-                    chunk = local_len // ssm_num_splits
-                    handle.append((addr + p_idx * chunk, chunk, dev))
-                else:
-                    handle.append((addr, local_len, dev))
+            handle = src_blocks_data.copy()
+            handle[fa_split, 1] //= fa_num_splits
+            handle[fa_split, 0] += fa_slot * handle[fa_split, 1]
+            if ssm_num_splits:
+                handle[ssm_split, 1] //= ssm_num_splits
+                handle[ssm_split, 0] += p_idx * handle[ssm_split, 1]
             yield handle
 
     def _needs_split_local_xfer_handles(self, tp_ratio: int, plan: TPMapping) -> bool:
@@ -346,25 +349,30 @@ class NixlBaseConnectorWorker:
         return tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1)
 
     def _fa_desc_replicated(
-        self, num_fa_descs: int, block_size_ratio: int = 1
-    ) -> list[bool]:
-        """Per-FA-descriptor replicate flag, in _build_fa_local emission order:
-        region-major, ``region_num_blocks[i] * block_size_ratio`` descs per
-        region. Length ``num_fa_descs``.
+        self,
+        desc_starts: np.ndarray,
+        num_fa_descs: int,
+        block_size_ratio: int = 1,
+    ) -> np.ndarray:
+        """Replicate flag per FA strided desc, given each one's flat block index.
+
+        _build_fa_local emits regions in order, with
+        ``region_num_blocks[i] * block_size_ratio`` blocks per region. A strided
+        desc never spans regions, so the region owning its first block decides.
         """
         assert self.transfer_topo is not None
         n_regions = len(self.block_len_per_layer)
         if n_regions == 0 or self.num_regions == 0:
-            return [False] * num_fa_descs
+            return np.zeros(len(desc_starts), dtype=bool)
         region_indices = self._transfer_layer_region_indices or range(n_regions)
-        flags: list[bool] = []
-        for i in region_indices:
-            num_descs = self.region_num_blocks[i] * block_size_ratio
-            flags.extend([self._is_region_replicated(i)] * num_descs)
-        assert len(flags) == num_fa_descs, (
-            f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}"
+        region_ends = np.cumsum(
+            [self.region_num_blocks[i] * block_size_ratio for i in region_indices]
         )
-        return flags
+        assert region_ends[-1] == num_fa_descs, (
+            f"FA region blocks {region_ends[-1]} != num_fa_descs {num_fa_descs}"
+        )
+        replicated = np.array([self._is_region_replicated(i) for i in region_indices])
+        return replicated[np.searchsorted(region_ends, desc_starts, side="right")]
 
     def _is_region_replicated(self, region_idx: int) -> bool:
         """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT.
@@ -512,6 +520,8 @@ class NixlBaseConnectorWorker:
         if nixl_wrapper_cls is None:
             logger.error("NIXL is not available")
             raise RuntimeError("NIXL is not available")
+        # Descriptor lists are prepared in NIXL's strided (Nx5) form.
+        check_nixl_version("1.5.0")
         logger.info("Initializing NIXL wrapper")
         logger.info("Initializing NIXL worker %s", engine_id)
 
@@ -1492,7 +1502,6 @@ class NixlBaseConnectorWorker:
                 if isinstance(layer_spec, MambaSpec)
                 else logical_num_blocks * self._physical_blocks_per_logical_kv_block
             )
-            base_addr = cache.data_ptr()
             is_mla_region = isinstance(
                 layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
             )
@@ -1827,7 +1836,7 @@ class NixlBaseConnectorWorker:
 
     def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray:
         """Build desc regions (conv sub-projections + ssm) per layer for
-        local mamba blocks with DS conv layout, as an Nx3 uint64 array.
+        local mamba blocks with DS conv layout, as Nx5 uint64 strided descs.
 
         A Mamba block interleaves conv and SSM state, which crucially differ in
         size, so the two are indexed as separate sub-regions. Attention blocks
@@ -1865,28 +1874,32 @@ class NixlBaseConnectorWorker:
         num_blocks = self._logical_num_blocks
         physical_per_logical = self._physical_blocks_per_logical_kv_block
         device_id = self.device_id
-        block_arange = np.arange(num_blocks, dtype=np.uint64)
-        parts: list[np.ndarray] = []
+        strided_descs: list[tuple[int, int, int, int, int]] = []
 
         region_indices = self._ssm_region_indices or range(len(base_addresses))
         for i in region_indices:
             base_addr = base_addresses[i]
             block_stride = self.block_stride_per_layer[i] * physical_per_logical
-            blk_addrs = base_addr + block_arange * block_stride
             for off, sz in conv_offsets:
-                parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
+                strided_descs.append(
+                    (base_addr + off, sz, device_id, block_stride, num_blocks)
+                )
             # SSM temporal state follows the conv state.
-            parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id))
+            strided_descs.append(
+                (base_addr + conv_size, ssm_size, device_id, block_stride, num_blocks)
+            )
 
         if (region_index := self._ple_region_index) is not None:
             block_len = self._ple_page_len() * physical_per_logical
             block_stride = (
                 self.block_stride_per_layer[region_index] * physical_per_logical
             )
-            block_addrs = base_addresses[region_index] + block_arange * block_stride
-            parts.append(self._stack_descs(block_addrs, block_len, self.device_id))
+            base_addr = base_addresses[region_index]
+            strided_descs.append(
+                (base_addr, block_len, device_id, block_stride, num_blocks)
+            )
 
-        return np.concatenate(parts)
+        return self._pack_strided_descs(strided_descs)
 
     def _build_mamba_remote(
         self,
@@ -1896,7 +1909,7 @@ class NixlBaseConnectorWorker:
     ) -> np.ndarray:
         """Build remote desc regions (conv sub-projections + ssm) per layer.
         For hetero-TP, each D rank reads only its sub-projection slice from
-        the P rank. Returns an Nx3 uint64 array."""
+        the P rank. Returns Nx5 uint64 strided descs."""
         assert nixl_agent_meta.kv_caches_base_addr, (
             "Remote KV cache base addresses must not be empty."
         )
@@ -1916,9 +1929,8 @@ class NixlBaseConnectorWorker:
         remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
         num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
         device_id = nixl_agent_meta.device_id
-        block_arange = np.arange(num_blocks, dtype=np.uint64)
 
-        parts: list[np.ndarray] = []
+        strided_descs: list[tuple[int, int, int, int, int]] = []
         # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
         # block lengths vary across layers (e.g. MLA).
         region_indices = self._ssm_region_indices or range(
@@ -1929,12 +1941,15 @@ class NixlBaseConnectorWorker:
             block_stride = (
                 nixl_agent_meta.block_strides[i] * remote_physical_per_logical
             )
-            blk_addrs = base_addr + block_arange * block_stride
             for off, sz in conv_offsets:
-                parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
+                strided_descs.append(
+                    (base_addr + off, sz, device_id, block_stride, num_blocks)
+                )
             # SSM temporal state is also TP-sharded on the heads dimension.
-            ssm_addrs = blk_addrs + conv_size_remote + local_offset * ssm_read_size
-            parts.append(self._stack_descs(ssm_addrs, ssm_read_size, device_id))
+            ssm_addr = base_addr + conv_size_remote + local_offset * ssm_read_size
+            strided_descs.append(
+                (ssm_addr, ssm_read_size, device_id, block_stride, num_blocks)
+            )
 
         if (region_index := self._ple_region_index) is not None:
             local_block_len = (
@@ -1954,52 +1969,85 @@ class NixlBaseConnectorWorker:
                 nixl_agent_meta.block_strides[region_index]
                 * remote_physical_per_logical
             )
-            block_addrs = (
-                nixl_agent_meta.kv_caches_base_addr[region_index]
-                + block_arange * remote_block_stride
+            addr = nixl_agent_meta.kv_caches_base_addr[region_index]
+            strided_descs.append(
+                (addr, remote_block_len, device_id, remote_block_stride, num_blocks)
             )
-            parts.append(self._stack_descs(block_addrs, remote_block_len, device_id))
 
-        return np.concatenate(parts)
+        return self._pack_strided_descs(strided_descs)
 
     @staticmethod
-    def _stack_descs(addrs: np.ndarray, length: int, device_id: int) -> np.ndarray:
-        out = np.empty((addrs.shape[0], 3), dtype=np.uint64)
-        out[:, 0] = addrs
-        out[:, 1] = length
-        out[:, 2] = device_id
-        return out
+    def _pack_strided_descs(
+        strided_descs: list[tuple[int, int, int, int, int]],
+    ) -> np.ndarray:
+        """Pack (addr, len, dev_id, stride, count) rows into an Nx5 uint64 array.
+
+        Each row describes `count` blocks of `len` bytes whose starts are
+        `stride` bytes apart; NIXL >= 1.5 prepares this layout directly.
+        """
+        return np.array(strided_descs, dtype=np.uint64).reshape(-1, 5)
+
+    @staticmethod
+    def _block_strided_descs(
+        base_addr: int,
+        block_len: int,
+        block_stride: int,
+        num_blocks: int,
+        block_size_ratio: int,
+        device_id: int,
+    ) -> np.ndarray:
+        """Nx5 strided descs for `num_blocks` blocks, each split into
+        `block_size_ratio` sub-blocks of `block_len` bytes.
+
+        Evenly spaced sub-blocks collapse into one strided desc. In a padded
+        block the sub-blocks abut only within the block, so each block gets
+        its own.
+        """
+        if block_size_ratio == 1 or block_stride == block_size_ratio * block_len:
+            stride = block_stride // block_size_ratio
+            count = num_blocks * block_size_ratio
+            return np.array(
+                [[base_addr, block_len, device_id, stride, count]], dtype=np.uint64
+            )
+        strided_descs = np.empty((num_blocks, 5), dtype=np.uint64)
+        strided_descs[:, 0] = (
+            base_addr + np.arange(num_blocks, dtype=np.uint64) * block_stride
+        )
+        strided_descs[:, 1] = block_len
+        strided_descs[:, 2] = device_id
+        strided_descs[:, 3] = block_len
+        strided_descs[:, 4] = block_size_ratio
+        return strided_descs
+
+    @staticmethod
+    def _strided_desc_starts(strided_descs: np.ndarray) -> np.ndarray:
+        """Flat block index of each strided desc's first block."""
+        counts = strided_descs[:, 4]
+        return np.cumsum(counts) - counts
 
     def _build_fa_local(
         self,
         base_addresses: list[int],
         block_size_ratio: int,
     ) -> np.ndarray:
-        """Build local FA descriptors as an Nx3 uint64 array."""
+        """Build local FA descriptors as Nx5 uint64 strided descs."""
         assert self.transfer_topo is not None
         assert base_addresses, "Local KV cache base addresses must not be empty."
         device_id = self.device_id
         region_indices = self._transfer_layer_region_indices or range(
             len(base_addresses)
         )
-        parts: list[np.ndarray] = []
-        for i in region_indices:
-            base_addr = base_addresses[i]
-            block_len = self.block_len_per_layer[i] // block_size_ratio
-            logical_blocks = np.repeat(
-                np.arange(self.region_num_blocks[i], dtype=np.uint64),
-                block_size_ratio,
-            )
-            split_offsets = np.tile(
-                np.arange(block_size_ratio, dtype=np.uint64),
+        parts = [
+            self._block_strided_descs(
+                base_addresses[i],
+                self.block_len_per_layer[i] // block_size_ratio,
+                self.block_stride_per_layer[i],
                 self.region_num_blocks[i],
+                block_size_ratio,
+                device_id,
             )
-            addrs = (
-                base_addr
-                + logical_blocks * self.block_stride_per_layer[i]
-                + split_offsets * block_len
-            )
-            parts.append(self._stack_descs(addrs, block_len, device_id))
+            for i in region_indices
+        ]
         return np.concatenate(parts)
 
     def _build_fa_remote(
@@ -2008,7 +2056,7 @@ class NixlBaseConnectorWorker:
         nixl_agent_meta: NixlAgentMetadata,
         block_size_ratio: int,
     ) -> np.ndarray:
-        """Build remote FA descriptors as an Nx3 uint64 array."""
+        """Build remote FA descriptors as Nx5 uint64 strided descs."""
         assert self.transfer_topo is not None
         assert nixl_agent_meta.kv_caches_base_addr, (
             "Remote KV cache base addresses must not be empty."
@@ -2024,7 +2072,7 @@ class NixlBaseConnectorWorker:
         ] * len(nixl_agent_meta.kv_caches_base_addr)
         device_id = nixl_agent_meta.device_id
         block_strides = nixl_agent_meta.block_strides
-        parts: list[np.ndarray] = []
+        strided_descs: list[tuple[int, int, int, int, int]] = []
         use_layer_name_routing = bool(self._transfer_layer_region_indices)
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
             local_region = (
@@ -2046,10 +2094,12 @@ class NixlBaseConnectorWorker:
             )
             local_block_len = local_block_len // num_reads
 
-            block_arange = np.arange(region_num_blocks[i], dtype=np.uint64)
-            addrs = base_addr + rank_offset + block_arange * block_strides[i]
-            parts.append(self._stack_descs(addrs, local_block_len, device_id))
-        return np.concatenate(parts)
+            addr = base_addr + rank_offset
+            count = region_num_blocks[i]
+            strided_descs.append(
+                (addr, local_block_len, device_id, block_strides[i], count)
+            )
+        return self._pack_strided_descs(strided_descs)
 
     def register_local_xfer_handler(
         self,
@@ -2072,13 +2122,13 @@ class NixlBaseConnectorWorker:
         blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio)
         logger.debug(
             "Created %s blocks for src engine %s and rank %s on device id %s",
-            len(blocks_data),
+            blocks_data[:, 4].sum(),
             self.engine_id,
             self.tp_rank,
             self.device_id,
         )
         if self._has_mamba:
-            assert self.num_descs * block_size_ratio == len(blocks_data)
+            assert self.num_descs * block_size_ratio == blocks_data[:, 4].sum()
             # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
             # is unnecessary — a single conv desc per block suffices.  Consider
             # adding a fast path that falls back to the standard 2-region
@@ -2104,7 +2154,7 @@ class NixlBaseConnectorWorker:
                     )
                 ]
             )
-            assert len(desc_is_dram) == len(blocks_data)
+            assert len(desc_is_dram) == blocks_data[:, 4].sum()
             desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
             dram_idx = np.where(desc_is_dram)[0]
             vram_idx = np.where(~desc_is_dram)[0]
@@ -2112,28 +2162,32 @@ class NixlBaseConnectorWorker:
             desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64)
             self._desc_is_dram_by_block_size[block_size] = desc_is_dram
             self._desc_pos_by_block_size[block_size] = desc_pos
-
             # DRAM descriptors are registered under CPU device 0.
-            blocks_data[desc_is_dram, 2] = 0
-            dram_blocks = blocks_data[dram_idx]
-            vram_blocks = blocks_data[vram_idx]
+            # A strided desc never spans regions, so its first block decides.
+            strided_is_dram = desc_is_dram[self._strided_desc_starts(blocks_data)]
+            blocks_data[strided_is_dram, 2] = 0
             if self._skip_dram_xfer:
                 self._dram_src_handles_by_block_size[block_size] = None
             else:
-                dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM")
                 self._dram_src_handles_by_block_size[block_size] = (
-                    self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
+                    self.nixl_wrapper.prep_xfer_dlist(
+                        "NIXL_INIT_AGENT",
+                        blocks_data[strided_is_dram],
+                        mem_type="DRAM",
+                    )
                 )
-            descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type)
-            return (
-                self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs),
-                blocks_data,
+            device_handle = self.nixl_wrapper.prep_xfer_dlist(
+                "NIXL_INIT_AGENT",
+                blocks_data[~strided_is_dram],
+                mem_type=self.nixl_memory_type,
             )
+            return device_handle, blocks_data
 
-        mem_type = self.region_mem_types[0]
-        descs = self.nixl_wrapper.get_xfer_descs(blocks_data, mem_type)
         # NIXL_INIT_AGENT to be used for preparations of local descs.
-        return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data
+        handle = self.nixl_wrapper.prep_xfer_dlist(
+            "NIXL_INIT_AGENT", blocks_data, mem_type=self.region_mem_types[0]
+        )
+        return handle, blocks_data
 
     def add_remote_agent(
         self,
@@ -2350,23 +2404,23 @@ class NixlBaseConnectorWorker:
                 block_size_ratio,
             ):
                 if self._mixed_mem_types:
-                    handle_data = np.asarray(handle_data, dtype=np.uint64)
                     desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size]
+                    strided_is_dram = desc_is_dram[
+                        self._strided_desc_starts(handle_data)
+                    ]
                     if self._skip_dram_xfer:
                         dram_handle = None
                     else:
-                        dram_descs = self.nixl_wrapper.get_xfer_descs(
-                            handle_data[desc_is_dram], "DRAM"
-                        )
                         dram_handle = self.nixl_wrapper.prep_xfer_dlist(
-                            "NIXL_INIT_AGENT", dram_descs
+                            "NIXL_INIT_AGENT",
+                            handle_data[strided_is_dram],
+                            mem_type="DRAM",
                         )
                     self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle)
-                    handle_data = handle_data[~desc_is_dram]
-                descs = self.nixl_wrapper.get_xfer_descs(
-                    handle_data, self.nixl_memory_type
+                    handle_data = handle_data[~strided_is_dram]
+                handle = self.nixl_wrapper.prep_xfer_dlist(
+                    "NIXL_INIT_AGENT", handle_data, mem_type=self.nixl_memory_type
                 )
-                handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
                 self.src_xfer_handles_by_tp_ratio[split_key].append(handle)
 
         ### Register remote agent memory regions
@@ -2383,7 +2437,7 @@ class NixlBaseConnectorWorker:
         )
         logger.debug(
             "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
-            len(blocks_data),
+            blocks_data[:, 4].sum(),
             engine_id,
             remote_tp_rank,
             self.tp_rank,
@@ -2402,9 +2456,10 @@ class NixlBaseConnectorWorker:
             raise NotImplementedError(
                 "NIXL pull does not yet support a mixed-memory producer"
             )
-        descs = self.nixl_wrapper.get_xfer_descs(blocks_data, remote_mem_types.pop())
         self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
-            self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
+            self.nixl_wrapper.prep_xfer_dlist(
+                remote_agent_name, blocks_data, mem_type=remote_mem_types.pop()
+            )
         )
 
         return remote_agent_name

vllm/distributed/nixl_utils.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import functools
 import importlib
 import os
 import sys
@@ -91,6 +92,24 @@ def is_nixl_available() -> bool:
     return pkg in sys.modules or importlib.util.find_spec(pkg) is not None
 
 
+@functools.cache
+def check_nixl_version(min_version: str) -> None:
+    """Raise if the installed NIXL package is older than ``min_version``."""
+    from importlib.metadata import packages_distributions, version
+
+    from packaging.version import Version
+
+    pkg = _get_nixl_package_name()
+    for dist in packages_distributions().get(pkg, ()):
+        installed = version(dist)
+        if Version(installed) < Version(min_version):
+            raise RuntimeError(
+                f"NIXL >= {min_version} is required, but {dist}=={installed} "
+                "is installed"
+            )
+        return
+
+
 def alias_nixl_for_ray() -> None:
     """Let Ray's ``nixl._api`` import resolve to the ROCm implementation."""
     if not current_platform.is_rocm() or "nixl._api" in sys.modules:
@@ -112,5 +131,6 @@ __all__ = [
     "nixl_agent_config",
     "nixlXferTelemetry",
     "is_nixl_available",
+    "check_nixl_version",
     "alias_nixl_for_ray",
 ]