diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 0db01f79e..6f12d50f7 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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: diff --git a/test/registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py b/test/registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py index 07f185c69..7598e09ef 100644 --- a/test/registered/models_e2e/test_gpt_oss_4gpu_mxfp4.py +++ b/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=220, stage="base-c", runner_config="4-gpu-b200") +register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100") +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()