[Perf][gfx1250] Add split-K and a decode config selector to the FP8 GEMM

[jhinpan]

Summary

For decode shapes (M <= 256), the gfx1250 FP8 GEMM leaves most CUs idle. This PR adds a split-K launcher, launch_gemm_a8w8_splitk, and a decode config selector, select_gemm_a8w8_decode_config. Over 94 decode rows, the selected configs are 2.23x faster in median than the best upstream-tested config.

Motivation

• The upstream-tested tiles have 128 or 256 rows. At M = 16, each tile computes only 16 valid rows.
• Few tiles give few workgroups. For example, q_a_proj 16 x 1536 x 7168 launches 6 workgroups, and o_proj 16 x 7168 x 2048 launches 28. The GPU has 256 CUs.
• Each workgroup walks the full K range alone, so a deep K is serial.

Changes

• launch_gemm_a8w8_splitk(..., split_k, i64_ws, i64_counters) splits K across grid z:
  • Split z computes K-tiles [z*T//S, (z+1)*T//S) of the T = K / tile_k K-tiles. S does not have to divide T.
  • Each workgroup writes its fp32 partial tile to a workspace and increments the tile's int32 arrival counter.
  • The last workgroup to arrive sums the partials in split order, so the result does not depend on the arrival order. It applies the per-token/per-channel scales to the sum, stores C, and resets the counter to zero.
  • The caller owns the workspace (split_k * ceil(M / tile_m) * ceil(N / tile_n) * tile_m * tile_n fp32 values) and the counters (ceil(M / tile_m) * ceil(N / tile_n) int32 values, zero before the first launch). Launches that can overlap need separate buffers. The docstring states these sizes and the shape requirements.
  • Cluster launch is not supported with split_k > 1.
• select_gemm_a8w8_decode_config(M, N, K) returns (tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, split_k) for M <= 256:
  • M <= 16: 16x64 tiles, or 16x128 tiles with 4 splits when K is deep and the N-tile count is small.
  • 16 < M <= 64: 16x64 tiles. 64 < M <= 256: 64x128 tiles (64x64 when N % 128 != 0).
  • M > 16 and N >= 16384: 128x256 tiles, or 256x256 tiles when M > 128 and the FP8 weights reach 128 MiB.
  • It raises ValueError for M > 256, K % 128 != 0, K < 256, and N % 64 != 0.
  • The thresholds come from measurements on one gfx1250 part with 256 CUs.
• launch_gemm_a8w8 keeps its signature. Its kernel gets two unused arguments.
• Tests:
  • Split-K in both scale modes: a split count that does not divide the K-tile count, ragged M, two launches on one workspace, and per-token/per-channel scales applied after the sum (partials that cancel).
  • One row per selector return. A GPU-free test checks the selected config, and the gfx1250 test runs it through its launcher.
  • Rejected shapes.

Performance (gfx1250, MI450-class, 256 CUs)

Before: the best upstream-tested config per row. After: the selected config. Blockscale, hot device time, median of 15 interleaved samples per side.

M x N x K	Layer	Before config	Before (µs)	After config	After (µs)	Speedup
16 x 1536 x 7168	DeepSeek-V3 q_a_proj	t128x256x128	21.93	t16x128x128, split 4	6.71	3.27x
16 x 7168 x 2048	DeepSeek-V3 o_proj, TP8	t128x256x128	8.57	t16x64x128	3.66	2.34x
16 x 7168 x 16384	DeepSeek-V3 o_proj, TP1	t128x256x128	48.72	t16x128x128, split 4	12.98	3.75x
16 x 7168 x 18432	DeepSeek-V3 dense down, TP1	t128x256x128	57.64	t16x128x128, split 4	14.02	4.11x
16 x 6144 x 16384	GLM-5 o_proj, TP1	t128x256x128	46.34	t16x128x128, split 4	11.99	3.87x
16 x 36864 x 7168	DeepSeek-V3 dense gate_up, TP1	t128x256x128	31.19	t16x64x128	17.28	1.81x
1 x 4096 x 7168	shared-expert gate_up	t128x256x128	22.07	t16x64x128	8.30	2.66x
64 x 1536 x 7168	q_a_proj	t128x256x128	21.31	t16x64x128	8.50	2.51x
256 x 7168 x 2048	o_proj, TP8	t128x256x128	8.47	t64x128x128	5.86	1.45x
256 x 36864 x 7168	dense gate_up, TP1	t256x256x128	36.64	t256x256x128	36.67	1.00x

• Over all 94 decode rows (M = 1, 16, 64, 256; DeepSeek-V3/R1, Kimi and GLM-5 shapes): median 2.23x, geomean 2.07x, and the summed time drops from 1793 µs to 879 µs.
• 12 rows select split 4. The other 82 rows use split 1, so their gain comes from the tile choice.
• Six rows select the baseline's config, so both sides run the same code. One of them, 256 x 16384 x 2048, read 0.80x in this run. It read 1.00x in an earlier run and in a separate check that runs two copies of each side, so the code does not cause it.
• The table was measured before the review fixes (34f5f1a and later). On the split rows, those fixes run at 0.99-1.02x of the measured code.

Testing

• Unit tests added/updated (run on gfx1250, which has no CI runner; the 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 and the new launcher.

Related PRs

• [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971 adds a different split-K to this launcher. There, K must divide evenly into the splits, the partials are bf16 or f16, and a second kernel reduces them. This PR keeps fp32 partials, allows any split count, and reduces them in the same launch. The two PRs conflict in this file.
• [Test][gfx1250] State the FP8 GEMM shape contract and check that stores stay inside N #1229 bounds the N tail of this launcher. The PR that merges second must apply the N bound in store_c and relax the N % tile_n == 0 line of the docstring.
• [Perf][gfx1250] Pipeline the FP8 GEMM K-loop and add a prefill config selector #1233 is stacked on this PR.

[zhiding512]

Benchmark on a gfx1250 box: the tile choices reproduce, the split-K rows run slower than without split

I measured this PR on a gfx1250 machine (4 × 256-CU parts, the 3 idle ones used; amdgpu 7.1.1, ROCm 10.1; FlyDSL built from main with the pinned LLVM e2a39f5). The base (1941889) and this PR (9f6019b) were loaded side by side in one process and sampled round-robin. Blockscale, bf16 output, hot kernel device time from torch.profiler, median of 15 rounds. Every config was checked against an fp32 reference first, and the split-K and selector tests in this PR pass here (20/20).

• The split_k == 1 rows reproduce. They run at 0.99–2.12x of the best pre-PR config, and the three M ≤ 64 rows at 2.0–2.1x (2.3–2.7x in your table). The selected configs run within 2% of the same configs on the base kernel, so these gains come from the tile choice alone.
• The split-K rows don't. The three deep-K rows run at 0.61–0.70x of the baseline (75–87 µs, vs 12–14 µs in your table). The same 16×128 tile without split is about 2x faster than with split 4.
• The cost is the device-scope release/acquire in the split-K epilogue. On this box it adds about 0.3 µs per workgroup; details below.
• [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971's split-K (bf16 partials plus a separate reduce kernel) runs the same 16×128 nb4 split-4 config in 14.5–15.6 µs on those three rows, 5.2–5.6x faster than this PR on this box.

M×N×K	your table: before → after (µs)	here: best pre-PR config (µs)	here: selected (µs)	here
16×1536×7168	21.93 → 6.71 (3.27x)	24.00 (128×256 nb4)	20.57 (16×128 nb4 split 4)	1.17x
16×7168×2048	8.57 → 3.66 (2.34x)	9.77 (128×256 nb4)	4.89 (16×64 nb6)	2.00x
16×7168×16384	48.72 → 12.98 (3.75x)	51.50 (128×256 nb4)	84.37 (16×128 nb4 split 4)	0.61x
16×7168×18432	57.64 → 14.02 (4.11x)	56.88 (128×256 nb4)	87.13 (16×128 nb4 split 4)	0.65x
16×6144×16384	46.34 → 11.99 (3.86x)	51.94 (128×256 nb4)	74.67 (16×128 nb4 split 4)	0.70x
16×36864×7168	31.19 → 17.28 (1.80x)	27.60 (128×256 nb4)	21.25 (16×64 nb6)	1.30x
1×4096×7168	22.07 → 8.30 (2.66x)	25.20 (128×256 nb4)	12.00 (16×64 nb6)	2.10x
64×1536×7168	21.31 → 8.50 (2.51x)	25.26 (128×256 nb4)	11.94 (16×64 nb6)	2.12x
256×7168×2048	8.47 → 5.86 (1.45x)	11.02 (128×256 nb4)	8.61 (64×128 nb4)	1.28x
256×36864×7168	36.64 → 36.67 (1.00x)	42.58 (256×256 nb4)	42.94 (256×256 nb4)	0.99x

The pre-PR candidates are the 128×256 (nb2/3/4) and 256×256 (nb4) configs from the base tests. Absolute times differ from yours (the baseline configs run at 0.88–1.30x of your table, the selected small tiles at 1.17–1.47x), so the ratios are what to compare. HIP-graph replay gives the same picture: 75–88 µs for the three deep-K split rows.

Where the split-K time goes

Device time in µs, 16×128×128 nb4, split 4 unless noted:

variant	16×7168×18432	16×7168×16384	16×1536×7168
no split	43.64	39.20	23.88
split 4, this PR (9f6019b)	88.57	85.63	22.25
split 4, first commit (afacf98)	87.75	84.97	19.53
without the release fence	36.37	37.09	11.58
acquire fence only, relaxed atomic	22.91	22.03	9.17
no fences, relaxed atomic	20.43	17.59	8.59
everything at workgroup scope	15.62	14.26	8.17
#971 split 4 (GEMM + reduce kernel)	15.74	14.46	8.25

The four rows in the middle are diagnostics only; they are not correct protocols (the workgroup-scope one returned NaN on 16×1536×7168). With the pinned LLVM, the epilogue compiles to:

global_store_b128 ... (x4)              ; fp32 partials -> workspace
s_wait_storecnt 0x0
global_wb scope:SCOPE_DEV               ; release fence, every lane
s_barrier_signal -1 / s_barrier_wait -1
global_wb scope:SCOPE_DEV               ; AcqRel atomic, one lane
global_atomic_add_u32 ... th:TH_ATOMIC_RETURN scope:SCOPE_DEV
global_inv scope:SCOPE_DEV
...
global_inv scope:SCOPE_DEV              ; acquire fence, last arrival

Compared with the workgroup-scope variant, the extra time grows with the workgroup count: +14 µs with 48 workgroups and +71–73 µs with 224, about 0.3 µs each, as if these operations serialize across the device. The compiler is the pinned one, so I suspect the platform (driver or firmware cache configuration) makes them cheap on your machine. Which driver and firmware was your table measured on?

Comparison with #971's split-K (same kernel file)

#971 writes bf16 partials to a split_k × M × N buffer and sums them in a second kernel; K must divide by tile_k × split_k. I ported #971 (e35fa3e) onto 1941889 and swept tiles 16×64 / 16×128 and split 2/4/8 over 21 decode shapes (M = 1/16/64):

M×N×K	best without split	this PR, best split	#971, best split
1×7168×18432	27.80 (16×64 nb6)	57.27 (16×128 split 2)	12.22 (16×64 split 8)
16×1536×7168	11.90 (16×64 nb6)	17.23 (16×128 split 2)	6.53 (16×64 split 8)
16×4096×7168	12.06 (16×64 nb6)	30.70 (16×128 split 2)	7.01 (16×128 split 8)
16×7168×18432	27.36 (16×64 nb6)	57.71 (16×128 split 2)	12.86 (16×64 split 8)
64×1536×7168	12.46 (16×64 nb6)	38.72 (16×128 split 2)	6.93 (16×128 split 8)
16×7168×2048	4.97 (16×64 nb6)	32.01 (16×128 split 2)	5.37 (16×128 split 4)
16×36864×7168	21.51 (16×64 nb6)	187.48 (16×128 split 2)	23.99 (16×64 split 2)

• At M ≤ 16 and K ≥ 7168, [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971's best split is 1.7–2.5x faster than the best config without split. At M = 64 it helps only for N ≤ 4096 (1.3–1.8x). At K = 2048 or N = 36864, no split wins.
• This PR's split is slower than no split on all 21 shapes; its fastest is always split 2, the fewest workgroups.
• What [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971 pays for it: rms error 6–27% higher than this PR's fp32 partials, a second launch (about 2 µs more under HIP-graph replay), the bf16 partial buffer, and the divisibility constraint.

The selector on shapes outside your table

22 shapes not in the table (Llama-3.1 8B, 70B TP8 and 405B TP8; Qwen3-32B and Qwen3-235B TP8; DeepSeek-V3 TP8 q_b_proj, kv_b_proj and dense MLP) × M = 1/16/32/64/128/256, 132 rows, against the best of 15 single-launch configs:

selector branch	rows	selected / best (geomean)	worst	slower than the best pre-PR config
M ≤ 64 → 16×64 nb6	72	1.00x	1.12x	0
M ≤ 16 and deep K → 16×128 split 4	8	1.81x	2.56x	2
M > 16 and N ≥ 16384 → 128×256 nb3 (or 256×256)	16	1.51x	2.55x	15
64 < M ≤ 256 → 64×128 (or 64×64)	36	1.21x	1.61x	0

Overall the selector gives 1.53x over the pre-PR configs (geomean over the 126 rows they cover) and runs 1.15x slower than the best config, within 10% of it on 87 of 132 rows. The 16×64 branch holds up on these shapes; the other three don't:

• The split branch picks split 4 for Llama-3.1-8B down, 16×4096×14336: 57.12 µs, vs 22.27 µs for 16×64 nb6 and 45.47 µs for the best pre-PR config.
• The N ≥ 16384 branch picks nb3. On all 15 rows where it regresses, 128×256 nb4 is 6–34% faster (32×16384×6656: 33.09 vs 24.64 µs).
• At M = 128, 64×128 is within 10% of the best on only 3 of 22 shapes.

Suggestions

1. Could you share the driver and firmware your table was measured on? If other gfx1250 machines behave like this one, the in-kernel reduction needs a cheaper synchronization there, or the selector should fall back to a separate reduce (as in [gfx1250] Add A8W8/A8W4/A4W4 compute-bound GEMM #971) or to no split.
2. In the 128×256 branch of the selector, min(3, k_tiles) → min(4, k_tiles).

[aoli26]

Our E2E path AITER already uses offline auto tuning, device/shape keyed tuned CSVs, and a legal heuristic fallback. See https://github.com/ROCm/aiter/blob/main/aiter/ops/gemm_op_a8w8.py#L1198

[aoli26]

We also already have a cluster-barrier fused split-K epilogue that has performed well in our tests. See https://github.com/ROCm/aiter/blob/main/aiter/ops/flydsl/kernels/splitk_fused_epilogue_gfx1250.py.
In our earlier tests, the last-arriving-workgroup approach was slower. Could you compare the two on the same workloads?

[aoli26]

Could these two launchers be consolidated?

[jhinpan]

Thanks for the detailed benchmark. Our table was measured on the box below. The numbers in this comment are from the same box today, with the same driver, firmware, and compiler.

• One 256-CU gfx1250 part (8 per node), SPX compute partition, NPS1 memory partition.
• amdgpu 7.1.0 (7.1.0.31300009, DKMS 7.1.0-2410483), ROCm 10.0.0.
• Firmware: CP_MEC1 2430, RLC 32, SDMA0/1 8992071, PM 02.125.13.03, TA_RAS 1B.44.00.06.
• amdgpu parameters: mtype_local=0 (MTYPE_RW), mtype_remote=-1 (ASIC default), user_partt_mode=-2 (auto).
• FlyDSL 1941889, LLVM e2a39f5 (the pinned one).

Your rows on this box. Base 1941889 and 9f6019b in one process; blockscale, bf16 output; hot device time from event-timed batches, median of 15 rounds; µs:

M×N×K	128×256 nb4 (base)	16×128 nb4, no split	16×128 nb4, split 4 (selected)	split 4 on your box
16×7168×18432	54.53	29.20	15.28	87.13
16×7168×16384	48.04	24.67	14.04	84.37
16×6144×16384	47.38	19.96	12.71	74.67
16×4096×14336	41.67	17.45	10.28	57.12
16×1536×7168	22.17	9.19	6.93	20.57

The per-workgroup cost you found shows up here too, but it is about 10x smaller. With many workgroups, split 4 also loses on this box:

• 256×7168×2048, 3584 workgroups: 114.82 µs, vs 11.74 µs without split.
• 256×36864×7168, 18432 workgroups: 741.93 µs, vs 157.25 µs without split.

That is about 0.03 µs per workgroup here, against about 0.3 µs on your box. The selector splits only for M ≤ 16 with at most 64 N tiles, so at most 256 workgroups. At 0.03 µs each, the split still wins here.

Could you share the same data for your box? That is cat /sys/module/amdgpu/version, /sys/module/amdgpu/parameters/{mtype_local,mtype_remote,user_partt_mode}, the CP_MEC1, RLC, and PM lines of amd-smi firmware, and the output of amd-smi partition. We will decide on the split branch and the selector changes once we know which setting differs.

[jhinpan]

@aoli26 Thanks for the review.

• Selector: noted. The E2E path through AITER uses offline tuning, shape-keyed tuned CSVs, and a heuristic fallback.
• Split-K: we will compare AITER's cluster-barrier fused split-K epilogue with the last-arriving-workgroup epilogue of this PR on the same decode workloads on our gfx1250 box, and post the numbers here.
• Launchers: noted.

We will decide on the selector, the split-K path, and the two launchers after that comparison.

[jhinpan]

@aoli26 Here is the comparison on our gfx1250 box (one 256-CU part, the box from the comment above). Both sides get the same A, B, and E8M0 block scales (AITER gets the 16x16-preshuffled B). All variants run in one process and match one fp64 reference. Hot device time, median of 15 interleaved samples, µs. AITER is at 2b8885a, built with our FlyDSL (1941889) instead of the pinned 0.3.4.1.

Decode, M = 16:

M×N×K	this PR, selected (16×128 split 4)	AITER t16x32x1024_nb4_sk1	AITER 128×256 split 4, fused	AITER 128×256 split 4, separate reduce	this PR, 128×256 split 4
16×7168×16384	13.59	9.86	13.60	14.61	34.16
16×7168×18432	15.32	10.94	14.80	14.38	34.73
16×6144×16384	12.62	9.49	13.28	11.85	30.90
16×4096×14336	10.34	7.91	10.70	11.58	26.89
16×1536×7168	6.95	4.36	7.43	11.49	24.59

M = 512, where your tuned CSV uses the fused split-K (compute_wmma_t128x256x128_mw2_nw2_nb3_sk2_cm4_cn2_fsk):

M×N×K	this PR, best config	AITER, fused	AITER, separate reduce
512×7168×16384	49.71 (128×256 nb4, split 2)	25.94	28.11
512×6144×7168	23.76 (128×256 nb3, no split)	14.22	15.27

• On every decode row, AITER's tuned no-split config is the fastest: 1.31–1.59x faster than the split-K config this PR selects. It needs no split, so it also avoids the per-workgroup synchronization cost that @zhiding512 measured.
• At the same 128×256 tile and split, AITER's kernel is 2.3–3.3x faster at M = 16 and 1.8–2.2x faster at M = 512. Most of the gap is in the main loop: with a separate reduce kernel, AITER is still 2.0–2.6x faster.
• The fused epilogue saves 7–8% against the separate reduce at M = 512. At M = 16 the result is mixed (0.89–1.55x).

So on this box the split-K and the selector of this PR do not beat what AITER already ships.

[jhinpan]

Closing in favor of AITER's gfx1250 FP8 GEMM and its tuned configs; see the comparison above. Thanks @zhiding512 and @aoli26 for the measurements and the review.