[ZLkanyo009]

Summary

Qwen3.8-Flash-Next-Quark-MXFP4's routed MoE (E=513 incl. the fused shared expert, topk 11, model_dim 2560, inter 640, a4w4 per_1x32) had no tuned fused_moe rows on gfx950 and ran the heuristic FlyDSL v1 kernels. This PR adds (513, 2560, 768, 11) to the MXFP4 MoE aux shapes so the MXMOE port (flydsl_mxmoe_g1_a4w4_* with inline sort/quant + flydsl_moe2_layout_*) can run it, and a tuned table for token 1–16384. MoE kernel time per step drops 24% at the decode working point and 29% at the 16k prefill working point.

The rows are keyed on inter_dim=640, the unpadded layout. ATOM pads this model's experts to 768 today and MXMOE does not take an intermediate pad, so the table only takes effect together with ROCm/ATOM#2649 (keeps 640 unpadded at TP1). Without that PR nothing changes (no rows match 768). Running the model with MTP also depends on ROCm/ATOM#2606.

Motivation

Qwen3.8-Flash-Next (amd/Qwen3.8-Flash-Next-Quark-MXFP4) optimization on MI355X, https://github.com/wuhuikx/LLM-Booster/issues/263 (item 3, MoE tuning: "prefill stage2 is slow", tuned_fmoe for this shape).

Changes

