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

[bingxche]

Motivation

Resolve issue ROCm/aiter#3227 surfaced after AITER commit ff9bf15cc5c5afcd527d5695cf5462ba0bc1cc2f (ROCm/aiter#3123)

Credit to @XiaobingSuper

Fix the nightly-accuracy-8-gpu-mi35x / nightly-accuracy-8-gpu-mi35x-rocm720 failure on test_gpt_oss_eval_mi35x.py (runs 76071145339 / 76071145701).

The crash chain (logs):

[aiter] [fused_moe] using 2stage default for
  (256, 128, 3072, 512, 32, 4, <ActivationType.Swiglu: 2>,
   'torch.bfloat16', 'torch.bfloat16', 'torch.float4_e2m1fn_x2',
   'QuantType.per_1x32', True, False)

[aiter] start build [module_moe_ck2stages_b16_fp4x2_preshuffle_off_b16_swiglu_per_1x32_mulWeightStage2]

gen_instances.py: error: argument -act/--activation:
    invalid choice: 'swiglu' (choose from 'silu', 'gelu')

/sgl-workspace/aiter/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu:6:10:
    fatal error: 'gemm_moe_ck2stages_lookup.h' file not found

ModuleNotFoundError: No module named
  'aiter.jit.module_moe_ck2stages_b16_fp4x2_preshuffle_off_b16_swiglu_per_1x32_mulWeightStage2'

Root cause:

1. Mxfp4MoEMethod hard-codes activation="swiglu" for the AITER backend (python/sglang/srt/layers/quantization/mxfp4.py:1019).
2. aiter.fused_moe(...) exposes gate_mode as a caller-controlled parameter (defaults to GateMode.SEPARATED). SGLang was not setting it, so AITER took the SEPARATED path.
3. For the gpt-oss MI35x workload no tuned config row matches, so AITER falls back to the default CK 2-stage kernel and synthesizes a JIT module name carrying swiglu (module_moe_ck2stages_..._swiglu_per_1x32_...).
4. The CK 2-stage codegen (ck_gemm_moe_2stages_codegen/gen_instances.py) only accepts silu or gelu for -act. Codegen aborts, gemm_moe_ck2stages_lookup.h is never produced, the subsequent .cu compile fails, and the scheduler dies with ModuleNotFoundError.

The correct path for MXFP4 + swiglu on AITER is the FlyDSL interleaved-gate stage1 kernels (flydsl_moe1_afp4_wfp4_bf16_*), which natively implement swiglu via fused silu_and_mul_fq. Those kernels are selected when the caller passes gate_mode=GateMode.INTERLEAVE.

Modifications

