[yangyuhuiling]

Builds on #6407's 8-wave mxfp8 kernel. Three changes, all to that kernel: a wider C store, a narrower K rule, and keeping the scale staging inside the scale tensors (which the K rule leans on).

1. A/B-swapped epilogue (2.7-8.5%)

StoreBf16 wrote one bf16 at a time: a lane holds four consecutive rows of one column, which cannot be coalesced. Computing B^T A^T instead of A B transposes every accumulator, so a lane holds four consecutive columns — eight contiguous bytes of a row-major C, one buffer_store_dwordx2.

Two cross-lane stages widen it further:

• permlane16_swap pairs the two N-tiles a lane holds into 16 bytes, giving 64-byte rows;
• a ds_bpermute lane transpose puts adjacent lanes on one row, so a 16-lane group covers 4 rows instead of 16.

The swap itself is a thin wrapper — the inner atom is constructed with the tile counts exchanged, so its asserts and idx() line up without further adjustment, and scales and opsels follow their operands.

Prior art. The first two stages are not new to FlyDSL: its MXFP4 4-wave GEMM already feeds the MFMA (B, A) to get C^T = B^T A^T and pairs N-tiles with permlane16_swap into 128-bit C stores (kernels/gemm/fp4_gemm_4wave.py, StoreCFp4 at L565-L607, from ROCm/FlyDSL#990). This PR brings the same epilogue to the mxfp8 8-wave kernel, where the swap goes through the scaled MxMfma (scales and opsels swapped with their operands), and adds the ds_bpermute lane transpose on top. Neither the mxfp8 nor the fp8 8-wave kernel, in aiter or in FlyDSL, has a wide C store today.

Measured on MI355X after merging current main (FlyDSL 0.4.0, as requirements.txt now pins), median of 40 launches after warmup, HIP_VISIBLE_DEVICES=0, min over three passes with the arms alternating (main, PR / PR, main / main, PR). Times in µs:

shape	main	this PR
1024×5120×2048	64.7	61.9	-4.3%
2048×5120×2048	66.3	64.5	-2.7%
4096×5120×2048	90.2	86.3	-4.3%
8192×5120×2048	123.1	115.5	-6.2%
16384×5120×2048	187.1	177.0	-5.4%
8192×7168×2048	152.3	142.9	-6.1%
16384×7168×2048	246.1	230.0	-6.6%
8192×8192×1280	128.3	117.4	-8.5%

The run is gated on an idle machine: it starts only after five consecutive samples show every GPU under 5%, and is discarded if anything is running when it finishes. Worst pass-to-pass spread within an arm was 4.0% (main) and 2.9% (this PR), so the small-M entries are near the noise and the large-M ones are not. An earlier run on FlyDSL 0.3.4.1, before the merge, gave -2.7% to -9.1% on the same shapes.

2. Narrower K rule

K % 256 was required unconditionally. Only ScalePreshuffledS2R needs it — it addresses scale words by K-pair (k1 = k // 2). ScaleRawLds stages a chunk of K steps at a time and never consumes the bytes of the last chunk that lie past K, so the row-major path needs only whole 128-deep tiles. That those bytes are read safely is section 3.

That admits odd multiples of BLOCK_K such as K = 640, 1152, 4224, which the kernel rejected before.

The preshuffled branch keeps its own assert. Nothing in the tree reaches that path today (shuffle_scale_w4 is named in comments but never defined, and _RUNNERS registers only ("gfx950", "1x32", "row", "row")), but keeping the check means a future host-side shuffle gets an error rather than a silently wrong result.

3. Scale chunk reads stay inside the scale tensors

ScaleRawLds.dma() passed each chunk's offset through soffset. The raw-buffer range check covers the VGPR and immediate offsets but not soffset, so the last chunk read past the end of the scale tensor rather than reading zero: the last row's bytes past K on the row-major form, and whole columns past K/128 on the 1x128 column-major A. Never consumed, but the read could leave the allocation. This is already on main whenever K is not a multiple of 512 (e.g. K=1792); section 2 makes it more common.

The offset now goes in the per-lane VGPR offset, which the check covers. The immediate offset is not an option: on an LDS-direct load it also moves the LDS destination. Same review point as on the mori port of this reader (ROCm/mori#806).

Cost, 6 alternating passes before/after on 15 cases (the 8 shapes above, K = 640/1152/4224, and three 128x128-form shapes): every case within -1.4% to +1.7%, against a 2-8% pass-to-pass spread on that run, so no measurable difference. All cases element-exact against the reference.

Tests

No new test file; the existing one covers the epilogue.

Epilogue. op_tests/test_gemm_a8w8_mxfp8_bpreshuffle.py over the tuned gfx950 CSV: 555 rows, of which 145 dispatch to this kernel, 59 of those with N not a multiple of 256 (a lane's four columns straddle the tile edge). All 145 report 0 mismatches.

K relaxation. No tuned row has K % 256 == 128, so dispatch cannot reach the new K values yet. I called the kernel directly instead:

M×N	K
1024×5120	640, 2944, 4224
4096×7168	1152, 7296
8192×7168	4224
320×5152	4224 (ragged N and M)

Each shape ran 6 launches with the output NaN-poisoned beforehand, so a skipped store shows up instead of being masked, and was compared element-wise against the torch reference at 5% relative tolerance. 0 elements off in every case. On main all of these are rejected by the K check. ue8m0 scales are exact powers of two, so the only error against the reference is fp8 quantisation plus bf16 rounding — a misplaced scale moves whole rows rather than hiding in tolerance.

After merging main, which brought #6603's 32x32 and 128x128 scale forms into the same reader, the same direct check across all three forms (w_scale_block = 1x32, 32x32, 128x128), K = 640 to 7296 on and off a multiple of 256, M = 320 and N = 5152 included: 24 cases, 4 launches each, 0 elements off. The dispatch test above gives the same result after the merge (145 of 555 rows on this kernel, all err 0). Re-run after the latest merge (#6380) on FlyDSL 0.4.0: the 24 direct cases, the 555-row dispatch test and test_gemm_group32_interface.py (159 passed, 10 skipped) all pass.

Not covered

gfx950 only, FlyDSL 0.4.0 (also checked on 0.3.4.1 before the merge). No config is re-tuned: the kernel name and tile are unchanged, so existing tuned rows keep selecting it. The preshuffled scale path is unreachable and therefore untested end to end.

🤖 Generated with Claude Code

[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
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 6605 --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.

[solinzby1]

this test file need to integrate into op_tests/*py(related op test file).