[johannes-graner]

Summary

Add the shape (NE=513, H=8192, D_INTER=256, TOPK=11) to the MXFP4 MoE aux kernel codegen list and the Python supported-shape set, so Qwen3.8 2.4T at TP8 has codegen'd aux instances.

Motivation

Target model: Qwen3.8 2.4T, TP8 (512 routed + 1 shared expert, hidden 8192, topk 11, D_INTER 256 per rank).

Online fmoeTuner logged:

moe_mxfp4_aux.cu:73 mxfp4_moe_scatter_reduce: no codegen'd instance for shape key 'aux_scatter_H8192_TOPK11_NT1'

aux_find uses AITER_CHECK, so a missing key aborts the process. 8 of 376 tuning tasks failed. All 8 are layout-v2 stage2 scatter candidates (flydsl_moe2_layout_afp4_wfp4_bf16_t128x{128,256}x{128,256}_scatter[_nt]_sbm128, blockM=128), which call aiter.mxfp4_moe_scatter_reduce(..., MB=128) after GEMM2. The link between the missing key and these 8 failures is inferred from the call path; it has not been confirmed by rerunning the tuner.

Changes

• csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py: add (513, 8192, 256, 11) to SHAPES.
• aiter/ops/moe_mxfp4_aux.py: add (513, 8192, 256, 11) to MXFP4_MOE_SUPPORTED_SHAPES.

The two lists must stay synchronized.

Testing

Run on gfx950 with ENABLE_CK=0 and AITER_REBUILD=1 (a stale module .so is otherwise reused).

• Baseline origin/main aborts with exit 134 and the message above, for both MB=128 (NT=1) and MB<128 (NT=0).
• With this change the module builds 343 instances. aux_scatter and aux_scatterq H8192_TOPK11_NT{0,1} all resolve, and a direct MB=128 call returns normally.
• mxfp4_moe_scatter_reduce verified against a torch reference at NE=513, H=8192, TOPK=11 for MB in {64, 128}, M in {1, 64, 1024, 8192}, with and without 25% EP-remote (-1) slots: 16/16 pass. Max abs error is 0 to 0.031 on outputs of magnitude ~10 (bf16 rounding and summation order). No LDS or launch failure. The scatter_reduce_q variant was not tested.
• op_tests/test_moe_2stage.py with AITER_SITUV2_A4W4=1, -dim 8192,256 -e 513 -k 11 -t 64 1024 8192: exit 0, no abort or NaN, logits_diff 1.3e-4 to 1.4e-4 (warn threshold 1e-3). These configs use atomic/reduce stage2 epilogues, so they exercise sort/quant/sort_scales but not scatter.

Test Results

Correctness tests pass.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[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 6468 --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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yifehuan]

LGTM