[Rubin] Enable NVLink all-reduce paths on SM107 (#49647)
Signed-off-by: Zachary Aristei <zaristei@nvidia.com> Co-authored-by: Sahithi Chigurupati <chigurupati.sahithi@gmail.com> Co-authored-by: hlu1 <14827759+hlu1@users.noreply.github.com>
zaristei committed
Jul 29, 2026 at 13:53 UTC
5fa01544480b60275edf913e1a620d229e40553d
3 files changed
+23
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+10
@@ -119,6 +119,11 @@ FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
119
8: 2, # 2MB
120
16: 64, # 64MB (mnnvl multi-node)
121
},
122
+ 107: {
123
+ 2: 64, # 64MB
124
+ 4: 64, # 64MB
125
+ 8: 2, # 2MB
126
+ },
127
}
128
129
# Max size of the input tensor per world size per device capability
@@ -140,6 +145,11 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
145
4: 4, # 4MB
146
8: 2, # 2MB
147
},
148
+ 107: {
149
+ 2: 32, # 32MB
150
+ 4: 4, # 4MB
151
+ 8: 2, # 2MB
152
+ },
153
}
154
155
MiB = 1024 * 1024
vllm/distributed/device_communicators/all_reduce_utils.py
+12
@@ -47,6 +47,12 @@ CUSTOM_ALL_REDUCE_MAX_SIZES = {
47
6: 8 * MiB, # 8 MB
48
8: 4 * MiB, # 4 MB
49
},
50
+ "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds
51
+ 2: 4 * MiB, # 4 MB
52
+ 4: 4 * MiB, # 4 MB
53
+ 6: 8 * MiB, # 8 MB
54
+ 8: 4 * MiB, # 4 MB
55
+ },
56
}
57
58
SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
@@ -68,6 +74,12 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
74
6: 32 * MiB, # 32 MB
75
8: 64 * MiB, # 64 MB
76
},
77
+ "10.7": { # sm_107 (Rubin): reuse 10.3 all-reduce thresholds
78
+ 2: 4 * MiB, # 4 MB
79
+ 4: 32 * MiB, # 32 MB
80
+ 6: 32 * MiB, # 32 MB
81
+ 8: 64 * MiB, # 64 MB
82
+ },
83
}
84
85
# NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks.
vllm/distributed/device_communicators/symm_mem.py
+1
@@ -27,6 +27,7 @@ class SymmMemCommunicator:
27
"9.0": [4, 6, 8],
28
"10.0": [6, 8],
29
"10.3": [6, 8],
30
+ "10.7": [6, 8], # sm_107 (Rubin): reuse 10.3 thresholds
31
}
32
33
def __init__(