diff --git a/python/sglang/srt/layers/moe/__init__.py b/python/sglang/srt/layers/moe/__init__.py index 3984a8322..b9bbcfad3 100644 --- a/python/sglang/srt/layers/moe/__init__.py +++ b/python/sglang/srt/layers/moe/__init__.py @@ -10,6 +10,7 @@ from sglang.srt.layers.moe.utils import ( get_tbo_token_distribution_threshold, initialize_moe_config, is_tbo_enabled, + should_skip_post_experts_all_reduce, should_use_dp_reduce_scatterv, should_use_flashinfer_cutlass_moe_fp4_allgather, ) @@ -24,6 +25,7 @@ __all__ = [ "get_moe_a2a_backend", "get_moe_runner_backend", "get_deepep_mode", + "should_skip_post_experts_all_reduce", "should_use_dp_reduce_scatterv", "should_use_flashinfer_cutlass_moe_fp4_allgather", "is_tbo_enabled", diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 73dde3166..59cdce5d1 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -346,6 +346,39 @@ def should_use_dp_reduce_scatterv(): ) +def should_skip_post_experts_all_reduce( + *, + is_tp_path: bool, + use_reduce_scatter: bool = False, + should_allreduce_fusion: bool = False, +) -> bool: + """Whether to skip the post-experts all-reduce (EP or TP) because a + downstream component will fuse, replace, or absorb it. + + Skip reasons, in order: + - ``should_allreduce_fusion``: LayerCommunicator will fuse the all-reduce + with the next layer's residual all-reduce. + - ``use_reduce_scatter``: LayerCommunicator's post-attention scatter will + do reduce-scatter, which would double-reduce on top of an all-reduce. + - ``should_use_dp_reduce_scatterv()``: the standard dispatcher's combine + path replaces the all-reduce with a reduce-scatterv. + - ``should_use_flashinfer_cutlass_moe_fp4_allgather()`` (TP path only): + the flashinfer cutlass FP4 kernel performs an all-gather that absorbs + the post-experts TP all-reduce. Not relevant to the EP all-reduce. + + The first two args are layer-context flags from ``LayerCommunicator`` and + default to ``False`` for models that don't use it. Pass ``is_tp_path=True`` + for the post-experts TP all-reduce, ``False`` for the EP all-reduce. + """ + if should_allreduce_fusion or use_reduce_scatter: + return True + if should_use_dp_reduce_scatterv(): + return True + if is_tp_path and should_use_flashinfer_cutlass_moe_fp4_allgather(): + return True + return False + + @contextmanager def speculative_moe_backend_context(): """ diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 4d9a400a4..6daafecf4 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -58,8 +58,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_deepep_mode, get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -382,12 +381,10 @@ class BailingMoESparseMoeBlock(nn.Module): if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index d8750da8f..3f47c038c 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -34,7 +34,7 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import should_use_dp_reduce_scatterv +from sglang.srt.layers.moe import should_skip_post_experts_all_reduce from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK @@ -348,11 +348,10 @@ class BailingMoE(nn.Module): if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if ( - self.tp_size > 1 - and not use_reduce_scatter - and not should_allreduce_fusion - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 9f2c82405..b226dc3b9 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -85,7 +85,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, get_moe_runner_backend, - should_use_dp_reduce_scatterv, + should_skip_post_experts_all_reduce, should_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -651,12 +651,10 @@ class DeepseekV2MoE(nn.Module): current_stream.wait_stream(self.alt_stream) final_hidden_states += shared_output - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states @@ -741,12 +739,10 @@ class DeepseekV2MoE(nn.Module): final_hidden_states *= self.routed_scaling_factor if shared_output is not None: final_hidden_states += shared_output - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 18012448a..ff0a02099 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -47,7 +47,10 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput -from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_skip_post_experts_all_reduce, +) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.topk import TopK @@ -300,10 +303,9 @@ class ExaoneMoESparseMoEBlock(nn.Module): if shared_output is not None: final_hidden_states = final_hidden_states + shared_output - if ( - self.tp_size > 1 - and not use_reduce_scatter - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 092b94b7e..a83a7b36d 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -61,7 +61,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, + should_skip_post_experts_all_reduce, should_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -594,12 +594,10 @@ class Glm4MoeSparseMoeBlock(nn.Module): current_stream.wait_stream(self.alt_stream) final_hidden_states += shared_output - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states @@ -629,12 +627,10 @@ class Glm4MoeSparseMoeBlock(nn.Module): final_hidden_states_out = torch.empty_like(final_hidden_states) torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) final_hidden_states = final_hidden_states_out - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index f11827b98..d44e3ef21 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/python/sglang/srt/models/hunyuan_v3.py @@ -34,7 +34,7 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import should_use_dp_reduce_scatterv +from sglang.srt.layers.moe import should_skip_post_experts_all_reduce from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -192,11 +192,14 @@ class HYV3MoEFused(nn.Module): hidden_states=hidden_states, topk_output=topk_output ) - skip_post_reduce = should_use_dp_reduce_scatterv() - if self.ep_size > 1 and not skip_post_reduce: + if self.ep_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=False, + ): final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) - if self.tp_size > 1 and not skip_post_reduce: + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + ): final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states ) @@ -224,11 +227,14 @@ class HYV3MoEFused(nn.Module): current_stream.wait_stream(self.alt_stream) final_hidden_states = final_hidden_states + shared_output - skip_post_reduce = should_use_dp_reduce_scatterv() - if self.ep_size > 1 and not skip_post_reduce: + if self.ep_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=False, + ): final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) - if self.tp_size > 1 and not skip_post_reduce: + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + ): final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states ) diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 194abf32b..7daf233d9 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -58,7 +58,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_deepep_mode, get_moe_a2a_backend, - should_use_dp_reduce_scatterv, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -383,10 +383,9 @@ class LLaDA2MoeSparseMoeBlock(nn.Module): if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if ( - self.tp_size > 1 - and not use_reduce_scatter - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index c2e4b12c9..8d29e0143 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -39,7 +39,7 @@ from sglang.srt.layers.linear import ( ReplicatedLinear, RowParallelLinear, ) -from sglang.srt.layers.moe import should_use_dp_reduce_scatterv +from sglang.srt.layers.moe import should_skip_post_experts_all_reduce from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -146,10 +146,9 @@ class Llama4MoE(nn.Module): out_aD = routed_out + shared_out - if ( - self.tp_size > 1 - and not use_reduce_scatter - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, ): out_aD = tensor_model_parallel_all_reduce(out_aD) diff --git a/python/sglang/srt/models/mimo_v2_flash.py b/python/sglang/srt/models/mimo_v2_flash.py index b6dd6c4b0..1ca4a4a2b 100644 --- a/python/sglang/srt/models/mimo_v2_flash.py +++ b/python/sglang/srt/models/mimo_v2_flash.py @@ -51,8 +51,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, get_moe_runner_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat @@ -298,12 +297,10 @@ class MiMoV2MoE(nn.Module): final_hidden_states = self.experts(hidden_states, topk_output) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index d6a38d082..5f7cab05e 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -61,8 +61,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -552,12 +551,10 @@ class MiniMaxM2MoE(nn.Module): topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts(hidden_states, topk_output) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 755fa70ac..2c6fd4da7 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -60,8 +60,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton import FusedMoE @@ -467,12 +466,10 @@ class Qwen2MoeSparseMoeBlock(nn.Module): # An out-of-place add would allocate a new tensor outside symm # memory, breaking subsequent symmetric collective operations. final_hidden_states += shared_output - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index dd9598c36..f255b90fd 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -50,8 +50,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -332,20 +331,17 @@ class Qwen3MoeSparseMoeBlock(nn.Module): topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) - if ( - self.ep_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_dp_reduce_scatterv() + if self.ep_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=False, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index faafd8c78..36ead547a 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -39,10 +39,7 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput -from sglang.srt.layers.moe import ( - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, -) +from sglang.srt.layers.moe import should_skip_post_experts_all_reduce from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK @@ -373,12 +370,10 @@ class SarvamMoESparseMoeBlock(nn.Module): final_hidden_states = final_hidden_states * self.routed_scaling_factor current_stream.wait_stream(self.alt_stream) final_hidden_states = final_hidden_states + shared_out - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) @@ -417,12 +412,10 @@ class SarvamMoESparseMoeBlock(nn.Module): elif self.routed_scaling_factor != 1.0: final_hidden_states = final_hidden_states * self.routed_scaling_factor - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index 55ccc04f2..c09bfeb17 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -34,8 +34,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -161,13 +160,10 @@ class SDARMoeSparseMoeBlock(nn.Module): topk_output = self.topk(hidden_states, router_logits) out = self.experts(hidden_states, topk_output) # (T, H) - # TP all-reduce (unless fused / reduce_scatter / fp4 allgather / dp reduce_scatterv path) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): out = tensor_model_parallel_all_reduce(out) diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index bcb978e0d..1f3a4d221 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -31,8 +31,7 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, - should_use_dp_reduce_scatterv, - should_use_flashinfer_cutlass_moe_fp4_allgather, + should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -233,12 +232,10 @@ class Step3p5MoEMLP(nn.Module): router_logits=topk_output.router_logits, ) final_hidden_states = self.experts(hidden_states, topk_output) - if ( - self.tp_size > 1 - and not should_allreduce_fusion - and not use_reduce_scatter - and not should_use_flashinfer_cutlass_moe_fp4_allgather() - and not should_use_dp_reduce_scatterv() + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)