Reenable MNNVL backend for FlashInfer allreduce fusion (#23402)

python/sglang/srt/layers/communicator.py

@@ -175,7 +175,7 @@ def apply_flashinfer_allreduce_fusion(batch_size: int):
         and batch_size > 0
         and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
         and not is_dp_attention_enabled()
-        and get_global_server_args().enable_flashinfer_allreduce_fusion
+        and get_global_server_args().flashinfer_allreduce_fusion_backend is not None
         and not is_flashinfer_allreduce_unavailable()
     )
 

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -3,6 +3,8 @@ import logging
 from typing import Optional, Tuple
 
 import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
 
 from sglang.srt.distributed import (
     get_attn_tensor_model_parallel_rank,
@@ -16,21 +18,69 @@ from sglang.srt.distributed import (
     get_moe_tp_group,
     get_tp_group,
 )
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     ceil_align,
     get_cuda_driver_bindings,
     is_flashinfer_available,
+    is_sm90_supported,
+    is_sm100_supported,
 )
 from sglang.srt.utils.custom_op import register_custom_op
 
 logger = logging.getLogger(__name__)
 
+# FlashInfer allreduce fusion: set when flashinfer is available (see block below)
 _flashinfer_comm = None
 _TorchDistBackend = None
+_mnnvl_comm_backend = None
+_create_allreduce_fusion_workspace = None
 _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.
