[HIP] [CK] [FlyDSL] [Kernel] Extend MXFP4 GEMM1 replacement to A4W4

[fsx950223]

Summary

• Extend the gfx950 FlyDSL MXFP4 GEMM1 replacement to A4W4 MoE workloads.
• Add activation-aware GEMM1 variants, inline BF16→FP4 quantization, hidden-load prefetch (_hpf), and expanded tuner coverage for BN / wave / K-wave / XCD choices.
• Integrate the replacement with the adaptive MoE auxiliary sort, cache-safe AOT/runtime dispatch, and layout-compatible stage2 paths.
• Remove legacy MXMOE-GEMM1 + generic-GEMM2 handling that is no longer supported.

GEMM2 scope

GEMM2 tuning is intentionally out of scope for this PR.

The final config set was audited against latest main e583a77c6:

Config file	Retained GEMM1-only rows	GEMM2 differences
glm5_fp4_tuned_fmoe.csv	4	0
kimik2_fp4_tuned_fmoe.csv	33	0
kimik3_a4w4_tuned_fmoe.csv	3	0
qwen3_5_397b_fp4_tuned_fmoe.csv	0	0
Total	40	0

For all 81 audited rows, kernelName2, us2, and err2 are byte-equivalent to latest main. Thirteen rows had also changed block_m; those were retuned with main's GEMM2 and block_m locked. Rows without a confirmed GEMM1 gain fell back to the main GEMM1. The Qwen3.5-397B FP4 config is fully restored to main.

Performance validation

Measured on one gfx950 MI355X system with warm FlyDSL caches. Main and PR were alternated on the same GPU; the 81-row set received three A/B rounds and every negative/borderline row received five-round confirmation.

Environment for Kimi-K3 A4W4:

AITER_SITUV2_A4W4=1 \
AITER_FLYDSL_STAGE2_FP8=1 \
AITER_MXFP4_INTERMEDIATE=1

Final result versus latest main:

• 35 rows improve by more than 2%
• No confirmed regression greater than 1%
• 3.39% geomean E2E improvement across the 81-row audit set
• 2.69% aggregate latency reduction
• Maximum logits_diff: 0.00733 (gate: <= 0.01)

The two small-token Kimi-K2 rows that regressed with the unchanged kernel were recovered with HPF GEMM1 variants:

• token 4, inter_dim=512: +0.82%
• token 8, inter_dim=256: +1.47%

Validation

• 702 locked-GEMM2 GEMM1 candidates profiled for the 18 rows that had changed GEMM2 selection.
• All 81 final rows passed strict fused-MoE accuracy validation.
• Config shape-collision guard: 16/16 passed.
• Black, Ruff, wheel build, standard tests, OPUS tests, Kimi-K2 accuracy/performance, and the current model accuracy checks pass in CI.

Test methodology

Config comparisons run the production fused_moe path. Main and PR use separate checkouts and caches but execute on the same GPU in alternating order. A row is retained only when repeated measurements show a real gain; otherwise its main config is preserved. Per-stage attribution uses rocprofv3 --kernel-trace, while final acceptance is based on end-to-end fused_moe latency and logits_diff <= 0.01.

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
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 4526 --add-label <label>

[Copilot]

🟡 Human review recommended

It changes Tier-2 MoE dispatch (aiter/fused_moe.py) and multiple FlyDSL kernel/tuning/AOT paths, so the blast radius is high despite local unit test coverage.

Pull request overview

Extends the FlyDSL MXFP4 “GEMM1 replacement” MoE pipeline beyond the prior A4W4-only port to also cover A8W4, adding cache-safe kernel naming/selection, runtime gating, AOT precompile support, and updated tuned model CSVs.

Changes:

• Add cache-safe GEMM1 kernel-name construction/parsing for new MXMOE variants (FP8 intermediates, SiTUv2 beta encoding, interleaved layouts).
• Update tuning + dispatch plumbing to couple new GEMM1 variants with mixed_moe_gemm_2stage GEMM2 kernels, plus profiling support to split stage timings.
• Expand supported shape/config coverage (aux-kernel shape set, AOT job enumeration, refreshed tuned CSVs for GLM-5 / DeepSeek-V4 / Kimi-K3).

File summaries

