[FlyDSL] Dev/a4w4 prefill v2 yadai

[yadaish]

Motivation

Improve FlyDSL grouped MXFP4 prefill performance on gfx1250, especially the A4W4 GEMM1 path, while preserving A8W4 support and numerical correctness. The previous implementation left substantial mainloop throughput on the table because LDS traffic, multicast topology, synchronization, and WMMA scheduling were not sufficiently overlapped.

Technical Details

• Reworks the gfx1250 A4W4/A8W4 TDM mainloop and LDS pipeline to improve overlap between data movement and WMMA compute, including planar/interleaved LDS loads.
• Adds configurable cluster geometry for GEMM1 and GEMM2 through cluster_m, cluster_n, cluster_m2, and cluster_n2 in the tuned CSV. Supported dimensions are cluster_m in {1, 2, 4}, cluster_n in {1, 2, 3, 4}, with at most 16 workgroups per cluster.
• Uses N-direction multicast without a full-cluster barrier for 1xN clusters. MxN clusters additionally use cross-M multicast and cluster synchronization.
• Fuses GEMM1 activation with MX requantization and scale generation, and fixes requant scale addressing for row-tile sizes that are not powers of two.
• Adds optional explicit VGPR placement when both EXPLICIT_VGPR_PARTITION is enabled and the installed FlyDSL exposes fx.set_register.
• Extends CSV-driven test coverage, supports deterministic zero initialization, and reorganizes the tuned CSV fields so epilogue and cluster controls are grouped consistently.

Test Plan

Run the representative 1536-token A4W4 kernel benchmark with zero initialization:

AITER_USE_GROUPED_GEMM=1 \
AITER_GROUPED_DEBUG=0 \
ENABLE_CK=0 \
FLYDSL_DUMP_IR=1 \
AITER_FLYDSL_DISABLE_GEMM1_REQUANT=0 \
AITER_LOG_MORE=1 \
AITER_MOE_EXPERT_BALANCE=true \
AITER_FLYDSL_MOE_EXPERT_SCHEDULING_MODE=1 \
AITER_FLYDSL_MMA_FIRST_GROUP=4 \
AITER_FLYDSL_MMA_GROUP=5 \
AITER_FLYDSL_DS_FIRST_N=1 \
AITER_FLYDSL_EXPLICIT_VGPR_PARTITION=1 \
AITER_FLYDSL_PLANAR_LDS=1 \
AITER_FLYDSL_INTERLEAVED_LDS_LOAD=1 \
python3 -u op_tests/flydsl_tests/test_flydsl_grouped_gemm.py \
  --scenario kernel \
  --data-format a4w4 \
  --experts 64 \
  --tokens 1536 \
  --topk 8 \
  --model-dim 7168 \
  --inter-dim 2048 \
  --act silu \
  --no-bias \
  --iters 32 \
  --no-check-aot-cache \
  --data-init zero \
  --scale-init zero

For nonzero correctness and performance, rerun the same command with:

--data-init uniform --scale-init auto

Run the complete single-GPU tuned-CSV sweep:

AITER_USE_GROUPED_GEMM=1 \
ENABLE_CK=0 \
python3 -u op_tests/flydsl_tests/test_flydsl_grouped_gemm.py \
  --scenario csv \
  --csv-path aiter/configs/tuned_grouped_fmoe.csv \
  --warmup 1 \
  --iters 2 \
  --no-check-aot-cache \
  --no-bias \
  --data-init uniform \
  --scale-init auto

Code-quality checks:

ruff check aiter/ops/flydsl/grouped_gemm_mxfp4.py \
  aiter/ops/flydsl/grouped_moe_gfx1250.py \
  aiter/ops/flydsl/kernels/gemm_common_gfx1250.py \
  aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py \
  op_tests/flydsl_tests/test_flydsl_grouped_gemm.py

black --check aiter/ops/flydsl/grouped_gemm_mxfp4.py \
  aiter/ops/flydsl/grouped_moe_gfx1250.py \
  aiter/ops/flydsl/kernels/gemm_common_gfx1250.py \
  aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py \
  op_tests/flydsl_tests/test_flydsl_grouped_gemm.py

Test Result

Representative A4W4 shape: 64 experts, 1536 tokens, topk=8, model dimension 7168, intermediate dimension 2048, 32 iterations. Each benchmark was run twice with the GPU idle before every run. The baseline is origin/main@475cf0f607; the PR result is d816736617.

Initialization	Revision	GEMM1 samples	GEMM1 mean	GEMM2 samples	GEMM2 mean	Accuracy
zero / zero	main	118.29, 117.97 us	118.13 us	74.34, 74.35 us	74.35 us	Pass (logits_diff=0, rel_l2=0)
zero / zero	PR	77.76, 76.54 us	77.15 us	58.31, 58.55 us	58.43 us	Pass (logits_diff=0, rel_l2=0)
uniform / auto	main	141.54, 141.02 us	141.28 us	81.37, 81.20 us	81.28 us	Pass (logits_diff=3.52166e-06, rel_l2=2.65393e-03)
uniform / auto	PR	89.79, 88.80 us	89.30 us	62.60, 63.47 us	63.04 us	Pass (logits_diff=3.52166e-06, rel_l2=2.65393e-03)

Compared with main, the PR reduces latency by 34.69% (1.53x speedup) for zero-init GEMM1 and 21.40% (1.27x) for zero-init GEMM2. With uniform inputs, it reduces latency by 36.79% (1.58x) for GEMM1 and 22.45% (1.29x) for GEMM2.

Main did not have a tuned CSV match for this shape and selected its default tile_m=64, n_warp=4, and num_buffers=2. The PR selected the new tuned A4W4 configuration.

• Tuned CSV: 116 rows total; all 88 executable single-GPU cases passed (56 A8W4 and 32 A4W4), with 0 failures. The remaining 28 EP rows use topk=-1 and were skipped because the single-GPU runner has no EP arena.
• CSV accuracy: maximum logits_diff was 4.16491e-06; maximum relative L2 error was 2.88615e-03, below the 1.0e-02 acceptance threshold.
• Ruff 0.16.0: passed.
• Black: passed.

Submission Checklist

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

[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; main branch always runs both MI35X and MI300X
multigpu	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: 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 5934 --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.