• python/sglang/srt/layers/moe/moe_runner/aiter.py
  • Add gate_mode: str = "separated" to AiterMoeQuantInfo (default preserves existing AITER MoE semantics — AITER's own default is also SEPARATED).
  • In AiterRunnerCore.run, resolve a single gate_mode and forward it into the aiter.fused_moe(...) call. This is the post-[MoE] Unify DeepEPMoE+MoriEPMoE through AITER MoeRunner pre/post-permute #23760 equivalent of the original patch's fused_experts_none_to_aiter edit, and now covers the unified standard / DeepEP normal / DeepEP LL dispatch paths through the AITER runner.
• python/sglang/srt/layers/quantization/mxfp4.py
  • In Mxfp4MoEMethod.apply AITER branch, set gate_mode=GateMode.INTERLEAVE.value when self.runner.config.activation == "swiglu", otherwise keep GateMode.SEPARATED.value. Import GateMode from aiter.ops.flydsl.moe_common locally inside the branch (kept lazy because aiter is ROCm-only).

Reconcile with main #26208 (added during the main merge)

After branching, main landed #26208 which, inside AiterRunnerCore.run, sets extra["gate_mode"] = GateMode.INTERLEAVE.value and extra["swiglu_limit"] whenever quant_info.swiglu_limit > 0 (the DSv4 FP8 path). A naive merge would then pass gate_mode twice to aiter.fused_moe(...) (once via this PR's explicit kwarg, once via **extra), raising TypeError: fused_moe() got multiple values for keyword argument 'gate_mode' on any swiglu_limit > 0 call.

The merge therefore reconciles both into a single source of truth for gate_mode in run:

gate_mode = quant_info.gate_mode          # MXFP4/GPT-OSS sets this from activation
if quant_info.swiglu_limit > 0:           # FP8/DSv4 path
    gate_mode = GateMode.INTERLEAVE.value  # swiglu_limit > 0 implies interleaved gate/up
    extra["swiglu_limit"] = quant_info.swiglu_limit
output = fused_moe(..., gate_mode=gate_mode, **extra)   # never via extra → no double-pass

This is byte-for-byte equivalent to #26208 on the DSv4 FP8 path (swiglu_limit > 0 ⇒ gate_mode="interleave" + swiglu_limit forwarded), so DSv4 behavior is unchanged, while the MXFP4/GPT-OSS path gets INTERLEAVE from quant_info.gate_mode as intended.

Required AITER version. This patch depends on two AITER changes:

1. aiter.fused_moe(...) exposing gate_mode as a caller-controlled kwarg.
2. The aiter.ops.flydsl.moe_common module that defines GateMode.

Both are present from ROCm/aiter PR #3123 (commit ff9bf15c). Older AITER pins (e.g. a6bb499... shipped in some rocm/sgl-dev images) will fail at runtime with either TypeError: fused_moe() got an unexpected keyword argument 'gate_mode' or ModuleNotFoundError: No module named 'aiter.ops.flydsl.moe_common'. CI images / Dockerfile pins must be bumped to at least ff9bf15c, or workflows must be dispatched with aiter_ref set to a sufficiently new commit, before this PR can land cleanly across all AITER MoE tests (LLaDA2, Qwen3 MoE, DeepEP, etc., not only MXFP4).

For AITER versions that meet the above bound, this change is a no-op on every non-MXFP4 / non-swiglu AITER path: quant_info.gate_mode stays at the "separated" default and matches AITER's own default behavior.

Tests (commit 134d74a)

Re-dispatched nightly-accuracy-8-gpu-mi35x on this PR branch after merging latest main (which required reconciling the gate_mode change with the new swiglu_limit handling from #26208). Both passed ✅:

• ROCm 7 — ✅ nightly-accuracy-8-gpu-mi35x (run)
• ROCm 7.2 — ✅ nightly-accuracy-8-gpu-mi35x-rocm720 (run)

Speed Tests and Profiling

The INTERLEAVE FlyDSL stage1 kernels are the ones the MXFP4 tuned CSVs target on gfx950, so this also takes the MXFP4 path off the un-tuned CK default and onto the intended high-performance path. No standalone microbench in this PR.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. (Covered by the existing AMD MI35x nightly accuracy job.)
• Update documentation according to Write documentations. (N/A — internal kernel selection.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (Will rely on the MI35x nightly re-run.)
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #26733545189
Latest PR Test (Extra): ❌ Run #26733545103

[cursor]

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
_{Learn more about Cursor Agents}

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[bingxche]

@amd-bot ci-status

[amd-bot]

@bingxche

CI Status for PR #25580

PR: [AMD] fix(mxfp4): route AITER MXFP4+swiglu through FlyDSL gate_mode=INTERLEAVE
Changed files: python/sglang/srt/layers/moe/moe_runner/aiter.py (+3/-0), python/sglang/srt/layers/quantization/mxfp4.py (+7/-0)

AMD: 3 test-job failures (0 likely related) | Others: 16 base-b job failures (0 related — all fast-failed on lint)

AMD CI Failures

Job	Test File	Test Function	Error	Related?	Explanation	Log
stage-b-test-1-gpu-large-amd (linux-mi300-1gpu-sglang, 0)	test/registered/perf/test_bench_serving_1gpu_part1.py	test_online_lora_latency_with_concurrent_adapter_updates	AssertionError: 15850.98 not less than 4000 (median E2E latency)	🟢 Unlikely	LoRA-adapter online-serving latency perf test on a Llama-3.1-8B BF16 path. PR only touches the MXFP4-AITER swiglu branch (Mxfp4MoEMethod.apply + AiterMoeQuantInfo.gate_mode default = "separated"). No LoRA / BF16 path passes through the modified gate_mode branch — extra["gate_mode"] is only injected when quant_info.gate_mode != "separated". Looks like a perf flake / runner-noise issue.	Log
stage-b-test-1-gpu-large-amd (linux-mi300-1gpu-sglang, 1)	test/registered/perf/test_bench_serving_1gpu_part2.py	test_score_api_batch_scaling	AssertionError: 80.09 not less than 70 (retry: 81.45 not less than 80)	🟢 Unlikely	Score-API latency perf test, marginal failure ~1.5% over threshold. Unrelated to MXFP4 — uses the score API path, not the MoE quant runner the PR touches.	Log
stage-b-test-2-gpu-large-amd (linux-mi300-2gpu-sglang, 0)	test/registered/quant/test_int4fp8_moe.py	N/A (job-level timeout)	Scheduler watchdog timeout (300s) then The action 'Run test' has timed out after 45 minutes. Hung in aiter JIT build for module_moe_ck2stages_f8_i4_preshuffle_off_b16_silu_per_token_mulWeightStage2 (ActivationType.Silu, QuantType.per_Token).	🟢 Unlikely	Int4-FP8 MoE path, not MXFP4. Uses Silu, not Swiglu, so the new gate_mode=INTERLEAVE branch is never hit; AiterMoeQuantInfo.gate_mode default stays "separated" and extra["gate_mode"] is not injected. Symptom is JIT build of a silu_per_token 2-stage kernel hanging — looks like an aiter-JIT/CK build / runner-disk issue, not anything this PR could affect.	Log

Other CI Failures

Job	Test File	Test Function	Error	Related?	Explanation	Log
lint	N/A	N/A	ruff F821 Undefined name 'BaseTpWorker' at python/sglang/srt/managers/scheduler.py:1004 (def get_draft_kv_pool(*, draft_worker: "BaseTpWorker", …))	🟢 Unlikely	The PR does not modify scheduler.py. The undefined BaseTpWorker annotation is in code that lives in the PR's base merge ref, not in this PR's diff. This is a pre-existing/transient base-branch lint issue — needs the missing BaseTpWorker import (TYPE_CHECKING) to be added on main.	Log
15× base-b-test-{1,2,4}-gpu-{small,large,b200} + wait-for-base-b + pr-test-finish	N/A	N/A	All show ##[error]Fast-fail: lint check failed — never executed any test.	🟢 Unlikely	Cascade of the lint failure above; the base-b stage is gated on lint. Once the unrelated BaseTpWorker ruff error is fixed on main and the PR is rebased, these jobs will run.	Run

Details

PR semantics (why none of the failures are PR-caused):

• python/sglang/srt/layers/moe/moe_runner/aiter.py: adds gate_mode: str = "separated" field to AiterMoeQuantInfo and only injects extra["gate_mode"] into the aiter.fused_moe(...) call when quant_info.gate_mode != "separated". For every existing caller (non-MXFP4, non-swiglu paths — including the failing int4fp8_moe, bf16 LoRA, score-API/perf benches), the default is "separated" ⇒ no behavior change vs. main.
• python/sglang/srt/layers/quantization/mxfp4.py: change is gated on the _use_aiter branch of Mxfp4MoEMethod.apply AND on self.runner.config.activation == "swiglu". None of the three failing AMD tests, nor any non-AMD test, exercise that combination.

Blocker / next step: the cascade of NV/base-b failures and the lint failure are caused by a pre-existing F821 BaseTpWorker ruff error in scheduler.py on the base ref — not by this PR. Rebasing onto a main that has the BaseTpWorker TYPE_CHECKING import fixed will unblock all 16 non-AMD failures. The 3 AMD failures appear to be perf flakes / JIT-build hang on the runner (none touch the MXFP4-swiglu path) — recommend /rerun-failed-ci after the rebase to confirm.

Generated by amd-bot using Claude Code CLI

[1am9trash]

LGTM