[HIP] [CK] [FlyDSL] FlyDSL split-K preshuffle decode GEMM for small-M shapes (FP8 + MX, gfx950/gfx942)

[johannes-graner]

Summary

Adds a FlyDSL split-K preshuffle GEMM tuned for decode (small-M) shapes, covering FP8 per-tensor and FP8 blockscale, plus a direct-output (split_k=1) epilogue for the whole scale-mode family. Adds MX datatype support (mxfp8, mxfp4) as correctness-validated, tunable dispatch candidates. Enables the blockscale and epilogue split-K kernel families on gfx942 in addition to gfx950. Rides along: a rework of the existing FP8 decode-GEMM dispatch/layout/AOT plumbing that this work shares infrastructure with.

What's included

• gemm_decode-family split-K preshuffle GEMM kernel (aiter/ops/flydsl/kernels/preshuffle_gemm_splitk.py + reduce/op helpers), wired into dispatch, the tuner, and AOT prebuild for both FP8 per-tensor and FP8 blockscale.
• Direct-output single-pass (split_k=1) epilogue, validated per scale mode against the FP32 oracle.
• mxfp8 (128-block E8M0, hardware-scaled MFMA) and mxfp4 (4-bit packed, 32-block E8M0) kernel support, wired as a third scale mode symmetric to existing per-tensor/blockscale dispatch.
• gfx942 enablement for the blockscale and epilogue split-K families via a software-dequant path (4x standard MFMA per 128-K scale block in place of the gfx950-only scaled MFMA); mx128 stays gfx950-only (needs a hardware scale operand gfx942 lacks).
• New op_tests/flydsl_tests coverage for the decode GEMM, split-K preshuffle GEMM, and small-M HGEMM kernels.
• Shared DRY infrastructure: a common split-K bpreshuffle dispatch helper, a shared tuner base class, and a shared kernel-fit check reused across the a8w8/a4w4 families.

Performance

Measured on gfx950, M in {1, 16}, N in {2048, 16384}, K=7168. "Best existing" is the best-tuned pre-existing CK/asm/FlyDSL incumbent for that shape, not a cold-dispatch heuristic fallback. Fuller benchmark data across more shapes is tracked separately.

