[Bug] ROCm images pin AITER below the FlyDSL MXFP4 MoE kernels, making tuned AITER_CONFIG_FMOE tables unusable on gfx950

[ThomasNing]

Summary

docker/rocm.Dockerfile pins AITER_COMMIT_DEFAULT=d9e5ef7ce08ee7045d583aed768cff41aa9210fe (ROCm/aiter, 2026-07-29) for all four arch stages. That pin predates ROCm/aiter db24351e — "[fmoe][gfx950] Flydsl mxmoe v2" (#4179), 2026-08-04, which introduces the flydsl_mxmoe_* / flydsl_moe2_layout_* FlyDSL MXFP4 fused-MoE kernel families.

Consequence on gfx950: SGLang cannot consume an AITER-tuned MXFP4 fused-MoE table (AITER_CONFIG_FMOE), and cannot produce one either, so the tuned-MoE kernel-selection path is unavailable on published ROCm images.

Repro

Load any AITER_CONFIG_FMOE table tuned with AITER ≥ db24351e (rows naming flydsl_mxmoe_* kernels) — the scheduler dies at model load on every TP rank:

ValueError: Invalid FlyDSL kernel name: flydsl_mxmoe_g1_a4w4_16x256x256_f16in_nt

raised from aiter/fused_moe.py after aiter.ops.flydsl.moe_kernels.get_flydsl_kernel_params() returns None.

Probing the registry inside the published images confirms the cause is the pin, not the table. Model GLM-5.2-MXFP4, TP4, gfx950 (MI350X); 4 distinct kernel names in the table:

image	published	bundled AITER	FlyDSL registry size	table's kernels resolved	--mxfp4-flydsl tuner mode
lmsysorg/sglang:v0.5.16-rocm720-mi35x	2026-07-24	9127c94a (06-25)	3393	0 / 4	no
lmsysorg/sglang:v0.5.17-rocm720-mi35x	2026-08-08	d9e5ef7c (07-29)	4257	0 / 4	no

Both registries contain only the older flydsl_moe1_* / flydsl_moe2_* families (1057 of them afp4_wfp4 in v0.5.16) — zero entries matching mxmoe, a4w4, or layout.

Note v0.5.17's image was published four days after the floor commit while its AITER pin is six days before it, so the image tag date is not a reliable proxy here.

Retuning locally is not a workaround: the bundled csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py at this pin has no MXFP4-FlyDSL mode (--grouped-gemm is gfx1250-only), so --mxfp4-flydsl is rejected with unrecognized arguments.

Impact

On MI350X/gfx950 with GLM-5.2-MXFP4 (TP4), the tuned MXFP4 MoE table measured +7.5–12.9% end-to-end output throughput at concurrency 16–64 on an engine that can load it (kernel-level: +21.6% at M=32, +19.1% at M=16, +7.7% at M=1; 6 of 11 shapes improved past a 3% gate). That gain is currently unreachable from any published SGLang ROCm image.

Suggested fix

Bump AITER_COMMIT_DEFAULT to a commit ≥ db24351e (or current AITER main) — at minimum for the gfx950 / gfx950-rocm720 stages. This follows the existing routine-bump precedent (#28248, #28249, #27555, #27376).

Users can already work around it today via the documented build arg (thanks to #21949):

docker build -f docker/rocm.Dockerfile --build-arg GPU_ARCH=gfx950-rocm720 \
  --build-arg AITER_COMMIT=<commit ≥ db24351e> .

Happy to submit the bump PR if a maintainer confirms which AITER commit you'd like to land on — I have a gfx950 (MI350X) node and can validate GLM-5.2-MXFP4 TP4 serving plus the tuned-table load before opening it.

Environment

• 8× MI350X (gfx950), ROCm 7.2, TP4
• Model: GLM-5.2-MXFP4
• Images as tabled above; probe via aiter.ops.flydsl.moe_kernels.get_flydsl_kernel_params(name) is not None

[ThomasNing]

Correction to my own analysis above — the conclusion holds, the stated mechanism was wrong

I validated the fix, and in doing so found that the evidence I gave in the issue body was based on the wrong lookup. Correcting it here rather than leaving it to mislead someone.

What I got wrong

I reported that the bundled AITER's FlyDSL registry "contains zero mxmoe / a4w4 / layout entries" and inferred the kernels were absent, identifying db24351e as the floor because that commit's diff contains the kernel name.

Both steps were faulty:

• My probe called aiter.ops.flydsl.moe_kernels.get_flydsl_kernel_params(). That registry (_KERNEL_PARAMS) is pre-populated at import from the stage1/stage2 builders only. flydsl_moe2_layout_* names come from get_flydsl_stage2_v2_kernels, and flydsl_mxmoe_* from mxmoe_dispatcher / parse_g2_kname_any. The string mxmoe does not appear in moe_kernels.py at any version. So those names report missing on every image — including images that serve them correctly. The 0/4 result was a false negative, not a floor.
• The db24351e citation came from that commit's diff touching aiter/aot/flydsl/mxfp4_moe.py — AOT codegen, not the registry. I matched the string without checking what consumed it.

What is actually true (validated by serving, not probing)

Replicating this Dockerfile's own AITER step inside the published lmsysorg/sglang:v0.5.17-rocm720-mi35x image — git checkout -f fc2e5d57 (AITER v0.1.20, 2026-08-18), then PREBUILD_KERNELS=1 GPU_ARCHS=gfx950 setup.py build_ext --inplace and the editable install:

	shipped pin d9e5ef7c	rebuilt at fc2e5d57
GLM-5.2-MXFP4 TP4 + AITER_CONFIG_FMOE (tuned MXFP4 MoE table)	scheduler dies on every TP rank	server READY, serves
Invalid FlyDSL kernel name in log	yes	zero occurrences
tuner --mxfp4-flydsl mode	absent (unrecognized arguments)	present (Mxfp4FlydslTuner, no arch gate)
registry probe (both)	0/4	0/4 — irrelevant, wrong lookup

So the real difference is the mxmoe front-end in aiter/fused_moe.py. On the older pin the tuned kernel name falls through to the generic FlyDSL wrapper, which raises ValueError: Invalid FlyDSL kernel name. At v0.1.20 the mxmoe path exists and handles it.

Note the registry probe reports 0/4 on both pins, including the one that serves — which is why I am reporting this as a serving result. A registry probe cannot settle this question in either direction.

Net effect on the request

Unchanged, and now validated: bumping AITER_COMMIT_DEFAULT to a pin containing the mxmoe front-end restores both consuming and producing tuned MXFP4 MoE tables on gfx950. AITER v0.1.20 (fc2e5d57fb5b8ad8e7e23f7103071dde798ea618) is a tagged release that works.

Caveats I would want a maintainer to weigh

• Transitive flydsl bump. Moving to v0.1.20 pulls flydsl 0.2.4 → 0.3.1 (pip flags the conflict during the switch). [Bug] ROCm diffusion fused-norm kernels depend on non-public FlyDSL APIs #33142 documents ROCm diffusion fused-norm kernels depending on non-public FlyDSL APIs, with fix(rocm): support flydsl 0.3.0 in the FlyDSL fused norm kernel #34536 / [Diffusion] Migrate FlyDSL fused norm kernels to the v0.3.0 stable API #33144 open to migrate them. A bump probably wants to land after that migration.
• Validation scope. LLM serving only (GLM-5.2-MXFP4, TP4, gfx950), by rebuilding AITER inside the published image — not a full docker build of every arch stage, and nothing on gfx942 or diffusion paths.
• Throughput attribution (tuned vs. untuned table on the same image) is being measured now; I will post it when the repeats are in rather than quote a single run.

Happy to open the bump PR whenever you want it, scoped to whichever stages you prefer.