From 24b408c2f03013b97ad155b372470540a140cb40 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 20 May 2026 11:24:38 +0000 Subject: [PATCH] fix(mxfp4): route AITER MXFP4+swiglu through FlyDSL gate_mode=INTERLEAVE AITER's CK 2-stage MoE codegen (gen_instances.py) only accepts silu/gelu as activation; passing swiglu makes the JIT build fail with 'invalid choice: swiglu' -> 'gemm_moe_ck2stages_lookup.h not found' -> ModuleNotFoundError when no tuned config row matches the workload (e.g. gpt-oss-20b/120b MXFP4 on mi35x gfx950 8-GPU). aiter.fused_moe exposes gate_mode as a caller-controlled parameter (default GateMode.SEPARATED). The MXFP4 path was not setting it, so swiglu fell back to the CK default and tried to JIT a non-existent ck2stages swiglu kernel. Thread gate_mode through AiterMoeQuantInfo and AiterRunnerCore.run, and select GateMode.INTERLEAVE in Mxfp4MoEMethod.apply when activation is swiglu. The FlyDSL interleaved-gate stage1 kernels natively support swiglu and are what the MXFP4 tuned CSV rows target. Note: AITER >= the commit that adds the gate_mode kwarg to aiter.fused_moe and ships aiter.ops.flydsl.moe_common is required. Co-authored-by: Bingxu Chen --- python/sglang/srt/layers/moe/moe_runner/aiter.py | 2 ++ python/sglang/srt/layers/quantization/mxfp4.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 0e4ab204c147..fe68341cc261 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -56,6 +56,7 @@ class AiterMoeQuantInfo(MoeQuantInfo): doweight_stage1: bool = False hidden_pad: int = 0 intermediate_pad: int = 0 + gate_mode: str = "separated" @dataclass @@ -147,6 +148,7 @@ def run( doweight_stage1=quant_info.doweight_stage1, hidden_pad=quant_info.hidden_pad, intermediate_pad=quant_info.intermediate_pad, + gate_mode=quant_info.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 fa1142b21411..7edd72458903 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -1194,6 +1194,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, @@ -1221,6 +1223,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