File	Description
csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py	Extends generated aux-kernel shape list (adds GLM-5/Kimi-K3 shapes) and minor formatting tweaks.
csrc/ck_gemm_moe_2stages_codegen/tune_mxfp4_flydsl.py	New CLI entrypoint for the updated FlyDSL MXFP4 GEMM1 + mixed GEMM2 coupled tuner.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py	Updates tuning logic to generate/lock candidates for A4W4 and A8W4 (FP8) replacement paths and record stage timings.
aiter/test_common.py	Adds optional kernel-time extraction from trace profiling for more detailed perf attribution.
aiter/ops/moe_mxfp4_aux.py	Adds an explicit supported-shape set + helper for MXFP4 aux-kernel coverage checks.
aiter/ops/flydsl/mxfp4_kname.py	Implements new cache-safe GEMM1 naming/parsing (including float-bit encoding for SiTUv2 betas) and selection helpers.
aiter/ops/flydsl/mxfp4_gemm1_kernels.py	Extends GEMM1 port wrapper to support dtype/activation variants, BN choices, and SiTUv2 parameters.
aiter/ops/flydsl/moe_kernels.py	Threads model-dim padding into reduction compilation and normalizes padded output in debug/raw-output mode.
aiter/ops/flydsl/kernels/mxfp4_gemm_common.py	Adds FP8/FP4 MMA atom/scaling helpers and activation math helpers used by new variants.
aiter/ops/flydsl/kernels/moe_gemm_2stage.py	Adds model-dim padding support to the reduction kernel path and updates comments/metadata.
aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py	Adds padding tags to module naming and bounds padded stores in stage2 epilog.
aiter/fused_moe.py	Extends runtime dispatch/gating to select MXMOE replacement kernels for both A4W4 and A8W4, including SiTUv2 beta threading and stage2 integration.
aiter/configs/model_configs/kimik3_a8w4_tuned_fmoe.csv	Updates tuned Kimi-K3 A8W4 configs to use new MXMOE GEMM1 kernel names.
aiter/configs/model_configs/glm5_fp4_newg1_mixedg2_stdsort_tuned_fmoe.csv	Adds a new tuned CSV for GLM-5 FP4 with new GEMM1 + mixed GEMM2 naming.
aiter/configs/model_configs/dsv4_fp8fp4_tuned_fmoe.csv	Updates DeepSeek-V4 FP8FP4 tuned rows to use MXMOE GEMM1 replacement kernels where intended.
aiter/aot/flydsl/mxfp4_moe.py	Expands AOT job parsing to cover fp4/fp8fp4/a8w4 tuned CSVs and generate correct compile jobs for replacement kernels.

Review details

Suppressed comments (1)

aiter/ops/moe_mxfp4_aux.py:19

• The shape glossary line uses ? {16, 32, 64, 128}, which is unclear. Use a clear set-membership phrasing (e.g., "in {..}") so readers don't interpret it as uncertainty.

#   MB       = block_m (sort/gemm block size, ? {16, 32, 64, 128})

• Files reviewed: 17/17 changed files
• Comments generated: 3
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The updated FlyDSL topk-decode/test gating can raise import-time failures when FlyDSL is unavailable, and the new --output-aux help text is now inconsistent with the runtime BM16 selection logic.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

op_tests/test_topk_per_row.py:512

• flydsl_available is now gated only on arch. If flydsl isn't installed/usable, the flydsl=True decode runs will fail at import time instead of being skipped.

flydsl_available = get_gfx() in ("gfx942", "gfx950")

op_tests/test_moe_2stage.py:784

• The --output-aux help text says "block_m=16 keeps its fused sort regardless", but fused_moe._aux_uses_opus() can select Opus for block_m=16 when routed_rows >= num_experts. This makes the CLI help misleading for BM16 cases.

    help="For configs whose sort emits the a4w4 aux arrays, pick which prologue "
    "does it: 'opus' (the default those configs already select) runs the Opus "
    "multi-phase sorter, one CTA per expert; 'threestage' forces the port's own "
    "3-stage sort instead. Same sorted output either way, and block_m=16 keeps "
    "its fused sort regardless. Passed straight through to "
    "fused_moe(output_aux=...); configs that do not use the aux sort are "
    "unaffected.",

• Files reviewed: 24/24 changed files
• Comments generated: 1
• Review effort level: Lite