| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for TP mapping and transfer plan utilities. |
| 4 | |
| 5 | These tests verify that TP mapping produces correct outputs |
| 6 | (source ranks, split handles, desc IDs). |
| 7 | No GPU or NIXL required. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | from types import SimpleNamespace |
| 13 | |
| 14 | import numpy as np |
| 15 | import pytest |
| 16 | |
| 17 | from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology |
| 18 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( |
| 19 | TPMapping, |
| 20 | compute_tp_mapping, |
| 21 | ) |
| 22 | from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( |
| 23 | NixlConnectorWorker, |
| 24 | ) |
| 25 | from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec |
| 26 | |
| 27 | # ====================================================================== |
| 28 | # Test fixtures / helpers |
| 29 | # ====================================================================== |
| 30 | |
| 31 | |
| 32 | def _compute_mapping( |
| 33 | tp_rank: int = 0, |
| 34 | tp_size: int = 1, |
| 35 | remote_tp_size: int = 1, |
| 36 | is_mla: bool = False, |
| 37 | num_kv_heads: int = 8, |
| 38 | group_spec_types: tuple[type, ...] = (FullAttentionSpec,), |
| 39 | dcp_size: int = 1, |
| 40 | remote_dcp_size: int = 1, |
| 41 | ) -> TPMapping: |
| 42 | transfer_topology = object.__new__(TransferTopology) |
| 43 | transfer_topology.tp_rank = tp_rank |
| 44 | transfer_topology.tp_size = tp_size |
| 45 | transfer_topology.is_mla = is_mla |
| 46 | transfer_topology.total_num_kv_heads = num_kv_heads |
| 47 | transfer_topology.dcp_size = dcp_size |
| 48 | return compute_tp_mapping( |
| 49 | transfer_topology=transfer_topology, |
| 50 | remote_tp_size=remote_tp_size, |
| 51 | group_spec_types=group_spec_types, |
| 52 | remote_dcp_size=remote_dcp_size, |
| 53 | ) |
| 54 | |
| 55 | |
| 56 | # ====================================================================== |
| 57 | # TP mapping structure tests |
| 58 | # ====================================================================== |
| 59 | |
| 60 | |
| 61 | class TestTPMappingStructure: |
| 62 | def test_source_ranks_homogeneous(self): |
| 63 | m = _compute_mapping(tp_size=2, tp_rank=1, remote_tp_size=2) |
| 64 | assert m.all_source_ranks == (1,) |
| 65 | |
| 66 | def test_source_ranks_d_gt_p(self): |
| 67 | m = _compute_mapping(tp_size=4, tp_rank=2, remote_tp_size=2) |
| 68 | assert m.all_source_ranks == (1,) |
| 69 | |
| 70 | def test_source_ranks_p_gt_d(self): |
| 71 | m = _compute_mapping(tp_size=1, tp_rank=0, remote_tp_size=2) |
| 72 | assert m.all_source_ranks == (0, 1) |
| 73 | |
| 74 | |
| 75 | @pytest.mark.parametrize( |
| 76 | "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks", |
| 77 | [ |
| 78 | (0, 1, 8, 1, 8, tuple(range(8))), |
| 79 | (0, 4, 4, 1, 4, (0, 1, 2, 3)), |
| 80 | (2, 4, 4, 4, 4, (2,)), |
| 81 | (0, 2, 4, 2, 4, (0, 2)), |
| 82 | (3, 4, 2, 4, 2, (1,)), |
| 83 | ], |
| 84 | ) |
| 85 | def test_mla_dcp_source_ranks( |
| 86 | tp_rank, |
| 87 | tp_size, |
| 88 | remote_tp_size, |
| 89 | dcp_size, |
| 90 | remote_dcp_size, |
| 91 | expected_ranks, |
| 92 | ): |
| 93 | mapping = _compute_mapping( |
| 94 | tp_rank=tp_rank, |
| 95 | tp_size=tp_size, |
| 96 | remote_tp_size=remote_tp_size, |
| 97 | is_mla=True, |
| 98 | num_kv_heads=1, |
| 99 | dcp_size=dcp_size, |
| 100 | remote_dcp_size=remote_dcp_size, |
| 101 | ) |
| 102 | |
| 103 | assert mapping.all_source_ranks == expected_ranks |
| 104 | |
| 105 | |
| 106 | @pytest.mark.parametrize( |
| 107 | "tp_size,remote_tp_size,dcp_size,remote_dcp_size", |
| 108 | [ |
| 109 | (4, 2, 4, 2), |
| 110 | (4, 4, 4, 4), |
| 111 | (4, 1, 4, 1), |
| 112 | (4, 4, 1, 4), |
| 113 | (2, 4, 2, 4), |
| 114 | ], |
| 115 | ) |
| 116 | def test_dcp_consumer_count_matches_readers( |
| 117 | tp_size, remote_tp_size, dcp_size, remote_dcp_size |
| 118 | ): |
| 119 | """Each producer waits for exactly the local ranks that read from it.""" |
| 120 | mappings = [ |
| 121 | _compute_mapping( |
| 122 | tp_rank=tp_rank, |
| 123 | tp_size=tp_size, |
| 124 | remote_tp_size=remote_tp_size, |
| 125 | is_mla=True, |
| 126 | num_kv_heads=1, |
| 127 | dcp_size=dcp_size, |
| 128 | remote_dcp_size=remote_dcp_size, |
| 129 | ) |
| 130 | for tp_rank in range(tp_size) |
| 131 | ] |
| 132 | |
| 133 | for remote_rank in range(remote_tp_size): |
| 134 | readers = [m for m in mappings if remote_rank in m.all_source_ranks] |
| 135 | for mapping in readers: |
| 136 | assert mapping.local_consumers == len(readers) |
| 137 | |
| 138 | |
| 139 | # ====================================================================== |
| 140 | # Split handle tests |
| 141 | # ====================================================================== |
| 142 | |
| 143 | |
| 144 | def _make_mock_worker_for_splits(group_spec_types): |
| 145 | """Build a mock NixlConnectorWorker with _group_spec_types for split tests. |
| 146 | |
| 147 | No per-region replicate flags are configured (``block_len_per_layer`` empty |
| 148 | and ``num_regions == 0``), so ``_fa_desc_replicated`` takes its early-return |
| 149 | path and treats every FA descriptor as SPLIT, matching the legacy behavior |
| 150 | these tests assert. |
| 151 | """ |
| 152 | worker = object.__new__(NixlConnectorWorker) |
| 153 | worker._group_spec_types = group_spec_types |
| 154 | worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=False) |
| 155 | worker.block_len_per_layer = [] |
| 156 | worker.num_regions = 0 |
| 157 | worker._region_is_mla = [] |
| 158 | worker._conv_decomp = SimpleNamespace(local_conv_offsets=()) |
| 159 | worker._ssm_region_indices = [0] if MambaSpec in group_spec_types else [] |
| 160 | worker._ple_region_index = None |
| 161 | return worker |
| 162 | |
| 163 | |
| 164 | class TestBuildSrcSplitHandles: |
| 165 | @pytest.mark.parametrize("remote_tp_size", [2, 4]) |
| 166 | def test_build_src_split_handles(self, remote_tp_size): |
| 167 | tp_rank = 0 |
| 168 | tp_size = 1 |
| 169 | |
| 170 | plan = _compute_mapping( |
| 171 | tp_rank=tp_rank, |
| 172 | tp_size=tp_size, |
| 173 | remote_tp_size=remote_tp_size, |
| 174 | ) |
| 175 | |
| 176 | worker = _make_mock_worker_for_splits((FullAttentionSpec,)) |
| 177 | src_blocks_data = np.array( |
| 178 | [(0x2000 + i * 1024, 1024, 0) for i in range(8)], |
| 179 | dtype=np.uint64, |
| 180 | ) |
| 181 | num_descs = len(src_blocks_data) |
| 182 | splits = list( |
| 183 | worker._build_local_splits_from_plan( |
| 184 | plan, |
| 185 | src_blocks_data, |
| 186 | num_descs, |
| 187 | ) |
| 188 | ) |
| 189 | |
| 190 | assert len(splits) == remote_tp_size |
| 191 | for handle in splits: |
| 192 | assert len(handle) == len(src_blocks_data) |
| 193 | for _, length, _ in handle: |
| 194 | assert length == 1024 // remote_tp_size |
| 195 | |
| 196 | |
| 197 | class TestMambaPlanSplitHandles: |
| 198 | """Verify split handles for Mamba with FA/SSM distinction.""" |
| 199 | |
| 200 | def test_fa_and_ssm_different_split_factors(self): |
| 201 | """Section 0 split by num_attn_reads, section 1 by abs_tp.""" |
| 202 | fa_readers = (0,) |
| 203 | ssm_readers = (0, 1) |
| 204 | plan = TPMapping( |
| 205 | source_ranks_per_group=(fa_readers, ssm_readers), |
| 206 | all_source_ranks=(0, 1), |
| 207 | rank_to_attention_slot={0: 0, 1: 0}, |
| 208 | rank_offset_factor=0, |
| 209 | ) |
| 210 | |
| 211 | worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) |
| 212 | # 2 FA descs + 1 SSM desc |
| 213 | src_blocks_data = np.array( |
| 214 | [ |
| 215 | (1000, 200, 0), # FA desc 0 |
| 216 | (2000, 200, 0), # FA desc 1 |
| 217 | (3000, 400, 0), # SSM desc 0 |
| 218 | ], |
| 219 | dtype=np.uint64, |
| 220 | ) |
| 221 | |
| 222 | splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2)) |
| 223 | |
| 224 | assert len(splits) == 2 # 2 source ranks |
| 225 | |
| 226 | # Rank 0 (FA source, p_idx=0): |
| 227 | # FA: chunk=200//1=200, slot=0 → (1000, 200, 0), (2000, 200, 0) |
| 228 | # SSM: chunk=400//2=200, idx=0 → (3000, 200, 0) |
| 229 | assert splits[0] == [(1000, 200, 0), (2000, 200, 0), (3000, 200, 0)] |
| 230 | |
| 231 | # Rank 1 (not FA source, p_idx=1): |
| 232 | # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0) |
| 233 | # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0) |
| 234 | assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)] |
| 235 | |
| 236 | def test_hetero_block_size_splits(self): |
| 237 | """With a block-size ratio, single-source FA sub-block descs pass |
| 238 | through whole; SSM descs are unexpanded and split per source.""" |
| 239 | plan = TPMapping( |
| 240 | source_ranks_per_group=((0,), (0, 1)), |
| 241 | all_source_ranks=(0, 1), |
| 242 | rank_to_attention_slot={0: 0, 1: 0}, |
| 243 | rank_offset_factor=0, |
| 244 | ) |
| 245 | |
| 246 | worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) |
| 247 | # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded). |
| 248 | src_blocks_data = np.array( |
| 249 | [ |
| 250 | (1000, 100, 0), |
| 251 | (1100, 100, 0), |
| 252 | (2000, 100, 0), |
| 253 | (2100, 100, 0), |
| 254 | (3000, 400, 0), |
| 255 | ], |
| 256 | dtype=np.uint64, |
| 257 | ) |
| 258 | |
| 259 | splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2)) |
| 260 | |
| 261 | assert len(splits) == 2 |
| 262 | fa_passthrough = [ |
| 263 | (1000, 100, 0), |
| 264 | (1100, 100, 0), |
| 265 | (2000, 100, 0), |
| 266 | (2100, 100, 0), |
| 267 | ] |
| 268 | assert splits[0] == fa_passthrough + [(3000, 200, 0)] |
| 269 | assert splits[1] == fa_passthrough + [(3200, 200, 0)] |
| 270 | |
| 271 | def test_hetero_block_size_head_sharded_asserts(self): |
| 272 | """Head-sharded FA reads (multiple FA sources) are incompatible with |
| 273 | a block-size mismatch and must fail loudly.""" |
| 274 | plan = TPMapping( |
| 275 | source_ranks_per_group=((0, 1), (0, 1)), |
| 276 | all_source_ranks=(0, 1), |
| 277 | rank_to_attention_slot={0: 0, 1: 1}, |
| 278 | rank_offset_factor=0, |
| 279 | ) |
| 280 | |
| 281 | worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) |
| 282 | src_blocks_data = np.array( |
| 283 | [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)], |
| 284 | dtype=np.uint64, |
| 285 | ) |
| 286 | |
| 287 | with pytest.raises(AssertionError, match="Head-sharded"): |
| 288 | list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2)) |
| 289 | |
| 290 | |
| 291 | @pytest.mark.parametrize( |
| 292 | ("total_kv_heads", "local_tp", "remote_tp", "compatible"), |
| 293 | [ |
| 294 | # 2 heads: TP1 packs both heads into a page, TP >= 2 packs one. |
| 295 | (2, 1, 2, False), |
| 296 | (2, 2, 1, False), |
| 297 | (2, 1, 1, True), |
| 298 | (2, 2, 4, True), |
| 299 | (2, 4, 2, True), |
| 300 | # 4 heads: the boundary moves with the head count. |
| 301 | (4, 2, 4, False), |
| 302 | (4, 2, 2, True), |
| 303 | (4, 4, 8, True), |
| 304 | ], |
| 305 | ) |
| 306 | def test_csa_linear_tp_layout_boundary(total_kv_heads, local_tp, remote_tp, compatible): |
| 307 | worker = object.__new__(NixlConnectorWorker) |
| 308 | worker.world_size = local_tp |
| 309 | worker._is_csa_linear = True |
| 310 | worker.transfer_topo = SimpleNamespace(total_num_kv_heads=total_kv_heads) |
| 311 | |
| 312 | if compatible: |
| 313 | worker._validate_csa_linear_tp_layout(remote_tp) |
| 314 | else: |
| 315 | with pytest.raises(ValueError, match="KV-head sharding boundary"): |
| 316 | worker._validate_csa_linear_tp_layout(remote_tp) |