• aiter/ops/moe_mxfp4_aux.py, csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py: add (513, 2560, 768, 11) (aux key = inter rounded up to 256, like MiniMax-M3's 384 → 512). Rebuilds module_moe_mxfp4_aux (~18 s JIT).
• aiter/configs/model_configs/qwen3_8_flash_fp4_{tuned,untuned}_fmoe.csv: 15 rows, token 1, 2, 4, …, 16384 (get_padded_M buckets), gfx950 / cu_num 256. BM16 inline-quant gemm1 for token ≤ 512, BM64 for 1024–8192, BM128 + scatter gemm2 for 16384.
• op_tests/test_mxfp4_moe_aux_codegen.py: CPU check that every entry of MXFP4_MOE_SUPPORTED_SHAPES has a generated aux instance (read via ast, no aiter import).

How the table was tuned: gemm_moe_tune.py's isolated stage timing did not match the production path for this shape (e.g. it reports 12 µs for a token-64 stage1 that must read ~600 MB of weights, and leaves out moe_reduction / scatter_reduce), so the candidates (all Mxfp4FlydslTuner._g1_variants gemm1 configs × layout gemm2 with matching tile_m) were timed end to end through aiter.fused_moe (sort + quant + both stages + reduce; CUDA graph replay for token ≤ 512), with ATOM's weight shuffle, the shared expert always selected, 6 rotating weight copies (so decode does not sit in MALL), and for token ≥ 1024 routing dumped from the real model on long coding prompts. Candidates with > 5% relative L2 error vs the default path were dropped. The us column is that end-to-end time.

Accuracy

GSM8K (lm-eval, real MTP verification), TP1, MTP 3:

config	strict-match	flexible-extract	MTP acceptance	avg tokens / forward
before (ATOM main: inter 768, heuristic FlyDSL)	0.9780 ± 0.0040	0.9780	0.721	3.163
ATOM PR only (inter 640, heuristic FlyDSL)	0.9757 ± 0.0042	0.9757	0.720	3.159
ATOM PR + this PR (MXMOE)	0.9757 ± 0.0042	0.9757	0.724	3.173

Performance

Test config: image rocm/atom-dev:latest (566aa3dc9d5d), ATOM 53e4fb524 (+ #2606 + the ATOM PR above), AITER 948dd267f + this PR's diff (the 13 commits between 948dd267f and this PR's base touch none of the files or kernels involved), MI355X, TP1, MTP 3, bf16 KV.

Server profile under AgentX C16 load (torch profiler, 15 s window after 420 s warm-up). Decode = all pure decode steps replaying the bs=16 graph (decode[bs=9/16 tok=36 d=9 spec=3] pads to it); prefill = prefill[bs=1 tok=16384 ctx=98304]. µs per step (48 MoE layers):

working point	config	kernel total	MoE total	stage1	stage2	reduce	sort+quant
decode bs=16 graph	before	18,211	7,916	3,935	2,307	–	1,118
	ATOM PR + this PR	16,387 (−10.0%)	6,030 (−23.8%)	3,304	1,747	–	438
prefill tok=16384 ctx=98304	before	473,463	74,177	30,624	38,341	–	4,021
	ATOM PR + this PR	429,626 (−9.3%)	52,592 (−29.1%)	17,532	21,410	8,275	4,216

(Prefill non-MoE kernels vary ±5% between servers; the MoE columns are the comparable part.)

Isolated fused_moe (same weights layout as ATOM, shared expert always selected + 10 uniform routed, 6 rotating weight copies, µs per call):

token	before (768, heuristic)	ATOM PR + this PR
4	41.8	26.3 (−37%)
36 (decode point, bucket 64)	146.9	113.7 (−23%)
512	276.4	245.8 (−11%)
4096	521.6	466.4 (−11%)
16384	1808.7	1101.6 (−39%)

(The heuristic kernels at 640 unpadded take the same time as at 768 with intermediate_pad=128 within noise, e.g. decode 76.6 / 42.9 vs 75.9 / 43.6 µs for stage1 / stage2, so the gain here comes from the table.)

AgentX C16:

run	tput/GPU	p90 norm intvty tok/s/user	p90 ITL intvty	realized prefix hit	theoretical hit	prompt tokens	generated	errors/reqs	Effective Conc
before (main, 768)	45,111	8.69	17.28	60.7%	95.9%	178,920,730	1,463,928	0/1644	11.24
before, repeat	48,221	8.68	17.05	63.9%	95.9%	190,279,227	1,549,957	0/1799	10.96
ATOM#2649 only (640, AITER heuristic)	64,635	29.13	43.22	83.7%	96.0%	249,262,024	2,087,631	0/2351	7.96
ATOM#2649 + aiter#6708 (640, MXMOE)	69,775	33.93	52.45	85.9%	96.1%	266,589,865	2,247,026	0/2468	7.39
ATOM#2649 + aiter#6708, repeat	68,047	31.25	48.76	84.9%	96.1%	261,268,100	2,206,191	0/2433	7.57

Cache = GPU prefix, TP1, concurrency 16 for all rows. Runs were done two at a time on separate GPUs of one node (the "ATOM#2649 only" run and the second "+ aiter#6708" run ran side by side).

• The two before runs differ by 7% in tput/GPU: that is the run-to-run noise.
• ATOM#2649 alone: tput/GPU +38.5%, p90 norm interactivity 3.35× vs the mean of the two before runs. The extra 25% of KV blocks lifts the realized prefix hit from 62% to 84% on these long-prefix agentic traces, so far fewer prompt tokens are recomputed.
• aiter#6708 on top: vs the side-by-side ATOM-only run, tput/GPU +5.3%, p90 norm interactivity +7.3%, p90 ITL interactivity +12.8% (vs the mean of both runs: +6.6% / +11.9% / +17.1%). The ITL gain (the two before runs differ by only 1.3% there) matches the −10% decode-step kernel time; the tput gain is of the same size as the noise.

How to use

No flags. AITER merges aiter/configs/model_configs/*tuned_fmoe*.csv automatically; the rows match once ATOM runs the model with unpadded inter 640.

Test plan

• black --check / ruff check on the changed Python files.
• pytest op_tests/test_mxfp4_moe_aux_codegen.py: 3 passed (CPU).
• MI355X: module_moe_mxfp4_aux rebuilt with the new shape; isolated fused_moe hits all 15 rows ([fused_moe] using 2stage (kernelName1='flydsl_mxmoe_g1_a4w4_…'), no heuristic fallback) and stays within the accuracy gate; full server: GSM8K and AgentX C16 above, server log shows the MXMOE configs for (…, 2560, 640, 513, 11, …).

Limitations / follow-ups

• TP1 only (the shape at TP > 1 differs and is not tuned).
• The us column is from the isolated benchmark. Absolute times differ from the server (e.g. the heuristic 16k stage1 is 860 µs isolated vs 638 µs in the server, while the tuned one is ~365 µs in both), so the server profile above is the end-to-end reference.

Update after #6769

Merged origin/main at 31729e8db into this branch (merge commit 722791eba, no rebase / force push).

• Conflict: only the MXFP4 MoE aux shape lists, MXFP4_MOE_SUPPORTED_SHAPES in aiter/ops/moe_mxfp4_aux.py and SHAPES in csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py. Resolved as the union: main's entries and order are kept as they are ([HIP] [MoE] Qwen3.8 TP=8 MXFP4 MoE aux kernel codegen #6468's (513, 8192, 256, 11) and [HIP] [GLM-5.3-Flash] MXFP4 fused-MoE tuned configs for MI355X #6435's six GLM-5.3-Flash entries), and this PR's (513, 2560, 768, 11) goes right after main's Qwen3.8 TP=8 entry in both lists, so the two Qwen3.8 entries sit together.
• fmoe csv: no conflict and no format change needed. qwen3_8_flash_fp4_{tuned,untuned}_fmoe.csv has the same header as the other fmoe tables on main; the extra nt column in [HIP] [GLM-5.3-Flash] MXFP4 fused-MoE tuned configs for MI355X #6435's GLM table is optional in fused_moe.py. The merged tuned_fmoe family (54 files, 3414 rows) has no duplicate key.

Validation (MI355X, rocm/atom-dev:nightly_202610221638, ROCm 7.2.4, the merged tree with JIT built from scratch)

• op_tests/test_mxfp4_moe_aux_codegen.py: 3 passed. is_mxfp4_moe_shape_supported(513, 2560, 640, 11) is True.
• op_tests/test_moe_2stage.py --no-legacy --csv-filter flydsl_ with AITER_CONFIG_FMOE set to this PR's table, which covers (2560, 640, 513, 11): 15/15 cases pass, logits_diff 4.9e-4 to 9.4e-4. Every case hits a tuned flydsl row, with no default fallback.
• End to end: ATOM main (5fea029d, includes fix(qwen4_exp): load unscaled Quark MTP weights as BF16 ATOM#2606) + perf(qwen4-exp): keep 128-aligned MXFP4 expert intermediate unpadded ATOM#2649, TP1, real MTP-3 verification. GSM8K (1319) 0.9788 ± 0.0040 (reference 0.9757–0.978), MTP acceptance 0.723 (3.17 tokens/fwd). With AITER_LOG_TUNED_CONFIG=1, all 30 MXFP4 MoE lookups in the server hit tuned flydsl rows of this table, with no default fallback.
  • Note: on this ROCm 7.2.4 image the server cannot start with main's bf16 GEMM tables. It exits at warmup in hipb_mm on (16384, 10240, 2560), because the hipBLASLt solution index from [Config] [Perf] Add gfx950 a8w8 bpreshuffle and BF16 GEMM tunings for Qwen3.8 Flash Next #6769 does not exist in hipBLASLt 1.2.70204. That is unrelated to this PR and is fixed in [Config] [Perf] Tune gfx950 BF16 GEMMs for Qwen3.8-Flash-Next (TP1, MTP 3) #6711. For this run, AITER_CONFIG_GEMM_BF16 pointed at the merged bf16 tables of [Config] [Perf] Tune gfx950 BF16 GEMMs for Qwen3.8-Flash-Next (TP1, MTP 3) #6711, and every other table came from this branch.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6708 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[mjkvaak-amd]

vLLM A/B on MI355X, Qwen3.8-Flash-Next-Quark-MXFP4 at TP1: +4.6% output tok/s at TP1 (geomean of 11 points, +1.3% to +8.8%) and GSM8K unchanged.

Scope. vLLM rounds the MXFP4 expert intermediate up to 128, so at TP1 the routed MoE reaches AITER as (2560, 640, 513, 11) and these rows match. At TP2 (the model card's recipe) each rank sees 384, which no row matches, so the PR is a no-op there; similarly TP>2.

Setup. vllm/vllm-openai-rocm:nightly-df3bd094 (AITER 0.1.25) vs the same image + this PR @28eb6c3 (aiter/ + csrc/, module_moe_mxfp4_aux rebuilt). VLLM_ROCM_USE_AITER=1, fused shared expert, bf16 KV, no MTP, no prefix caching. 2 runs per arm, 1 GPU each, all 4 concurrent. Each server log on the PR arm selects flydsl_mxmoe_g1_a4w4_* (30 lines); stock selects none.

ISL/OSL, conc	stock tok/s	PR tok/s	Δ	TPOT stock (ms)	TPOT PR (ms)
8K/1K, 4	387	404	+4.4%	9.5	8.6
8K/1K, 8	612	666	+8.8%	12.0	11.0
8K/1K, 16	854	901	+5.4%	17.1	16.2
8K/1K, 32	1074	1126	+4.8%	26.3	25.0
8K/1K, 64	1275	1323	+3.8%	44.9	43.2
8K/1K, 128	1437	1498	+4.2%	76.7	73.4
8K/1K, 256	1596	1640	+2.8%	137.0	133.3
60K/600, 4	126	132	+4.9%	23.0	22.0
60K/600, 64	163	165	+1.3%	328.3	320.8
128K/1K, 4	108	111	+3.4%	27.3	26.2
128K/1K, 16	114	122	+6.5%	117.6	110.3

GSM8K (5-shot, n = 1319, flexible-extract): stock 0.869 / 0.870, PR 0.877 / 0.879 (± 0.009).

[yzhou103]

please resolve conflicts

[ZLkanyo009]

@yzhou103 Conflicts are resolved. I merged origin/main at 31729e8db (merge commit 722791eba, no force push).

• The only conflict was in the MXFP4 MoE aux shape lists (moe_mxfp4_aux.py and gen_instances.py). I took the union: main's entries stay as they were, including (513, 8192, 256, 11) from [HIP] [MoE] Qwen3.8 TP=8 MXFP4 MoE aux kernel codegen #6468 and the GLM entries from [HIP] [GLM-5.3-Flash] MXFP4 fused-MoE tuned configs for MI355X #6435, and (513, 2560, 768, 11) goes right after the existing Qwen3.8 entry.
• The fmoe csv merged cleanly. Its format already matches main, and the tuned_fmoe family has no duplicate key.
• I verified on MI355X with ROCm 7.2.4 and a JIT build from scratch:
  • test_mxfp4_moe_aux_codegen.py passes.
  • test_moe_2stage.py --csv-filter flydsl_ passes 15/15 cases on this table, which includes (2560, 640, 513, 11).
  • End to end with ATOM main + perf(qwen4-exp): keep 128-aligned MXFP4 expert intermediate unpadded ATOM#2649 and real MTP-3, GSM8K is 0.9788. All MXFP4 MoE lookups hit the tuned rows.

Details are under "Update after #6769" in the description.

[yzhou103]

LGTM