[NVIDIA] Add flashinfer MNNVL backend for allreduce only (#30700)

python/sglang/srt/arg_groups/overrides.py

@@ -1955,6 +1955,7 @@ _FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset(
     {
         "DeepseekV3ForCausalLM",
         "DeepseekV32ForCausalLM",
+        "DeepseekV4ForCausalLM",
         "GptOssForCausalLM",
         "GlmMoeDsaForCausalLM",
         "Glm4MoeForCausalLM",

python/sglang/srt/distributed/bootstrap.py

@@ -16,9 +16,13 @@ from sglang.srt.distributed import (
     init_distributed_environment,
     initialize_model_parallel,
     set_custom_all_reduce,
+    set_flashinfer_allreduce_only,
     set_mscclpp_all_reduce,
     set_torch_symm_mem_all_reduce,
 )
+from sglang.srt.distributed.parallel_state import (
+    _tag_groups_for_flashinfer_allreduce_only,
+)
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
@@ -164,6 +168,9 @@ def _set_all_reduce_flags(*, server_args: ServerArgs) -> None:
     set_custom_all_reduce(not server_args.disable_custom_all_reduce)
     set_mscclpp_all_reduce(server_args.enable_mscclpp)
     set_torch_symm_mem_all_reduce(server_args.enable_torch_symm_mem)
+    set_flashinfer_allreduce_only(
+        server_args.flashinfer_allreduce_fusion_backend is not None
+    )
 
 
 def _init_cpu_threads_env(
@@ -233,6 +240,7 @@ def _init_parallel_groups(
         rank_offset=rank_offset,
         max_world_size=server_args.max_ep_size,
     )
+    _tag_groups_for_flashinfer_allreduce_only()
     initialize_dp_attention(
         server_args=server_args,
         model_config=model_config,

python/sglang/srt/distributed/parallel_state.py

@@ -185,6 +185,22 @@ def outplace_all_reduce(
     return group._all_reduce_out_place(tensor, outplace_all_reduce_method)
 
 
+@register_custom_op(out_shape="tensor")
+def flashinfer_allreduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    """FlashInfer kAllReduce over ``group_name``.
+
+    Registered as a custom op so it stays opaque under Dynamo and can run inside
+    piecewise CUDA graphs. Applicability is decided by
+    ``GroupCoordinator._can_use_flashinfer_allreduce`` before the call -- this op
+    has no fallback of its own.
+    """
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    return group._flashinfer_allreduce(tensor)
+
+
 @register_custom_op(mutates_args=["output"])
 def reg_all_gather_into_tensor(
     output: torch.Tensor, input: torch.Tensor, group_name: str
@@ -291,6 +307,10 @@ class GroupCoordinator:
         self.local_rank = local_rank
         self.device_group = None
         self.cpu_group = None
+        # Which FlashInfer fusion workspace this group owns, or None when the
+        # group is not eligible for the allreduce-only kAllReduce path. Stamped
+        # by _tag_groups_for_flashinfer_allreduce_only() after group init.
+        self._fi_workspace_hint: Optional[str] = None
         self.local_size = get_int_env_var("LOCAL_SIZE", 0)
 
         if is_cuda_alike():
@@ -672,6 +692,9 @@ class GroupCoordinator:
             return self.npu_communicator.all_reduce(input_)
 
         if torch.compiler.is_compiling():
+            if self._can_use_flashinfer_allreduce(input_):
+                return flashinfer_allreduce(input_, group_name=self.unique_name)
+
             # Byte-size thresholds in method selection (e.g. `_pick_algo` or
             # `should_mscclpp_allreduce`) would guard on the symbolic token dim
             # and recompile per shape; defer the selection to runtime inside
@@ -723,6 +746,9 @@ class GroupCoordinator:
                 self.pynccl_comm.all_reduce(input_)
                 return input_
 
+        if self._can_use_flashinfer_allreduce(input_):
+            return flashinfer_allreduce(input_, group_name=self.unique_name)
+
         outplace_all_reduce_method = self._resolve_outplace_all_reduce_method(
             input_=input_,
             should_use_pymscclpp_allreduce=should_use_pymscclpp_allreduce,
@@ -919,6 +945,29 @@ class GroupCoordinator:
             return "pynccl"
         return None
 
+    def _can_use_flashinfer_allreduce(self, input_: torch.Tensor) -> bool:
+        if self._fi_workspace_hint is None:
+            return False
+        from sglang.srt.layers.flashinfer_comm_fusion import (
+            can_use_flashinfer_allreduce,
+        )
+
+        return can_use_flashinfer_allreduce(
+            input_,
+            use_attn_tp_group=(self._fi_workspace_hint == "attn_tp"),
+            expected_world_size=self.world_size,
+            expected_group_key=(self.device_group, self.cpu_group),
+        )
+
+    def _flashinfer_allreduce(self, input_: torch.Tensor) -> torch.Tensor:
+        from sglang.srt.layers.flashinfer_comm_fusion import (
+            flashinfer_allreduce as _flashinfer_allreduce_impl,
+        )
+
+        return _flashinfer_allreduce_impl(
+            input_, use_attn_tp_group=(self._fi_workspace_hint == "attn_tp")
+        )
+
     def _all_reduce_out_place(
         self, input_: torch.Tensor, outplace_all_reduce_method: str
     ) -> torch.Tensor:
@@ -2008,6 +2057,7 @@ _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = False
 # Read once at import: whether CustomAllReduceV2 is opted in on a multi-node
 # (MNNVL) group. Used on the all_reduce hot path (see GroupCoordinator).
 _CA_V2_MULTINODE = envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()
+_ENABLE_FLASHINFER_ALLREDUCE_ONLY = False
 
 
 def set_custom_all_reduce(enable: bool):
@@ -2025,6 +2075,41 @@ def set_torch_symm_mem_all_reduce(enable: bool):
     _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = enable
 
 
+def set_flashinfer_allreduce_only(enable: bool):
+    global _ENABLE_FLASHINFER_ALLREDUCE_ONLY
+    _ENABLE_FLASHINFER_ALLREDUCE_ONLY = enable
+
+
+def _tag_groups_for_flashinfer_allreduce_only():
+    """Stamp _fi_workspace_hint on the group coordinators that own a FlashInfer
+    fusion workspace, so all_reduce() can dispatch to flashinfer_allreduce()
+    without touching the call sites.
+
+    Only two workspaces exist (see ``_get_workspace_manager``): one for
+    attention TP and one for MoE. A group may only be tagged for the workspace
+    that was rendezvoused on its own peers -- reducing over a workspace built
+    for a different set of peers silently returns wrong data.
+
+    - ``_TP`` is deliberately absent: it *is* ``_ATTN_TP`` when
+      ``attn_tp_size == tp_size``, and a strict superset of it otherwise (DP
+      attention), where the attention workspace addresses the wrong peers.
+    - The MoE workspace rendezvouses on the EP group when ``moe_ep_size > 1``
+      and on the MoE-TP group otherwise, so exactly one of ``_MOE_EP`` /
+      ``_MOE_TP`` is eligible. Tagging both makes a MoE-TP allreduce reduce
+      across the EP peers under hybrid EP+TP (e.g. tp=4, ep=2).
+    """
+    if not _ENABLE_FLASHINFER_ALLREDUCE_ONLY:
+        return
+
+    moe_group = _MOE_EP if (_MOE_EP is not None and _MOE_EP.world_size > 1) else _MOE_TP
+    # Attention is tagged last on purpose: when a coordinator backs both roles
+    # (e.g. _ATTN_TP is _MOE_EP is _TP at tp=4, ep=4) either workspace spans the
+    # same peers and is correct, so we just pick one deterministically.
+    for group, hint in ((moe_group, "moe"), (_ATTN_TP, "attn_tp")):
+        if group is not None:
+            group._fi_workspace_hint = hint
+
+
 # TODO: refactor in-tree platforms to get rid of this wrapper
 def get_default_distributed_backend(device: str) -> str:
     # We deliberately go through ``platforms.current_platform`` (rather than

python/sglang/srt/layers/communicator.py

@@ -816,6 +816,18 @@ class LayerCommunicator:
         if is_enable_moe_cp_allgather():
             return False
 
+        # Fusing makes the next layer's residual+LN absorb the post-experts
+        # all-reduce, and that fused kernel reduces over a single group. Under
+        # hybrid EP+TP the post-experts reduction spans two disjoint groups
+        # (moe_expert_parallel_all_reduce over _MOE_EP, then
+        # moe_tensor_model_parallel_all_reduce over _MOE_TP), and
+        # should_skip_post_experts_all_reduce() skips *both* once fusion is
+        # published -- so the fused reduce would cover only half the peers and
+        # silently return under-reduced activations.
+        parallel = get_parallel()
+        if parallel.moe_ep_size > 1 and parallel.moe_tp_size > 1:
+            return False
+
         if (
             is_dp_attention_enabled()
             and self._speculative_algo is not None

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -860,6 +860,105 @@ def flashinfer_allreduce_residual_rmsnorm(
     return norm_out, residual_out
 
 
+def can_use_flashinfer_allreduce(
+    input_: torch.Tensor,
+    *,
+    use_attn_tp_group: bool,
+    expected_world_size: int,
+    expected_group_key: Tuple[Optional[ProcessGroup], Optional[ProcessGroup]],
+) -> bool:
+    """Whether ``flashinfer_allreduce`` can service this all-reduce.
+
+    Split out from the kernel call so the decision happens in plain Python,
+    outside the custom op: the op is opaque to Dynamo and has to return a
+    tensor, so it cannot carry a data-dependent fallback of its own.
+
+    ``expected_world_size`` / ``expected_group_key`` describe the calling group;
+    the workspace is only usable when it was rendezvoused on exactly those peers.
+
+    Every check here is rank-invariant by construction, and must stay that way:
+    a rank that quietly falls back to NCCL while its peers enter the kernel
+    mismatches and hangs. The unavailable flag and workspace initialization are
+    cross-rank synced at init time (``_sync_allreduce_unavailable_across_tp``);
+    the rest are pure functions of the group identity and of tensor metadata,
+    which is identical on every rank of the group.
+    """
+    if _flashinfer_allreduce_unavailable or _flashinfer_comm is None:
+        return False
+
+    if input_.ndim != 2 or not input_.is_contiguous():
+        return False
+
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+    if not workspace_manager.initialized or workspace_manager.workspace is None:
+        return False
+
+    # The two workspaces are keyed by attention-TP vs MoE, but the MoE one
+    # rendezvouses on either the EP or the MoE-TP group depending on topology.
+    # Under hybrid EP+TP those groups have equal world size but pair different
+    # ranks, so a mismatch here reduces across the wrong peers and silently
+    # produces garbage rather than failing. Require an exact match.
+    if (
+        workspace_manager.world_size != expected_world_size
+        or workspace_manager.group != expected_group_key
+    ):
+        return False
+
+    # Size checks stay last: they read the token dim, which is symbolic under
+    # Dynamo, so statically-off configs must short-circuit before reaching them
+    # (same ordering rule as apply_flashinfer_allreduce_fusion).
+    token_num, hidden_dim = input_.shape
+    if torch.compiler.is_compiling():
+        # Don't call into the flashinfer workspace object while tracing. The
+        # workspace was allocated for (max_token_num, hidden_dim, dtype) and
+        # vetted by is_buffer_size_sufficient() at init; the requirement is
+        # monotone in token_num/hidden_dim, so staying within the allocation
+        # (including dtype) is a conservative stand-in here.
+        return (
+            workspace_manager.max_token_num is not None
+            and workspace_manager.hidden_dim is not None
+            and workspace_manager.dtype is not None
+            and token_num <= workspace_manager.max_token_num
+            and hidden_dim <= workspace_manager.hidden_dim
+            and workspace_manager.dtype == input_.dtype
+        )
+
+    return workspace_manager.is_buffer_size_sufficient(
+        token_num=token_num,
+        hidden_dim=hidden_dim,
+        dtype=input_.dtype,
+    )
+
+
+def flashinfer_allreduce(
+    input_: torch.Tensor,
+    *,
+    use_attn_tp_group: bool,
+) -> torch.Tensor:
+    """Allreduce-only FlashInfer kAllReduce.
+
+    Assumes ``can_use_flashinfer_allreduce`` returned True for this call; there
+    is no fallback here. Kernel errors are deliberately not caught -- swallowing
+    one would put this rank on NCCL while its peers stay in the kernel, which
+    mismatch-hangs instead of failing.
+    """
+    workspace_manager = _get_workspace_manager(use_attn_tp_group)
+
+    output = torch.empty_like(input_)
+    kwargs = dict(
+        input=input_,
+        workspace=workspace_manager.workspace,
+        pattern=_flashinfer_comm.AllReduceFusionPattern.kAllReduce,
+        launch_with_pdl=True,
+        fp32_acc=False,
+        output=output,
+    )
+    if _flashinfer_allreduce_supports_trigger_completion:
+        kwargs["trigger_completion_at_end"] = False
+    _flashinfer_comm.allreduce_fusion(**kwargs)
+    return output
+
+
 def pre_initialize_workspaces(
     max_token_num: int,
     hidden_dim: int,