From c7c03ec53b1e664c2d415db4f02e43f86661f31d Mon Sep 17 00:00:00 2001 From: Shu Wang Date: Tue, 11 Aug 2026 18:46:46 -0500 Subject: [PATCH] [NVIDIA] Add flashinfer MNNVL backend for allreduce only (#30700) --- python/sglang/srt/arg_groups/overrides.py | 1 + python/sglang/srt/distributed/bootstrap.py | 8 + .../sglang/srt/distributed/parallel_state.py | 85 +++++++ python/sglang/srt/layers/communicator.py | 12 + .../srt/layers/flashinfer_comm_fusion.py | 99 ++++++++ .../layers/test_flashinfer_comm_fusion.py | 212 +++++++++++++++++- .../test_layer_communicator_fusion_gate.py | 65 ++++++ 7 files changed, 476 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/layers/test_layer_communicator_fusion_gate.py diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 62f9fdb21..3ea7ab042 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1955,6 +1955,7 @@ _FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset( { "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", + "DeepseekV4ForCausalLM", "GptOssForCausalLM", "GlmMoeDsaForCausalLM", "Glm4MoeForCausalLM", diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index f21844415..173b5ee61 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/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, diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index f1d134e91..903fdff79 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index bb99bdf1f..e92806fd4 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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 diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 194d97e05..43c512df4 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/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, diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index c2ad459fd..6f176931c 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/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, - norm_out, - residual_in, - rms_gamma, - rms_eps, + pattern, + output=None, + residual_out=None, + norm_out=None, + 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() diff --git a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py new file mode 100644 index 000000000..16f49cf15 --- /dev/null +++ b/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()