fix(gpt-oss): avoid duplicate MoE reduction with DP attention (#37199)

python/sglang/srt/models/gpt_oss.py

@@ -44,7 +44,10 @@ from sglang.srt.layers.linear import (
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
-from sglang.srt.layers.moe import get_moe_a2a_backend
+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.layer import FusedMoE
 from sglang.srt.layers.moe.topk import TopK
@@ -329,7 +332,9 @@ class GptOssSparseMoeBlock(nn.Module):
             topk_output = self.topk(router_input, router_logits)
             final_hidden_states = self.experts(hidden_states, topk_output)
 
-        if self.tp_size > 1 and not get_forward().fuse_mlp_allreduce:
+        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
+            is_tp_path=True,
+        ):
             final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
         # When input was pre-padded, FusedMoE.forward_impl captured the
@@ -606,6 +611,7 @@ class GptOssDecoderLayer(nn.Module):
             layer_scatter_modes=self.layer_scatter_modes,
             input_layernorm=self.input_layernorm,
             post_attention_layernorm=self.post_attention_layernorm,
+            allow_reduce_scatter=True,
             is_last_layer=(
                 self.is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
             ),
@@ -639,7 +645,14 @@ class GptOssDecoderLayer(nn.Module):
             )
         )
 
-        with get_forward().scoped(fuse_mlp_allreduce=fuse_mlp_allreduce):
+        mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
+            forward_batch
+        )
+
+        with get_forward().scoped(
+            fuse_mlp_allreduce=fuse_mlp_allreduce,
+            mlp_reduce_scatter=mlp_reduce_scatter,
+        ):
             hidden_states = self.mlp(hidden_states, forward_batch)
 
         if fuse_mlp_allreduce:

test/registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py

@@ -3,8 +3,8 @@ import unittest
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.gpt_oss_common import BaseTestGptOss
 
-register_cuda_ci(est_time=220, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100")
-register_cuda_ci(est_time=220, stage="base-c", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-b200")
 
 
 class TestGptOss4GpuMxfp4(BaseTestGptOss):
@@ -23,6 +23,28 @@ class TestGptOss4GpuMxfp4(BaseTestGptOss):
             ],
         )
 
+    def test_mxfp4_120b_dpa(self):
+        self.run_test(
+            model_variant="120b",
+            quantization="mxfp4",
+            expected_score_of_reasoning_effort={
+                "low": 0.50,
+            },
+            other_args=[
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--ep",
+                "4",
+                "--enable-dp-attention",
+                "--moe-a2a-backend",
+                "none",
+                "--cuda-graph-max-bs-decode",
+                "200",
+            ],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()