diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index e478d8e5a9..ed66b8e85e 100644 --- a/docker/Dockerfile.rock +++ b/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 \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 8ba942346b..36aa0aa4df 100644 --- a/docker/Dockerfile.rocm +++ b/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 \ diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index dcbfa7eb04..b7dd3ce70f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/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] ] diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 469e0f5d82..e431beba69 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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, 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 4e7a7b8431..a437a5ee97 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/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 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 a36c7e5b2a..48cf893210 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index 5ae235c24d..2a10e7557f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/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: diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 3ddc568a52..e3c73fb2c8 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/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, ) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 59b2e7ad03..1e7bfeb14f 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/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 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 778eb0f756..7e3698f0ac 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 @@ -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 diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py index c063b3ea74..c2f8d6be5a 100644 --- a/vllm/distributed/nixl_utils.py +++ b/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", ]