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 7, 16:30 UTC
[Kernel][perf] optimize NCCL symm_mem vs custom_AR selection thresholds (#33839)
Signed-off-by: <> Signed-off-by: pkousha <43781676+pkousha@users.noreply.github.com> Co-authored-by: Pouya Kousha <pkousha@login-eos01.eos.clusters.nvidia.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
pkousha committed
Feb 26, 2026 at 14:35 UTC
31fb6f43dac735369851c1d908d3d6ed5d6dc1c2
1 file changed
+43
-6
vllm/distributed/device_communicators/all_reduce_utils.py
+43
-6
@@ -27,6 +27,7 @@ from vllm.utils.torch_utils import cuda_device_count_stateless
27
28
logger = init_logger(__name__)
29
30
+KiB = 1024
31
MiB = 1024 * 1024
32
# Max size for each world size in case symmetric memory is available
33
# For different SM architectures
@@ -60,17 +61,44 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
61
},
62
}
63
64
+# NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks.
65
+# PyNCCL-symm outperforms custom_AR for small and large tensor sizes,
66
+# while custom_AR wins for mid-range sizes.
67
+#
68
+# Benchmark results (8 GPUs):
69
+# 2K - 16K: PyNCCL-symm wins (1.35x - 1.48x faster)
70
+# 32K - 64K: custom_AR wins
71
+# 128K - 1G: PyNCCL-symm wins (1.12x - 6.14x faster)
72
+#
73
+# Benchmark results (4 GPUs):
74
+# 2K - 16K: PyNCCL-symm wins (1.21x - 1.30x faster)
75
+# 32K - 256K: custom_AR wins (1.07x - 1.35x faster)
76
+# 512K - 1G: PyNCCL-symm wins (1.10x - 2.32x faster)
77
+#
78
+# The config defines ranges where custom_AR is preferred (symm_mem disabled).
79
NCCL_SYMM_MEM_ALL_REDUCE_CONFIG: dict[str, Any] = {
80
"min_world_size": 4,
65
- "thresholds": {
66
- 4: 2 * MiB, # 2 MB
67
- 8: 1 * MiB, # 1 MB
81
+ # Ranges where custom_AR outperforms NCCL symm_mem: (lower_bound, upper_bound)
82
+ # NCCL symm_mem will NOT be used for sizes in range: lower < size < upper
83
+ "custom_ar_preferred_ranges": {
84
+ 4: (16 * KiB, 512 * KiB), # custom_AR wins for 32K-256K
85
+ 8: (16 * KiB, 128 * KiB), # custom_AR wins for 32K-64K
86
},
87
"always_use_above_world_size": 8, # Always use symm mem for world_size > 8
88
}
89
90
91
def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool:
92
+ """
93
+ Determine if NCCL symmetric memory allreduce should be used.
94
+
95
+ Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for:
96
+ - Small tensors (≤16K): Lower latency than custom_AR
97
+ - Large tensors (≥128K for 8 GPUs, ≥512K for 4 GPUs): Better bandwidth
98
+
99
+ Custom_AR is preferred for mid-range sizes where its P2P approach
100
+ has lower overhead than the symm_mem copy-in/copy-out pattern.
101
+ """
102
from vllm.distributed.device_communicators.pynccl_allocator import (
103
is_symmetric_memory_enabled,
104
)
@@ -80,11 +108,20 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor)
108
109
if not is_symmetric_memory_enabled():
110
return False
111
+
112
if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]:
113
return False
85
- threshold = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["thresholds"].get(world_size)
86
- if threshold is not None and input_tensor.nbytes >= threshold:
87
- return True
114
+
115
+ tensor_size = input_tensor.nbytes
116
+ custom_ar_range = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"].get(
117
+ world_size
118
+ )
119
+
120
+ if custom_ar_range is not None:
121
+ lower_bound, upper_bound = custom_ar_range
122
+ # Use symm_mem for small sizes (≤ lower_bound) and large sizes (≥ upper_bound)
123
+ # Use custom_AR (not symm_mem) for mid-range sizes
124
+ return tensor_size <= lower_bound or tensor_size >= upper_bound
125
return world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
126
127