FlyDSL decode GEMM (M=1–5) and gfx942 small-M VGPR HGEMM

[samremes]

Why

Skinny decode GEMM (exact M=1–5) and tiny-K small-M GEMM on gfx942 leave a gap versus the backends the A16W16 tuner already had: hipBLASLt/torch.mm, AITER HIP wvSpltK (tuner libtype skinny), generic FlyDSL HGEMM, and Triton. The gap is small M (decode/skinny GEMM): hipBLASLt and generic HGEMM use large GEMM tiles/policies that are a poor fit when M is 1–5 (and, for small-M, when K is tiny). HIP wvSpltK is the existing skinny option and still loses on the measured cells. Decode wins at tiny K (1536,128) and at large K (896,7168) / (2112,7168); vocab (20480,7168) is a smaller win. Small-M wins only tiny-K (1536,128) at M=8/16 and loses at (20480,7168). These are packed-graph wall/N comparisons, not a host-launch-overhead story.

This PR adds two kernel families for that gap:

• FlyDSL decode GEMM — exact-shape Wave and BlockMFMA BF16 kernels for M=1–5 (no padded M). Wave is a one-wave/no-LDS policy; BlockMFMA is a multi-wave 4×4×4 MFMA policy with optional N-persistent full-A LDS.
• FlyDSL small-M HGEMM — a dedicated M=1–16 family with gfx942 VGPR staging (gfx942 cannot issue 16-byte global-to-LDS DMA, so A is staged through VGPRs by architecture rather than a catalog knob). The strongest measured island is tiny-K (N,K)=(1536,128) at M=6–16, not a general M=6–16 winner.

What this PR does

• Public exact-shape decode identity (gemm_decode_bf16 + Wave/BlockMFMA configs) and flydsl_small_m_hgemm.
• Runtime pytest coverage in two files: op_tests/flydsl_tests/test_flydsl_decode_gemm.py and test_flydsl_small_m_hgemm.py.
• CI lists those two FlyDSL pytest files (split_tests.sh / aiter_test.sh).
• AOT overlay compiles CSV wins for decode; kernel names carry the exact (arch,M,N,K) identity. No shipping tuned-CSV changes.
• Tuner grows flydsl_decode tasks (bounded candidate policy by default). Comparison-only vLLM timing stays opt-in and is never promoted.

gfx942 matched microbench (this revision)

Hardware: MI325X, gfx942, 304 CU, HIP_VISIBLE_DEVICES=0. BF16, no bias.

Method: packed same-stream HIP graph wall/N for every backend (50 launches, 7 samples, median µs). Tiny kernels also N-swept at 20/50/100 launches. This is not decode-graph vs eager-torch. rocprofv3 exclusive duration was not used (it floors near ~4 µs on this GPU). Decode numbers are the best of a few representative Wave/BlockMFMA configs, not a full autotune. torch is torch.mm with hipBLASLt preferred. skinny is AITER HIP wvSpltK (decode) / wv_splitk_small (small-M) from csrc/kernels/custom_kernels.cu (tuner libtype skinny). vLLM was not timed.

N-sweep sanity on the two sub-10 µs cells (median µs): decode M=1 (1536,128) 2.09 / 1.85 / 1.77 at N=20/50/100; small-M M=8 same shape 2.18 / 1.95 / 1.87. Torch stays ~4.2–4.5 µs across the sweep.

Decode (M=1 and M=4)

M	(N,K)	decode	torch	skinny	HGEMM	triton	vs torch	vs skinny
1	(1536,128)	1.85	4.23	2.09	3.10	3.90	2.3×	1.1×
4	(1536,128)	2.47	4.28	5.15	3.06	3.86	1.7×	2.1×
1	(896,7168)	2.81	6.73	5.87	65.1	80.1	2.4×	2.1×
4	(896,7168)	5.82	6.79	20.5	65.1	80.2	1.2×	3.5×
1	(2112,7168)	3.86	10.3	9.12	49.0	89.0	2.7×	2.4×
4	(2112,7168)	6.97	10.3	27.8	48.6	88.8	1.5×	4.0×
1	(20480,7168)	58.8	74.4	59.8	132	109	1.3×	1.02×
4	(20480,7168)	70.1	75.0	84.7	133	108	1.07×	1.2×

Wide-N M=4 vs torch is modest (1.07–1.2×). The clear decode wins are tiny-K and mid-N × large-K at M=1–4.

Small-M VGPR 16×16×64 (not a general M=6–16 winner)

