[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,

test/registered/unit/layers/test_flashinfer_comm_fusion.py

@@ -1,3 +1,4 @@
+import contextlib
 import types
 import unittest
 from unittest.mock import patch
@@ -7,6 +8,7 @@ import torch
 from sglang.srt.layers import flashinfer_comm_fusion as fusion
 from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
 
 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")
@@ -24,6 +26,7 @@ class _FakeWorkspace:
 
 class _FakeFlashInferComm:
     class AllReduceFusionPattern:
+        kAllReduce = object()
         kARResidualRMSNorm = object()
 
     def __init__(self):
@@ -38,13 +41,25 @@ class _FakeFlashInferComm:
         *,
         input,
         workspace,
-        residual_out,
+        pattern,
-        norm_out,
+        output=None,
-        residual_in,
+        residual_out=None,
-        rms_gamma,
+        norm_out=None,
-        rms_eps,
+        residual_in=None,
+        rms_gamma=None,
+        rms_eps=None,
         **_kwargs,
     ):
+        if pattern is self.AllReduceFusionPattern.kAllReduce:
+            allreduced = input * workspace.world_size
+            if output is None:
+                return allreduced
+            output.copy_(allreduced)
+            return output
+
+        if pattern is not self.AllReduceFusionPattern.kARResidualRMSNorm:
+            raise ValueError(f"Unexpected pattern: {pattern}")
+
         allreduced = input * workspace.world_size
         expected_residual = allreduced + residual_in
         variance = expected_residual.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
