Skip to content

[ROCm][RDNA] Skinny GEMMs: no VGPR spills on gfx11/gfx12, wvSplitK launch rows for gfx1150/gfx1151 #1401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
roberteg16 merged 4 commits into from
Oct 26, 2026
Merged

[ROCm][RDNA] Skinny GEMMs: no VGPR spills on gfx11/gfx12, wvSplitK launch rows for gfx1150/gfx1151 #1401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cc26d00
[ROCm][RDNA] skinny GEMMs: no VGPR spills in wvSplitK on gfx11 and gfx12
roberteg16 Oct 26, 2026
9c15ea7
[ROCm][RDNA] int4 skinny GEMMs: no spilling tile on 1024-VGPR archs, …
roberteg16 Oct 26, 2026
f993e6b
[ROCm][RDNA] skinny GEMMs: wvSplitK launch rows of their own for gfx1…
roberteg16 Oct 26, 2026
79aa663
[ROCm][RDNA] skinny GEMMs: drop the wvSplitK VGPR-budget fallback
roberteg16 Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions benchmarks/kernels/sweep_bf16_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,14 @@
(unquantized) wvSplitK path used by every Linear that hasn't been replaced
by an int4/int8/AWQ wrapper.

Limitation: the in-tree ``wvSplitK_sweep`` op (csrc/rocm/skinny_gemms.cu)
currently exposes only ``(ytile, unrl)`` as runtime args; ``WvPrGrp`` and
``A_CHUNK`` are macro-baked at 16 and 8 respectively. The script therefore
sweeps only YTILE x UNRL. For 4-axis sweeps (the (W, AC) corner of the
search space) use the out-of-tree HIP harness in
``rdna35-asm-expert/microbench/wv_splitk/`` until ``wvSplitK_sweep`` is
extended to match ``wvSplitK_int4g_sweep`` (which already takes all 4 axes).
The in-tree ``wvSplitK_sweep`` op (csrc/rocm/skinny_gemms.cu, built with
``VLLM_SKINNY_GEMM_SWEEP_BF16=1``) takes every launch knob of the dispatcher:
YTILE x UNRL x A_CHUNK x WvPrGrp, for N = 1..5. Combinations that spill on
the current arch are worth excluding from any conclusion (their speed depends
on the process's scratch history), and A_CHUNK > 8 needs K % A_CHUNK == 0.
Weights are padded as vLLM pads them at load on gfx11 (a row stride that is a
multiple of 2048 B gets one more 128 B line); unpadded, a weight that lands on
an aliasing address can run 2x slower.

For each (M, K, N) cell the script also calls the production ``wvSplitK``
op once to record the dispatcher's chosen wall-clock — useful both as a
Expand Down Expand Up @@ -69,9 +70,9 @@
(1024, 4096, "K=4096 small-M sanity"),
]

YTILES = [1, 2]
UNRLS = [1, 2, 4]
ACHUNKS = [8, 16]
YTILES = [1, 2, 3, 4]
UNRLS = [1, 2, 4, 8]
ACHUNKS = [8, 16, 32]
WVPRGRPS = [16, 32]

# Working-set guards used by the dispatcher's variant selection
Expand Down Expand Up @@ -141,7 +142,9 @@ def run_sweep(shapes, batch_sizes, warmup, rep, dtype, csv_path):
# three variants under the hood. Just noting it for the user.
pass

A = torch.randn(M, K, dtype=dtype, device="cuda") * 0.01
pad = 64 if (K * 2) % 2048 == 0 else 0
A = torch.empty(M, K + pad, dtype=dtype, device="cuda")[:, :K]
A.copy_(torch.randn(M, K, dtype=dtype, device="cuda") * 0.01)
B = torch.randn(N, K, dtype=dtype, device="cuda") * 0.01
weight_bytes = M * K * A.element_size()

Expand All @@ -157,7 +160,7 @@ def run_sweep(shapes, batch_sizes, warmup, rep, dtype, csv_path):
for ytile, unrl, ac, w in itertools.product(
YTILES, UNRLS, ACHUNKS, WVPRGRPS
):
if M % ytile != 0:
if K % ac != 0:
skipped += 1
continue
try:
Expand Down
Loading
Loading