FlyDSL A4W4 decode GEMM for M<=16 on gfx950

[johannes-graner]

Motivation

Qwen3.8-27B MXFP4 decode on MI355X runs small-M dense GEMMs through the asm f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128 and _64x128 kernels via gemm_a4w4. This adds a FlyDSL A4W4 (MXFP4, per_1x32 e8m0) decode GEMM that is 1.7-2.3x faster on the tuned shapes with N<=8192.

Technical Details

• New kernel kernels/decode_gemm_mxfp4.py plus optional split-K reduce kernel kernels/decode_gemm_mxfp4_reduce.py, with wrapper aiter/ops/flydsl/decode_gemm_mxfp4.py. Targets M<=16 on gfx950.
• Constraints: N % tile_n == 0, K % 64 == 0. Valid (tile_n, k_waves) pairs: (32, 2) and (64, 1).
• Dispatch: gemm_a4w4 routes kernelName prefix flydsl_decode_t16x{tile_n}x{256*k_waves}_kw{k_waves}_nb{num_buffers}_sk{split_k} to the kernel. Unsupported shapes and unparsable names fall back to the asm heuristic.
• 25 rows added to model_configs/qwen38_27b_a4w4_blockscale_tuned_gemm.csv.
• Tuner: gemm_a4w4_blockscale_tune.py gains --libtype flydsl. Split-K candidates require --splitK and splitK is stored as log2(partitions), as for CK and asm.
• CK gen_instances.py excludes flydsl_* rows.
• AOT: decode jobs are prebuilt as kind decode under OpKind.GEMM, for every padded M served by a tuned row (80 jobs).
• CSV timing method: committed rows use decode-accurate cold-ring HIP-graph medians, not tuner CLI output. The tuner's run_perftest path reads 2-19% higher on 20 of 25 shapes, most on short K=3072 kernels (+15-19%), because of launch overhead; graph replay reproduces the CSV values. At 3 points it picks a different num_buffers (7168x5120 M=16; 8192x5120 M=1 and M=16). Rows were kept because dispatch reads only kernelName, the cold-ring method better represents decode with non-cache-resident weights, and the tuner still picks FlyDSL over asm at all ten checked points. A plain re-tune therefore yields slightly different rows.
• Known gap: FlyDSL loses to asm on wide-N shapes at K=16384; no rows ship there.

Design notes for reviewers:

• The A swizzle and the A/B operand thread mappings stay hand-written. Structured swizzled views and MMA-derived tiled copies were prototyped and lowered to extra address arithmetic (+2 VGPRs, more s_waitcnt). Code comments state the constraint.

Test Plan

• op_tests/flydsl_tests/test_decode_gemm_mxfp4.py: 29 oracle-based tests; each also checks that the call reached the FlyDSL kernel, so a silent fallback to asm fails the test.
• op_tests/test_gemm_a4w4.py: now includes the 25 decode shapes.
• aiter/aot/test/test_flydsl_gemm.py and an AOT prebuild from the shipped CSV.
• End-to-end timing against the replaced asm kernels (run_perftest, 10 warmup, 101 iters).

Test Result

• Decode suite 29/29 passed; test_gemm_a4w4.py err 0 on all 94 rows; AOT test passed; AOT prebuild 80/80.
• CI runs test_gemm_a4w4.py on gfx950 only (the test skips other archs).
• End to end vs asm on the 25 tuned shapes: 1.71-2.30x for N<=8192, 1.01-1.06x for N=17408 (bandwidth-bound). All 25 dispatch FlyDSL; output is bit-identical to asm.

N	K	M	speedup vs asm
5120	3072	1-16	1.76-1.79x
7168	5120	1-16	1.79-1.83x
8192	5120	1-16	1.71-2.30x
5120	8704	1-16	2.07-2.28x
17408	5120	1-16	1.01-1.06x

N=8192 K=5120 at M=1,2,4,8 has no asm row at the merge-base; asm there used the default heuristic.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
• flydsl-kernel-code-cleanup: the cleanup pass and two review rounds are addressed.

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