Re-enable SM90 FlashInfer allreduce fusion with safe backend defaults (#28789)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user