Skip to content

[WIP DO NOT MERGE] [AMD] fix(mxfp4): route AITER MXFP4+swiglu through FlyDSL gate_mode=INTERLEAVE #25580

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bingxche wants to merge 3 commits into from
+15 −1
Closed

[WIP DO NOT MERGE] [AMD] fix(mxfp4): route AITER MXFP4+swiglu through FlyDSL gate_mode=INTERLEAVE #25580

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
24b408c
fix(mxfp4): route AITER MXFP4+swiglu through FlyDSL gate_mode=INTERLEAVE
cursoragent May 20, 2026
3d3937d
Merge branch 'main' into bingxche/mxfp4-aiter-swiglu-gate-mode-5279
bingxche May 22, 2026
134d74a
Merge main into bingxche/mxfp4-aiter-swiglu-gate-mode-5279
bingxche Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion python/sglang/srt/layers/moe/moe_runner/aiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand Down
7 changes: 7 additions & 0 deletions python/sglang/srt/layers/quantization/mxfp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
Loading