Re-enable SM90 FlashInfer allreduce fusion with safe backend defaults (#28789)

This commit is contained in:
Mohammad Miadh Angkad
2026-06-23 01:29:19 -07:00
committed by GitHub
parent 854c688121
commit 7b1a20344c
3 changed files with 117 additions and 47 deletions
@@ -35,16 +35,10 @@ _flashinfer_allreduce_unavailable = False
_flashinfer_create_workspace_supports_group = False
_flashinfer_create_workspace_supports_comm_backend = False
_flashinfer_allreduce_supports_trigger_completion = False
_mnnvl_non_blackwell_fallback_logged = False
def _mnnvl_supported(is_multi_node: bool) -> bool:
"""Whether the mnnvl backend is usable on the current system.
mnnvl runs on Blackwell (SM10x) for both single- and multi-node, and on
SM90 for single-node only. Multi-node mnnvl on non-Blackwell is not
supported and must fall back to trtllm.
"""
"""Whether the mnnvl backend is usable on the current system."""
if is_sm100_supported():
return True
return is_sm90_supported() and not is_multi_node
@@ -52,21 +46,31 @@ def _mnnvl_supported(is_multi_node: bool) -> bool:
def _resolve_backend(backend: str, is_multi_node: bool = False) -> str:
"""Resolve the requested FlashInfer allreduce fusion backend."""
global _mnnvl_non_blackwell_fallback_logged
if not (is_sm90_supported() or is_sm100_supported()):
raise ValueError(
"FlashInfer allreduce fusion requires SM90 or SM10X NVIDIA GPUs."
)
if backend == "auto":
# Prefer mnnvl wherever it is supported (any Blackwell system, or SM90
# single-node); fall back to trtllm otherwise.
return "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm"
if is_multi_node:
if is_sm100_supported():
return "mnnvl"
raise ValueError(
"FlashInfer allreduce fusion does not support multi-node on "
"non-Blackwell systems."
)
return "trtllm"
if backend == "trtllm" and is_multi_node:
raise ValueError(
"FlashInfer allreduce fusion trtllm backend supports single-node only."
)
if backend == "mnnvl" and not _mnnvl_supported(is_multi_node):
if not _mnnvl_non_blackwell_fallback_logged:
logger.info(
"FlashInfer allreduce fusion: forcing trtllm backend "
"(mnnvl requires a Blackwell system, or SM90 single-node)."
)
_mnnvl_non_blackwell_fallback_logged = True
return "trtllm"
raise ValueError(
"FlashInfer allreduce fusion mnnvl backend requires a Blackwell "
"system, or SM90 single-node."
)
return backend
@@ -198,11 +202,10 @@ if is_flashinfer_available():
# trtllm | Yes | Yes | Yes | Yes | No |
# mnnvl | Yes | Yes | Single-node | Yes | Blackwell |
#
# mnnvl runs on any Blackwell GPU (SM10x) for both single- and multi-node, and
# on SM90 for single-node only. auto resolves to mnnvl wherever it is supported
# and to trtllm otherwise. An explicit mnnvl request on an unsupported
# configuration (e.g. SM90 multi-node) falls back to trtllm (see
# _resolve_backend).
# FlashInfer allreduce fusion requires SM90 or SM10X. auto resolves to trtllm
# on single-node systems and to mnnvl on Blackwell multi-node systems.
# Non-Blackwell multi-node allreduce fusion is rejected. Explicit mnnvl remains
# available on SM90 single-node systems.
def is_flashinfer_allreduce_unavailable() -> bool:
+17 -9
View File
@@ -2057,7 +2057,17 @@ class ServerArgs:
flashinfer_allreduce_fusion_backend: A[
Optional[Literal["auto", "trtllm", "mnnvl"]],
Arg(
help="Enable FlashInfer allreduce fusion and choose backend. Defaults to auto. 'auto': choose mnnvl on SM90 single-node systems and SM100/SM103 single-node or multi-node systems; choose trtllm otherwise. 'trtllm': available on single-node systems only. 'mnnvl': available on SM90 single-node systems and SM100/SM103 single-node or multi-node systems via MNNVL fabric. Fuses allreduce with Residual + RMSNorm for supported MoE models.",
help=(
"Enable FlashInfer allreduce fusion and choose backend. "
"Requires SM90 or SM10X NVIDIA GPUs. "
"Defaults to auto. "
"'auto': choose trtllm on single-node systems and mnnvl on "
"SM100/SM103 multi-node systems. "
"'trtllm': available on single-node systems only. "
"'mnnvl': available on SM90 single-node systems and SM100/SM103 "
"single-node or multi-node systems via MNNVL fabric. "
"Fuses allreduce with Residual + RMSNorm for supported MoE models."
),
),
] = None
enable_aiter_allreduce_fusion: A[bool, "Enable Aiter AllReduce Fusion."] = False
@@ -4291,13 +4301,10 @@ class ServerArgs:
"Overlap scheduler is disabled when using sparse head for embedding model."
)
# Auto-enable FlashInfer AllReduce Fusion on SM100 only, for models with
# Auto-enable FlashInfer AllReduce Fusion on SM90/SM100, for models with
# explicit support (DeepseekV3, GptOss, Glm4Moe, MistralLarge3,
# Qwen3/Qwen3-VL/Qwen3Next/Qwen3.5 MoE families). SM90 is not
# auto-enabled because auto resolves to mnnvl, which requires a working
# NVLink multicast fabric that SM90 nodes do not reliably have; SM90
# users can opt in explicitly via
# --flashinfer-allreduce-fusion-backend.
# Qwen3/Qwen3-VL/Qwen3Next/Qwen3.5 MoE families). auto resolves to trtllm on
# single-node systems and mnnvl on Blackwell multi-node systems.
if (
self.flashinfer_allreduce_fusion_backend is None
and model_arch
@@ -4319,14 +4326,15 @@ class ServerArgs:
"NemotronHForCausalLM",
"NemotronHPuzzleForCausalLM",
]
and is_sm100_supported()
and (is_sm90_supported() or is_sm100_supported())
and self.tp_size > 1
and not self.enable_dp_attention
and (self.nnodes == 1 or is_sm100_supported())
and self.moe_a2a_backend == "none"
):
self.flashinfer_allreduce_fusion_backend = "auto"
logger.info(
f"Auto-enabling FlashInfer AllReduce Fusion on SM10X for {model_arch}"
f"Auto-enabling FlashInfer AllReduce Fusion on SM90/SM10X for {model_arch}"
)
# Apply enforce_disable_flashinfer_allreduce_fusion after all model-specific adjustments