FP8 per-tensor (this PR's split-K decode GEMM vs best-tuned existing kernel):

M	N	GB/s (this PR)	% of HBM bandwidth wall	speedup vs best existing
1	2048	1974	55%	1.54x
1	16384	5760	98%	1.05x
16	2048	1728	48%	1.50x
16	16384	5691	97%	1.03x

FP8 blockscale (best existing -> this PR):

M	N	best existing GB/s	this PR GB/s	speedup
1	2048	632	2106	3.33x
1	16384	4394	4620	1.05x
16	2048	636	1917	3.01x
16	16384	4247	4531	1.07x

Gains are largest at small N, the latency-bound decode regime; at large N both tables converge toward parity as the kernel approaches the HBM-bandwidth wall.

Correctness

Measured on gfx950, current main + FlyDSL 0.3.2. The tracked FlyDSL pytest suite passes (17/17, 0 skipped) against a torch FP32 reference. mxfp8 is exact across the tile/split sweep. mxfp4 is bit-exact on the identity case and reaches ~3e-3 max relative error on real data, with a negative control that correctly diverges.

Notes

• MX (mxfp8, mxfp4) is correctness-complete and wired as tunable dispatch candidates; performance tuning is deferred and no tuned configuration rows are shipped in this PR.
• gfx942 support is newly enabled in this PR, validated with a single measured winning shape out of the decode shape set tested — an enablement, not a performance headline.

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

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 opt this PR out.

[Copilot]

Pull request overview

Adds new FlyDSL GEMM capabilities focused on decode/small‑M shapes: a two-pass fp32-workspace split‑K preshuffle GEMM (including blockscale + MX variants), plus a unified exact-shape BF16 decode GEMM family. The PR also wires these into tuning/AOT infrastructure, expands arch coverage (gfx942/gfx950 as described), and introduces new correctness tests.

Changes:

• Add split‑K preshuffle GEMM (workspace partials + fp32 reduce) and shared dispatch/tuner scaffolding for bpreshuffle families (a8w8 + a4w4/mxfp4).
• Add unified exact-shape BF16 decode GEMM kernels (Wave + BlockMFMA policies) and integrate decode candidates into the a16w16 tuner + AOT compilation.
• Add new FlyDSL correctness tests and update CI sharding scripts to include FlyDSL GEMM pytest-based tests (with one currently missing from the shard list).

Reviewed changes

Copilot reviewed 34 out of 35 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
op_tests/test_pretune.py	Updates pretune expectations for the new blockscale bpreshuffle tuner script path.
op_tests/flydsl_tests/test_flydsl_small_m_hgemm.py	Adds runtime correctness tests (and optional perf sweep entrypoint) for small‑M BF16 HGEMM.
op_tests/flydsl_tests/test_flydsl_preshuffle_gemm_splitk.py	Adds correctness tests for split‑K preshuffle GEMM (epilogue + blockscale coverage).
op_tests/flydsl_tests/test_flydsl_decode_gemm.py	Adds correctness tests for unified BF16 decode GEMM + per-tensor FP8 decode path.
csrc/gemm_a16w16/gemm_a16w16_tune.py	Adds flydsl_decode tuner libtype, decode candidate policy, and optional vLLM wvSplitK comparison timing.
csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py	Adds split‑K pipeline runner support for the a8w8 bpreshuffle tuner.
csrc/ck_gemm_a8w8_blockscale_bpreshuffle/gemm_a8w8_blockscale_bpreshuffle_tune.py	New dedicated tuner driver for FlyDSL split‑K blockscale/mx128 bpreshuffle variants.
csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_bpreshuffle_tune.py	New dedicated tuner driver for FlyDSL a4w4 (mxfp4) split‑K bpreshuffle variants.
aiter/utility/pretune.py	Updates pretune script mapping/comments for the new blockscale bpreshuffle tuner.
aiter/tuned_gemm.py	Adds flydsl_decode dispatch path and strengthens FlyDSL HGEMM kernel identity checks.
aiter/ops/shuffle.py	Adds CDNA4 preshuffle helpers for mxfp4 weights/scales.
aiter/ops/gemm_op_a8w8.py	Wires FlyDSL split‑K bpreshuffle kernels into a8w8 dispatch via shared helper.
aiter/ops/gemm_op_a4w4.py	Adds FlyDSL split‑K dispatch branch for a4w4 (mxfp4), including inline operand preshuffle hook.
aiter/ops/flydsl/splitk_bpreshuffle_common.py	New shared dispatch helper for split‑K bpreshuffle kernels across a8w8/a4w4.
aiter/ops/flydsl/kernels/tensor_shim.py	Adds unused_tensor_arg helper and plumbs cache_modifier through buffer loads.
aiter/ops/flydsl/kernels/splitk_hgemm.py	Extends split‑K HGEMM kernel naming to include arch + exact N/K identity.
aiter/ops/flydsl/kernels/preshuffle_gemm.py	Updates preshuffle GEMM to support smaller tile_n by varying waves/threads.
aiter/ops/flydsl/kernels/preshuffle_gemm_splitk_reduce.py	New fp32 split‑K reduction kernel (downcast once at final store).
aiter/ops/flydsl/kernels/preshuffle_gemm_splitk_op.py	New driver API for the two-pass split‑K preshuffle GEMM (now used by dispatch).
aiter/ops/flydsl/kernels/gemm_decode_wave.py	New Wave-policy exact-shape decode GEMM implementation (BF16 + per-tensor FP8).
aiter/ops/flydsl/kernels/gemm_decode_block_mfma.py	New BlockMFMA-policy exact-shape BF16 decode GEMM implementation.
aiter/ops/flydsl/gemm_tune/flydsl_splitk_bpreshuffle_tuner_common.py	New shared tuner scaffolding for split‑K bpreshuffle tuners.
aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py	Expands candidate space (adds split‑K pipeline + narrow decode tiles) and exposes split‑K blockscale/mx128 pipelines.
aiter/ops/flydsl/gemm_tune/flydsl_gemm_a4w4_bpreshuffle_common.py	New split‑K candidate space for a4w4/mxfp4 (gfx950-only).
aiter/ops/flydsl/gemm_kernels.py	Broadens GEMM APIs: adds decode GEMM entrypoints, small‑M HGEMM entrypoint, and refactors split‑K HGEMM metadata handling.
aiter/ops/flydsl/__init__.py	Exposes new decode GEMM/small‑M HGEMM symbols from the FlyDSL package surface.
aiter/jit/optCompilerConfig.json	Adds new tuner module entry for a4w4 blockscale bpreshuffle tune.
aiter/jit/core.py	Makes duplicate-row merge handling read-only (dedup in merged view without rewriting source CSVs).
aiter/aot/flydsl/gemm.py	Extends FlyDSL AOT precompile to include split‑K preshuffle GEMM, a4w4 split‑K, and decode GEMM jobs.
.github/scripts/split_tests.sh	Adds explicit FlyDSL GEMM pytest files to the sharded aiter test list (currently missing the split‑K preshuffle test).
.github/scripts/aiter_test.sh	Runs selected FlyDSL GEMM test files via pytest instead of python for correct collection.

---

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

[Copilot]

Pull request overview

Copilot reviewed 34 out of 35 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 33 out of 34 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

csrc/ck_gemm_a8w8_blockscale_bpreshuffle/gemm_a8w8_blockscale_bpreshuffle_tune.py:86

• run_torch_blockscale() uses view(m, k // block_k, block_k) but computes scale_k with ceil-division and does not validate that K is a multiple of the block size. For an invalid shape this will fail with a low-signal view error (or silently mis-handle scale shapes). Add an explicit K%block_k check and compute scale_k from the exact divisor so failures are actionable.

    block_n, block_k = BLOCK_SHAPE
    m, k = x.shape
    n = weight.shape[0]
    scale_n = (n + block_n - 1) // block_n
    scale_k = (k + block_k - 1) // block_k

[Copilot]

Pull request overview

Copilot reviewed 32 out of 33 changed files in this pull request and generated 2 comments.

[andyluo7]

Hi @johannes-graner — while validating this PR against current main (0598b04f, which contains #5116), I found a concrete rebase integration issue.

Several files introduced or changed by #5148 still import aiter.ops.flydsl.utils, but #5116 removed aiter/ops/flydsl/utils.py. In an exact integration tree, import aiter succeeds, but this fails:

python3 -c "from aiter.ops.flydsl.utils import is_flydsl_available"
# ModuleNotFoundError: No module named 'aiter.ops.flydsl.utils'

This also affects downstream capability checks and some of the new small-M/split-K paths. The PR head still contains utils.py; current main does not. For local gfx950 Hy4 profiling I restored the PR-head helper as an isolated, hashed integration repair, but I have not pushed anything to your branch.

When rebasing, could you either retain the helper or update all #5148-added imports to the post-#5116 chip_info / direct FlyDSL availability pattern? I can share the exact integration diff and runtime probe results if useful.

[johannes-graner]

Hi @andyluo7,
Thanks for the heads-up, I'll refactor to use the new patterns.

[Copilot]

🟡 Changes recommended

There are confirmed issues in newly added FlyDSL test scripts (argparse defaults) and CI test runner timeout behavior that should be fixed before approval.

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

Review details

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

[Copilot]

🔵 Needs a closer look

It changes multiple GPU kernel families plus dispatch/tuning/AOT plumbing across many files, so it needs careful manual validation (and the review found at least one correctness guard to add).

Review details

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

[andyluo7]

Hi @johannes-graner — one more downstream integration point from Hy4-preview FP8 TP8/C1 validation on gfx950.

vLLM's correctness-preserving path splits the routed and shared experts because they use different activation/clamp semantics. That changes the small-M MXFP8 FMoE lookups from the existing combined MiniMax-M3 rows (expert=129, topk=5, inter_dim=768/384, Swiglu) to these two generic call shapes:

• routed: model_dim=6144, inter_dim=256, expert=256, topk=8, Silu
• shared: model_dim=6144, inter_dim=256, expert=1, topk=1, Silu
• token tiers: 1, 2, 4
• gfx950, cu_num=256, bf16 output, MXFP8 per_1x32

I tuned those six rows against the #5148 kernel family. The resulting CSV is SHA-256 1f05eadfb3654b7d4adb059f069fe1d0de7d4ce24d9e8cb3e1b6f23699b92534. Preliminary end-to-end vLLM TP8/C1 graph-mode runs using it are about 19.29-19.39 ms mean ITL, versus 25.562 ms on matched gfx942 (~1.32x); the clean three-repeat rerun is queued behind another user's workload.

Current #5148 head f8500b854 has the existing combined rows in aiter/configs/model_configs/minimax_m3_mxfp8_tuned_fmoe.csv, but not these split-call rows, so downstream performance currently requires an external AITER_CONFIG_FMOE file.

Would you be comfortable including the six tuned rows in that existing per-model table after I send the clean repeat/profile evidence, or do you prefer another AITER-owned configuration location? I will not hard-code these shapes in vLLM. I can provide the exact rows and per-kernel measurements once the uncontaminated rerun completes.

[johannes-graner]

Hi @andyluo7,

Of course I can include the tuned rows. Either send them to me, or add them to the tuned csv and push directly to this branch, either is fine by me.

[Copilot]

🟡 Changes recommended

The new FlyDSL a4w4 split‑K dispatch path double-shuffles already-bpreshuffled weights when bpreshuffle=True, which can break existing callers and produce incorrect results.

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

Review details

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