[Perf] Enable + tune FlashInfer fused allreduce at world_size=16 on SM 10.3 (GB300) (#46392)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jeff (Junze) Ma committed
Jun 24, 2026 at 23:04 UTC
fc61c6fc26a273e9e0fb3d5054ca1e7c64f645fc
3 files changed
+13
-4
benchmarks/kernels/benchmark_fused_collective.py
+10
-3
@@ -80,13 +80,17 @@ _FI_MAX_SIZES = {
80
2: 64 * MiB, # 64MB
81
4: 64 * MiB, # 64MB
82
8: 64 * MiB, # 64MB
83
+ 16: 64 * MiB, # 64MB (multi-node)
84
}
85
86
# Global workspace tensors for FlashInfer (keyed by backend name)
87
_FI_WORKSPACES: dict = {}
88
88
-# Backends to benchmark
89
-FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
89
+# Backends to benchmark. trtllm is single-node only and can hang cross-node, so
90
+# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl.
91
+FLASHINFER_BACKENDS = [
92
+ b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b
93
+]
94
95
96
def setup_flashinfer_workspace(
@@ -995,7 +999,10 @@ def main():
999
rank = int(os.environ["RANK"])
1000
world_size = int(os.environ["WORLD_SIZE"])
1001
998
- device = torch.device(f"cuda:{rank}")
1002
+ # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per
1003
+ # node) map to a valid local GPU; falls back to global rank single-node.
1004
+ local_rank = int(os.environ.get("LOCAL_RANK", rank))
1005
+ device = torch.device(f"cuda:{local_rank}")
1006
torch.accelerator.set_device_index(device)
1007
torch.set_default_device(device)
1008
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+2
@@ -98,11 +98,13 @@ FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
98
2: 64, # 64MB
99
4: 32, # 32MB
100
8: 1, # 1MB
101
+ 16: 64, # 64MB (mnnvl multi-node)
102
},
103
103: {
104
2: 64, # 64MB
105
4: 64, # 64MB
106
8: 2, # 2MB
107
+ 16: 64, # 64MB (mnnvl multi-node)
108
},
109
}
110
vllm/config/compilation.py
+1
-1
@@ -186,7 +186,7 @@ class PassConfig:
186
"""
187
188
MiB = 1024 * 1024
189
- FI_SUPPORTED_WORLD_SIZES = [2, 4, 8]
189
+ FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16]
190
if world_size not in FI_SUPPORTED_WORLD_SIZES:
191
return None
192
max_size_mb = self.fi_allreduce_fusion_max_size_mb