@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a047e2543d

Fix MNNVL Lamport mailbox publication and cleanup (#53000)

Signed-off-by: Ace Eldeib <aeldeib@coreweave.com> Co-authored-by: Codex <codex@openai.com>

Ace Eldeib committed Aug 24, 2026 at 02:56 UTC a047e2543da570a64d1bbfeac4fe44eff3e87a81
3 files changed +372 -6
.buildkite/test_areas/distributed.yaml
+1
@@ -266,6 +266,7 @@ steps:
266 working_dir: "/vllm-workspace/"
267 num_devices: 2
268 commands:
269 + - pytest -v -s tests/distributed/test_custom_all_gather_reduce_scatter.py
270 - pytest -v -s tests/distributed/test_context_parallel.py
271 - pytest -v -s tests/distributed/test_nccl_symm_mem.py
272 - pytest -v -s tests/v1/distributed/test_dbo.py
csrc/custom_all_gather_reduce_scatter.cuh
+45 -6
@@ -104,6 +104,38 @@ DINLINE P sanitize_lamport_payload(P packed) {
104 return value.packed;
105 }
106
107 +__host__ __device__ constexpr int mnnvl_lamport_dirty_stage(int current_stage) {
108 + return (current_stage + 2) % 3;
109 +}
110 +
111 +__host__ __device__ constexpr int mnnvl_lamport_next_stage(int current_stage) {
112 + return (current_stage + 1) % 3;
113 +}
114 +
115 +template <typename P>
116 +DINLINE void store_multimem_lamport_payload(P* ptr, P packed) {
117 + static_assert(sizeof(P) == 16);
118 + // multimem was introduced in PTX 8.1 (CUDA 12.1) for SM90 and newer.
119 +#if !defined(USE_ROCM) && CUDA_VERSION >= 12010 && defined(__CUDA_ARCH__) && \
120 + (__CUDA_ARCH__ >= 900)
121 + LamportPack<P> value{.packed = packed};
122 + // The local alias remains sentinel until the multicast payload becomes
123 + // observable; readers reject prior or partial values and retry.
124 + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};"
125 + :
126 + : "l"(ptr), "r"(value.words[0]), "r"(value.words[1]),
127 + "r"(value.words[2]), "r"(value.words[3])
128 + : "memory");
129 +#elif defined(USE_ROCM)
130 + __builtin_trap();
131 +#else
132 + // Multicast mappings do not exist before SM90. Fail closed if this kernel is
133 + // ever dispatched for an unsupported target instead of issuing an undefined
134 + // ordinary store to a multicast address.
135 + asm volatile("trap;");
136 +#endif
137 +}
138 +
139 template <typename P>
140 DINLINE P wait_lamport_payload(const P* ptr) {
141 auto value = load_lamport_pack(ptr);
@@ -199,7 +231,9 @@ __global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
231 int stride = gridDim.x * blockDim.x;
232 uint32_t epoch = epochs[0];
233 int current_stage = epoch % 3;
202 - int dirty_stage = (epoch + 1) % 3;
234 + // A peer may start the next epoch after we publish but before we finish.
235 + // Clean the previous stage, which cannot be reused until two epochs later.
236 + int dirty_stage = mnnvl_lamport_dirty_stage(current_stage);
237 int dirty_size = epochs[2 + dirty_stage];
238 auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
239 auto current_local = local_buffer + current_stage * stage_size;
@@ -213,8 +247,11 @@ __global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
247 P local_value;
248 if (tid < size_per_rank) {
249 local_value = packed_input[tid];
216 - current_multicast[rank * size_per_rank + tid] =
217 - sanitize_lamport_payload(local_value);
250 + // A CUDA multicast mapping may only be accessed with multimem PTX;
251 + // ordinary global loads and stores have undefined behavior.
252 + store_multimem_lamport_payload(
253 + current_multicast + rank * size_per_rank + tid,
254 + sanitize_lamport_payload(local_value));
255 }
256 #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
257 (__CUDA_ARCH__ >= 900)
@@ -241,7 +278,7 @@ __global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
278 if (tid == 0) {
279 while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
280 epochs[2 + current_stage] = total_size;
244 - epochs[0] = epoch + 1;
281 + epochs[0] = mnnvl_lamport_next_stage(current_stage);
282 epochs[1] = 0;
283 }
284 }
@@ -267,7 +304,9 @@ __global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
304 int stride = gridDim.x * blockDim.x;
305 uint32_t epoch = epochs[0];
306 int current_stage = epoch % 3;
270 - int dirty_stage = (epoch + 1) % 3;
307 + // A peer may start the next epoch after we publish but before we finish.
308 + // Clean the previous stage, which cannot be reused until two epochs later.
309 + int dirty_stage = mnnvl_lamport_dirty_stage(current_stage);
310 int dirty_size = epochs[2 + dirty_stage];
311 auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
312 auto current_local = local_buffer + current_stage * stage_size;
@@ -319,7 +358,7 @@ __global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
358 if (tid == 0) {
359 while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
360 epochs[2 + current_stage] = size_per_rank * ngpus;
322 - epochs[0] = epoch + 1;
361 + epochs[0] = mnnvl_lamport_next_stage(current_stage);
362 epochs[1] = 0;
363 }
364 }
tests/distributed/test_custom_all_gather_reduce_scatter.py new
+326
@@ -0,0 +1,326 @@
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.platforms import current_platform
27 +
28 +from ..utils import (
29 + ensure_model_parallel_initialized,
30 + init_test_distributed_environment,
31 + multi_process_parallel,
32 +)
33 +
34 +_SENTINEL = -2147483648
35 +_ROWS = 32
36 +_HIDDEN = 8
37 +_NUM_STAGES = 3
38 +_NUM_RANKS = 2
39 +_FAST_RANK = 0
40 +_SLOW_RANK = 1
41 +_PACK_BYTES = 16
42 +_UINT32_MODULUS = 1 << 32
43 +_UINT32_MAX = _UINT32_MODULUS - 1
44 +_ALL_GATHER_EPOCH_ROW = 0
45 +_REDUCE_SCATTER_EPOCH_ROW = 1
46 +_REDUCE_SCATTER_REGION_STAGE_OFFSET = 3
47 +
48 +
49 +def _supports_multimem():
50 + capability = current_platform.get_device_capability()
51 + return (
52 + current_platform.is_cuda() and capability is not None and capability.major >= 9
53 + )
54 +
55 +
56 +def _payload(value, device):
57 + return torch.full((_ROWS, _HIDDEN), value, dtype=torch.bfloat16, device=device)
58 +
59 +
60 +def _payload_bytes(value):
61 + return value.contiguous().view(torch.uint8).flatten()
62 +
63 +
64 +def _int32_bits(value):
65 + """Represent an unsigned 32-bit value in the signed epoch tensor."""
66 + value %= _UINT32_MODULUS
67 + return value if value < (1 << 31) else value - _UINT32_MODULUS
68 +
69 +
70 +def _epoch_value(epoch):
71 + return int(epoch.item()) % _UINT32_MODULUS
72 +
73 +
74 +def _stage_slot(buffer, stage_size, region_offset, stage, src_rank, slot_size):
75 + offset = region_offset + stage * stage_size + src_rank * slot_size
76 + return buffer.narrow(0, offset, slot_size)
77 +
78 +
79 +def _prepare_rank_state(
80 + fa, epoch_row, region_offset, epoch, stage, peer_rank, peer_payload
81 +):
82 + """Reset this rank's replica and seed the peer payload needed to finish."""
83 + buffer = fa.mnnvl_buffer
84 + all_epochs = fa.mnnvl_lamport_epochs
85 + assert buffer is not None and all_epochs is not None
86 + slot_size = peer_payload.nbytes
87 + slot = _stage_slot(
88 + buffer, fa.mnnvl_buffer_size, region_offset, stage, peer_rank, slot_size
89 + )
90 + buffer.view(torch.int32).fill_(_SENTINEL)
91 + epochs = all_epochs[epoch_row]
92 + # epochs[0] selects the stage, epochs[1] is the CTA completion counter,
93 + # and epochs[2:5] record the valid pack count for each stage.
94 + epochs.zero_()
95 + epochs[0] = epoch
96 + epochs[2:5].fill_(slot_size * _NUM_RANKS // _PACK_BYTES)
97 + slot.copy_(_payload_bytes(peer_payload))
98 + return buffer, slot_size
99 +
100 +
101 +def _all_ranks_agree(checks, device):
102 + """Give every rank the same result so one failure cannot strand a peer."""
103 + values = torch.tensor(checks, dtype=torch.int32, device=device)
104 + dist.all_reduce(values, op=dist.ReduceOp.MIN)
105 + return [bool(value) for value in values.cpu().tolist()]
106 +
107 +
108 +def _run_all_gather_retention_case(fa, rank, local_rank, base_epoch, device):
109 + slow_epoch = base_epoch
110 + fast_epoch = base_epoch + 1
111 + slow_stage = slow_epoch % _NUM_STAGES
112 + fast_stage = fast_epoch % _NUM_STAGES
113 + fast_local = _payload(11, device)
114 + fast_peer = _payload(21, device)
115 + slow_peer = _payload(31, device)
116 + slow_local = _payload(41, device)
117 +
118 + if rank == _FAST_RANK:
119 + buffer, slot_size = _prepare_rank_state(
120 + fa,
121 + _ALL_GATHER_EPOCH_ROW,
122 + 0,
123 + fast_epoch,
124 + fast_stage,
125 + _SLOW_RANK,
126 + fast_peer,
127 + )
128 + else:
129 + buffer, slot_size = _prepare_rank_state(
130 + fa,
131 + _ALL_GATHER_EPOCH_ROW,
132 + 0,
133 + slow_epoch,
134 + slow_stage,
135 + _FAST_RANK,
136 + slow_peer,
137 + )
138 + torch.accelerator.synchronize()
139 + dist.barrier(device_ids=[local_rank])
140 +
141 + # Fast rank publishes N+1 through the shipping all-gather kernel.
142 + fast_output_ok = True
143 + if rank == _FAST_RANK:
144 + output = fa.custom_all_gather(fast_local)
145 + assert output is not None
146 + torch.accelerator.synchronize()
147 + fast_output_ok = torch.equal(output, torch.cat((fast_local, fast_peer)))
148 + dist.barrier(device_ids=[local_rank])
149 +
150 + # Slow rank verifies that publication, then runs shipping epoch-N cleanup.
151 + publication_ok = True
152 + slow_output_ok = True
153 + retained = True
154 + if rank == _SLOW_RANK:
155 + target = _stage_slot(
156 + buffer,
157 + fa.mnnvl_buffer_size,
158 + 0,
159 + fast_stage,
160 + _FAST_RANK,
161 + slot_size,
162 + )
163 + publication_ok = torch.equal(target, _payload_bytes(fast_local))
164 + output = fa.custom_all_gather(slow_local)
165 + assert output is not None
166 + torch.accelerator.synchronize()
167 + slow_output_ok = torch.equal(output, torch.cat((slow_peer, slow_local)))
168 + retained = torch.equal(target, _payload_bytes(fast_local))
169 +
170 + checks = _all_ranks_agree(
171 + [fast_output_ok, publication_ok, slow_output_ok, retained], device
172 + )
173 + assert all(checks[:3]), "production all-gather setup did not reach target state"
174 + assert checks[3], "all-gather cleanup erased the next-stage publication"
175 +
176 +
177 +def _run_reduce_scatter_retention_case(fa, rank, local_rank, base_epoch, device):
178 + slow_epoch = base_epoch
179 + fast_epoch = base_epoch + 1
180 + slow_stage = slow_epoch % _NUM_STAGES
181 + fast_stage = fast_epoch % _NUM_STAGES
182 + region_offset = _REDUCE_SCATTER_REGION_STAGE_OFFSET * fa.mnnvl_buffer_size
183 +
184 + fast_chunk_0 = _payload(11, device)
185 + fast_chunk_1 = _payload(12, device)
186 + fast_peer = _payload(21, device)
187 + slow_chunk_0 = _payload(31, device)
188 + slow_chunk_1 = _payload(32, device)
189 + slow_peer = _payload(41, device)
190 + fast_input = torch.cat((fast_chunk_0, fast_chunk_1))
191 + slow_input = torch.cat((slow_chunk_0, slow_chunk_1))
192 +
193 + if rank == _FAST_RANK:
194 + buffer, slot_size = _prepare_rank_state(
195 + fa,
196 + _REDUCE_SCATTER_EPOCH_ROW,
197 + region_offset,
198 + fast_epoch,
199 + fast_stage,
200 + _SLOW_RANK,
201 + fast_peer,
202 + )
203 + else:
204 + buffer, slot_size = _prepare_rank_state(
205 + fa,
206 + _REDUCE_SCATTER_EPOCH_ROW,
207 + region_offset,
208 + slow_epoch,
209 + slow_stage,
210 + _FAST_RANK,
211 + slow_peer,
212 + )
213 + torch.accelerator.synchronize()
214 + dist.barrier(device_ids=[local_rank])
215 +
216 + # Fast rank publishes its N+1 contribution through shipping reduce-scatter.
217 + fast_output_ok = True
218 + if rank == _FAST_RANK:
219 + output = fa.custom_reduce_scatter(fast_input)
220 + assert output is not None
221 + torch.accelerator.synchronize()
222 + fast_output_ok = torch.equal(output, fast_chunk_0 + fast_peer)
223 + dist.barrier(device_ids=[local_rank])
224 +
225 + # Slow rank verifies that contribution, then runs shipping epoch-N cleanup.
226 + publication_ok = True
227 + slow_output_ok = True
228 + retained = True
229 + if rank == _SLOW_RANK:
230 + target = _stage_slot(
231 + buffer,
232 + fa.mnnvl_buffer_size,
233 + region_offset,
234 + fast_stage,
235 + _FAST_RANK,
236 + slot_size,
237 + )
238 + publication_ok = torch.equal(target, _payload_bytes(fast_chunk_1))
239 + output = fa.custom_reduce_scatter(slow_input)
240 + assert output is not None
241 + torch.accelerator.synchronize()
242 + slow_output_ok = torch.equal(output, slow_peer + slow_chunk_1)
243 + retained = torch.equal(target, _payload_bytes(fast_chunk_1))
244 +
245 + checks = _all_ranks_agree(
246 + [fast_output_ok, publication_ok, slow_output_ok, retained], device
247 + )
248 + assert all(checks[:3]), "production reduce-scatter setup did not reach target state"
249 + assert checks[3], "reduce-scatter cleanup erased the next-stage publication"
250 +
251 +
252 +def _run_epoch_wrap_case(fa, rank, local_rank, seed_epoch, device):
253 + """Run both shipping collectives across the uint32 wrap boundary."""
254 + buffer = fa.mnnvl_buffer
255 + all_epochs = fa.mnnvl_lamport_epochs
256 + assert buffer is not None and all_epochs is not None
257 + buffer.view(torch.int32).fill_(_SENTINEL)
258 + all_epochs.zero_()
259 + all_epochs[:, 0] = _int32_bits(seed_epoch)
260 + torch.accelerator.synchronize()
261 + dist.barrier(device_ids=[local_rank])
262 +
263 + ag_input = _payload(rank + 1, device)
264 + ag_output = fa.custom_all_gather(ag_input)
265 + assert ag_output is not None
266 +
267 + rs_input = torch.cat(
268 + (_payload(rank + 1, device), _payload((rank + 1) * 10, device))
269 + )
270 + rs_output = fa.custom_reduce_scatter(rs_input)
271 + assert rs_output is not None
272 + torch.accelerator.synchronize()
273 +
274 + next_stage = (seed_epoch % _NUM_STAGES + 1) % _NUM_STAGES
275 + checks = _all_ranks_agree(
276 + [
277 + torch.equal(
278 + ag_output,
279 + torch.cat((_payload(1, device), _payload(2, device))),
280 + ),
281 + torch.equal(rs_output, _payload(3 if rank == 0 else 30, device)),
282 + _epoch_value(all_epochs[_ALL_GATHER_EPOCH_ROW, 0]) == next_stage,
283 + _epoch_value(all_epochs[_REDUCE_SCATTER_EPOCH_ROW, 0]) == next_stage,
284 + ],
285 + device,
286 + )
287 + assert all(checks), "Lamport stages did not advance safely across uint32 wrap"
288 +
289 +
290 +@ray.remote(num_gpus=1, max_calls=1)
291 +def _run_stage_cleanup_test(
292 + monkeypatch: pytest.MonkeyPatch,
293 + tp_size,
294 + pp_size,
295 + rank,
296 + distributed_init_port,
297 +):
298 + with monkeypatch.context() as m:
299 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
300 + m.delenv("HIP_VISIBLE_DEVICES", raising=False)
301 + m.setenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1")
302 + device = torch.device(f"cuda:{rank}")
303 + torch.accelerator.set_device_index(device)
304 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
305 + ensure_model_parallel_initialized(tp_size, pp_size)
306 +
307 + fa = get_tp_group().device_communicator.ca_comm
308 + assert fa is not None and not fa.disabled
309 + assert fa.mnnvl_multicast_ptr
310 +
311 + for base_epoch in range(_NUM_STAGES):
312 + _run_all_gather_retention_case(fa, rank, rank, base_epoch, device)
313 + _run_reduce_scatter_retention_case(fa, rank, rank, base_epoch, device)
314 + # Old unbounded metadata stores UINT32_MAX and then repeats stage 0.
315 + for seed_epoch in (_UINT32_MAX - 1, _UINT32_MAX):
316 + _run_epoch_wrap_case(fa, rank, rank, seed_epoch, device)
317 +
318 +
319 +@pytest.mark.skipif(
320 + not _supports_multimem(),
321 + reason="MNNVL Lamport collectives require an SM90 or newer NVIDIA GPU.",
322 +)
323 +def test_mnnvl_lamport_stage_cleanup(monkeypatch: pytest.MonkeyPatch):
324 + if torch.accelerator.device_count() < 2:
325 + pytest.skip("Need at least two GPUs to run the test.")
326 + multi_process_parallel(monkeypatch, 2, 1, _run_stage_cleanup_test)