+    """
+    if is_sm100_supported():
+        return True
+    return is_sm90_supported() and not is_multi_node
+
+
+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 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 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"
+    return backend
+
+
+def resolve_flashinfer_allreduce_fusion_backend(server_args) -> Optional[str]:
+    backend = getattr(server_args, "flashinfer_allreduce_fusion_backend", None)
+    if backend is None:
+        return None
+    is_multi_node = getattr(server_args, "nnodes", 1) > 1
+    return _resolve_backend(backend, is_multi_node)
 
 
 if is_flashinfer_available():
@@ -41,6 +91,7 @@ if is_flashinfer_available():
             comm, "create_allreduce_fusion_workspace"
         ):
             _flashinfer_comm = comm
+            _create_allreduce_fusion_workspace = comm.create_allreduce_fusion_workspace
             workspace_params = inspect.signature(
                 comm.create_allreduce_fusion_workspace
             ).parameters
@@ -58,11 +109,12 @@ if is_flashinfer_available():
                 "flashinfer.comm unified allreduce_fusion API is not available, "
                 "falling back to standard implementation"
             )
-    except ImportError:
+    except (ImportError, AttributeError) as e:
         _flashinfer_allreduce_unavailable = True
         logger.warning(
-            "flashinfer.comm is not available, falling back to standard "
+            "flashinfer.comm allreduce_fusion API is not available (%s), "
-            "implementation"
+            "falling back to standard implementation",
+            e,
         )
 
     try:
@@ -102,6 +154,61 @@ if is_flashinfer_available():
             "allreduce fusion will use the default process group"
         )
 
+    try:
+        from flashinfer.comm.mnnvl import CommBackend
+
+        class TorchDistributedCommBackend(CommBackend):
+            """
+            Use torch distributed instead of MPI to set up flashinfer MNNVL
+            workspaces during initialization.
+            """
+
+            def __init__(self, group: ProcessGroup):
+                self._group = group
+
+            def Get_rank(self) -> int:
+                return self._group.rank()
+
+            def Get_size(self) -> int:
+                return self._group.size()
+
+            def allgather(self, data: int):
+                gathered = [None] * self.Get_size()
+                dist.all_gather_object(gathered, data, group=self._group)
+                return gathered
+
+            def bcast(self, data, root: int = 0):
+                """Broadcast a picklable Python object from root to all ranks."""
+                obj_list = [data]
+                dist.broadcast_object_list(obj_list, src=root, group=self._group)
+                return obj_list[0]
+
+            def barrier(self):
+                dist.barrier(group=self._group)
+
+            def Split(self, color: int, key: int):
+                # No need to split; we already use the proper group.
+                return self._group
+
+        _mnnvl_comm_backend = TorchDistributedCommBackend
+    except ImportError:
+        _mnnvl_comm_backend = None
+
+
+# FlashInfer allreduce fusion backend support matrix for
+# --flashinfer-allreduce-fusion-backend:
+#
+#   Backend   | SM103 | SM100 | SM90        | Single-Node | Multi-Node |
+#   --------- | ----- | ----- | ----------- | ----------- | ---------- |
+#   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).
+
 
 def is_flashinfer_allreduce_unavailable() -> bool:
     return _flashinfer_allreduce_unavailable
@@ -269,6 +376,11 @@ def _preflight_check_workspace_memory(
 
 
 class FlashInferWorkspaceManager:
+    """
+    Manages FlashInfer's unified allreduce workspace.
+    Supports trtllm and mnnvl backends via create_allreduce_fusion_workspace().
+    """
+
     def __init__(self):
         self.workspace = None
         self.world_size = None
@@ -278,6 +390,10 @@ class FlashInferWorkspaceManager:
         self.hidden_dim = None
         self.dtype = None
         self.initialized = False
+        # Track max sizes ever requested so the workspace only grows (fewer recreates)
+        self._max_token_num_seen: Optional[int] = None
+        self._max_hidden_dim_seen: Optional[int] = None
+        self._logged_init = False
 
     def initialize(
         self,
@@ -285,13 +401,39 @@ class FlashInferWorkspaceManager:
         rank: int,
         max_token_num: int,
         hidden_dim: int,
-        dtype: torch.dtype,
+        backend: str = "auto",
+        group: Optional[ProcessGroup] = None,
+        use_fp32_lamport: bool = False,
+        dtype: Optional[torch.dtype] = None,
         use_oneshot: Optional[bool] = None,
         device_group: Optional["torch.distributed.ProcessGroup"] = None,
         cpu_group: Optional["torch.distributed.ProcessGroup"] = None,
     ):
-        """Initialize workspace"""
+        """Initialize workspace using FlashInfer's unified API."""
-        if _flashinfer_comm is None:
+        global _flashinfer_allreduce_unavailable
+
+        # Track the high-water mark so allocations only grow
+        self._max_token_num_seen = max(max_token_num, self._max_token_num_seen or 0)
+        self._max_hidden_dim_seen = max(hidden_dim, self._max_hidden_dim_seen or 0)
+
+        # Reuse existing workspace if it already covers this problem size
+        if (
+            self.initialized
+            and self.world_size == world_size
+            and self.is_buffer_size_sufficient(
+                token_num=max_token_num,
+                hidden_dim=hidden_dim,
+                dtype=dtype or torch.bfloat16,
+                use_oneshot=use_oneshot,
+            )
+        ):
+            return
+
+        # Same world_size but buffer too small: free old workspace before creating new
+        if self.initialized and self.world_size == world_size:
+            self.cleanup()
+
+        if _flashinfer_comm is None or _create_allreduce_fusion_workspace is None:
             logger.warning(
                 "FlashInfer comm not available, skipping workspace initialization"
             )
@@ -299,7 +441,6 @@ class FlashInferWorkspaceManager:
 
         self.cleanup()
 
