[Perf] Enable + tune FlashInfer fused allreduce at world_size=16 on SM 10.3 (GB300) (#46392)

Signed-off-by: Jeff Ma <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

benchmarks/kernels/benchmark_fused_collective.py

@@ -80,13 +80,17 @@ _FI_MAX_SIZES = {
     2: 64 * MiB,  # 64MB
     4: 64 * MiB,  # 64MB
     8: 64 * MiB,  # 64MB
+    16: 64 * MiB,  # 64MB (multi-node)
 }
 
 # Global workspace tensors for FlashInfer (keyed by backend name)
 _FI_WORKSPACES: dict = {}
 
-# Backends to benchmark
-FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
+# Backends to benchmark. trtllm is single-node only and can hang cross-node, so
+# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl.
+FLASHINFER_BACKENDS = [
+    b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b
+]
 
 
 def setup_flashinfer_workspace(
@@ -995,7 +999,10 @@ def main():
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
 
-    device = torch.device(f"cuda:{rank}")
+    # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per
+    # node) map to a valid local GPU; falls back to global rank single-node.
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    device = torch.device(f"cuda:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
 

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -98,11 +98,13 @@ FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
         2: 64,  # 64MB
         4: 32,  # 32MB
         8: 1,  # 1MB
+        16: 64,  # 64MB (mnnvl multi-node)
     },
     103: {
         2: 64,  # 64MB
         4: 64,  # 64MB
         8: 2,  # 2MB
+        16: 64,  # 64MB (mnnvl multi-node)
     },
 }
 

vllm/config/compilation.py

@@ -186,7 +186,7 @@ class PassConfig:
         """
 
         MiB = 1024 * 1024
-        FI_SUPPORTED_WORLD_SIZES = [2, 4, 8]
+        FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16]
         if world_size not in FI_SUPPORTED_WORLD_SIZES:
             return None
         max_size_mb = self.fi_allreduce_fusion_max_size_mb