[Perf][gfx1250] Pipeline the FP8 GEMM K-loop and add a prefill config selector

[jhinpan]

Summary

Large FP8 GEMMs (prefill, M > 256) on gfx1250 wait for the L2-to-LDS tile traffic. This PR pipelines the K-loop across K-tiles and adds a prefill config selector that uses 2x2 clusters. On 59 blockscale prefill rows, the selected configs are 1.19x faster in geomean (1.045-1.348x) than the best upstream-tested config, and every row is faster. On the 5 per-token/per-channel rows that the selector takes, they are 1.19-1.26x faster. Outputs are bit-identical.

Stacked on #1228. The diff against the #1228 branch shows only this PR's changes. After #1228 merges, this PR will be retargeted to main.

Motivation

• A 256x256 tile reads its A row-block and B column-block from L2 into LDS once per workgroup. On the large rows, this traffic runs at 22-24 TB/s, near the TDM plateau of 24-28 TB/s for cache-resident data. So the traffic, not the WMMA rate, limits these GEMMs.
• A 2x2 cluster sends each A tile and each B tile to two workgroups with one multicast load, so it halves the traffic. The launcher supports clusters, but only the tests use them.
• Each K-tile also starts with a stall. In an ATT trace (rocprofv3) of 8192 x 36864 x 7168, about 390 of the 1036 cycles of each K-tile come before its first WMMA: the fence, 57 LDS reads, and 21 VALU instructions.

Changes

• Cross-K-tile pipeline for the quadrant schedule (warp tiles with an even number of WMMA rows and columns, and at least 8 accumulators):
  • In the last K-step of K-tile kt, the kernel waits for K-tile kt+1 and loads its first operands (scales, the left half of B, the top half of A) into carry registers.
  • The barrier of that wait is split around the third quadrant (signal before it, wait after it), so its WMMAs absorb the arrival skew between waves.
• LDS loads are interleaved with the WMMAs through scheduling group barriers. Their addresses are computed before each interleaved run, so expert scheduling waits once per run, not at each load.
• The pipeline also works with cluster launch. A multicast load returns data only to the workgroups that requested it (CDNA5 ISA §10.7), and each workgroup requests it after its own fence.
• E8M0 block scales are broadcast with v_perm_b32 at the full VALU rate, instead of a quarter-rate multiply by 0x01010101.
• SCHED_MODE.DISABLE_XDL_ARB_STALL (bit 2, CDNA5 ISA §5.7.2.1) is set at kernel start, so a wave can issue other work while its WMMAs run. In an ablation, it gave 0.9-1.2% on 17 rows (blockscale 128x256 and per-token/per-channel 256x256, one wave per SIMD). It was neutral on the blockscale 256x256 rows and on the decode configs, which run several waves per SIMD.
• select_gemm_a8w8_prefill_config(M, N, K, is_bsc=True, num_cus=256) returns (tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, cluster_m, cluster_n) for launch_gemm_a8w8:
  • 256x256 tiles with 4 buffers and a 2x2 cluster (cluster_n = 1 for an odd N-tile count);
  • for blockscale grids that 256x256 tiles leave idle: 128x256 tiles with 3 buffers, if they still fit in one wave of workgroups;
  • ValueError for per-token/per-channel grids that 256x256 tiles do not fill, and for M <= 256, N % 256 != 0, K % 128 != 0, and K < 512. Callers keep their own config there.
  • The thresholds come from measurements on one gfx1250 part with 256 CUs.
• The previous quadrant K-tile schedule is removed. The row schedule (16-row tiles) does not change.
• Tests: both tiles, the odd N-tile count, ragged M, the rejected shapes, and a GPU-free check of the selected configs.

Resources

VGPRs go up on the quadrant configs (compile-only, gfx1250). No config spills.

Config	Scales	VGPRs before -> after	Waves per SIMD before -> after
256x256x128, 2x2 warps, 4 buffers	blockscale / per-token/per-channel	746 -> 862 / 793 -> 848	1 -> 1
128x256x128, 2x2 warps, 3 buffers	blockscale / per-token/per-channel	425 -> 468 / 472 -> 524	2 -> 2 / 2 -> 1

The per-token/per-channel 128x256 config drops to one wave per SIMD. #1228's decode selector returns it for per-token/per-channel rows with N >= 16384 and 16 < M <= 128. Those grids have one workgroup per CU at N <= 65536; larger N with this config was not measured.

Performance (gfx1250, MI450-class, 256 CUs)

Before: the best upstream-tested config per row, without a cluster. After: the selected config. Hot device time, median of 20 interleaved samples.