-        global _flashinfer_allreduce_unavailable
         if not _preflight_check_workspace_memory(
             world_size=world_size,
             max_token_num=max_token_num,
@@ -312,56 +453,82 @@ class FlashInferWorkspaceManager:
             self.initialized = False
             return
 
-        try:
+        # Determine GPUs per node for MNNVL topology detection
-            kwargs = dict(
+        gpus_per_node = None
-                backend="trtllm",
+        node_pg = cpu_group if cpu_group is not None else group
-                world_size=world_size,
+        if node_pg is not None:
-                rank=rank,
+            gpus_per_node = sum(in_the_same_node_as(node_pg, source_rank=0))
-                max_token_num=max_token_num,
+        comm_backend = None
-                hidden_dim=hidden_dim,
-                dtype=dtype,
-                force_oneshot_support=bool(use_oneshot),
-            )
-            create_workspace = _flashinfer_comm.create_allreduce_fusion_workspace
-            if _flashinfer_create_workspace_supports_group:
-                # Pin the symmetric-memory rendezvous to the actual subgroup.
-                # Older FlashInfer releases only support comm_backend.
-                kwargs["group"] = device_group
         if (
             _TorchDistBackend is not None
-                and _flashinfer_create_workspace_supports_comm_backend
             and device_group is not None
             and cpu_group is not None
         ):
-                kwargs["comm_backend"] = _TorchDistBackend(
+            comm_backend = _TorchDistBackend(
                 device_group=device_group, cpu_group=cpu_group
             )
-            self.workspace = create_workspace(**kwargs)
+        elif _mnnvl_comm_backend is not None and group is not None:
+            comm_backend = _mnnvl_comm_backend(group)
+
+        try:
+            alloc_token_num = max(max_token_num, self._max_token_num_seen or 0)
+            alloc_hidden_dim = max(hidden_dim, self._max_hidden_dim_seen or 0)
+            create_kw = dict(
+                backend=backend,
+                world_size=world_size,
+                rank=rank,
+                max_token_num=alloc_token_num,
+                hidden_dim=alloc_hidden_dim,
+                dtype=dtype or torch.bfloat16,
+                gpus_per_node=gpus_per_node,
+            )
+            if (
+                _flashinfer_create_workspace_supports_comm_backend
+                and comm_backend is not None
+            ):
+                create_kw["comm_backend"] = comm_backend
+            if _flashinfer_create_workspace_supports_group:
+                # Pin the symmetric-memory rendezvous to the actual
+                # subgroup. Without this, flashinfer >=0.6.10 falls back
+                # to WORLD and TP/EP/CP subgroup peers get addressed
+                # incorrectly (kernel hangs in cuda-graph warmup).
+                create_kw["group"] = device_group
+            if use_oneshot is not None:
+                create_kw["force_oneshot_support"] = bool(use_oneshot)
+            if use_fp32_lamport:
+                create_kw["use_fp32_lamport"] = True
+            self.workspace = _create_allreduce_fusion_workspace(**create_kw)
+            self.world_size = world_size
+            self.rank = rank
+            self.group = (device_group, cpu_group)
+            self.max_token_num = alloc_token_num
+            self.hidden_dim = alloc_hidden_dim
+            self.dtype = dtype or torch.bfloat16
+            self.initialized = True
+
+            backend_name = getattr(self.workspace, "backend", "unknown")
+            if not self._logged_init:
+                logger.info(
+                    f"FlashInfer AllReduce Fusion enabled and workspace initialized: "
+                    f"backend={backend_name}, rank={rank}, world_size={world_size}, "
+                    f"max_token_num={self.max_token_num}, hidden_dim={self.hidden_dim}"
+                )
+                self._logged_init = True
+            else:
+                logger.debug(
+                    f"FlashInfer workspace re-initialized: backend={backend_name}, "
+                    f"rank={rank}, world_size={world_size}"
+                )
         except Exception as e:
             _flashinfer_allreduce_unavailable = True
             logger.warning(
-                f"Failed to initialize FlashInfer workspace: {e}. "
+                f"Failed to initialize FlashInfer workspace (backend={backend}): {e}. "
                 "Disabling flashinfer allreduce fusion permanently."
             )
             self.workspace = None
             self.initialized = False
             return
 
-        self.world_size = world_size
-        self.rank = rank
-        self.group = (device_group, cpu_group)
-        self.max_token_num = max_token_num
-        self.hidden_dim = hidden_dim
-        self.dtype = dtype
-        self.initialized = True
-
-        backend = getattr(self.workspace, "backend", "unknown")
-        logger.info(
-            f"FlashInfer workspace initialized for rank {rank}, "
-            f"world_size {world_size}, backend {backend}, "
-            f"max_token_num {max_token_num}, hidden_dim {hidden_dim}"
-        )
-
     def is_buffer_size_sufficient(
         self,
         token_num: int,
@@ -381,12 +548,22 @@ class FlashInferWorkspaceManager:
             )
         except Exception as e:
             logger.debug(f"FlashInfer workspace size check failed: {e}")
+            # Fallback: some backends may not implement is_buffer_size_sufficient;
+            # reuse if within our allocated dimensions.
+            if (
+                self.max_token_num is not None
+                and self.hidden_dim is not None
+                and token_num <= self.max_token_num
+                and hidden_dim <= self.hidden_dim
+            ):
+                return True
             return False
 
     def cleanup(self):
-        """Clean up workspace"""
+        """Clean up workspace."""
         if self.workspace is not None:
             try:
+                if hasattr(self.workspace, "destroy"):
                     self.workspace.destroy()
             except Exception as e:
                 logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
@@ -399,6 +576,7 @@ class FlashInferWorkspaceManager:
                 self.max_token_num = None
                 self.hidden_dim = None
                 self.dtype = None
+                self._logged_init = False
 
 
 _attn_tp_workspace_manager = FlashInferWorkspaceManager()
@@ -445,12 +623,13 @@ def _sync_allreduce_unavailable_across_tp():
 def ensure_workspace_initialized(
     max_token_num: int = 2048,
     hidden_dim: int = 4096,
-    dtype: torch.dtype = torch.float16,
+    use_fp32_lamport: bool = False,
+    dtype: Optional[torch.dtype] = None,
     token_num: Optional[int] = None,
     use_oneshot: Optional[bool] = None,
     use_attn_tp_group: bool = True,
 ):
-    """Ensure workspace is initialized"""
+    """Ensure workspace is initialized."""
     if _flashinfer_allreduce_unavailable:
         return False
 
@@ -484,6 +663,11 @@ def ensure_workspace_initialized(
     workspace_manager = _get_workspace_manager(use_attn_tp_group)
     token_num = token_num or max_token_num
     group_key = (device_group, cpu_group)
+    effective_dtype = dtype or torch.bfloat16
+    server_args = get_global_server_args()
+    backend = resolve_flashinfer_allreduce_fusion_backend(server_args)
+    if backend is None:
+        return False
 
     if (
         not workspace_manager.initialized
@@ -493,7 +677,7 @@ def ensure_workspace_initialized(
         or not workspace_manager.is_buffer_size_sufficient(
             token_num=token_num,
             hidden_dim=hidden_dim,
-            dtype=dtype,
+            dtype=effective_dtype,
             use_oneshot=use_oneshot,
         )
     ):
@@ -502,6 +686,9 @@ def ensure_workspace_initialized(
             rank=rank,
             max_token_num=max_token_num,
             hidden_dim=hidden_dim,
+            backend=backend,
+            group=cpu_group,
+            use_fp32_lamport=use_fp32_lamport,
             dtype=dtype,
             use_oneshot=use_oneshot,
             device_group=device_group,
@@ -545,7 +732,9 @@ def flashinfer_allreduce_residual_rmsnorm(
     use_attn_tp_group: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Use FlashInfer's fused allreduce + residual + RMS norm operation
+    Use FlashInfer's unified fused allreduce + residual + RMS norm operation.
+    Automatically selects between trtllm and mnnvl backends based on topology
+    and hardware (controlled by --flashinfer-allreduce-fusion-backend).
 
     Args:
         input_tensor: Input tensor that needs allreduce
@@ -570,9 +759,6 @@ def flashinfer_allreduce_residual_rmsnorm(
     if use_attn_tp_group:
         world_size = get_attn_tensor_model_parallel_world_size()
     else:
-        # If MoE expert parallel world size > 1, use expert parallel group
-        # Otherwise, use tensor parallel group
-        # The two values cannot be larger than 1 at the same time
         if get_moe_expert_parallel_world_size() > 1:
             world_size = get_moe_expert_parallel_world_size()
         else:
@@ -594,6 +780,7 @@ def flashinfer_allreduce_residual_rmsnorm(
     if not ensure_workspace_initialized(
         max_token_num=max_token_num,
         hidden_dim=input_tensor.shape[-1],
+        use_fp32_lamport=(input_tensor.dtype == torch.float32),
         dtype=input_tensor.dtype,
         token_num=input_tensor.shape[0],
         use_oneshot=use_oneshot,
@@ -602,10 +789,14 @@ def flashinfer_allreduce_residual_rmsnorm(
         logger.debug("FlashInfer workspace not available")
         return None, None
 
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+    if workspace_manager.workspace is None:
+        logger.debug("FlashInfer workspace is None")
+        return None, None
+
     residual_out = torch.empty_like(residual)
     norm_out = torch.empty_like(input_tensor)
 
-    workspace_manager = _get_workspace_manager(use_attn_tp_group)
     kwargs = dict(
         input=input_tensor,
         workspace=workspace_manager.workspace,

python/sglang/srt/model_executor/model_runner.py

@@ -2463,7 +2463,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         (broadcasts, barriers) inside the graph capture context, which can
         deadlock with custom_all_reduce.register_graph_buffers.
         """
-        if not self.server_args.enable_flashinfer_allreduce_fusion:
+        if self.server_args.flashinfer_allreduce_fusion_backend is None:
             return
 
         from sglang.srt.layers.communicator import FUSE_ALLREDUCE_MAX_BATCH_SIZE

python/sglang/srt/server_args.py

@@ -659,6 +659,9 @@ class ServerArgs:
     flashinfer_mxfp4_moe_precision: Literal["default", "bf16"] = "default"
     enable_flashinfer_allreduce_fusion: bool = False
     enforce_disable_flashinfer_allreduce_fusion: bool = False
+    flashinfer_allreduce_fusion_backend: Optional[
+        Literal["auto", "trtllm", "mnnvl"]
+    ] = None
     enable_aiter_allreduce_fusion: bool = False
     deepep_mode: Literal["auto", "normal", "low_latency"] = "auto"
     deepep_dispatcher_output_dtype: Literal["auto", "bf16", "fp8", "int8", "nvfp4"] = (
@@ -1200,6 +1203,17 @@ class ServerArgs:
             )
             self.tool_call_parser = deprecated_tool_call_parsers[self.tool_call_parser]
 
+        # When user passes --enable-flashinfer-allreduce-fusion, enable with auto backend
+        if (
+            self.enable_flashinfer_allreduce_fusion
+            and self.flashinfer_allreduce_fusion_backend is None
+        ):
+            logger.warning(
+                "--enable-flashinfer-allreduce-fusion is deprecated. "
+                "Please use --flashinfer-allreduce-fusion-backend=auto instead."
+            )
+            self.flashinfer_allreduce_fusion_backend = "auto"
+        self.enable_flashinfer_allreduce_fusion = False
         # Deprecated attention-backend alias: "compressed" -> "dsv4".
         for attr in (
             "attention_backend",
@@ -2777,13 +2791,14 @@ class ServerArgs:
                 "Overlap scheduler is disabled when using sparse head for embedding model."
             )
 
-        # TRTLLM AllReduce Fusion supports SM90/100, enable it by default
+        # Auto-enable FlashInfer AllReduce Fusion on SM100 only, for models with
-        # for models with explicit support (DeepseekV3, GptOss, Glm4Moe,
+        # explicit support (DeepseekV3, GptOss, Glm4Moe, MistralLarge3,
-        # MistralLarge3, Qwen3/Qwen3Next/Qwen3.5 MoE families)
+        # Qwen3/Qwen3Next/Qwen3.5 MoE families). SM90 is not auto-enabled because
-        # TODO: currently, it is only supported in the single node scenario. https://github.com/flashinfer-ai/flashinfer/issues/2006
+        # 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.
         if (
-            not self.enable_flashinfer_allreduce_fusion
+            self.flashinfer_allreduce_fusion_backend is None
             and model_arch
             in [
                 "DeepseekV3ForCausalLM",
@@ -2800,20 +2815,19 @@ class ServerArgs:
                 "InternS2PreviewForConditionalGeneration",
                 "Qwen3_5ForConditionalGeneration",
             ]
-            and (is_sm90_supported() or is_sm100_supported())
+            and is_sm100_supported()
             and self.tp_size > 1
             and not self.enable_dp_attention
-            and self.nnodes == 1
             and self.moe_a2a_backend == "none"
         ):
-            self.enable_flashinfer_allreduce_fusion = True
+            self.flashinfer_allreduce_fusion_backend = "auto"
             logger.info(
-                f"Auto-enabling FlashInfer AllReduce Fusion on SM90/SM10X for {model_arch}"
+                f"Auto-enabling FlashInfer AllReduce Fusion on SM10X for {model_arch}"
             )
 
         # Apply enforce_disable_flashinfer_allreduce_fusion after all model-specific adjustments
         if self.enforce_disable_flashinfer_allreduce_fusion:
-            self.enable_flashinfer_allreduce_fusion = False
+            self.flashinfer_allreduce_fusion_backend = None
             logger.info(
                 "FlashInfer allreduce fusion is forcibly disabled "
                 "via --enforce-disable-flashinfer-allreduce-fusion."
@@ -4402,11 +4416,11 @@ class ServerArgs:
                 )
                 self.enable_aiter_allreduce_fusion = False
 
-            if self.enable_flashinfer_allreduce_fusion:
+            if self.flashinfer_allreduce_fusion_backend is not None:
                 logger.warning(
-                    "Disable --enable-flashinfer-allreduce-fusion because deterministic inference is enabled."
+                    "Disable --flashinfer-allreduce-fusion-backend because deterministic inference is enabled."
                 )
-                self.enable_flashinfer_allreduce_fusion = False
+                self.flashinfer_allreduce_fusion_backend = None
 
             # Check sampling backend
             if self.sampling_backend != "ascend":
@@ -6256,10 +6270,27 @@ class ServerArgs:
             default=ServerArgs.flashinfer_mxfp4_moe_precision,
             help="Choose the computation precision of flashinfer mxfp4 moe",
         )
+        parser.add_argument(
+            "--flashinfer-allreduce-fusion-backend",
+            type=str,
+            choices=["auto", "trtllm", "mnnvl"],
+            default=None,
+            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."
+            ),
+        )
         parser.add_argument(
             "--enable-flashinfer-allreduce-fusion",
             action="store_true",
-            help="Enable FlashInfer allreduce fusion with Residual RMSNorm.",
+            help="(Deprecated: use --flashinfer-allreduce-fusion-backend=auto) "
+            "Enable FlashInfer allreduce fusion with Residual RMSNorm.",
         )
         parser.add_argument(
             "--enforce-disable-flashinfer-allreduce-fusion",

test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py

@@ -82,6 +82,7 @@ class TestLoRAQwen3_30B_A3B_Instruct_2507_LogprobDiff(CustomTestCase):
             lora_paths={"my_lora": adapter_path},
             lora_backend=LORA_BACKEND,
             attention_backend="flashinfer",
+            flashinfer_allreduce_fusion_backend="trtllm",
             moe_runner_backend=MOE_RUNNER_BACKEND,
             experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
             prefill_attention_backend=PREFILL_ATTENTION_BACKEND,

test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py

@@ -46,7 +46,15 @@ class TestDeepseekV32FP4TPSpec(GSM8KMixin, DefaultServerBase):
     gsm8k_accept_length_thres = 2.7
 
     def test_z_bs_1_speed(self):
-        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        args = BenchArgs(
+            port=int(self.base_url.split(":")[-1]),
+            max_new_tokens=2048,
+            prompt=(
+                "Human: Think carefully before answering. Build a fully functional FastAPI todo server. "
+                "Start with a short design plan, then output the complete Python code, then show how to run it "
+                "and test three endpoints.\n\nAssistant:"
+            ),
+        )
         acc_length, speed = send_one_prompt(args)
 
         print(f"{acc_length=:.2f} {speed=:.2f}")

test/registered/unit/layers/test_flashinfer_comm_fusion.py

@@ -0,0 +1,181 @@
+import types
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers import flashinfer_comm_fusion as fusion
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300")
+
+
+class _FakeWorkspace:
+    def __init__(self, backend, world_size):
+        self.backend = backend
+        self.world_size = world_size
+
+    def is_buffer_size_sufficient(self, **_kwargs):
+        return True
+
+
+class _FakeFlashInferComm:
+    class AllReduceFusionPattern:
+        kARResidualRMSNorm = object()
+
+    def __init__(self):
+        self.calls = []
+
+    def create_allreduce_fusion_workspace(self, **kwargs):
+        self.calls.append(kwargs)
+        return _FakeWorkspace(kwargs["backend"], kwargs["world_size"])
+
+    def allreduce_fusion(
+        self,
+        *,
+        input,
+        workspace,
+        residual_out,
+        norm_out,
+        residual_in,
+        rms_gamma,
+        rms_eps,
+        **_kwargs,
+    ):
+        allreduced = input * workspace.world_size
+        expected_residual = allreduced + residual_in
+        variance = expected_residual.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+        expected_norm = (
+            expected_residual.to(torch.float32)
+            * torch.rsqrt(variance + rms_eps)
+            * rms_gamma.to(torch.float32)
+        ).to(input.dtype)
+        residual_out.copy_(expected_residual)
+        norm_out.copy_(expected_norm)
+
+
+def _torch_allreduce_residual_rmsnorm_baseline(
+    input_tensor, residual, weight, world_size, eps
+):
+    allreduced = input_tensor * world_size
+    residual_out = allreduced + residual
+    variance = residual_out.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+    norm_out = (
+        residual_out.to(torch.float32)
+        * torch.rsqrt(variance + eps)
+        * weight.to(torch.float32)
+    ).to(input_tensor.dtype)
+    return norm_out, residual_out
+
+
+class TestFlashInferCommFusion(unittest.TestCase):
+    def test_auto_backend_resolves_by_arch(self):
+        single_node = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="auto", nnodes=1
+        )
+        multi_node = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="auto", nnodes=2
+        )
+
+        # Blackwell: mnnvl regardless of node count.
+        with patch.object(fusion, "is_sm100_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), "mnnvl"
+            )
+
+        # SM90: mnnvl on single-node, trtllm fallback on multi-node.
+        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",
+            )
+
+    def test_allreduce_fusion_backends_match_torch_baseline(self):
+        fake_comm = _FakeFlashInferComm()
+        original_comm = fusion._flashinfer_comm
+        original_create = fusion._create_allreduce_fusion_workspace
+        original_manager = fusion._attn_tp_workspace_manager
+        original_unavailable = fusion._flashinfer_allreduce_unavailable
+        try:
+            fusion._flashinfer_comm = fake_comm
+            fusion._create_allreduce_fusion_workspace = (
+                fake_comm.create_allreduce_fusion_workspace
+            )
+            fusion._flashinfer_allreduce_unavailable = False
+
+            for backend in ("trtllm", "mnnvl"):
+                with self.subTest(backend=backend):
+                    world_size = 4
+                    manager = fusion.FlashInferWorkspaceManager()
+                    manager.workspace = _FakeWorkspace(backend, world_size)
+                    manager.initialized = True
+                    fusion._attn_tp_workspace_manager = manager
+                    if not torch.cuda.is_available():
+                        self.skipTest("FlashInfer allreduce custom op is CUDA-only")
+                    device = torch.device("cuda")
+                    torch.manual_seed(0)
+                    input_tensor = torch.randn(4, 8, dtype=torch.float32, device=device)
+                    residual = torch.randn(4, 8, dtype=torch.float32, device=device)
+                    weight = torch.randn(8, dtype=torch.float32, device=device)
+                    eps = 1e-6
+
+                    expected_norm, expected_residual = (
+                        _torch_allreduce_residual_rmsnorm_baseline(
+                            input_tensor, residual, weight, world_size, eps
+                        )
+                    )
+
+                    with (
+                        patch.object(
+                            fusion, "is_flashinfer_available", return_value=True
+                        ),
+                        patch.object(
+                            fusion,
+                            "get_attn_tensor_model_parallel_world_size",
+                            return_value=world_size,
+                        ),
+                        patch.object(
+                            fusion, "ensure_workspace_initialized", return_value=True
+                        ),
+                    ):
+                        norm_out, residual_out = (
+                            fusion.flashinfer_allreduce_residual_rmsnorm(
+                                input_tensor=input_tensor,
+                                residual=residual,
+                                weight=weight,
+                                eps=eps,
+                                max_token_num=8,
+                            )
+                        )
+
+                    torch.testing.assert_close(norm_out, expected_norm)
+                    torch.testing.assert_close(residual_out, expected_residual)
+        finally:
+            fusion._flashinfer_comm = original_comm
+            fusion._create_allreduce_fusion_workspace = original_create
+            fusion._attn_tp_workspace_manager = original_manager
+            fusion._flashinfer_allreduce_unavailable = original_unavailable
+
+
+if __name__ == "__main__":
+    unittest.main()