diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index faaadc90a..4d9a400a4 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -58,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -386,6 +387,7 @@ class BailingMoESparseMoeBlock(nn.Module): 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() ): 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 0833a68c2..d8750da8f 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -34,6 +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.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 @@ -347,7 +348,12 @@ 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: + if ( + self.tp_size > 1 + and not use_reduce_scatter + and not should_allreduce_fusion + and not should_use_dp_reduce_scatterv() + ): 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 6b23c2c88..9f2c82405 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -85,6 +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_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -655,6 +656,7 @@ class DeepseekV2MoE(nn.Module): 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() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states @@ -744,6 +746,7 @@ class DeepseekV2MoE(nn.Module): 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() ): 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 ab31c0a59..18012448a 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -47,7 +47,7 @@ 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 +from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv 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,7 +300,11 @@ 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: + if ( + self.tp_size > 1 + and not use_reduce_scatter + and not should_use_dp_reduce_scatterv() + ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 7e8dd9605..092b94b7e 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -61,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -598,6 +599,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): 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() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states @@ -632,6 +634,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): 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() ): 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 b0f87dc32..f11827b98 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/python/sglang/srt/models/hunyuan_v3.py @@ -34,6 +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.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -191,10 +192,11 @@ class HYV3MoEFused(nn.Module): hidden_states=hidden_states, topk_output=topk_output ) - if self.ep_size > 1: + skip_post_reduce = should_use_dp_reduce_scatterv() + if self.ep_size > 1 and not skip_post_reduce: final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) - if self.tp_size > 1: + if self.tp_size > 1 and not skip_post_reduce: final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states ) @@ -222,10 +224,11 @@ class HYV3MoEFused(nn.Module): current_stream.wait_stream(self.alt_stream) final_hidden_states = final_hidden_states + shared_output - if self.ep_size > 1: + skip_post_reduce = should_use_dp_reduce_scatterv() + if self.ep_size > 1 and not skip_post_reduce: final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) - if self.tp_size > 1: + if self.tp_size > 1 and not skip_post_reduce: 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 85c9e6a32..194abf32b 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -55,7 +55,11 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import get_deepep_mode, get_moe_a2a_backend +from sglang.srt.layers.moe import ( + get_deepep_mode, + get_moe_a2a_backend, + should_use_dp_reduce_scatterv, +) 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.token_dispatcher import DeepEPDispatcher @@ -379,7 +383,11 @@ 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: + if ( + self.tp_size > 1 + and not use_reduce_scatter + and not should_use_dp_reduce_scatterv() + ): 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 f0a237aa6..c2e4b12c9 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -39,6 +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.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -145,7 +146,11 @@ class Llama4MoE(nn.Module): out_aD = routed_out + shared_out - if self.tp_size > 1 and not use_reduce_scatter: + if ( + self.tp_size > 1 + and not use_reduce_scatter + and not should_use_dp_reduce_scatterv() + ): out_aD = tensor_model_parallel_all_reduce(out_aD) return out_aD diff --git a/python/sglang/srt/models/mimo_v2_flash.py b/python/sglang/srt/models/mimo_v2_flash.py index 9902b63d4..b6dd6c4b0 100644 --- a/python/sglang/srt/models/mimo_v2_flash.py +++ b/python/sglang/srt/models/mimo_v2_flash.py @@ -51,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class @@ -302,6 +303,7 @@ class MiMoV2MoE(nn.Module): 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() ): 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 b99c2471b..d6a38d082 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -61,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -556,6 +557,7 @@ class MiniMaxM2MoE(nn.Module): 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() ): final_hidden_states = 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 9c4179362..faafd8c78 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -39,7 +39,10 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput -from sglang.srt.layers.moe import should_use_flashinfer_cutlass_moe_fp4_allgather +from sglang.srt.layers.moe import ( + should_use_dp_reduce_scatterv, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) 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 @@ -375,6 +378,7 @@ class SarvamMoESparseMoeBlock(nn.Module): 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() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) @@ -418,6 +422,7 @@ class SarvamMoESparseMoeBlock(nn.Module): 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() ): 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 a1b20f67c..55ccc04f2 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -34,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -160,12 +161,13 @@ 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 path) + # 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() ): out = tensor_model_parallel_all_reduce(out) diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index e1f13e604..bcb978e0d 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -31,6 +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, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -237,6 +238,7 @@ class Step3p5MoEMLP(nn.Module): 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() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)