M	(N,K)	small-M	torch	skinny	HGEMM	triton	vs torch	vs HGEMM
8	(1536,128)	2.06	4.30	8.21	3.11	4.12	2.1×	1.5×
16	(1536,128)	1.99	4.29	15.3	3.22	4.85	2.2×	1.6×
8	(20480,7168)	103	75.3	—	—	—	0.73× (loss)	—

Tiny-K is the island. The same VGPR 16×16×64 config loses to hipBLASLt at M=8 (20480,7168). Do not treat small-M as a default for all M=6–16 shapes.

[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
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 1 --add-label <label>

[Copilot]

Pull request overview

Adds exact-shape BF16 FlyDSL decode GEMM support, small-M kernel improvements, tuning integration, AOT compilation, and runtime tests.

Changes:

• Adds Wave and BlockMFMA decode kernels and public dispatch APIs.
• Integrates decode candidates into tuning, runtime selection, AOT, and tuned configurations.
• Extends gfx942 small-M support and CI coverage.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
op_tests/flydsl_tests/test_flydsl_small_m_hgemm.py	Adds small-M correctness tests.
op_tests/flydsl_tests/test_flydsl_decode_gemm.py	Adds decode kernel and graph replay tests.
csrc/gemm_a16w16/gemm_a16w16_tune.py	Adds decode and vLLM tuning candidates.
aiter/utility/mp_tuner.py	Adds timing divisors and candidate timeout handling.
aiter/tuned_gemm.py	Adds tuned decode dispatch.
aiter/ops/flydsl/kernels/tensor_shim.py	Adds optional tensor placeholders and cache controls.
aiter/ops/flydsl/kernels/splitk_hgemm.py	Makes kernel identities shape-specific.
aiter/ops/flydsl/kernels/small_m_hgemm.py	Extends small-M support to gfx942.
aiter/ops/flydsl/kernels/gemm_decode_wave.py	Implements Wave decode GEMM.
aiter/ops/flydsl/kernels/gemm_decode_common.py	Defines decode configurations and primitives.
aiter/ops/flydsl/kernels/gemm_decode_block_mfma.py	Implements BlockMFMA decode GEMM.
aiter/ops/flydsl/gemm_kernels.py	Exposes and dispatches new GEMM APIs.
aiter/ops/flydsl/__init__.py	Exports decode and small-M APIs.
aiter/jit/core.py	Makes merged-config deduplication read-only.
aiter/configs/bf16_tuned_gemm.csv	Adds gfx942 decode tuning rows.
aiter/aot/flydsl/gemm.py	Adds decode AOT parsing and compilation.
aiter/aot/flydsl/common.py	Supports selecting AOT operation kinds.
.github/scripts/update_split_test_times.py	Includes new tests in timing updates.
.github/scripts/split_tests.sh	Adds new tests to CI shards.
.github/scripts/aiter_test.sh	Runs new files through pytest.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[vedenev-amd]

gfx950 (MI355X) microbench — BF16 decode GEMM, N=16384 K=7168

Hardware: MI355X gfx950, 1 XCD (32 CUs), HBM ~1332 GB/s · ROCm: 7.2.0 · PyTorch: 2.7.1+rocm7.2.0 · FlyDSL: 0.3.1 · aiter: samremes/flydsl-gemm-decode f5d97b9

Method: CUDA graph 200 launches × 5 trials, median. Python launch overhead eliminated. CK = bench_gemm_decode_msweep best swept config (warmup=200, repeat=200).
wvSplitK: M≤4 → wvSpltK · M≥5 → wv_splitk_small_fp16_bf16.

M	Sami kernel	Sami µs	wvSplitK µs	CK µs	vs wvSplitK	vs CK
1	Wave mp=1 np=2	32.38	35.32	33.01	0.917×	0.981×
2	Wave mp=2 np=2	32.46	37.70	33.07	0.861×	0.982×
3	Wave mp=3 np=4	33.62	40.30	33.02	0.834×	1.018×
4	Wave mp=4 np=4	33.81	44.29	33.05	0.763×	1.023×
5	Wave mp=5 np=4	34.96	359.94	33.07	0.097×	1.057×
8	small_m_hgemm	34.10	481.92	33.05	0.071×	1.032×

ratio < 1.0× = Sami faster · wv_splitk_small at M=5,8 is 360–480 µs — targets small-N shapes, not competitive at N=16384.
M=4 improved from 1.067× to 1.023× after full config sweep (np=4 > np=2 at M=4).