Repository navigation
[ROCm][RDNA] Skinny GEMMs: no VGPR spills on gfx11/gfx12, wvSplitK launch rows for gfx1150/gfx1151 - #1401
Merged
Conversation
roberteg16
force-pushed
the
rogarcia.skinny-gemm-no-spill
branch
from
October 24, 2026 20:21
52f36e7 to
9eb172c
Compare
mgehre-amd
reviewed
Oct 26, 2026
mgehre-amd
approved these changes
Oct 26, 2026
wvSplitK_hf_sml_, wvSplitK_hf_ and wvSplitK_hf_big_ held every activation chunk of an unrolled step in registers next to the weights, and bf16 widened the weights to f32. With 1024-thread workgroups a wave gets 128 VGPRs on gfx1150 (192 on gfx1151), and 102 of the 388 kernels spilled on gfx1150, 78 on gfx1151. A spilling GEMM is slow, and its speed depends on the largest scratch an earlier kernel of the process used (1.3x slower after a 256 B/lane kernel, 2.1x after 4 KiB). On gfx1x: - bf16 multiplies with v_dot2_f32_bf16, as fp16 with v_dot2_f32_f16. - A tile holds the activation chunks while (N + YTILE) * UNRL * A_CHUNK / 2 operand VGPRs plus 64 fit the wave's VGPRs; otherwise it reads each chunk from LDS where it is used. gfx9 code is unchanged. No VGPR spills in the 388 kernels on gfx1150, gfx1151 and gfx1201. bench_rocm_skinny_gemm.py, N 1..4, against gfx11: gfx1150 bf16 1.050x, fp16 1.008x; gfx1151 1.00x. The 4096x2048 bf16 GEMM keeps 186-210 us whatever scratch an earlier kernel used (242-651 us before). test_rocm_skinny_gemms.py: 172 passed on gfx1150 and gfx1151. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
…only reachable tiles built The int4 dispatch macros expand every tile for N = 1..5, so tiles only picked for N = 1 were also built for N = 2..5, where they are never launched: 480 of the 1684 kernels, and most of the spilling ones. Those arms are now `if constexpr (__N == 1)`; the dispatch is unchanged. The MoE (YTILE 4, W 32, AC 32, UNRL 4) tile for N = 1, K = 2048 spills with the 128 VGPRs a wave gets on a 1024-VGPR arch: gfx1102, gfx1103 and gfx1150..gfx1153 take (4, 4, 16, 16) there; gfx1151 keeps it. Kernel code is unchanged. Spilling kernels: 334 -> 24 on gfx1150 (the tile above, not launched there), 293 -> 8 on gfx1151 (at the VGPR limit, as before). W4A16 group 128, M = 1..4, against gfx11: gfx1150 MoE 1.19x, dense 0.99x-1.00x; gfx1151 1.00x. test_rdna_hybrid_w4a16.py (102 passed), test_hybrid_w4a16_moe.py (70 passed) and test_rocm_skinny_gemms.py (172 passed) on gfx1150 and gfx1151. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
…151 and gfx1150 The wave32 wvSplitK dispatch was one chain for every gfx11 arch, tuned shape by shape on gfx1151. A sweep of every non-spilling launch configuration (YTILE 1-4, UNRL 1/2/4/8, A_CHUNK 8/16/32, WvPrGrp 16/32, N = 1..5, bf16 and fp16) finds one tile per arch within ~1% of the best for nearly every shape: - gfx1151: (YTILE 1, UNRL 4, A_CHUNK 16); (2, 4, 8) for K < 1024. - gfx1150, gfx1152, gfx1153: (2, 1, 16); (1, 2, 16) when M <= 32 * CUs. - A_CHUNK 8 when K % 16 != 0: wider chunks give wrong results there. Other archs take the generic rows. The sweep op covers every knob the dispatcher uses, and sweep_bf16_kernel.py pads the weights as they are padded at load. Interleaved A/B against the previous rows, 43 shapes (13 not used for tuning, odd M, K 768/1000/2880), N = 1..5, bf16/fp16 and bias: gfx1151 +0.2% to +1.3%, gfx1150 +0.6% to +1.1%; against upstream's rows +1.8% to +4.2% and +0.8% to +3.9%. 99.3-99.8% of the per-shape best. test_rocm_skinny_gemms.py: 190 passed on gfx1150 and gfx1151, with new K % 16 != 0 and N = 5 cases. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
roberteg16
force-pushed
the
rogarcia.skinny-gemm-no-spill
branch
from
October 26, 2026 19:17
9eb172c to
f993e6b
Compare
roberteg16
changed the title
The fallback that reads each activation chunk from LDS when a tile's operands exceed the wave's VGPRs only served the large tiles of the previous rows. No launched tile needs it now (the tuned rows, the generic rows and the fused epilogues all fit), so the kernels go back to holding the activation chunks, as before. 0 VGPR spills on gfx1150 and gfx1151; the launched kernels compile to the same instructions as the same tiles upstream. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The skinny GEMMs (
wvSplitK*) spilled VGPRs in many instantiations on gfx11/gfx12. A spilling GEMM is slower, and its speed depends on earlier kernels: the runtime sizes every wave's scratch slot to the largest scratch used so far in the process, and with large slots the spills go to DRAM.v_dot2_f32_bf16, and the launch rows below only use tiles that fit the wave's VGPRs. 0 VGPR spills on gfx1150, gfx1151, gfx1200, gfx1201 (78-102 before). The kernels compile to the same instructions as the same tiles upstream.gfx9 code is unchanged.
Results
Spill fixes against gfx11 (
bench_rocm_skinny_gemm.py, 95 shape x N cells per dtype, N = 1..5; int4 W4A16 group 128, M = 1..5):Launch rows against the previous rows (interleaved A/B, cold caches, 43 shapes x N = 1..5, bf16/fp16, 13 shapes not used for tuning, odd M, K 768/1000/2880, bias): gfx1150 +0.6% to +1.1%, gfx1151 +0.2% to +1.3%. They reach 99.3-99.8% of the per-shape best tile (previous rows 98-99%, upstream's 96-98%).
The 4096x2048 bf16 GEMM on gfx1150 takes 188 us whatever scratch earlier kernels used (236-525 us before).
Tests
On gfx1150 and gfx1151:
test_rocm_skinny_gemms.py: 190 passed (new K % 16 != 0 and N = 5 cases).test_wvsplitk_fused_silu.py: 124 passed.test_rdna_hybrid_w4a16.py: 102 passed.test_hybrid_w4a16_moe.py: 70 passed.gfx1201 (spill commits):
test_rocm_skinny_gemms.py311 passed; its fp8 cases run out of memory with and without this PR.pre-commitpasses.Overlap
#1246 (draft) also changes
wvSplitK_hf_sml_and the gfx11 dispatch; the two will conflict textually.AI assistance was used; the submitter reviews every line.