@@ -71,7 +86,7 @@ def _torch_allreduce_residual_rmsnorm_baseline(
     return norm_out, residual_out
 
 
-class TestFlashInferCommFusion(unittest.TestCase):
+class TestFlashInferCommFusion(CustomTestCase):
     def test_auto_backend_resolves_by_arch(self):
         single_node = types.SimpleNamespace(
             flashinfer_allreduce_fusion_backend="auto", nnodes=1
@@ -240,5 +255,190 @@ class TestFlashInferCommFusion(unittest.TestCase):
             fusion._flashinfer_allreduce_unavailable = original_unavailable
 
 
+_GROUP_KEY = ("device_group", "cpu_group")
+_OTHER_GROUP_KEY = ("other_device_group", "other_cpu_group")
+
+
+class TestFlashInferAllReduceOnly(CustomTestCase):
+    def _make_manager(self, world_size, group_key=_GROUP_KEY):
+        manager = fusion.FlashInferWorkspaceManager()
+        manager.workspace = _FakeWorkspace(None, world_size)
+        manager.initialized = True
+        manager.world_size = world_size
+        manager.group = group_key
+        manager.max_token_num = 2048
+        manager.hidden_dim = 4096
+        manager.dtype = torch.float32
+        return manager
+
+    @contextlib.contextmanager
+    def _patched_attn_workspace(self, manager):
+        from sglang.srt.runtime_context import get_resources
+
+        buffers = get_resources().buffers
+        manager_key = "flashinfer_fusion_attn_tp_workspace"
+        original_manager = buffers.get(manager_key)
+        original_comm = fusion._flashinfer_comm
+        original_unavailable = fusion._flashinfer_allreduce_unavailable
+
+        buffers[manager_key] = manager
+        fusion._flashinfer_comm = _FakeFlashInferComm()
+        fusion._flashinfer_allreduce_unavailable = False
+        try:
+            yield
+        finally:
+            fusion._flashinfer_comm = original_comm
+            fusion._flashinfer_allreduce_unavailable = original_unavailable
+            if original_manager is None:
+                buffers.pop(manager_key, None)
+            else:
+                buffers[manager_key] = original_manager
+
+    def _can_use(self, input_, world_size=4, group_key=_GROUP_KEY):
+        return fusion.can_use_flashinfer_allreduce(
+            input_,
+            use_attn_tp_group=True,
+            expected_world_size=world_size,
+            expected_group_key=group_key,
+        )
+
+    def test_allreduce_output_equals_input_times_world_size(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required for flashinfer custom op")
+        world_size = 4
+        with self._patched_attn_workspace(self._make_manager(world_size)):
+            input_ = torch.randn(8, 16, dtype=torch.bfloat16, device="cuda")
+            expected = input_ * world_size
+
+            with get_parallel().override(attn_tp_size=world_size):
+                self.assertTrue(self._can_use(input_, world_size=world_size))
+                result = fusion.flashinfer_allreduce(input_, use_attn_tp_group=True)
+
+            torch.testing.assert_close(result, expected)
+
+    def test_shape_guard_rejects_non_2d(self):
+        with self._patched_attn_workspace(self._make_manager(4)):
+            self.assertFalse(self._can_use(torch.randn(16)))
+            self.assertFalse(self._can_use(torch.randn(2, 8, 16)))
+
+    def test_shape_guard_rejects_non_contiguous(self):
+        with self._patched_attn_workspace(self._make_manager(4)):
+            non_contiguous = torch.randn(16, 8).t()
+            self.assertFalse(non_contiguous.is_contiguous())
+            self.assertFalse(self._can_use(non_contiguous))
+
+    def test_rejects_when_unavailable(self):
+        original_unavailable = fusion._flashinfer_allreduce_unavailable
+        try:
+            fusion._flashinfer_allreduce_unavailable = True
+            self.assertFalse(self._can_use(torch.randn(8, 16)))
+        finally:
+            fusion._flashinfer_allreduce_unavailable = original_unavailable
+
+    def test_rejects_when_workspace_uninitialized(self):
+        with self._patched_attn_workspace(fusion.FlashInferWorkspaceManager()):
+            with get_parallel().override(attn_tp_size=4):
+                self.assertFalse(self._can_use(torch.randn(8, 16)))
+
+    def test_rejects_when_workspace_group_differs(self):
+        """A workspace rendezvoused on other peers must not be reused.
+
+        Under hybrid EP+TP (e.g. tp=4, ep=2) the MoE-TP and MoE-EP groups have
+        the same world size but pair different ranks, so a workspace built for
+        one silently reduces across the wrong peers when used by the other --
+        wrong output rather than a crash.
+        """
+        with self._patched_attn_workspace(self._make_manager(2)):
+            self.assertFalse(
+                self._can_use(
+                    torch.randn(8, 16), world_size=2, group_key=_OTHER_GROUP_KEY
+                )
+            )
+
+    def test_rejects_when_workspace_world_size_differs(self):
+        with self._patched_attn_workspace(self._make_manager(4)):
+            self.assertFalse(self._can_use(torch.randn(8, 16), world_size=2))
+
+    def test_rejects_when_token_num_exceeds_workspace_capacity(self):
+        """Under Dynamo the capacity check replaces is_buffer_size_sufficient().
+
+        _FakeWorkspace.is_buffer_size_sufficient() always says yes, so this only
+        passes if the compiling branch consults the manager's own allocation.
+        """
+        manager = self._make_manager(4)
+        manager.max_token_num = 8
+        with self._patched_attn_workspace(manager):
+            with patch.object(torch.compiler, "is_compiling", return_value=True):
+                self.assertTrue(self._can_use(torch.randn(8, 16)))
+                self.assertFalse(self._can_use(torch.randn(9, 16)))
+
+    def test_rejects_when_hidden_dim_exceeds_workspace_capacity(self):
+        manager = self._make_manager(4)
+        manager.hidden_dim = 16
+        with self._patched_attn_workspace(manager):
+            with patch.object(torch.compiler, "is_compiling", return_value=True):
+                self.assertTrue(self._can_use(torch.randn(8, 16)))
+                self.assertFalse(self._can_use(torch.randn(8, 17)))
+
+    def test_rejects_when_dtype_mismatches_workspace(self):
+        manager = self._make_manager(4)
+        manager.dtype = torch.bfloat16
+        with self._patched_attn_workspace(manager):
+            with patch.object(torch.compiler, "is_compiling", return_value=True):
+                self.assertTrue(self._can_use(torch.randn(8, 16, dtype=torch.bfloat16)))
+                self.assertFalse(self._can_use(torch.randn(8, 16, dtype=torch.float32)))
+
+
+class _FakeGroupCoordinator:
+    def __init__(self, world_size):
+        self.world_size = world_size
+        self._fi_workspace_hint = None
+
+
+class TestTagGroupsForFlashInferAllReduceOnly(CustomTestCase):
+    """The MoE workspace rendezvouses on the EP group when moe_ep_size > 1 and
+    on the MoE-TP group otherwise, so only that one group may be tagged."""
+
+    def _tag(self, *, attn_tp, moe_ep, moe_tp):
+        from sglang.srt.distributed import parallel_state as ps
+
+        with patch.object(ps, "_ENABLE_FLASHINFER_ALLREDUCE_ONLY", True), patch.object(
+            ps, "_ATTN_TP", attn_tp
+        ), patch.object(ps, "_MOE_EP", moe_ep), patch.object(ps, "_MOE_TP", moe_tp):
+            ps._tag_groups_for_flashinfer_allreduce_only()
+
+    def test_hybrid_ep_tp_tags_only_the_ep_group(self):
+        attn_tp = _FakeGroupCoordinator(4)
+        moe_ep = _FakeGroupCoordinator(2)
+        moe_tp = _FakeGroupCoordinator(2)
+
+        self._tag(attn_tp=attn_tp, moe_ep=moe_ep, moe_tp=moe_tp)
+
+        self.assertEqual(attn_tp._fi_workspace_hint, "attn_tp")
+        self.assertEqual(moe_ep._fi_workspace_hint, "moe")
+        self.assertIsNone(moe_tp._fi_workspace_hint)
+
+    def test_pure_moe_tp_tags_only_the_moe_tp_group(self):
+        attn_tp = _FakeGroupCoordinator(4)
+        moe_ep = _FakeGroupCoordinator(1)
+        moe_tp = _FakeGroupCoordinator(4)
+
+        self._tag(attn_tp=attn_tp, moe_ep=moe_ep, moe_tp=moe_tp)
+
+        self.assertEqual(moe_tp._fi_workspace_hint, "moe")
+        self.assertIsNone(moe_ep._fi_workspace_hint)
+
+    def test_shared_coordinator_prefers_attn_tp(self):
+        # tp=4, ep=4: _ATTN_TP is _MOE_EP is _TP. Either workspace spans the
+        # same peers, but the choice must be deterministic.
+        shared = _FakeGroupCoordinator(4)
+        moe_tp = _FakeGroupCoordinator(1)
+
+        self._tag(attn_tp=shared, moe_ep=shared, moe_tp=moe_tp)
+
+        self.assertEqual(shared._fi_workspace_hint, "attn_tp")
+        self.assertIsNone(moe_tp._fi_workspace_hint)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/layers/test_layer_communicator_fusion_gate.py

@@ -0,0 +1,65 @@
+import types
+import unittest
+from unittest.mock import patch
+
+from sglang.srt.layers import communicator as comm
+from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode
+from sglang.srt.runtime_context import get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _fake_communicator():
+    return types.SimpleNamespace(
+        _speculative_algo=None,
+        layer_scatter_modes=types.SimpleNamespace(mlp_mode=ScatterMode.TP_ATTN_FULL),
+        is_last_layer=False,
+        _context=types.SimpleNamespace(tp_size=4),
+    )
+
+
+class TestFuseMlpAllReduceGate(CustomTestCase):
+    """Hybrid EP+TP must not fuse the post-experts all-reduce away.
+
+    The fused residual+LN reduces over a single group, but with moe_ep_size > 1
+    and moe_tp_size > 1 the post-experts reduction spans two disjoint groups
+    (_MOE_EP then _MOE_TP) and should_skip_post_experts_all_reduce() drops both
+    once fusion is published. The result is activations reduced over only half
+    the peers -- wrong output, no crash. Observed as garbage completions on
+    Qwen3-30B-A3B with --tp-size 4 --ep-size 2.
+    """
+
+    def _should_fuse(self, *, moe_ep_size, moe_tp_size):
+        forward_batch = types.SimpleNamespace(
+            input_ids=types.SimpleNamespace(shape=(8,))
+        )
+        with (
+            patch.object(comm, "is_enable_moe_cp_allgather", return_value=False),
+            patch.object(comm, "apply_flashinfer_allreduce_fusion", return_value=True),
+            patch.object(
+                comm,
+                "get_attn_tp_context",
+                return_value=types.SimpleNamespace(input_scattered=False),
+            ),
+            get_parallel().override(
+                moe_ep_size=moe_ep_size, moe_tp_size=moe_tp_size, tp_size=4
+            ),
+        ):
+            return LayerCommunicator.should_fuse_mlp_allreduce_with_next_layer(
+                _fake_communicator(), forward_batch
+            )
+
+    def test_hybrid_ep_tp_does_not_fuse(self):
+        self.assertFalse(self._should_fuse(moe_ep_size=2, moe_tp_size=2))
+
+    def test_pure_tp_still_fuses(self):
+        self.assertTrue(self._should_fuse(moe_ep_size=1, moe_tp_size=4))
+
+    def test_pure_ep_still_fuses(self):
+        self.assertTrue(self._should_fuse(moe_ep_size=4, moe_tp_size=1))
+
+
+if __name__ == "__main__":
+    unittest.main()