M x N x K	Scales	Before config	Before (µs)	After config	After (µs)	Speedup
8192 x 36864 x 7168	blockscale	t256x256x128_w2x2_nb4	730.09	t256x256x128_w2x2_nb4, cluster 2x2	627.48	1.16x
8192 x 7168 x 18432	blockscale	t256x256x128_w2x2_nb4	374.89	t256x256x128_w2x2_nb4, cluster 2x2	315.32	1.19x
2048 x 7168 x 16384	blockscale	t256x256x128_w2x2_nb4	90.64	t256x256x128_w2x2_nb4, cluster 2x2	71.20	1.27x
8192 x 4096 x 6144	blockscale	t256x256x128_w2x2_nb4	75.63	t256x256x128_w2x2_nb4, cluster 2x2	56.09	1.35x
2048 x 24576 x 1536	blockscale	t256x256x128_w2x2_nb4	41.50	t256x256x128_w2x2_nb4, cluster 2x2	31.64	1.31x
2048 x 4096 x 7168	blockscale	t128x256x128_w2x2_nb3	27.15	t128x256x128_w2x2_nb3, cluster 2x2	23.21	1.17x
2048 x 2048 x 512	blockscale	t128x256x128_w2x2_nb3	4.99	t128x256x128_w2x2_nb3, cluster 2x2	4.77	1.05x
16384 x 16384 x 16384	blockscale	t256x256x128_w2x2_nb4	1558.73	t256x256x128_w2x2_nb4, cluster 2x2	1273.08	1.22x
4096 x 4096 x 4096	per-token/per-channel	t256x256x128_w2x2_nb4	24.86	t256x256x128_w2x2_nb4, cluster 2x2	20.23	1.23x
8192 x 8192 x 8192	per-token/per-channel	t256x256x128_w2x2_nb4	173.94	t256x256x128_w2x2_nb4, cluster 2x2	138.44	1.26x
9728 x 8192 x 8320	per-token/per-channel	t256x256x128_w2x2_nb4	265.93	t256x256x128_w2x2_nb4, cluster 2x2	223.19	1.19x

• Blockscale, all 59 prefill rows: 47 rows on 256x256 tiles at 1.21x geomean (1.093-1.348x), and 12 rows on 128x256 tiles at 1.12x (1.045-1.230x).
• Per-token/per-channel, 8 rows: the 5 selected rows run at 1.192-1.256x. On the 3 rows where the selector raises, both sides run the same config (0.995-1.012x).
• Decode rows (M <= 256, with [Perf][gfx1250] Add split-K and a decode config selector to the FP8 GEMM #1228's selector on both sides, so only the kernel changes): blockscale 94 rows at 1.02x geomean (0.990-1.134x), per-token/per-channel 6 rows at 1.01x (0.997-1.030x).
• Where the gain comes from, on the 59 blockscale rows: the 2x2 cluster alone on the base kernel gives 1.084x; the pipeline alone gives 1.020x; both give 1.158x; the final version gives 1.188x.
• A separately loaded copy of the baseline measured 0.978-1.010x against it.

These numbers were measured before the rebase onto #1228's review fixes. The rebase does not change the generated code of the selected configs (LLVM IR and ISA compared).

Testing

• Unit tests added/updated (run on gfx1250, which has no CI runner; the GPU-free selector tests also run on CI)
• Performance benchmarks run (gfx1250, table above)
• Tested on MI300X: not applicable (gfx1250-only kernel)

Dependencies

• No new third-party dependencies added
• If new dependencies added, justification provided

Breaking Changes

None. Callers opt in through the selector. launch_gemm_a8w8 and launch_gemm_a8w8_splitk keep their signatures.

Related PRs

• [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971 adds a hand-scheduled 256x256 kernel, for MX scales only, for the tiles that this selector picks, and makes the same v_perm_b32 change. It also changes this launcher, so the two PRs conflict. They have not been compared on the same rows.
• [Test][gfx1250] State the FP8 GEMM shape contract and check that stores stay inside N #1229 bounds the N tail. This PR keeps [Perf][gfx1250] Add split-K and a decode config selector to the FP8 GEMM #1228's B and C copies, which are not bounded in N. The selector requires N % 256 == 0, so its configs are safe; direct calls with an N tail need [Test][gfx1250] State the FP8 GEMM shape contract and check that stores stay inside N #1229.
• [Perf] gemm_bf16 gfx1250: cross-tile carry to hide K-tile prologue ds… #1125 uses a different cross-K-tile scheme in the BF16 kernel.

[jhinpan]

Closing together with #1228, which this PR is stacked on. On our gfx1250 box, AITER's compute kernel with the fused split-K ran 1.7–1.9x faster than the #1228 kernel at M = 512 (#1228 (comment)). We did not measure the head of this PR against AITER.