| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Validate production Lamport cleanup from a reachable one-epoch skew. |
| 4 | |
| 5 | The fast rank publishes epoch N+1 before the slow rank runs epoch N cleanup. |
| 6 | Peer slots are prefilled only so each shipping kernel can finish independently; |
| 7 | the test then checks that slow-rank cleanup retained the real N+1 publication. |
| 8 | The production kernels are also seeded at the final two uint32 values to verify |
| 9 | that stage selection remains bounded and never reuses a stage across wrap. |
| 10 | |
| 11 | Each rank's local buffer contains three all-gather stages followed by three |
| 12 | reduce-scatter stages:: |
| 13 | |
| 14 | all-gather: [stage 0][stage 1][stage 2] |
| 15 | reduce-scatter: [stage 0][stage 1][stage 2] |
| 16 | |
| 17 | Each stage contains one payload slot per source rank. |
| 18 | """ |
| 19 | |
| 20 | import pytest |
| 21 | import ray |
| 22 | import torch |
| 23 | import torch.distributed as dist |
| 24 | |
| 25 | from vllm.distributed.parallel_state import get_tp_group |
| 26 | from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter |
| 27 | from vllm.platforms import current_platform |
| 28 | |
| 29 | from ..utils import ( |
| 30 | ensure_model_parallel_initialized, |
| 31 | init_test_distributed_environment, |
| 32 | multi_process_parallel, |
| 33 | ) |
| 34 | |
| 35 | _SENTINEL = -2147483648 |
| 36 | _ROWS = 32 |
| 37 | _HIDDEN = 8 |
| 38 | _NUM_STAGES = 3 |
| 39 | _NUM_RANKS = 2 |
| 40 | _FAST_RANK = 0 |
| 41 | _SLOW_RANK = 1 |
| 42 | _PACK_BYTES = 16 |
| 43 | _UINT32_MODULUS = 1 << 32 |
| 44 | _UINT32_MAX = _UINT32_MODULUS - 1 |
| 45 | _ALL_GATHER_EPOCH_ROW = 0 |
| 46 | _REDUCE_SCATTER_EPOCH_ROW = 1 |
| 47 | _REDUCE_SCATTER_REGION_STAGE_OFFSET = 3 |
| 48 | |
| 49 | |
| 50 | def _supports_multimem(): |
| 51 | capability = current_platform.get_device_capability() |
| 52 | return ( |
| 53 | current_platform.is_cuda() and capability is not None and capability.major >= 9 |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def _supports_sm100_or_sm103(): |
| 58 | capability = current_platform.get_device_capability() |
| 59 | return current_platform.is_cuda() and capability in ((10, 0), (10, 3)) |
| 60 | |
| 61 | |
| 62 | def _payload(value, device): |
| 63 | return torch.full((_ROWS, _HIDDEN), value, dtype=torch.bfloat16, device=device) |
| 64 | |
| 65 | |
| 66 | def _payload_bytes(value): |
| 67 | return value.contiguous().view(torch.uint8).flatten() |
| 68 | |
| 69 | |
| 70 | def _int32_bits(value): |
| 71 | """Represent an unsigned 32-bit value in the signed epoch tensor.""" |
| 72 | value %= _UINT32_MODULUS |
| 73 | return value if value < (1 << 31) else value - _UINT32_MODULUS |
| 74 | |
| 75 | |
| 76 | def _epoch_value(epoch): |
| 77 | return int(epoch.item()) % _UINT32_MODULUS |
| 78 | |
| 79 | |
| 80 | def _stage_slot(buffer, stage_size, region_offset, stage, src_rank, slot_size): |
| 81 | offset = region_offset + stage * stage_size + src_rank * slot_size |
| 82 | return buffer.narrow(0, offset, slot_size) |
| 83 | |
| 84 | |
| 85 | def _prepare_rank_state( |
| 86 | fa, epoch_row, region_offset, epoch, stage, peer_rank, peer_payload |
| 87 | ): |
| 88 | """Reset this rank's replica and seed the peer payload needed to finish.""" |
| 89 | buffer = fa.mnnvl_buffer |
| 90 | all_epochs = fa.mnnvl_lamport_epochs |
| 91 | assert buffer is not None and all_epochs is not None |
| 92 | slot_size = peer_payload.nbytes |
| 93 | slot = _stage_slot( |
| 94 | buffer, fa.mnnvl_buffer_size, region_offset, stage, peer_rank, slot_size |
| 95 | ) |
| 96 | buffer.view(torch.int32).fill_(_SENTINEL) |
| 97 | epochs = all_epochs[epoch_row] |
| 98 | # epochs[0] selects the stage, epochs[1] is the CTA completion counter, |
| 99 | # and epochs[2:5] record the valid pack count for each stage. |
| 100 | epochs.zero_() |
| 101 | epochs[0] = epoch |
| 102 | epochs[2:5].fill_(slot_size * _NUM_RANKS // _PACK_BYTES) |
| 103 | slot.copy_(_payload_bytes(peer_payload)) |
| 104 | return buffer, slot_size |
| 105 | |
| 106 | |
| 107 | def _all_ranks_agree(checks, device): |
| 108 | """Give every rank the same result so one failure cannot strand a peer.""" |
| 109 | values = torch.tensor(checks, dtype=torch.int32, device=device) |
| 110 | dist.all_reduce(values, op=dist.ReduceOp.MIN) |
| 111 | return [bool(value) for value in values.cpu().tolist()] |
| 112 | |
| 113 | |
| 114 | def _run_all_gather_retention_case(fa, rank, local_rank, base_epoch, device): |
| 115 | slow_epoch = base_epoch |
| 116 | fast_epoch = base_epoch + 1 |
| 117 | slow_stage = slow_epoch % _NUM_STAGES |
| 118 | fast_stage = fast_epoch % _NUM_STAGES |
| 119 | fast_local = _payload(11, device) |
| 120 | fast_peer = _payload(21, device) |
| 121 | slow_peer = _payload(31, device) |
| 122 | slow_local = _payload(41, device) |
| 123 | |
| 124 | if rank == _FAST_RANK: |
| 125 | buffer, slot_size = _prepare_rank_state( |
| 126 | fa, |
| 127 | _ALL_GATHER_EPOCH_ROW, |
| 128 | 0, |
| 129 | fast_epoch, |
| 130 | fast_stage, |
| 131 | _SLOW_RANK, |
| 132 | fast_peer, |
| 133 | ) |
| 134 | else: |
| 135 | buffer, slot_size = _prepare_rank_state( |
| 136 | fa, |
| 137 | _ALL_GATHER_EPOCH_ROW, |
| 138 | 0, |
| 139 | slow_epoch, |
| 140 | slow_stage, |
| 141 | _FAST_RANK, |
| 142 | slow_peer, |
| 143 | ) |
| 144 | torch.accelerator.synchronize() |
| 145 | dist.barrier(device_ids=[local_rank]) |
| 146 | |
| 147 | # Fast rank publishes N+1 through the shipping all-gather kernel. |
| 148 | fast_output_ok = True |
| 149 | if rank == _FAST_RANK: |
| 150 | output = fa.custom_all_gather(fast_local) |
| 151 | assert output is not None |
| 152 | torch.accelerator.synchronize() |
| 153 | fast_output_ok = torch.equal(output, torch.cat((fast_local, fast_peer))) |
| 154 | dist.barrier(device_ids=[local_rank]) |
| 155 | |
| 156 | # Slow rank verifies that publication, then runs shipping epoch-N cleanup. |
| 157 | publication_ok = True |
| 158 | slow_output_ok = True |
| 159 | retained = True |
| 160 | if rank == _SLOW_RANK: |
| 161 | target = _stage_slot( |
| 162 | buffer, |
| 163 | fa.mnnvl_buffer_size, |
| 164 | 0, |
| 165 | fast_stage, |
| 166 | _FAST_RANK, |
| 167 | slot_size, |
| 168 | ) |
| 169 | publication_ok = torch.equal(target, _payload_bytes(fast_local)) |
| 170 | output = fa.custom_all_gather(slow_local) |
| 171 | assert output is not None |
| 172 | torch.accelerator.synchronize() |
| 173 | slow_output_ok = torch.equal(output, torch.cat((slow_peer, slow_local))) |
| 174 | retained = torch.equal(target, _payload_bytes(fast_local)) |
| 175 | |
| 176 | checks = _all_ranks_agree( |
| 177 | [fast_output_ok, publication_ok, slow_output_ok, retained], device |
| 178 | ) |
| 179 | assert all(checks[:3]), "production all-gather setup did not reach target state" |
| 180 | assert checks[3], "all-gather cleanup erased the next-stage publication" |
| 181 | |
| 182 | |
| 183 | def _run_reduce_scatter_retention_case(fa, rank, local_rank, base_epoch, device): |
| 184 | slow_epoch = base_epoch |
| 185 | fast_epoch = base_epoch + 1 |
| 186 | slow_stage = slow_epoch % _NUM_STAGES |
| 187 | fast_stage = fast_epoch % _NUM_STAGES |
| 188 | region_offset = _REDUCE_SCATTER_REGION_STAGE_OFFSET * fa.mnnvl_buffer_size |
| 189 | |
| 190 | fast_chunk_0 = _payload(11, device) |
| 191 | fast_chunk_1 = _payload(12, device) |
| 192 | fast_peer = _payload(21, device) |
| 193 | slow_chunk_0 = _payload(31, device) |
| 194 | slow_chunk_1 = _payload(32, device) |
| 195 | slow_peer = _payload(41, device) |
| 196 | fast_input = torch.cat((fast_chunk_0, fast_chunk_1)) |
| 197 | slow_input = torch.cat((slow_chunk_0, slow_chunk_1)) |
| 198 | |
| 199 | if rank == _FAST_RANK: |
| 200 | buffer, slot_size = _prepare_rank_state( |
| 201 | fa, |
| 202 | _REDUCE_SCATTER_EPOCH_ROW, |
| 203 | region_offset, |
| 204 | fast_epoch, |
| 205 | fast_stage, |
| 206 | _SLOW_RANK, |
| 207 | fast_peer, |
| 208 | ) |
| 209 | else: |
| 210 | buffer, slot_size = _prepare_rank_state( |
| 211 | fa, |
| 212 | _REDUCE_SCATTER_EPOCH_ROW, |
| 213 | region_offset, |
| 214 | slow_epoch, |
| 215 | slow_stage, |
| 216 | _FAST_RANK, |
| 217 | slow_peer, |
| 218 | ) |
| 219 | torch.accelerator.synchronize() |
| 220 | dist.barrier(device_ids=[local_rank]) |
| 221 | |
| 222 | # Fast rank publishes its N+1 contribution through shipping reduce-scatter. |
| 223 | fast_output_ok = True |
| 224 | if rank == _FAST_RANK: |
| 225 | output = fa.custom_reduce_scatter(fast_input) |
| 226 | assert output is not None |
| 227 | torch.accelerator.synchronize() |
| 228 | fast_output_ok = torch.equal(output, fast_chunk_0 + fast_peer) |
| 229 | dist.barrier(device_ids=[local_rank]) |
| 230 | |
| 231 | # Slow rank verifies that contribution, then runs shipping epoch-N cleanup. |
| 232 | publication_ok = True |
| 233 | slow_output_ok = True |
| 234 | retained = True |
| 235 | if rank == _SLOW_RANK: |
| 236 | target = _stage_slot( |
| 237 | buffer, |
| 238 | fa.mnnvl_buffer_size, |
| 239 | region_offset, |
| 240 | fast_stage, |
| 241 | _FAST_RANK, |
| 242 | slot_size, |
| 243 | ) |
| 244 | publication_ok = torch.equal(target, _payload_bytes(fast_chunk_1)) |
| 245 | output = fa.custom_reduce_scatter(slow_input) |
| 246 | assert output is not None |
| 247 | torch.accelerator.synchronize() |
| 248 | slow_output_ok = torch.equal(output, slow_peer + slow_chunk_1) |
| 249 | retained = torch.equal(target, _payload_bytes(fast_chunk_1)) |
| 250 | |
| 251 | checks = _all_ranks_agree( |
| 252 | [fast_output_ok, publication_ok, slow_output_ok, retained], device |
| 253 | ) |
| 254 | assert all(checks[:3]), "production reduce-scatter setup did not reach target state" |
| 255 | assert checks[3], "reduce-scatter cleanup erased the next-stage publication" |
| 256 | |
| 257 | |
| 258 | def _run_epoch_wrap_case(fa, rank, local_rank, seed_epoch, device): |
| 259 | """Run both shipping collectives across the uint32 wrap boundary.""" |
| 260 | buffer = fa.mnnvl_buffer |
| 261 | all_epochs = fa.mnnvl_lamport_epochs |
| 262 | assert buffer is not None and all_epochs is not None |
| 263 | buffer.view(torch.int32).fill_(_SENTINEL) |
| 264 | all_epochs.zero_() |
| 265 | all_epochs[:, 0] = _int32_bits(seed_epoch) |
| 266 | torch.accelerator.synchronize() |
| 267 | dist.barrier(device_ids=[local_rank]) |
| 268 | |
| 269 | ag_input = _payload(rank + 1, device) |
| 270 | ag_output = fa.custom_all_gather(ag_input) |
| 271 | assert ag_output is not None |
| 272 | |
| 273 | rs_input = torch.cat( |
| 274 | (_payload(rank + 1, device), _payload((rank + 1) * 10, device)) |
| 275 | ) |
| 276 | rs_output = fa.custom_reduce_scatter(rs_input) |
| 277 | assert rs_output is not None |
| 278 | torch.accelerator.synchronize() |
| 279 | |
| 280 | next_stage = (seed_epoch % _NUM_STAGES + 1) % _NUM_STAGES |
| 281 | checks = _all_ranks_agree( |
| 282 | [ |
| 283 | torch.equal( |
| 284 | ag_output, |
| 285 | torch.cat((_payload(1, device), _payload(2, device))), |
| 286 | ), |
| 287 | torch.equal(rs_output, _payload(3 if rank == 0 else 30, device)), |
| 288 | _epoch_value(all_epochs[_ALL_GATHER_EPOCH_ROW, 0]) == next_stage, |
| 289 | _epoch_value(all_epochs[_REDUCE_SCATTER_EPOCH_ROW, 0]) == next_stage, |
| 290 | ], |
| 291 | device, |
| 292 | ) |
| 293 | assert all(checks), "Lamport stages did not advance safely across uint32 wrap" |
| 294 | |
| 295 | |
| 296 | @ray.remote(num_gpus=1, max_calls=1) |
| 297 | def _run_stage_cleanup_test( |
| 298 | monkeypatch: pytest.MonkeyPatch, |
| 299 | tp_size, |
| 300 | pp_size, |
| 301 | rank, |
| 302 | distributed_init_port, |
| 303 | ): |
| 304 | with monkeypatch.context() as m: |
| 305 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 306 | m.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 307 | m.setenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1") |
| 308 | device = torch.device(f"cuda:{rank}") |
| 309 | torch.accelerator.set_device_index(device) |
| 310 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 311 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 312 | |
| 313 | fa = get_tp_group().device_communicator.ca_comm |
| 314 | assert fa is not None and not fa.disabled |
| 315 | assert fa.mnnvl_multicast_ptr |
| 316 | |
| 317 | for base_epoch in range(_NUM_STAGES): |
| 318 | _run_all_gather_retention_case(fa, rank, rank, base_epoch, device) |
| 319 | _run_reduce_scatter_retention_case(fa, rank, rank, base_epoch, device) |
| 320 | # Old unbounded metadata stores UINT32_MAX and then repeats stage 0. |
| 321 | for seed_epoch in (_UINT32_MAX - 1, _UINT32_MAX): |
| 322 | _run_epoch_wrap_case(fa, rank, rank, seed_epoch, device) |
| 323 | |
| 324 | |
| 325 | @pytest.mark.skipif( |
| 326 | not _supports_multimem(), |
| 327 | reason="MNNVL Lamport collectives require an SM90 or newer NVIDIA GPU.", |
| 328 | ) |
| 329 | def test_mnnvl_lamport_stage_cleanup(monkeypatch: pytest.MonkeyPatch): |
| 330 | if torch.accelerator.device_count() < 2: |
| 331 | pytest.skip("Need at least two GPUs to run the test.") |
| 332 | multi_process_parallel(monkeypatch, 2, 1, _run_stage_cleanup_test) |
| 333 | |
| 334 | |
| 335 | @ray.remote(num_gpus=1, max_calls=1) |
| 336 | def _run_multimem_reduce_scatter_test( |
| 337 | monkeypatch: pytest.MonkeyPatch, |
| 338 | tp_size, |
| 339 | pp_size, |
| 340 | rank, |
| 341 | distributed_init_port, |
| 342 | ): |
| 343 | with monkeypatch.context() as m: |
| 344 | m.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 345 | m.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 346 | device = torch.device(f"cuda:{rank}") |
| 347 | torch.accelerator.set_device_index(device) |
| 348 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 349 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 350 | |
| 351 | fa = get_tp_group().device_communicator.ca_comm |
| 352 | assert fa is not None and not fa.disabled |
| 353 | assert fa.mnnvl_multimem_rs_supported |
| 354 | assert not fa.mnnvl_multimem_rs_initialized |
| 355 | assert not fa.mnnvl_multimem_rs_multicast_ptr |
| 356 | |
| 357 | message_bytes = { |
| 358 | "mnnvl_lamport": 1 * 1024 * 1024, |
| 359 | "mnnvl_multimem": 20 * 1024 * 1024, |
| 360 | "nccl": 64 * 1024 * 1024 + 128, |
| 361 | } |
| 362 | for _ in range(2): |
| 363 | for backend, size_bytes in message_bytes.items(): |
| 364 | shard_elements = size_bytes // torch.bfloat16.itemsize // tp_size |
| 365 | inp = torch.full( |
| 366 | (tp_size, shard_elements), |
| 367 | rank + 1, |
| 368 | dtype=torch.bfloat16, |
| 369 | device=device, |
| 370 | ) |
| 371 | selected = fa._select_reduce_scatter_backend(inp) |
| 372 | if backend == "nccl": |
| 373 | assert selected is None |
| 374 | out = sp_reduce_scatter(inp) |
| 375 | else: |
| 376 | assert selected == backend |
| 377 | out = fa.custom_reduce_scatter(inp) |
| 378 | assert out is not None |
| 379 | if backend == "mnnvl_multimem": |
| 380 | assert fa.mnnvl_multimem_rs_initialized |
| 381 | assert fa.mnnvl_multimem_rs_multicast_ptr |
| 382 | expected = tp_size * (tp_size + 1) // 2 |
| 383 | torch.testing.assert_close(out, torch.full_like(out, expected)) |
| 384 | |
| 385 | |
| 386 | @pytest.mark.skipif( |
| 387 | not _supports_sm100_or_sm103(), |
| 388 | reason="The low-SM MNNVL reduce-scatter path requires SM100 or SM103.", |
| 389 | ) |
| 390 | @pytest.mark.parametrize("tp_size", [2, 4, 8]) |
| 391 | def test_mnnvl_multimem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, tp_size: int): |
| 392 | if torch.accelerator.device_count() < tp_size: |
| 393 | pytest.skip(f"Need at least {tp_size} GPUs to run the test.") |
| 394 | multi_process_parallel(monkeypatch, tp_size, 1, _run_multimem_reduce_scatter_test) |