@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 442bd00d9f

[Perf][Distributed] Add low-SM multimem reduce-scatter for SM100/SM103 (#55072)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com>

Yongye Zhu committed Sep 24, 2026 at 20:27 UTC 442bd00d9feff7d6a0cebf74248674bad74ec391
10 files changed +1243 -31
benchmarks/kernels/benchmark_multimem_reduce_scatter.py new
+480
@@ -0,0 +1,480 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark production MNNVL reduce-scatter backends on SM100/SM103.
4 +
5 +The multimem candidate includes the device-to-device copy into its persistent
6 +symmetric input buffer. Timings use pointer-distinct CUDA graphs and report the
7 +median of the per-sample maximum across ranks.
8 +
9 +Run the default BF16 message-size sweep with:
10 +
11 + .venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=8 \
12 + benchmarks/kernels/benchmark_multimem_reduce_scatter.py
13 +"""
14 +
15 +import argparse
16 +import json
17 +import os
18 +import statistics
19 +from collections.abc import Callable
20 +from dataclasses import dataclass
21 +
22 +import torch
23 +import torch.distributed as dist
24 +
25 +import vllm._custom_ops as ops
26 +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
27 +
28 +_DTYPES = {
29 + "bf16": torch.bfloat16,
30 + "f16": torch.float16,
31 + "f32": torch.float32,
32 +}
33 +
34 +
35 +@dataclass
36 +class Candidate:
37 + name: str
38 + runs: list[Callable[[], None]]
39 + outputs: list[torch.Tensor]
40 +
41 +
42 +def parse_size(value: str) -> int:
43 + suffixes = {"k": 1 << 10, "m": 1 << 20, "g": 1 << 30}
44 + value = value.strip().lower()
45 + multiplier = suffixes.get(value[-1], 1)
46 + number = value[:-1] if multiplier != 1 else value
47 + return int(float(number) * multiplier)
48 +
49 +
50 +def parse_args() -> argparse.Namespace:
51 + parser = argparse.ArgumentParser()
52 + parser.add_argument(
53 + "--message-bytes",
54 + type=parse_size,
55 + nargs="+",
56 + default=[
57 + 16 << 20,
58 + 32 << 20,
59 + 64 << 20,
60 + 68 << 20,
61 + 72 << 20,
62 + 76 << 20,
63 + 96 << 20,
64 + ],
65 + help="Full input bytes on each rank. K/M/G suffixes are accepted.",
66 + )
67 + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16")
68 + parser.add_argument(
69 + "--blocks",
70 + type=int,
71 + nargs="+",
72 + default=[8],
73 + help="Maximum CTAs used by each multimem candidate.",
74 + )
75 + parser.add_argument("--num-workspaces", type=int, default=4)
76 + parser.add_argument("--warmup-replays", type=int, default=8)
77 + parser.add_argument("--samples", type=int, default=31)
78 + parser.add_argument(
79 + "--cold-l2",
80 + action="store_true",
81 + help="Flush approximately twice L2 capacity before each replay.",
82 + )
83 + return parser.parse_args()
84 +
85 +
86 +def make_lamport_run(
87 + comm: CustomAllreduce,
88 + input_tensor: torch.Tensor,
89 + output: torch.Tensor,
90 +) -> Callable[[], None]:
91 + def run() -> None:
92 + ops.mnnvl_lamport_reduce_scatter(
93 + comm._ptr,
94 + input_tensor,
95 + output,
96 + comm.mnnvl_lamport_rs_local_ptr,
97 + comm.mnnvl_lamport_rs_epoch_ptr,
98 + comm.mnnvl_buffer_size,
99 + )
100 +
101 + return run
102 +
103 +
104 +def make_multimem_run(
105 + comm: CustomAllreduce,
106 + input_tensor: torch.Tensor,
107 + output: torch.Tensor,
108 + blocks: int,
109 +) -> Callable[[], None]:
110 + def run() -> None:
111 + ops.mnnvl_multimem_reduce_scatter(
112 + comm._ptr,
113 + input_tensor,
114 + output,
115 + comm.mnnvl_multimem_rs_local_ptr,
116 + comm.mnnvl_multimem_rs_multicast_ptr,
117 + comm.mnnvl_multimem_rs_buffer_size,
118 + blocks,
119 + )
120 +
121 + return run
122 +
123 +
124 +def make_nccl_run(
125 + input_tensor: torch.Tensor,
126 + output: torch.Tensor,
127 + device_group: dist.ProcessGroup,
128 +) -> Callable[[], None]:
129 + def run() -> None:
130 + dist.reduce_scatter_single(output, input_tensor, group=device_group)
131 +
132 + return run
133 +
134 +
135 +def build_candidates(
136 + comm: CustomAllreduce,
137 + inputs: list[torch.Tensor],
138 + local_elements: int,
139 + blocks_values: list[int],
140 + device_group: dist.ProcessGroup,
141 +) -> list[Candidate]:
142 + candidates: list[Candidate] = []
143 +
144 + lamport_outputs = [
145 + input_tensor.new_empty(local_elements) for input_tensor in inputs
146 + ]
147 + candidates.append(
148 + Candidate(
149 + "lamport",
150 + [
151 + make_lamport_run(comm, input_tensor, output)
152 + for input_tensor, output in zip(inputs, lamport_outputs)
153 + ],
154 + lamport_outputs,
155 + )
156 + )
157 +
158 + for blocks in blocks_values:
159 + outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs]
160 + candidates.append(
161 + Candidate(
162 + f"multimem_b{blocks}",
163 + [
164 + make_multimem_run(comm, input_tensor, output, blocks)
165 + for input_tensor, output in zip(inputs, outputs)
166 + ],
167 + outputs,
168 + )
169 + )
170 +
171 + nccl_outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs]
172 + candidates.append(
173 + Candidate(
174 + "nccl",
175 + [
176 + make_nccl_run(input_tensor, output, device_group)
177 + for input_tensor, output in zip(inputs, nccl_outputs)
178 + ],
179 + nccl_outputs,
180 + )
181 + )
182 + return candidates
183 +
184 +
185 +def check_correctness(
186 + candidates: list[Candidate],
187 + inputs: list[torch.Tensor],
188 + rank: int,
189 + world_size: int,
190 + cpu_group: dist.ProcessGroup,
191 + device_group: dist.ProcessGroup,
192 + graphs: dict[str, list[torch.cuda.CUDAGraph]] | None = None,
193 +) -> None:
194 + expected_outputs = []
195 + for input_tensor in inputs:
196 + expected = input_tensor.float().clone()
197 + dist.all_reduce(expected, group=device_group)
198 + expected_outputs.append(expected.to(input_tensor.dtype).chunk(world_size)[rank])
199 +
200 + if inputs[0].dtype == torch.float32:
201 + rtol, atol = 1e-5, 1e-5
202 + elif inputs[0].dtype == torch.float16:
203 + rtol, atol = 1e-2, 1e-2
204 + else:
205 + rtol, atol = 5e-2, 1.25e-1
206 +
207 + for candidate in candidates:
208 + for workspace_index, expected in enumerate(expected_outputs):
209 + dist.barrier(group=cpu_group)
210 + if graphs is None:
211 + candidate.runs[workspace_index]()
212 + else:
213 + graphs[candidate.name][workspace_index].replay()
214 + torch.accelerator.synchronize()
215 + error = None
216 + try:
217 + torch.testing.assert_close(
218 + candidate.outputs[workspace_index],
219 + expected,
220 + rtol=rtol,
221 + atol=atol,
222 + )
223 + except AssertionError as assertion_error:
224 + error = assertion_error
225 + correct = torch.tensor(
226 + error is None,
227 + dtype=torch.int32,
228 + device=candidate.outputs[workspace_index].device,
229 + )
230 + dist.all_reduce(correct, op=dist.ReduceOp.MIN, group=device_group)
231 + if not correct.item():
232 + location = f"{candidate.name} workspace {workspace_index}"
233 + detail = str(error) if error is not None else "failed on another rank"
234 + raise AssertionError(
235 + f"{location} produced an incorrect shard: {detail}"
236 + )
237 + dist.barrier(group=cpu_group)
238 +
239 +
240 +def capture_graph(
241 + run: Callable[[], None],
242 + cpu_group: dist.ProcessGroup,
243 +) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]:
244 + stream = torch.cuda.Stream()
245 + stream.wait_stream(torch.cuda.current_stream())
246 + dist.barrier(group=cpu_group)
247 + with torch.cuda.stream(stream):
248 + for _ in range(3):
249 + run()
250 + stream.synchronize()
251 + dist.barrier(group=cpu_group)
252 +
253 + graph = torch.cuda.CUDAGraph()
254 + with torch.cuda.graph(graph, stream=stream):
255 + run()
256 + torch.cuda.current_stream().wait_stream(stream)
257 + dist.barrier(group=cpu_group)
258 + return graph, stream
259 +
260 +
261 +def capture_candidates(
262 + candidates: list[Candidate],
263 + cpu_group: dist.ProcessGroup,
264 +) -> tuple[dict[str, list[torch.cuda.CUDAGraph]], list[object]]:
265 + graphs = {}
266 + keepalive: list[object] = []
267 + for candidate in candidates:
268 + bundles = [capture_graph(run, cpu_group) for run in candidate.runs]
269 + graphs[candidate.name] = [graph for graph, _ in bundles]
270 + keepalive.extend(bundles)
271 + return graphs, keepalive
272 +
273 +
274 +def benchmark_graphs(
275 + graphs: dict[str, list[torch.cuda.CUDAGraph]],
276 + warmup_replays: int,
277 + samples: int,
278 + device_group: dist.ProcessGroup,
279 + device_barrier: Callable[[], None],
280 + flush: torch.Tensor | None,
281 +) -> dict[str, float]:
282 + names = list(graphs)
283 + for round_index in range(warmup_replays):
284 + for candidate_index in range(len(names)):
285 + name = names[(round_index + candidate_index) % len(names)]
286 + device_barrier()
287 + graphs[name][round_index % len(graphs[name])].replay()
288 + torch.accelerator.synchronize()
289 +
290 + timings = {name: [] for name in names}
291 + start = torch.cuda.Event(enable_timing=True)
292 + end = torch.cuda.Event(enable_timing=True)
293 + for sample_index in range(samples):
294 + for candidate_index in range(len(names)):
295 + name = names[(sample_index + candidate_index) % len(names)]
296 + if flush is not None:
297 + flush.zero_()
298 + device_barrier()
299 + start.record()
300 + graphs[name][sample_index % len(graphs[name])].replay()
301 + end.record()
302 + end.synchronize()
303 +
304 + elapsed = torch.tensor(
305 + start.elapsed_time(end) * 1000,
306 + dtype=torch.float64,
307 + device=torch.accelerator.current_device_index(),
308 + )
309 + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
310 + timings[name].append(elapsed.item())
311 +
312 + return {name: statistics.median(values) for name, values in timings.items()}
313 +
314 +
315 +def benchmark_size(
316 + args: argparse.Namespace,
317 + comm: CustomAllreduce,
318 + message_bytes: int,
319 + dtype: torch.dtype,
320 + device: torch.device,
321 + rank: int,
322 + world_size: int,
323 + device_group: dist.ProcessGroup,
324 + cpu_group: dist.ProcessGroup,
325 + device_barrier: Callable[[], None],
326 +) -> dict[str, object]:
327 + element_size = dtype.itemsize
328 + if message_bytes % (element_size * world_size) != 0:
329 + raise ValueError("Message bytes must divide evenly into TP shards")
330 + num_elements = message_bytes // element_size
331 + local_elements = num_elements // world_size
332 + if local_elements * element_size % 16 != 0:
333 + raise ValueError("Each output shard must be a multiple of 16 bytes")
334 +
335 + inputs = []
336 + for workspace_index in range(args.num_workspaces):
337 + generator = torch.Generator(device=device)
338 + generator.manual_seed(1000 + rank * 101 + workspace_index)
339 + input_tensor = torch.empty(num_elements, dtype=dtype, device=device)
340 + input_tensor.uniform_(-1.0, 1.0, generator=generator)
341 + inputs.append(input_tensor)
342 +
343 + candidates = build_candidates(
344 + comm,
345 + inputs,
346 + local_elements,
347 + args.blocks,
348 + device_group,
349 + )
350 + torch.accelerator.synchronize()
351 + dist.barrier(group=cpu_group)
352 + check_correctness(
353 + candidates,
354 + inputs,
355 + rank,
356 + world_size,
357 + cpu_group,
358 + device_group,
359 + )
360 + graphs, graph_keepalive = capture_candidates(candidates, cpu_group)
361 + check_correctness(
362 + candidates,
363 + inputs,
364 + rank,
365 + world_size,
366 + cpu_group,
367 + device_group,
368 + graphs,
369 + )
370 +
371 + flush = None
372 + if args.cold_l2:
373 + l2_bytes = torch.cuda.get_device_properties(device).L2_cache_size
374 + flush = torch.empty(2 * l2_bytes, dtype=torch.uint8, device=device)
375 + times = benchmark_graphs(
376 + graphs,
377 + args.warmup_replays,
378 + args.samples,
379 + device_group,
380 + device_barrier,
381 + flush,
382 + )
383 + graph_keepalive.clear()
384 +
385 + bus_bytes = message_bytes * (world_size - 1) / world_size
386 + fastest = min(times, key=times.__getitem__)
387 + return {
388 + "message_bytes": message_bytes,
389 + "shard_bytes": message_bytes // world_size,
390 + "latency_us": times,
391 + "logical_bus_gbps": {
392 + name: bus_bytes / (latency_us * 1e3) for name, latency_us in times.items()
393 + },
394 + "fastest": fastest,
395 + }
396 +
397 +
398 +def main() -> None:
399 + args = parse_args()
400 + if not args.message_bytes or min(args.message_bytes) <= 0:
401 + raise ValueError("--message-bytes values must be positive")
402 + if not args.blocks or min(args.blocks) <= 0 or max(args.blocks) > 8:
403 + raise ValueError("--blocks values must be in [1, 8]")
404 + if min(args.num_workspaces, args.samples) <= 0:
405 + raise ValueError("Workspace and sample counts must be positive")
406 +
407 + local_rank = int(os.environ["LOCAL_RANK"])
408 + torch.accelerator.set_device_index(local_rank)
409 + dist.init_process_group("nccl")
410 + device_group = dist.group.WORLD
411 + cpu_group = dist.new_group(backend="gloo")
412 + rank = dist.get_rank(device_group)
413 + world_size = dist.get_world_size(device_group)
414 + device = torch.device("cuda", local_rank)
415 + capability = torch.cuda.get_device_capability(device)
416 + if world_size not in (2, 4, 8) or capability not in ((10, 0), (10, 3)):
417 + raise RuntimeError(
418 + "This benchmark targets TP2/TP4/TP8 SM100/SM103, "
419 + f"got TP{world_size} SM{capability}"
420 + )
421 +
422 + max_message_bytes = max(args.message_bytes)
423 + comm = CustomAllreduce(
424 + group=cpu_group,
425 + device=device,
426 + max_mnnvl_reduce_scatter_size=max_message_bytes,
427 + max_mnnvl_multimem_reduce_scatter_size=max_message_bytes,
428 + )
429 + if comm.disabled:
430 + raise RuntimeError("The production custom all-reduce backend is unavailable")
431 + comm._init_mnnvl_multimem_reduce_scatter_buffer()
432 + if not comm.mnnvl_multimem_rs_multicast_ptr:
433 + raise RuntimeError("The production MNNVL multimem backend is unavailable")
434 +
435 + dtype = _DTYPES[args.dtype]
436 + warmup = torch.zeros(1, dtype=torch.float32, device=device)
437 + dist.all_reduce(warmup, group=device_group)
438 + sync = torch.zeros_like(warmup)
439 +
440 + def device_barrier() -> None:
441 + dist.all_reduce(sync, group=device_group)
442 +
443 + results = [
444 + benchmark_size(
445 + args,
446 + comm,
447 + message_bytes,
448 + dtype,
449 + device,
450 + rank,
451 + world_size,
452 + device_group,
453 + cpu_group,
454 + device_barrier,
455 + )
456 + for message_bytes in args.message_bytes
457 + ]
458 + if rank == 0:
459 + metadata = {
460 + "world_size": world_size,
461 + "gpu": torch.cuda.get_device_name(device),
462 + "sm": f"{capability[0]}{capability[1]}",
463 + "num_sms": torch.cuda.get_device_properties(device).multi_processor_count,
464 + "dtype": args.dtype,
465 + "torch": torch.__version__,
466 + "cuda": torch.version.cuda,
467 + "blocks": args.blocks,
468 + "num_workspaces": args.num_workspaces,
469 + "cold_l2": args.cold_l2,
470 + }
471 + print(json.dumps({"metadata": metadata, "results": results}, indent=2))
472 +
473 + comm.close()
474 + dist.barrier(group=cpu_group)
475 + dist.destroy_process_group(cpu_group)
476 + dist.destroy_process_group()
477 +
478 +
479 +if __name__ == "__main__":
480 + main()
csrc/custom_all_gather_reduce_scatter.cuh
+125
@@ -7,6 +7,10 @@ namespace vllm {
7 constexpr int kMnnvlLamportAgThreads = 128;
8 constexpr int kMnnvlLamportRsThreads = 256;
9 constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
10 +constexpr int kMnnvlMultimemRsThreads = 1024;
11 +constexpr int kMnnvlMultimemRsBlockLimit = 8;
12 +constexpr int kMnnvlMultimemRsVectorBytes = 16;
13 +constexpr int kMnnvlMultimemRsUnroll = 8;
14
15 using CopyPack = array_t<uint64_t, 2>;
16
@@ -48,6 +52,127 @@ __global__ void __launch_bounds__(512, 1)
52 barrier_at_end<ngpus, true>(sg, self_sg, rank);
53 }
54
55 +// Multimem reduce-scatter reduces through an NVLS multicast mapping, which
56 +// has no AMD equivalent. The whole path is compiled out on ROCm and the host
57 +// entry point rejects the call instead of launching anything.
58 +#if !defined(USE_ROCM)
59 +
60 +template <typename T>
61 +DINLINE void multimem_load_reduce_16(uint32_t (&result)[4], const T* address) {
62 + #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
63 + if constexpr (std::is_same<T, nv_bfloat16>::value) {
64 + asm volatile(
65 + "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.bf16x2 "
66 + "{%0,%1,%2,%3}, [%4];"
67 + : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
68 + : "l"(address)
69 + : "memory");
70 + } else if constexpr (std::is_same<T, half>::value) {
71 + asm volatile(
72 + "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.f16x2 "
73 + "{%0,%1,%2,%3}, [%4];"
74 + : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
75 + : "l"(address)
76 + : "memory");
77 + } else {
78 + static_assert(std::is_same<T, float>::value);
79 + asm volatile(
80 + "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 "
81 + "{%0,%1,%2,%3}, [%4];"
82 + : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
83 + : "l"(address)
84 + : "memory");
85 + }
86 + #else
87 + asm volatile("trap;");
88 + #endif
89 +}
90 +
91 +DINLINE void store_global_16(void* address, const uint32_t (&value)[4]) {
92 + asm volatile("st.global.v4.u32 [%0], {%1,%2,%3,%4};"
93 + :
94 + : "l"(address), "r"(value[0]), "r"(value[1]), "r"(value[2]),
95 + "r"(value[3])
96 + : "memory");
97 +}
98 +
99 +DINLINE void mnnvl_multimem_publish_flag(FlagType* flag_addr, FlagType flag) {
100 + #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
101 + asm volatile("multimem.st.release.sys.global.u32 [%0], %1;"
102 + :
103 + : "l"(flag_addr), "r"(flag)
104 + : "memory");
105 + // The flag is stored through the multicast alias and polled through the
106 + // local unicast alias in the same grid.
107 + asm volatile("fence.proxy.alias;" ::: "memory");
108 + #else
109 + asm volatile("trap;");
110 + #endif
111 +}
112 +
113 +template <bool ready, int ngpus>
114 +__global__ void mnnvl_multimem_barrier_kernel(Signal* local_signal,
115 + Signal* multicast_signal,
116 + int rank) {
117 + FlagType flag = local_signal->_flag[0] + 1;
118 + FlagType* multicast_counters =
119 + ready ? multicast_signal->start[0] : multicast_signal->end[0];
120 + FlagType* local_counters =
121 + ready ? local_signal->start[0] : local_signal->end[0];
122 + // Replicate this rank's counter into every backing allocation.
123 + mnnvl_multimem_publish_flag(&multicast_counters[rank], flag);
124 + #pragma unroll
125 + for (int peer = 0; peer < ngpus; ++peer) {
126 + while (ld_flag_acquire(&local_counters[peer]) != flag);
127 + }
128 + local_signal->_flag[0] = flag;
129 +}
130 +
131 +template <typename T, int ngpus>
132 +__global__ void __launch_bounds__(kMnnvlMultimemRsThreads, 1)
133 + mnnvl_multimem_reduce_scatter_kernel(const T* __restrict__ multicast_input,
134 + T* __restrict__ result, int rank,
135 + int packs_per_rank) {
136 + constexpr int kThreadsPerWarp = 32;
137 + constexpr int kWarpsPerCta = kMnnvlMultimemRsThreads / kThreadsPerWarp;
138 + constexpr int kPacksPerWarpIteration =
139 + kThreadsPerWarp * kMnnvlMultimemRsUnroll;
140 +
141 + int lane;
142 + asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane));
143 + int warp = blockIdx.x * kWarpsPerCta + threadIdx.x / kThreadsPerWarp;
144 + int num_warps = gridDim.x * kWarpsPerCta;
145 + int pack_offset = warp * kPacksPerWarpIteration + lane;
146 + int pack_stride = num_warps * kPacksPerWarpIteration;
147 + auto* rank_input = reinterpret_cast<const char*>(multicast_input) +
148 + rank * packs_per_rank * kMnnvlMultimemRsVectorBytes;
149 + auto* rank_output = reinterpret_cast<char*>(result);
150 +
151 + while (pack_offset < packs_per_rank) {
152 + uint32_t reduced[kMnnvlMultimemRsUnroll][4];
153 + #pragma unroll
154 + for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) {
155 + int pack = pack_offset + u * kThreadsPerWarp;
156 + if (pack < packs_per_rank) {
157 + multimem_load_reduce_16<T>(
158 + reduced[u], reinterpret_cast<const T*>(
159 + rank_input + pack * kMnnvlMultimemRsVectorBytes));
160 + }
161 + }
162 + #pragma unroll
163 + for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) {
164 + int pack = pack_offset + u * kThreadsPerWarp;
165 + if (pack < packs_per_rank) {
166 + store_global_16(rank_output + pack * kMnnvlMultimemRsVectorBytes,
167 + reduced[u]);
168 + }
169 + }
170 + pack_offset += pack_stride;
171 + }
172 +}
173 +
174 +#endif // !defined(USE_ROCM)
175 +
176 template <typename P>
177 union LamportPack {
178 P packed;
csrc/custom_all_reduce.cuh
+6
@@ -339,6 +339,12 @@ class CustomAllreduce {
339 void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output,
340 void* local_buffer, uint32_t* epochs,
341 int size, int stage_size_bytes);
342 + template <typename T>
343 + void mnnvl_multimem_reduce_scatter(cudaStream_t stream,
344 + const T* multicast_input, T* output,
345 + Signal* local_signal,
346 + Signal* multicast_signal, int size,
347 + int block_limit);
348
349 ~CustomAllreduce() {
350 for (auto [_, ptr] : ipc_handles_) {
csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu
+108
@@ -191,6 +191,55 @@ void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream,
191 #undef MNNVL_LAMPORT_RS_LAUNCH
192 }
193
194 +template <typename T>
195 +void CustomAllreduce::mnnvl_multimem_reduce_scatter(
196 + cudaStream_t stream, const T* multicast_input, T* output,
197 + Signal* local_signal, Signal* multicast_signal, int size, int block_limit) {
198 +#if defined(USE_ROCM)
199 + throw std::runtime_error(
200 + "MNNVL multimem reduce-scatter requires NVLS multicast and is not "
201 + "supported on ROCm");
202 +#else
203 + if (size <= 0)
204 + throw std::runtime_error(
205 + "MNNVL multimem reduce-scatter requires a non-empty input");
206 + int size_bytes = size * sizeof(T);
207 + if (size_bytes % (kMnnvlMultimemRsVectorBytes * world_size_) != 0)
208 + throw std::runtime_error(
209 + "MNNVL multimem reduce-scatter requires each output shard byte size "
210 + "to be a multiple of 16");
211 + if (block_limit <= 0 || block_limit > kMnnvlMultimemRsBlockLimit)
212 + throw std::runtime_error(
213 + "MNNVL multimem reduce-scatter block limit is out of range");
214 +
215 + int packs_per_rank = size_bytes / kMnnvlMultimemRsVectorBytes / world_size_;
216 + int packs_per_block = kMnnvlMultimemRsThreads * kMnnvlMultimemRsUnroll;
217 + int blocks = std::min(
218 + block_limit, (packs_per_rank + packs_per_block - 1) / packs_per_block);
219 +
220 + #define MNNVL_MULTIMEM_RS_CASE(ngpus) \
221 + case ngpus: \
222 + mnnvl_multimem_barrier_kernel<true, ngpus> \
223 + <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \
224 + mnnvl_multimem_reduce_scatter_kernel<T, ngpus> \
225 + <<<blocks, kMnnvlMultimemRsThreads, 0, stream>>>( \
226 + multicast_input, output, rank_, packs_per_rank); \
227 + mnnvl_multimem_barrier_kernel<false, ngpus> \
228 + <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \
229 + break;
230 +
231 + switch (world_size_) {
232 + MNNVL_MULTIMEM_RS_CASE(2)
233 + MNNVL_MULTIMEM_RS_CASE(4)
234 + MNNVL_MULTIMEM_RS_CASE(8)
235 + default:
236 + throw std::runtime_error(
237 + "MNNVL multimem reduce-scatter only supports num gpus in (2,4,8)");
238 + }
239 + #undef MNNVL_MULTIMEM_RS_CASE
240 +#endif
241 +}
242 +
243 } // namespace vllm
244
245 using fptr_t = int64_t;
@@ -360,3 +409,62 @@ void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
409 "bfloat16");
410 }
411 }
412 +
413 +void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
414 + torch::stable::Tensor& out,
415 + fptr_t _local_buffer,
416 + fptr_t _multicast_buffer,
417 + int64_t stage_sz_bytes,
418 + int64_t block_limit) {
419 + auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
420 + const torch::stable::accelerator::DeviceGuard device_guard(
421 + inp.get_device_index());
422 + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
423 +
424 + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
425 + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
426 + STD_TORCH_CHECK(inp.numel() > 0);
427 + STD_TORCH_CHECK(_is_weak_contiguous(out));
428 + STD_TORCH_CHECK(_is_weak_contiguous(inp));
429 + auto input_size = inp.numel() * inp.element_size();
430 + STD_TORCH_CHECK(input_size <= stage_sz_bytes);
431 + auto local_buffer = reinterpret_cast<void*>(_local_buffer);
432 + auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
433 + STD_TORCH_CHECK(local_buffer != nullptr);
434 + STD_TORCH_CHECK(multicast_buffer != nullptr);
435 + auto* local_signal = reinterpret_cast<vllm::Signal*>(
436 + static_cast<char*>(local_buffer) - sizeof(vllm::Signal));
437 + auto* multicast_signal = reinterpret_cast<vllm::Signal*>(
438 + static_cast<char*>(multicast_buffer) - sizeof(vllm::Signal));
439 + STD_CUDA_CHECK(cudaMemcpyAsync(local_buffer, inp.const_data_ptr(), input_size,
440 + cudaMemcpyDeviceToDevice, stream));
441 + switch (out.scalar_type()) {
442 + case torch::headeronly::ScalarType::Float: {
443 + fa->mnnvl_multimem_reduce_scatter<float>(
444 + stream, reinterpret_cast<float*>(multicast_buffer),
445 + reinterpret_cast<float*>(out.mutable_data_ptr()), local_signal,
446 + multicast_signal, inp.numel(), block_limit);
447 + break;
448 + }
449 + case torch::headeronly::ScalarType::Half: {
450 + fa->mnnvl_multimem_reduce_scatter<half>(
451 + stream, reinterpret_cast<half*>(multicast_buffer),
452 + reinterpret_cast<half*>(out.mutable_data_ptr()), local_signal,
453 + multicast_signal, inp.numel(), block_limit);
454 + break;
455 + }
456 +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
457 + case torch::headeronly::ScalarType::BFloat16: {
458 + fa->mnnvl_multimem_reduce_scatter<nv_bfloat16>(
459 + stream, reinterpret_cast<nv_bfloat16*>(multicast_buffer),
460 + reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_signal,
461 + multicast_signal, inp.numel(), block_limit);
462 + break;
463 + }
464 +#endif
465 + default:
466 + throw std::runtime_error(
467 + "MNNVL multimem reduce-scatter only supports float32, float16 and "
468 + "bfloat16");
469 + }
470 +}
csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp
+6
@@ -17,6 +17,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) {
17 custom_ag_rs.def(
18 "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
19 "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()");
20 + custom_ag_rs.def(
21 + "mnnvl_multimem_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
22 + "local_buffer, int multicast_buffer, int stage_sz_bytes, int "
23 + "block_limit) -> ()");
24 }
25
26 STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
@@ -26,4 +30,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
30 custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter));
31 custom_ag_rs.impl("mnnvl_lamport_reduce_scatter",
32 TORCH_BOX(&mnnvl_lamport_reduce_scatter));
33 + custom_ag_rs.impl("mnnvl_multimem_reduce_scatter",
34 + TORCH_BOX(&mnnvl_multimem_reduce_scatter));
35 }
csrc/libtorch_stable/ops.h
+4
@@ -527,6 +527,10 @@ void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
527 torch::stable::Tensor& out,
528 fptr_t local_buffer, fptr_t epoch_buffer,
529 int64_t stage_sz_bytes);
530 +void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
531 + torch::stable::Tensor& out,
532 + fptr_t local_buffer, fptr_t multicast_buffer,
533 + int64_t stage_sz_bytes, int64_t block_limit);
534 void dispose(fptr_t _fa);
535 int64_t meta_size();
536 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
tests/distributed/test_custom_all_gather_reduce_scatter.py
+68
@@ -23,6 +23,7 @@ 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 (
@@ -53,6 +54,11 @@ def _supports_multimem():
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
@@ -324,3 +330,65 @@ 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)
tests/distributed/test_custom_all_reduce.py
+263 -1
@@ -151,9 +151,9 @@ def test_custom_allreduce_filters_dtype(
151 ) -> None:
152 communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
153 communicator.disabled = False
154 + communicator._ptr = 0
155 communicator.world_size = 2
156 communicator.max_size = 1024
156 - communicator._ptr = 0
157
158 assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected
159
@@ -219,6 +219,268 @@ def test_local_multicast_support_rejects_non_cuda(monkeypatch):
219 assert not car._has_local_multicast_support(torch.device("cuda:0"))
220
221
222 +@pytest.mark.parametrize(
223 + ("world_size", "device_capability", "local_multicast", "expected"),
224 + [
225 + (2, (10, 0), True, True),
226 + (4, (10, 3), True, True),
227 + (8, (10, 0), True, True),
228 + (8, (10, 3), True, True),
229 + (6, (10, 3), True, False),
230 + (8, (10, 1), True, False),
231 + (8, (9, 0), True, False),
232 + (8, (10, 3), False, False),
233 + ],
234 +)
235 +def test_mnnvl_multimem_reduce_scatter_platform_gate(
236 + monkeypatch,
237 + world_size,
238 + device_capability,
239 + local_multicast,
240 + expected,
241 +):
242 + def is_device_capability(capability, device_id):
243 + assert capability in ((10, 0), (10, 3))
244 + assert device_id == 3
245 + return device_capability == capability
246 +
247 + monkeypatch.setattr(
248 + car.current_platform,
249 + "is_device_capability",
250 + is_device_capability,
251 + )
252 + monkeypatch.setattr(
253 + car,
254 + "_has_local_multicast_support",
255 + lambda _device: local_multicast,
256 + )
257 +
258 + supported = car._supports_mnnvl_multimem_reduce_scatter(
259 + torch.device("cuda:3"), world_size
260 + )
261 + assert supported is expected
262 +
263 +
264 +@pytest.mark.parametrize(
265 + (
266 + "message_bytes",
267 + "multimem_ptr",
268 + "multimem_initialized",
269 + "batch_invariant",
270 + "expected",
271 + ),
272 + [
273 + (16 * 1024 * 1024, 1, True, False, "mnnvl_lamport"),
274 + (16 * 1024 * 1024 + 128, 1, True, False, "mnnvl_multimem"),
275 + (64 * 1024 * 1024, 1, True, False, "mnnvl_multimem"),
276 + (64 * 1024 * 1024 + 128, 1, True, False, None),
277 + (32 * 1024 * 1024, 0, True, False, None),
278 + (32 * 1024 * 1024, 0, False, False, "mnnvl_multimem"),
279 + (8 * 1024 * 1024, 1, True, True, "mnnvl_lamport"),
280 + (32 * 1024 * 1024, 1, True, True, None),
281 + ],
282 +)
283 +@pytest.mark.parametrize("world_size", [2, 4, 8])
284 +def test_mnnvl_reduce_scatter_backend_gate(
285 + monkeypatch,
286 + world_size,
287 + message_bytes,
288 + multimem_ptr,
289 + multimem_initialized,
290 + batch_invariant,
291 + expected,
292 +):
293 + monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True)
294 + monkeypatch.setattr(car.envs, "VLLM_BATCH_INVARIANT", batch_invariant)
295 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
296 + communicator.disabled = False
297 + communicator._ptr = 0
298 + communicator.world_size = world_size
299 + communicator.mnnvl_only = False
300 + communicator.fully_connected = True
301 + communicator.mnnvl_multicast_ptr = 1
302 + communicator.mnnvl_multimem_rs_supported = True
303 + communicator.mnnvl_multimem_rs_initialized = multimem_initialized
304 + communicator.mnnvl_multimem_rs_multicast_ptr = multimem_ptr
305 + communicator.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024
306 + communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
307 + communicator.max_reduce_scatter_size = 16 * 1024 * 1024
308 + inp = torch.empty(
309 + (world_size, message_bytes // torch.bfloat16.itemsize // world_size),
310 + dtype=torch.bfloat16,
311 + )
312 +
313 + assert inp.nbytes == message_bytes
314 + assert communicator._select_reduce_scatter_backend(inp) == expected
315 + assert communicator.should_custom_reduce_scatter(inp) is (expected is not None)
316 + assert communicator.should_mnnvl_multimem_reduce_scatter(inp) is (
317 + expected == "mnnvl_multimem"
318 + )
319 +
320 +
321 +def test_mnnvl_multimem_reduce_scatter_skips_rendezvous_after_peer_alloc_failure(
322 + monkeypatch,
323 +):
324 + events = []
325 +
326 + class FakeSymmMem:
327 + @staticmethod
328 + def empty(*_args, **_kwargs):
329 + events.append("empty")
330 + return torch.empty(1, dtype=torch.uint8)
331 +
332 + @staticmethod
333 + def rendezvous(*_args, **_kwargs):
334 + events.append("rendezvous")
335 + return None
336 +
337 + def report_peer_allocation_failure(group_value, **_kwargs):
338 + events.append("all_reduce")
339 + assert group_value.item() == 1
340 + group_value.zero_()
341 +
342 + monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
343 + monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
344 + monkeypatch.setattr(car.dist, "all_reduce", report_peer_allocation_failure)
345 + warnings = []
346 + monkeypatch.setattr(
347 + car.logger,
348 + "warning_once",
349 + lambda message, *_args, **_kwargs: warnings.append(message),
350 + )
351 +
352 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
353 + communicator.disabled = True
354 + communicator._ptr = 0
355 + communicator.group = object()
356 + communicator.device = torch.device("cuda:0")
357 + communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
358 + communicator.mnnvl_multimem_rs_supported = True
359 + communicator.mnnvl_multimem_rs_initialized = False
360 + communicator.mnnvl_multimem_rs_buffer = None
361 + communicator.mnnvl_multimem_rs_multicast_ptr = 0
362 +
363 + communicator._init_mnnvl_multimem_reduce_scatter_buffer()
364 +
365 + assert events == ["empty", "all_reduce"]
366 + assert communicator.mnnvl_multimem_rs_initialized
367 + assert communicator.mnnvl_multimem_rs_buffer is None
368 + assert communicator.mnnvl_multimem_rs_multicast_ptr == 0
369 + assert warnings == [
370 + "MNNVL multimem reduce-scatter symmetric-memory allocation "
371 + "failed on at least one rank; falling back to NCCL."
372 + ]
373 +
374 +
375 +def test_mnnvl_multimem_reduce_scatter_warns_on_rendezvous_failure(monkeypatch):
376 + events = []
377 +
378 + class FakeSymmMem:
379 + @staticmethod
380 + def empty(*_args, **_kwargs):
381 + events.append("empty")
382 + return torch.empty(1, dtype=torch.uint8)
383 +
384 + @staticmethod
385 + def rendezvous(*_args, **_kwargs):
386 + events.append("rendezvous")
387 + raise RuntimeError("rendezvous failed")
388 +
389 + def preserve_local_result(_group_value, **_kwargs):
390 + events.append("all_reduce")
391 +
392 + warnings = []
393 + monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
394 + monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
395 + monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result)
396 + monkeypatch.setattr(
397 + car.logger,
398 + "warning_once",
399 + lambda message, *_args, **_kwargs: warnings.append(message),
400 + )
401 +
402 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
403 + communicator.disabled = True
404 + communicator._ptr = 0
405 + communicator.group = type("Group", (), {"group_name": "test"})()
406 + communicator.device = torch.device("cuda:0")
407 + communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
408 + communicator.mnnvl_multimem_rs_supported = True
409 + communicator.mnnvl_multimem_rs_initialized = False
410 + communicator.mnnvl_multimem_rs_buffer = None
411 + communicator.mnnvl_multimem_rs_multicast_ptr = 0
412 +
413 + communicator._init_mnnvl_multimem_reduce_scatter_buffer()
414 +
415 + assert events == ["empty", "all_reduce", "rendezvous", "all_reduce"]
416 + assert communicator.mnnvl_multimem_rs_initialized
417 + assert communicator.mnnvl_multimem_rs_buffer is None
418 + assert communicator.mnnvl_multimem_rs_multicast_ptr == 0
419 + assert warnings == [
420 + "MNNVL multimem reduce-scatter symmetric-memory rendezvous "
421 + "failed on at least one rank; falling back to NCCL."
422 + ]
423 +
424 +
425 +def test_mnnvl_multimem_reduce_scatter_initializes_signals(monkeypatch):
426 + events = []
427 + buffers = []
428 +
429 + class FakeHandle:
430 + multicast_ptr = 0x3000
431 +
432 + class FakeSymmMem:
433 + @staticmethod
434 + def empty(size, **_kwargs):
435 + events.append(("empty", size))
436 + buffer = torch.ones(size, dtype=torch.uint8)
437 + buffers.append(buffer)
438 + return buffer
439 +
440 + @staticmethod
441 + def rendezvous(*_args, **_kwargs):
442 + events.append(("rendezvous", None))
443 + return FakeHandle()
444 +
445 + def preserve_local_result(*_args, **_kwargs):
446 + events.append(("all_reduce", None))
447 +
448 + monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
449 + monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
450 + monkeypatch.setattr(
451 + car.torch.accelerator,
452 + "synchronize",
453 + lambda: events.append(("synchronize", None)),
454 + )
455 + monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result)
456 +
457 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
458 + communicator.disabled = True
459 + communicator._ptr = 0
460 + communicator.group = type("Group", (), {"group_name": "test"})()
461 + communicator.device = torch.device("cpu")
462 + communicator.max_mnnvl_multimem_reduce_scatter_size = 129
463 + communicator.mnnvl_multimem_rs_supported = True
464 + communicator.mnnvl_multimem_rs_initialized = False
465 + communicator.mnnvl_multimem_rs_buffer = None
466 + communicator.mnnvl_multimem_rs_multicast_ptr = 0
467 +
468 + communicator._init_mnnvl_multimem_reduce_scatter_buffer()
469 +
470 + assert events == [
471 + ("empty", 257),
472 + ("all_reduce", None),
473 + ("rendezvous", None),
474 + ("synchronize", None),
475 + ("all_reduce", None),
476 + ]
477 + assert torch.all(buffers[0][:128] == 0)
478 + assert torch.all(buffers[0][128:] == 1)
479 + assert communicator.mnnvl_multimem_rs_buffer_size == 129
480 + assert communicator.mnnvl_multimem_rs_local_ptr == buffers[0].data_ptr() + 128
481 + assert communicator.mnnvl_multimem_rs_multicast_ptr == 0x3080
482 +
483 +
484 @ray.remote(num_gpus=1, max_calls=1)
485 def graph_allreduce(
486 monkeypatch: pytest.MonkeyPatch,
vllm/_custom_ops.py
+20
@@ -3159,6 +3159,26 @@ def mnnvl_lamport_reduce_scatter(
3159 )
3160
3161
3162 +def mnnvl_multimem_reduce_scatter(
3163 + fa: int,
3164 + inp: torch.Tensor,
3165 + out: torch.Tensor,
3166 + local_buffer: int,
3167 + multicast_buffer: int,
3168 + stage_sz_bytes: int,
3169 + block_limit: int,
3170 +) -> None:
3171 + torch.ops._C_custom_ar.mnnvl_multimem_reduce_scatter(
3172 + fa,
3173 + inp,
3174 + out,
3175 + local_buffer,
3176 + multicast_buffer,
3177 + stage_sz_bytes,
3178 + block_limit,
3179 + )
3180 +
3181 +
3182 def dispose(fa: int) -> None:
3183 torch.ops._C_custom_ar.dispose(fa)
3184
vllm/distributed/device_communicators/custom_all_reduce.py
+163 -30
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from contextlib import contextmanager
5 -from typing import cast
5 +from typing import Literal, cast
6
7 import torch
8 import torch.distributed as dist
@@ -32,6 +32,9 @@ except ImportError:
32
33 logger = init_logger(__name__)
34
35 +_ReduceScatterBackend = Literal["legacy", "mnnvl_lamport", "mnnvl_multimem"]
36 +_MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES = (2, 4, 8)
37 +
38
39 def _has_local_multicast_support(device: torch.device) -> bool:
40 """Return whether this CUDA device can allocate multicast memory."""
@@ -53,6 +56,12 @@ def _has_local_multicast_support(device: torch.device) -> bool:
56 return False
57
58
59 +def _all_ranks_true(group: ProcessGroup, local_value: bool) -> bool:
60 + group_value = torch.tensor(int(local_value), dtype=torch.int32, device="cpu")
61 + dist.all_reduce(group_value, op=dist.ReduceOp.MIN, group=group)
62 + return bool(group_value.item())
63 +
64 +
65 def _group_can_attempt_mnnvl(
66 group: ProcessGroup,
67 device: torch.device,
@@ -71,17 +80,7 @@ def _group_can_attempt_mnnvl(
80 and current_platform.has_device_capability(100, device_index)
81 and _has_local_multicast_support(device)
82 )
74 - group_support = torch.tensor(
75 - int(local_support),
76 - dtype=torch.int32,
77 - device="cpu",
78 - )
79 - dist.all_reduce(
80 - group_support,
81 - op=dist.ReduceOp.MIN,
82 - group=group,
83 - )
84 - return bool(group_support.item())
83 + return _all_ranks_true(group, local_support)
84
85
86 def _can_p2p(rank: int, world_size: int) -> bool:
@@ -101,6 +100,20 @@ def _can_p2p(rank: int, world_size: int) -> bool:
100 return True
101
102
103 +def _supports_mnnvl_multimem_reduce_scatter(
104 + device: torch.device, world_size: int
105 +) -> bool:
106 + return (
107 + world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES
108 + and device.index is not None
109 + and (
110 + current_platform.is_device_capability((10, 0), device.index)
111 + or current_platform.is_device_capability((10, 3), device.index)
112 + )
113 + and _has_local_multicast_support(device)
114 + )
115 +
116 +
117 from vllm.distributed.utils import is_weak_contiguous # noqa: E402
118
119
@@ -116,6 +129,8 @@ class CustomAllreduce:
129 }
130 _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
131 _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
132 + _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE = 64 * 1024 * 1024
133 + _MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS = 8
134
135 # max_size: max supported allreduce size
136 def __init__(
@@ -127,6 +142,9 @@ class CustomAllreduce:
142 max_mnnvl_all_gather_size=None,
143 max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE,
144 max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE,
145 + max_mnnvl_multimem_reduce_scatter_size=(
146 + _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE
147 + ),
148 symm_mem_enabled=False,
149 ) -> None:
150 """Args:
@@ -153,6 +171,13 @@ class CustomAllreduce:
171 self.mnnvl_lamport_epochs = None
172 self.mnnvl_lamport_ag_epoch_ptr = 0
173 self.mnnvl_lamport_rs_epoch_ptr = 0
174 + self.mnnvl_multimem_rs_supported = False
175 + self.mnnvl_multimem_rs_initialized = False
176 + self.mnnvl_multimem_rs_buffer = None
177 + self.mnnvl_multimem_rs_handle = None
178 + self.mnnvl_multimem_rs_buffer_size = 0
179 + self.mnnvl_multimem_rs_local_ptr = 0
180 + self.mnnvl_multimem_rs_multicast_ptr = 0
181 self.mnnvl_only = False
182
183 if not custom_ar:
@@ -197,6 +222,9 @@ class CustomAllreduce:
222 # now `device` is a `torch.device` object
223 assert isinstance(device, torch.device)
224 self.device = device
225 + mnnvl_multimem_rs_supported = _supports_mnnvl_multimem_reduce_scatter(
226 + device, world_size
227 + )
228 if not same_node and not _group_can_attempt_mnnvl(group, device):
229 logger.warning(
230 "Custom collectives are disabled because this multi-node "
@@ -294,6 +322,9 @@ class CustomAllreduce:
322 self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size
323 self.max_reduce_scatter_size = max_reduce_scatter_size
324 self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size
325 + self.max_mnnvl_multimem_reduce_scatter_size = (
326 + max_mnnvl_multimem_reduce_scatter_size
327 + )
328 self.rank = rank
329 self.world_size = world_size
330 self.fully_connected = fully_connected
@@ -307,6 +338,11 @@ class CustomAllreduce:
338 max_mnnvl_reduce_scatter_size,
339 )
340 )
341 + if world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES:
342 + self.mnnvl_multimem_rs_supported = _all_ranks_true(
343 + self.group,
344 + mnnvl_multimem_rs_supported and bool(self.mnnvl_multicast_ptr),
345 + )
346 if not same_node and not self.mnnvl_multicast_ptr:
347 logger.warning(
348 "Custom collectives are disabled because this multi-node "
@@ -368,6 +404,59 @@ class CustomAllreduce:
404 except RuntimeError as error:
405 logger.debug("MNNVL AG/RS initialization failed: %s", error)
406
407 + def _init_mnnvl_multimem_reduce_scatter_buffer(self) -> None:
408 + if self.mnnvl_multimem_rs_initialized:
409 + return
410 + self.mnnvl_multimem_rs_initialized = True
411 + if not self.mnnvl_multimem_rs_supported:
412 + return
413 +
414 + assert torch_symm_mem is not None
415 + buffer = None
416 + signal_size = ops.meta_size()
417 + try:
418 + buffer = torch_symm_mem.empty(
419 + signal_size + self.max_mnnvl_multimem_reduce_scatter_size,
420 + dtype=torch.uint8,
421 + device=self.device,
422 + )
423 + except RuntimeError as error:
424 + logger.debug("MNNVL multimem RS allocation failed: %s", error)
425 + if not _all_ranks_true(self.group, buffer is not None):
426 + logger.warning_once(
427 + "MNNVL multimem reduce-scatter symmetric-memory allocation "
428 + "failed on at least one rank; falling back to NCCL.",
429 + scope="global",
430 + )
431 + return
432 + assert buffer is not None
433 +
434 + handle = None
435 + try:
436 + handle = torch_symm_mem.rendezvous(buffer, self.group.group_name)
437 + if handle is not None:
438 + buffer[:signal_size].zero_()
439 + torch.accelerator.synchronize()
440 + except RuntimeError as error:
441 + logger.debug("MNNVL multimem RS rendezvous failed: %s", error)
442 + if not _all_ranks_true(
443 + self.group,
444 + handle is not None and bool(handle.multicast_ptr),
445 + ):
446 + logger.warning_once(
447 + "MNNVL multimem reduce-scatter symmetric-memory rendezvous "
448 + "failed on at least one rank; falling back to NCCL.",
449 + scope="global",
450 + )
451 + return
452 + assert handle is not None
453 +
454 + self.mnnvl_multimem_rs_buffer = buffer
455 + self.mnnvl_multimem_rs_handle = handle
456 + self.mnnvl_multimem_rs_buffer_size = self.max_mnnvl_multimem_reduce_scatter_size
457 + self.mnnvl_multimem_rs_local_ptr = buffer.data_ptr() + signal_size
458 + self.mnnvl_multimem_rs_multicast_ptr = handle.multicast_ptr + signal_size
459 +
460 @contextmanager
461 def capture(self):
462 """The main responsibility of this context manager is the
@@ -509,35 +598,77 @@ class CustomAllreduce:
598 )
599 return out
600
512 - def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
601 + def _select_reduce_scatter_backend(
602 + self, inp: torch.Tensor
603 + ) -> _ReduceScatterBackend | None:
604 if self.disabled or not current_platform.is_cuda():
514 - return False
605 + return None
606 if self.world_size == 16 and not self.mnnvl_only:
516 - return False
607 + return None
608 inp_size = inp.nbytes
609 if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
519 - return False
610 + return None
611 if inp.shape[0] % self.world_size != 0:
521 - return False
612 + return None
613 output_size = inp_size // self.world_size
523 - max_size = (
524 - self.max_mnnvl_reduce_scatter_size
525 - if self.mnnvl_multicast_ptr
526 - else self.max_reduce_scatter_size
527 - )
528 - return (
529 - 0 < inp_size <= max_size
530 - and output_size % 16 == 0
531 - and is_weak_contiguous(inp)
532 - and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
533 - )
614 + if inp_size <= 0 or output_size % 16 != 0 or not is_weak_contiguous(inp):
615 + return None
616 +
617 + if self.mnnvl_multicast_ptr:
618 + if inp_size <= self.max_mnnvl_reduce_scatter_size:
619 + return "mnnvl_lamport"
620 + if (
621 + (
622 + self.mnnvl_multimem_rs_multicast_ptr
623 + or (
624 + self.mnnvl_multimem_rs_supported
625 + and not self.mnnvl_multimem_rs_initialized
626 + )
627 + )
628 + and not envs.VLLM_BATCH_INVARIANT
629 + and inp_size <= self.max_mnnvl_multimem_reduce_scatter_size
630 + ):
631 + return "mnnvl_multimem"
632 + return None
633 +
634 + if self.fully_connected and inp_size <= self.max_reduce_scatter_size:
635 + return "legacy"
636 + return None
637 +
638 + def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
639 + return self._select_reduce_scatter_backend(inp) is not None
640 +
641 + def should_mnnvl_multimem_reduce_scatter(self, inp: torch.Tensor) -> bool:
642 + return self._select_reduce_scatter_backend(inp) == "mnnvl_multimem"
643
644 def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None:
536 - if not self.should_custom_reduce_scatter(inp):
645 + backend = self._select_reduce_scatter_backend(inp)
646 + if backend is None:
647 return None
648 + if backend == "mnnvl_multimem" and not self.mnnvl_multimem_rs_multicast_ptr:
649 + if self._IS_CAPTURING:
650 + return None
651 + self._init_mnnvl_multimem_reduce_scatter_buffer()
652 + if not self.mnnvl_multimem_rs_multicast_ptr:
653 + return None
654 +
655 out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:]
656 out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
540 - if self.mnnvl_multicast_ptr:
657 + if backend == "mnnvl_multimem":
658 + logger.info_once(
659 + "Using the low-SM MNNVL multimem reduce-scatter kernel.",
660 + scope="global",
661 + )
662 + ops.mnnvl_multimem_reduce_scatter(
663 + self._ptr,
664 + inp,
665 + out,
666 + self.mnnvl_multimem_rs_local_ptr,
667 + self.mnnvl_multimem_rs_multicast_ptr,
668 + self.mnnvl_multimem_rs_buffer_size,
669 + self._MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS,
670 + )
671 + elif backend == "mnnvl_lamport":
672 logger.info_once(
673 "Using the MNNVL Lamport reduce-scatter kernel.",
674 scope="global",
@@ -571,6 +702,8 @@ class CustomAllreduce:
702 self.mnnvl_handle = None
703 self.mnnvl_buffer = None
704 self.mnnvl_lamport_epochs = None
705 + self.mnnvl_multimem_rs_handle = None
706 + self.mnnvl_multimem_rs_buffer = None
707
708 def __del__(self):
709 self.close()