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

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -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.
+    """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.
-    """
     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
+        if is_multi_node:
-        # single-node); fall back to trtllm otherwise.
+            if is_sm100_supported():
-        return "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm"
+                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:
+        raise ValueError(
-            logger.info(
+            "FlashInfer allreduce fusion mnnvl backend requires a Blackwell "
-                "FlashInfer allreduce fusion: forcing trtllm backend "
+            "system, or SM90 single-node."
-                "(mnnvl requires a Blackwell system, or SM90 single-node)."
         )
-            _mnnvl_non_blackwell_fallback_logged = True
-        return "trtllm"
     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
+# FlashInfer allreduce fusion requires SM90 or SM10X. auto resolves to trtllm
-# on SM90 for single-node only. auto resolves to mnnvl wherever it is supported
+# on single-node systems and to mnnvl on Blackwell multi-node systems.
-# and to trtllm otherwise. An explicit mnnvl request on an unsupported
+# Non-Blackwell multi-node allreduce fusion is rejected. Explicit mnnvl remains
-# configuration (e.g. SM90 multi-node) falls back to trtllm (see
+# available on SM90 single-node systems.
-# _resolve_backend).
 
 
 def is_flashinfer_allreduce_unavailable() -> bool:

python/sglang/srt/server_args.py

@@ -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
+        # Qwen3/Qwen3-VL/Qwen3Next/Qwen3.5 MoE families). auto resolves to trtllm on
-        # auto-enabled because auto resolves to mnnvl, which requires a working
+        # single-node systems and mnnvl on Blackwell multi-node systems.
-        # NVLink multicast fabric that SM90 nodes do not reliably have; SM90
-        # users can opt in explicitly via
-        # --flashinfer-allreduce-fusion-backend.
         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

test/registered/unit/layers/test_flashinfer_comm_fusion.py

@@ -80,36 +80,95 @@ class TestFlashInferCommFusion(unittest.TestCase):
             flashinfer_allreduce_fusion_backend="auto", nnodes=2
         )
 
-        # Blackwell: mnnvl regardless of node count.
+        # Blackwell: trtllm on single-node, mnnvl on multi-node.
         with patch.object(fusion, "is_sm100_supported", return_value=True):
             self.assertEqual(
-                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl"
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node),
+                "trtllm",
             )
             self.assertEqual(
                 fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "mnnvl"
             )
 
-        # SM90: mnnvl on single-node, trtllm fallback on multi-node.
+        # SM90: auto uses trtllm on single-node, multi-node is unsupported.
         with (
             patch.object(fusion, "is_sm100_supported", return_value=False),
             patch.object(fusion, "is_sm90_supported", return_value=True),
-        ):
-            self.assertEqual(
-                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl"
-            )
-            self.assertEqual(
-                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "trtllm"
-            )
-
-        # Pre-SM90: trtllm everywhere.
-        with (
-            patch.object(fusion, "is_sm100_supported", return_value=False),
-            patch.object(fusion, "is_sm90_supported", return_value=False),
         ):
             self.assertEqual(
                 fusion.resolve_flashinfer_allreduce_fusion_backend(single_node),
                 "trtllm",
             )
+            with self.assertRaises(ValueError):
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node)
+
+        # Architectures outside SM90/SM10X are unsupported. Both pre-SM90
+        # and post-SM10X devices (e.g. SM120) must fail closed.
+        for arch in ("pre_sm90", "post_sm10x"):
+            with (
+                self.subTest(arch=arch),
+                patch.object(fusion, "is_sm100_supported", return_value=False),
+                patch.object(fusion, "is_sm90_supported", return_value=False),
+            ):
+                with self.assertRaises(ValueError):
+                    fusion.resolve_flashinfer_allreduce_fusion_backend(single_node)
+                with self.assertRaises(ValueError):
+                    fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node)
+
+    def test_explicit_backend_validation(self):
+        single_node_mnnvl = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="mnnvl", nnodes=1
+        )
+        multi_node_mnnvl = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="mnnvl", nnodes=2
+        )
+        single_node_trtllm = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="trtllm", nnodes=1
+        )
+        multi_node_trtllm = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="trtllm", nnodes=2
+        )
+
+        with (
+            patch.object(fusion, "is_sm100_supported", return_value=False),
+            patch.object(fusion, "is_sm90_supported", return_value=True),
+        ):
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node_mnnvl),
+                "mnnvl",
+            )
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node_trtllm),
+                "trtllm",
+            )
+            with self.assertRaises(ValueError):
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node_mnnvl)
+            with self.assertRaises(ValueError):
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node_trtllm)
+
+        with patch.object(fusion, "is_sm100_supported", return_value=True):
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node_mnnvl),
+                "mnnvl",
+            )
+            with self.assertRaises(ValueError):
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node_trtllm)
+
+        for arch in ("pre_sm90", "post_sm10x"):
+            with (
+                self.subTest(arch=arch),
+                patch.object(fusion, "is_sm100_supported", return_value=False),
+                patch.object(fusion, "is_sm90_supported", return_value=False),
+            ):
+                for args in (
+                    single_node_mnnvl,
+                    multi_node_mnnvl,
+                    single_node_trtllm,
+                    multi_node_trtllm,
+                ):
+                    with self.subTest(backend=args.flashinfer_allreduce_fusion_backend):
+                        with self.assertRaises(ValueError):
+                            fusion.resolve_flashinfer_allreduce_fusion_backend(args)
 
     def test_allreduce_fusion_backends_match_torch_baseline(self):
         fake_comm = _FakeFlashInferComm()