@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1fdd5c42d7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Kernels] Enable Torch Symmetric Memory All-Reduce By Default (#24111)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Ilya Markov committed Sep 11, 2025 at 18:45 UTC 1fdd5c42d7c59d6c73e969bc2c3820ead0528eaf
7 files changed +572 -30
benchmarks/kernels/benchmark_device_communicators.py new
+486
@@ -0,0 +1,486 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +"""
6 +Benchmark script for device communicators:
7 +CustomAllreduce (oneshot, twoshot), PyNcclCommunicator,
8 +and SymmMemCommunicator (multimem, two-shot).
9 +
10 +Usage:
11 + torchrun --nproc_per_node=<N> benchmark_device_communicators.py [options]
12 +
13 +Example:
14 + torchrun --nproc_per_node=2 benchmark_device_communicators.py
15 + --sequence-lengths 512 1024 2048 --num-warmup 10 --num-trials 100
16 +"""
17 +
18 +import json
19 +import os
20 +import time
21 +from contextlib import nullcontext
22 +from typing import Callable, Optional
23 +
24 +import torch
25 +import torch.distributed as dist
26 +from torch.distributed import ProcessGroup
27 +
28 +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
29 +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
30 +from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
31 +from vllm.logger import init_logger
32 +from vllm.utils import FlexibleArgumentParser
33 +
34 +logger = init_logger(__name__)
35 +
36 +# Default sequence lengths to benchmark
37 +DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192]
38 +
39 +# Fixed hidden size and dtype for all benchmarks
40 +HIDDEN_SIZE = 8192
41 +BENCHMARK_DTYPE = torch.bfloat16
42 +
43 +# CUDA graph settings
44 +CUDA_GRAPH_CAPTURE_CYCLES = 10
45 +
46 +
47 +class CommunicatorBenchmark:
48 + """Benchmark class for testing device communicators."""
49 +
50 + def __init__(
51 + self,
52 + rank: int,
53 + world_size: int,
54 + device: torch.device,
55 + cpu_group: ProcessGroup,
56 + sequence_lengths: list[int],
57 + ):
58 + self.rank = rank
59 + self.world_size = world_size
60 + self.device = device
61 + self.cpu_group = cpu_group
62 +
63 + # Calculate max_size_override based on largest sequence length
64 + max_seq_len = max(sequence_lengths)
65 + max_tensor_elements = max_seq_len * HIDDEN_SIZE
66 + self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1
67 +
68 + # Initialize communicators
69 + self.custom_allreduce = None
70 + self.pynccl_comm = None
71 + self.symm_mem_comm = None
72 + self.symm_mem_comm_multimem = None
73 + self.symm_mem_comm_two_shot = None
74 +
75 + self._init_communicators()
76 +
77 + def _init_communicators(self):
78 + """Initialize all available communicators."""
79 + try:
80 + self.custom_allreduce = CustomAllreduce(
81 + group=self.cpu_group,
82 + device=self.device,
83 + max_size=self.max_size_override,
84 + )
85 + if not self.custom_allreduce.disabled:
86 + logger.info("Rank %s: CustomAllreduce initialized", self.rank)
87 + else:
88 + logger.info("Rank %s: CustomAllreduce disabled", self.rank)
89 + except Exception as e:
90 + logger.warning(
91 + "Rank %s: Failed to initialize CustomAllreduce: %s", self.rank, e
92 + )
93 + self.custom_allreduce = None
94 +
95 + try:
96 + self.pynccl_comm = PyNcclCommunicator(
97 + group=self.cpu_group, device=self.device
98 + )
99 + if not self.pynccl_comm.disabled:
100 + logger.info("Rank %s: PyNcclCommunicator initialized", self.rank)
101 + else:
102 + logger.info("Rank %s: PyNcclCommunicator disabled", self.rank)
103 + self.pynccl_comm = None
104 + except Exception as e:
105 + logger.warning(
106 + "Rank %s: Failed to initialize PyNcclCommunicator: %s", self.rank, e
107 + )
108 + self.pynccl_comm = None
109 +
110 + # Initialize variants for SymmMemCommunicator
111 + try:
112 + self.symm_mem_comm_multimem = SymmMemCommunicator(
113 + group=self.cpu_group,
114 + device=self.device,
115 + force_multimem=True,
116 + max_size_override=self.max_size_override,
117 + )
118 + if not self.symm_mem_comm_multimem.disabled:
119 + logger.info(
120 + "Rank %s: SymmMemCommunicator (multimem) initialized", self.rank
121 + )
122 + else:
123 + self.symm_mem_comm_multimem = None
124 + except Exception as e:
125 + logger.warning(
126 + "Rank %s: Failed to initialize SymmMemCommunicator (multimem): %s",
127 + self.rank,
128 + e,
129 + )
130 + self.symm_mem_comm_multimem = None
131 +
132 + try:
133 + self.symm_mem_comm_two_shot = SymmMemCommunicator(
134 + group=self.cpu_group,
135 + device=self.device,
136 + force_multimem=False,
137 + max_size_override=self.max_size_override,
138 + )
139 + if not self.symm_mem_comm_two_shot.disabled:
140 + logger.info(
141 + "Rank %s: SymmMemCommunicator (two_shot) initialized", self.rank
142 + )
143 + else:
144 + self.symm_mem_comm_two_shot = None
145 + except Exception as e:
146 + logger.warning(
147 + "Rank %s: Failed to initialize SymmMemCommunicator (two_shot): %s",
148 + self.rank,
149 + e,
150 + )
151 + self.symm_mem_comm_two_shot = None
152 +
153 + def benchmark_allreduce(
154 + self, sequence_length: int, num_warmup: int, num_trials: int
155 + ) -> dict[str, float]:
156 + """Benchmark allreduce operations for all available communicators."""
157 +
158 + results = {}
159 +
160 + # Define communicators with their benchmark functions
161 + communicators = []
162 +
163 + if self.custom_allreduce is not None:
164 + comm = self.custom_allreduce
165 + # CustomAllreduce one-shot
166 + communicators.append(
167 + (
168 + "ca_1stage",
169 + lambda t, c=comm: c.custom_all_reduce(t),
170 + lambda t, c=comm: c.should_custom_ar(t),
171 + comm.capture(),
172 + "1stage", # env variable value
173 + )
174 + )
175 + # CustomAllreduce two-shot
176 + communicators.append(
177 + (
178 + "ca_2stage",
179 + lambda t, c=comm: c.custom_all_reduce(t),
180 + lambda t, c=comm: c.should_custom_ar(t),
181 + comm.capture(),
182 + "2stage", # env variable value
183 + )
184 + )
185 +
186 + if self.pynccl_comm is not None:
187 + comm = self.pynccl_comm
188 + communicators.append(
189 + (
190 + "pynccl",
191 + lambda t, c=comm: c.all_reduce(t),
192 + lambda t: True, # Always available if initialized
193 + nullcontext(),
194 + None, # no env variable needed
195 + )
196 + )
197 +
198 + if self.symm_mem_comm_multimem is not None:
199 + comm = self.symm_mem_comm_multimem
200 + communicators.append(
201 + (
202 + "symm_mem_multimem",
203 + lambda t, c=comm: c.all_reduce(t),
204 + lambda t, c=comm: c.should_use_symm_mem(t),
205 + nullcontext(),
206 + None, # no env variable needed
207 + )
208 + )
209 +
210 + if self.symm_mem_comm_two_shot is not None:
211 + comm = self.symm_mem_comm_two_shot
212 + communicators.append(
213 + (
214 + "symm_mem_two_shot",
215 + lambda t, c=comm: c.all_reduce(t),
216 + lambda t, c=comm: c.should_use_symm_mem(t),
217 + nullcontext(),
218 + None, # no env variable needed
219 + )
220 + )
221 +
222 + # Benchmark each communicator
223 + for name, allreduce_fn, should_use_fn, context, env_var in communicators:
224 + # Set environment variable if needed
225 + if env_var is not None:
226 + os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var
227 + else:
228 + # Clear the environment variable to avoid interference
229 + os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None)
230 +
231 + latency = self.benchmark_allreduce_single(
232 + sequence_length,
233 + allreduce_fn,
234 + should_use_fn,
235 + context,
236 + num_warmup,
237 + num_trials,
238 + )
239 + if latency is not None:
240 + results[name] = latency
241 +
242 + return results
243 +
244 + def benchmark_allreduce_single(
245 + self,
246 + sequence_length: int,
247 + allreduce_fn: Callable[[torch.Tensor], Optional[torch.Tensor]],
248 + should_use_fn: Callable[[torch.Tensor], bool],
249 + context,
250 + num_warmup: int,
251 + num_trials: int,
252 + ) -> Optional[float]:
253 + """Benchmark method with CUDA graph optimization."""
254 + try:
255 + # Create test tensor (2D: sequence_length x hidden_size)
256 + tensor = torch.randn(
257 + sequence_length, HIDDEN_SIZE, dtype=BENCHMARK_DTYPE, device=self.device
258 + )
259 + if not should_use_fn(tensor):
260 + return None
261 +
262 + torch.cuda.synchronize()
263 + stream = torch.cuda.Stream()
264 + with torch.cuda.stream(stream):
265 + graph_input = tensor.clone()
266 +
267 + # Warmup before capture
268 + for _ in range(3):
269 + allreduce_fn(graph_input)
270 +
271 + # Capture the graph using context manager
272 + with context:
273 + graph = torch.cuda.CUDAGraph()
274 + with torch.cuda.graph(graph):
275 + for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
276 + allreduce_fn(graph_input)
277 +
278 + torch.cuda.synchronize()
279 + for _ in range(num_warmup):
280 + graph.replay()
281 + torch.cuda.synchronize()
282 +
283 + torch.cuda.synchronize()
284 + start_time = time.perf_counter()
285 +
286 + for _ in range(num_trials):
287 + graph.replay()
288 + torch.cuda.synchronize()
289 +
290 + end_time = time.perf_counter()
291 +
292 + # Convert to ms and divide by CUDA_GRAPH_CAPTURE_CYCLES
293 + return (
294 + (end_time - start_time) / num_trials / CUDA_GRAPH_CAPTURE_CYCLES * 1000
295 + )
296 +
297 + except Exception as e:
298 + logger.error("CUDA graph benchmark failed: %s", e)
299 + raise RuntimeError(
300 + f"CUDA graph benchmark failed for communicator: {e}"
301 + ) from e
302 +
303 +
304 +def _calculate_speedup_info(comm_results: dict[str, float]) -> str:
305 + """Calculate speedup information for a single tensor size."""
306 + if not comm_results:
307 + return "N/A"
308 +
309 + # Find the fastest communicator
310 + fastest_comm = min(comm_results.keys(), key=lambda k: comm_results[k])
311 + fastest_time = comm_results[fastest_comm]
312 +
313 + # Calculate speedup vs PyNccl if available
314 + if "pynccl" in comm_results:
315 + pynccl_time = comm_results["pynccl"]
316 + speedup = pynccl_time / fastest_time
317 + return f"{fastest_comm} ({speedup:.2f}x)"
318 + else:
319 + return f"{fastest_comm} (N/A)"
320 +
321 +
322 +def print_results(
323 + results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int
324 +):
325 + """Print benchmark results in a formatted table."""
326 +
327 + print(f"\n{'=' * 130}")
328 + print("Device Communicator Benchmark Results")
329 + print(
330 + f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, "
331 + f"Hidden Size: {HIDDEN_SIZE}"
332 + )
333 + print(f"{'=' * 130}")
334 +
335 + # Get all communicator names
336 + all_comms = set()
337 + for size_results in results.values():
338 + all_comms.update(size_results.keys())
339 +
340 + all_comms = sorted(list(all_comms))
341 +
342 + # Print header
343 + header = f"{'Tensor Shape':<20}{'Tensor Size':<15}"
344 + for comm in all_comms:
345 + header += f"{comm:<20}"
346 + header += f"{'Best (Speedup vs PyNccl)':<30}"
347 + print(header)
348 + print("-" * len(header))
349 +
350 + # Print results for each sequence length
351 + for seq_len in sequence_lengths:
352 + if seq_len in results:
353 + # Calculate tensor size in elements and bytes
354 + tensor_elements = seq_len * HIDDEN_SIZE
355 + tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize
356 +
357 + # Format tensor size (MB)
358 + tensor_size_mb = tensor_bytes / (1024 * 1024)
359 + tensor_size_str = f"{tensor_size_mb:.2f} MB"
360 +
361 + # Format tensor shape
362 + tensor_shape = f"({seq_len}, {HIDDEN_SIZE})"
363 +
364 + row = f"{tensor_shape:<20}{tensor_size_str:<15}"
365 + for comm in all_comms:
366 + if comm in results[seq_len]:
367 + row += f"{results[seq_len][comm]:<20.3f}"
368 + else:
369 + row += f"{'N/A':<20}"
370 +
371 + # Calculate speedup information
372 + speedup_info = _calculate_speedup_info(results[seq_len])
373 + row += f"{speedup_info:<30}"
374 +
375 + print(row)
376 +
377 + print(f"{'=' * 130}")
378 + print("All times are in milliseconds (ms) per allreduce operation")
379 + print("Speedup column shows: fastest_algorithm (speedup_vs_pynccl)")
380 +
381 +
382 +def main():
383 + parser = FlexibleArgumentParser(description="Benchmark device communicators")
384 +
385 + parser.add_argument(
386 + "--sequence-lengths",
387 + type=int,
388 + nargs="+",
389 + default=DEFAULT_SEQUENCE_LENGTHS,
390 + help="Sequence lengths to benchmark (tensor shape: seq_len x hidden_size)",
391 + )
392 +
393 + parser.add_argument(
394 + "--num-warmup", type=int, default=5, help="Number of warmup iterations"
395 + )
396 +
397 + parser.add_argument(
398 + "--num-trials", type=int, default=50, help="Number of benchmark trials"
399 + )
400 +
401 + parser.add_argument("--output-json", type=str, help="Output results to JSON file")
402 +
403 + args = parser.parse_args()
404 +
405 + # Initialize distributed
406 + if not dist.is_initialized():
407 + dist.init_process_group(backend="gloo")
408 + rank = dist.get_rank()
409 + world_size = dist.get_world_size()
410 +
411 + # Set device
412 + device = torch.device(f"cuda:{rank}")
413 + torch.cuda.set_device(device)
414 +
415 + # Get CPU process group
416 + cpu_group = dist.new_group(backend="gloo")
417 +
418 + # Disable USE_SYMM_MEM to avoid affecting the max_sizes
419 + # in symm_mem and custom_all_reduce for benchmark
420 + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
421 +
422 + # Initialize benchmark
423 + benchmark = CommunicatorBenchmark(
424 + rank, world_size, device, cpu_group, args.sequence_lengths
425 + )
426 +
427 + # Run benchmarks
428 + all_results = {}
429 +
430 + for seq_len in args.sequence_lengths:
431 + if rank == 0:
432 + logger.info(
433 + "Benchmarking sequence length: %s (tensor shape: %s x %s)",
434 + seq_len,
435 + seq_len,
436 + HIDDEN_SIZE,
437 + )
438 +
439 + results = benchmark.benchmark_allreduce(
440 + sequence_length=seq_len,
441 + num_warmup=args.num_warmup,
442 + num_trials=args.num_trials,
443 + )
444 +
445 + all_results[seq_len] = results
446 +
447 + # Synchronize between ranks
448 + dist.barrier()
449 +
450 + # Print results (only rank 0)
451 + if rank == 0:
452 + print_results(all_results, args.sequence_lengths, world_size)
453 +
454 + # Save to JSON if requested
455 + if args.output_json:
456 + # Add speedup information to results
457 + enhanced_results = {}
458 + for seq_len, comm_results in all_results.items():
459 + enhanced_results[seq_len] = {
460 + "timings": comm_results,
461 + "speedup_info": _calculate_speedup_info(comm_results),
462 + }
463 +
464 + output_data = {
465 + "world_size": world_size,
466 + "dtype": str(BENCHMARK_DTYPE),
467 + "hidden_size": HIDDEN_SIZE,
468 + "sequence_lengths": args.sequence_lengths,
469 + "num_warmup": args.num_warmup,
470 + "num_trials": args.num_trials,
471 + "cuda_graph_capture_cycles": CUDA_GRAPH_CAPTURE_CYCLES,
472 + "results": enhanced_results,
473 + }
474 +
475 + with open(args.output_json, "w") as f:
476 + json.dump(output_data, f, indent=2)
477 +
478 + logger.info("Results saved to %s", args.output_json)
479 +
480 + # Cleanup
481 + if cpu_group != dist.group.WORLD:
482 + dist.destroy_process_group(cpu_group)
483 +
484 +
485 +if __name__ == "__main__":
486 + main()
csrc/custom_all_reduce.cuh
+40 -13
@@ -15,6 +15,8 @@ typedef __hip_bfloat16 nv_bfloat16;
15 #include <map>
16 #include <unordered_map>
17 #include <vector>
18 +#include <cstdlib>
19 +#include <cstring>
20
21 namespace vllm {
22 #define CUDACHECK(cmd) \
@@ -555,22 +557,47 @@ class CustomAllreduce {
557 size /= d;
558 auto bytes = size * sizeof(typename packed_t<T>::P);
559 int blocks = std::min(block_limit, (size + threads - 1) / threads);
560 +
561 + // Check environment variable once
562 + const char* env_algo = std::getenv("VLLM_CUSTOM_ALLREDUCE_ALGO");
563 + bool force_1stage = false;
564 + bool force_2stage = false;
565 + if (env_algo != nullptr) {
566 + if (std::strcmp(env_algo, "1stage") == 0 ||
567 + std::strcmp(env_algo, "oneshot") == 0) {
568 + force_1stage = true;
569 + } else if (std::strcmp(env_algo, "2stage") == 0 ||
570 + std::strcmp(env_algo, "twoshot") == 0) {
571 + force_2stage = true;
572 + } else {
573 + throw std::runtime_error(
574 + "Invalid VLLM_CUSTOM_ALLREDUCE_ALGO: " + std::string(env_algo) +
575 + ". Valid values: 1stage, oneshot, 2stage, twoshot");
576 + }
577 + }
578 +
579 #define KL(ngpus, name) \
580 name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, \
581 rank_, size);
561 -#define REDUCE_CASE(ngpus) \
562 - case ngpus: { \
563 - if (world_size_ == 2) { \
564 - KL(ngpus, cross_device_reduce_1stage); \
565 - } else if (fully_connected_) { \
566 - if ((world_size_ <= 4 && bytes < 512 * 1024) || \
567 - (world_size_ <= 8 && bytes < 256 * 1024)) { \
568 - KL(ngpus, cross_device_reduce_1stage); \
569 - } else { \
570 - KL(ngpus, cross_device_reduce_2stage); \
571 - } \
572 - } \
573 - break; \
582 +#define REDUCE_CASE(ngpus) \
583 + case ngpus: { \
584 + if (force_1stage) { \
585 + KL(ngpus, cross_device_reduce_1stage); \
586 + } else if (force_2stage) { \
587 + KL(ngpus, cross_device_reduce_2stage); \
588 + } else { \
589 + if (world_size_ == 2) { \
590 + KL(ngpus, cross_device_reduce_1stage); \
591 + } else if (fully_connected_) { \
592 + if ((world_size_ <= 4 && bytes < 512 * 1024) || \
593 + (world_size_ <= 8 && bytes < 256 * 1024)) { \
594 + KL(ngpus, cross_device_reduce_1stage); \
595 + } else { \
596 + KL(ngpus, cross_device_reduce_2stage); \
597 + } \
598 + } \
599 + } \
600 + break; \
601 }
602
603 switch (world_size_) {
vllm/distributed/device_communicators/all_reduce_utils.py
+2 -2
@@ -36,8 +36,8 @@ CUSTOM_ALL_REDUCE_MAX_SIZES = {
36 "10.0": {
37 2: 2 * MiB, # 2 MB
38 4: 2 * MiB, # 2 MB
39 - 6: 2 * MiB, # 2 MB
40 - 8: 2 * MiB, # 2 MB
39 + 6: 1 * MiB, # 1 MB
40 + 8: 1 * MiB, # 1 MB
41 }
42 }
43
vllm/distributed/device_communicators/cuda_communicator.py
+8 -5
@@ -57,11 +57,19 @@ class CudaCommunicator(DeviceCommunicatorBase):
57 self.ca_comm: Optional[CustomAllreduce] = None
58 self.qr_comm: Optional[QuickAllReduce] = None
59 self.symm_mem_comm: Optional[SymmMemCommunicator] = None
60 + if envs.VLLM_ALLREDUCE_USE_SYMM_MEM and current_platform.is_cuda():
61 + self.symm_mem_comm = SymmMemCommunicator(
62 + group=self.cpu_group,
63 + device=self.device,
64 + )
65 +
66 if use_custom_allreduce and self.world_size > 1:
67 # Initialize a custom fast all-reduce implementation.
68 self.ca_comm = CustomAllreduce(
69 group=self.cpu_group,
70 device=self.device,
71 + symm_mem_enabled=(self.symm_mem_comm is not None
72 + and not self.symm_mem_comm.disabled),
73 )
74
75 if current_platform.is_rocm():
@@ -72,11 +80,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
80 # currently be an MI300 series.
81 self.qr_comm = QuickAllReduce(group=self.cpu_group,
82 device=self.device)
75 - if envs.VLLM_ALLREDUCE_USE_SYMM_MEM and current_platform.is_cuda():
76 - self.symm_mem_comm = SymmMemCommunicator(
77 - group=self.cpu_group,
78 - device=self.device,
79 - )
83
84 if self.use_all2all:
85 all2all_backend = envs.VLLM_ALL2ALL_BACKEND
vllm/distributed/device_communicators/custom_all_reduce.py
+3 -2
@@ -54,7 +54,8 @@ class CustomAllreduce:
54 def __init__(self,
55 group: ProcessGroup,
56 device: Union[int, str, torch.device],
57 - max_size=8192 * 1024) -> None:
57 + max_size=8192 * 1024,
58 + symm_mem_enabled=False) -> None:
59 """
60 Args:
61 group: the process group to work on. If None, it will use the
@@ -111,7 +112,7 @@ class CustomAllreduce:
112 self.device = device
113 device_capability = current_platform.get_device_capability(
114 ).as_version_str()
114 - if (current_platform.is_cuda() and envs.VLLM_ALLREDUCE_USE_SYMM_MEM
115 + if (current_platform.is_cuda() and symm_mem_enabled
116 and device_capability in CUSTOM_ALL_REDUCE_MAX_SIZES):
117 max_size = min(
118 CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability][world_size],
vllm/distributed/device_communicators/symm_mem.py
+31 -6
@@ -27,8 +27,13 @@ class SymmMemCommunicator:
27 "10.0": [6, 8],
28 }
29
30 - def __init__(self, group: ProcessGroup, device: Union[int, str,
31 - torch.device]):
30 + def __init__(
31 + self,
32 + group: ProcessGroup,
33 + device: Union[int, str, torch.device],
34 + # add options for testing
35 + force_multimem: Optional[bool] = None,
36 + max_size_override: Optional[int] = None):
37 self.disabled = True
38
39 if not symm_mem_available:
@@ -64,8 +69,17 @@ class SymmMemCommunicator:
69 self.world_size,
70 )
71 return
67 - self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability][
68 - self.world_size]
72 + # Use override max_size if provided, otherwise use default
73 + if max_size_override is not None:
74 + self.max_size = max_size_override
75 + logger.info(
76 + "SymmMemCommunicator: Using override max_size: %s bytes",
77 + self.max_size,
78 + )
79 + else:
80 + self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[
81 + self.device_capability][self.world_size]
82 +
83 self.buffer = torch_symm_mem.empty(
84 self.max_size // self.dtype.itemsize,
85 device=self.device,
@@ -76,6 +90,7 @@ class SymmMemCommunicator:
90 logger.warning("SymmMemCommunicator: symmetric memory "
91 "multicast operations are not supported.")
92 return
93 + self.force_multimem = force_multimem
94 self.disabled = False
95
96 def should_use_symm_mem(self, inp: torch.Tensor):
@@ -98,8 +113,18 @@ class SymmMemCommunicator:
113 if out is None:
114 out = torch.empty_like(inp)
115 self.buffer[:inp.numel()].copy_(inp.view(-1))
101 - if self.world_size in self._WORLD_SIZES_MULTIMEM[
102 - self.device_capability]:
116 +
117 + # Determine which algorithm to use
118 + use_multimem = False
119 + if self.force_multimem is not None:
120 + # Test override: use forced setting
121 + use_multimem = self.force_multimem
122 + else:
123 + # Normal logic: use multimem for supported world sizes
124 + use_multimem = self.world_size in self._WORLD_SIZES_MULTIMEM[
125 + self.device_capability]
126 +
127 + if use_multimem:
128 torch.ops.symm_mem.multimem_all_reduce_(self.buffer[:inp.numel()],
129 "sum",
130 self.group.group_name)
vllm/envs.py
+2 -2
@@ -166,7 +166,7 @@ if TYPE_CHECKING:
166 VLLM_HAS_FLASHINFER_CUBIN: bool = False
167 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False
168 VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
169 - VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False
169 + VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
170 VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
171 VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False
172 VLLM_GPT_OSS_USE_CONTAINER_TOOL: bool = False
@@ -1203,7 +1203,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
1203
1204 # Whether to use pytorch symmetric memory for allreduce
1205 "VLLM_ALLREDUCE_USE_SYMM_MEM":
1206 - lambda: bool(int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "0"))),
1206 + lambda: bool(int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1"))),
1207
1208 # Allows vllm to find tuned config under customized folder
1209 "VLLM_TUNED_CONFIG_FOLDER":