diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index 36cbd715f18..c999c16021b 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/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) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index d1470029216..ee706037abb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/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) }, } diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index bc38ec6a8a8..4a392a7e3bd 100644 --- a/vllm/config/compilation.py +++ b/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