diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index ed402a283ed7..9bc7486eba65 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -57,6 +57,7 @@ class AiterMoeQuantInfo(MoeQuantInfo): hidden_pad: int = 0 intermediate_pad: int = 0 swiglu_limit: float = 0.0 + gate_mode: str = "separated" @dataclass @@ -130,8 +131,13 @@ def run( extra["num_local_tokens"] = runner_input.num_local_tokens if runner_input.output_dtype is not None: extra["dtype"] = runner_input.output_dtype + # Single source of truth for gate_mode. The MXFP4 (GPT-OSS) path sets + # quant_info.gate_mode from the activation type; the FP8/DSv4 path drives + # it via swiglu_limit. swiglu_limit > 0 implies the interleaved gate/up + # layout, so force INTERLEAVE there to preserve the original behavior. + gate_mode = quant_info.gate_mode if quant_info.swiglu_limit > 0: - extra["gate_mode"] = GateMode.INTERLEAVE.value + gate_mode = GateMode.INTERLEAVE.value extra["swiglu_limit"] = quant_info.swiglu_limit output = fused_moe( @@ -152,6 +158,7 @@ def run( doweight_stage1=quant_info.doweight_stage1, hidden_pad=quant_info.hidden_pad, intermediate_pad=quant_info.intermediate_pad, + gate_mode=gate_mode, **extra, ) return AiterRunnerOutput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 6ac793d0270f..0fed7c9005e0 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -1221,6 +1221,8 @@ def apply( )[0] return StandardCombineInput(hidden_states=trtllm_gen_output) if _use_aiter: + from aiter.ops.flydsl.moe_common import GateMode + from sglang.srt.layers.moe.moe_runner.aiter import ( AiterMoeQuantInfo, AiterQuantType, @@ -1248,6 +1250,11 @@ def apply( doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, hidden_pad=self.hidden_pad, intermediate_pad=self.intermediate_pad, + gate_mode=( + GateMode.INTERLEAVE.value + if self.runner.config.activation == "swiglu" + else GateMode.SEPARATED.value + ), ) return self.runner.run( dispatch_output._replace(hidden_states=x_padded), quant_info