feat(rocm): wvSplitK skinny GEMM for decode M<=4 — the #487 decode GEMM lever

[VikashLoomba]

Row

BACKEND-ROCM — M5-adjacent decode perf, the RDNA3 skinny-GEMM path. Issue #487 (decode M=1 GEMMs on 128-tile rocBLAS), coordinating with @joral (gfx1200 generic line). This board is gfx1100 (RDNA3) — the only arch that can test the RDNA3-specific config.

What changed

Ports vLLM's wvSplitK_hf_sml_ (csrc/rocm/skinny_gemms.cu, de-torched) into NEW src/vt/rocm/rocm_skinny_gemm.hip and routes decode-skinny MatmulBT (M in 1..4, bf16, K%8==0, activation fits the 64KB LDS stage) to it in rocm_matmul_hipblaslt.hip, ahead of the default-off naive GEMV and the rocBLAS tile path. VT_ROCM_SKINNY=0 rolls back to BLAS for A/B. New cross-device case (decode-skinny MatmulBT, bf16, M∈{1,4}) gates it vs the CPU oracle.

Evidence (4× RX 7900 XTX gfx1100, ROCm 7.14, Release)

Kernel-level A/B vs the current rocBLAS path (same buffers, back-to-back, many iters):

shape	rocBLAS	wvSplitK	speedup
qkv 5120×1024	33.1 us	13.2 us	2.52x
o_proj 1024×2048	18.3 us	5.3 us	3.47x
mlp gate/up 3072×1024	10.9 us	6.1 us	1.78x
lm_head 151936×1024	1238.2 us	340.1 us	3.64x

In-engine decode (Qwen3-0.6B, 128-token steady state): 88.1 vs 70.4 tok/s (+25%) with VT_ROCM_SKINNY on vs off. Qwen3.5-0.8B GDN model output unchanged.

Correctness: ported kernel standalone-validated vs CPU reference across real decode shapes (NMSE ~3e-6, zero bad outputs); new in-tree cross-device case 8/8 green; full ctest zero-delta vs base (same 7 pre-existing host/lane failures).

Speed claims

• The numbers above were run on this board under $GPU_LOCK; recorded here with the repro recipe. (Kernel bench + in-engine A/B, same binary.)

Honest gaps

• gfx1100/RDNA3 only. gfx1200 (RDNA4) and gfx9 (wave64) paths in the donor (MFMA variants) are NOT ported — those need their own boards. joral's gfx1200 line is the generic one.
• The donor's bf16 path has no HW dot on gfx1x (unpacks to f32 mul-add), mirrored exactly; a dot2-f16 path exists for fp16 but our decode is bf16.
• CuCount is the device MP count (donor passes it in); the sweep found CU=40-96 all correct and near-best on this board, but a per-shape autotune is a follow-on.

[localai-bot]

Reviewed as part of a sweep over the open external PRs. The port itself is careful — the upstream anchor is cited correctly and I verified csrc/rocm/skinny_gemms.cu:351-573 and :1169 are both exact at the pinned oracle, and the arithmetic matches the donor. What needs work is not the kernel body but the dispatch preconditions around it, three of which did not come across.

Flagging up front: this is the same shape as findings on #523 and #509, so it is worth treating as one lesson rather than three separate reviews. In each case the calculation was ported faithfully and the guards vLLM performs before choosing that calculation were not.

1. Out-of-bounds device write on an odd output dimension.

Upstream launches wvSplitK_hf_sml_ only under if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) (skinny_gemms.cu:1217). Odd M_in falls through to wvSplitK_hf_. The port carried the LDS-fit half of that predicate across as K * M <= 32768 and dropped the % YTILE half.

The store at rocm_skinny_gemm.hip:122 is unguarded:

if (threadIdx.x == (kThrds - 1)) {
  for (int n = 0; n < N; n++)
    for (int y = 0; y < kYtile; y++) C[m + y + n * M] = __float2bfloat16(sum[n][y]);
}

m is always even and the loop condition is m < M. The read is clamped (min__(y + m, M - 1), line 86); the write is not. With an odd output dim the last active wave has m == N-1, and y == 1 writes C[N + n*N] — for n == M-1 that is two bytes past the end of the output buffer. On a discrete card this is silent memory corruption, not a fault. vt::MatmulBT is a general shared-seam op with no documented parity precondition, and the gate at rocm_matmul_hipblaslt.hip:487 routes every qualifying call, so a GPT-2-family lm_head at vocab 50257 reaches it.

2. No arch guard, and wave32 is hardcoded while the path defaults ON.

Upstream guards twice — on_gfx9() or on_gfx1x() at dispatch (utils.py:174) and #if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) at compile, with an UNREACHABLE_CODE fallback. It also branches the final reduction on wave width: gfx9 needs ROW_BCAST15/ROW_BCAST31 rather than __shfl_xor(sum, 16) (skinny_gemms.cu:489-496).

The port has neither guard and takes only the wave32 arm, while SkinnyGemmEnabled() defaults true. Built with VLLM_CPP_HIP_ARCHITECTURES=gfx942, every decode MatmulBT with M<=4 silently produces wrong numbers with no error — __shfl_xor(x, 16) does not complete a reduction on a 64-wide wave. rocm_backend.hip:75 already reads props.gcnArchName into caps.gcn_arch, so the information is there and unused. The PR body's "Honest gaps" says gfx1100-only; nothing in the code does.

3. The m > 8 lower bound was dropped (utils.py:181: if m > 8 and 0 < n <= 5). The port gates only on M and never bounds N below. This compounds #1: N == 1 gives m = 0, y = 1 → C[1], out of bounds on the very first wave.

4. The upstream test exists and was not ported — and it contains exactly these cases.

tests/kernels/quantization/test_rocm_skinny_gemms.py::test_rocm_wvsplitk_kernel is present at the pin, and its NKM_FACTORS_WVSPLITK list deliberately includes (4, 4096, 4096 + 1), (4, 16384 * 2, 8192 + 1) and a block commented # Minimum M constraint validation (m >= 8). AGENTS.md asks for the upstream test in the same change with parameters preserved; doing that here would have caught findings 1 and 3 before the PR was opened. It also sweeps n ∈ {1,2,3,4} and both bf16 and fp16.

Two smaller things while you are in there. The shipped test bounds aggregate NMSE where upstream bounds elementwise (atol = eps * sqrt(k), rtol = 1e-2) — at the largest shape one arbitrarily-wrong element contributes ~4.9e-5 against a 5e-4 budget, so roughly ten completely wrong outputs still pass, and the tail defect above touches 1–4. And every K in the test is an exact multiple of the 512 stride, so if (k_ >= K) break never fires and the entire K-tail path is uncovered while real shapes hit it every call.

On the merge conflict: only scripts/env-doc-allowlist.txt, but do not reapply it mechanically — main re-sorted and de-duplicated that whole file (235 → 221 lines) while this branch inserted VT_ROCM_SKINNY three times to match the old three-section layout. Take main's version and insert once, in sorted position after VT_ROCM_HIPBLASLT. (That file being a shared must-write surface is our problem, not yours; it is on the list to fix.)

We have no AMD hardware here, so none of your measurements could be reproduced and I am not disputing them — findings 1–3 are read from the source against the donor at the pin, and are independent of any run. Happy to look again once the guards are back.

[VikashLoomba]

All four findings accepted and reworked (commit 96764fe, rebased onto current main):

1. Odd-N OOB: the M_in % _YTILE == 0 half of the dispatch predicate is restored as N > 8 && (N % 2) == 0 (our naming: N=features). The odd-features case now routes to the BLAS path.
2. Arch guard: dispatch refuses non-wave32 arches via CapabilityFromGcnArch(DeviceArchName()) (gfx11xx/gfx12xx only); the gfx9 ROW_BCAST arm is explicitly recorded as owed, not silently miscompiled.
3. m > 8 lower bound: restored (N > 8) — the N==1 first-wave OOB you identified is covered by the boundary case.
4. Upstream test ported: test_rocm_wvsplitk_kernel's applicable NKM factors, xavier on/off, and the elementwise tolerance (atol = eps_bf16·sqrt(k), rtol = 1e-2) replace the aggregate NMSE. Added boundary shapes (features≤8, odd features, odd K, K%512≠0 for the K-tail). Every case writes into a 0xDEAD sentinel guard band so any residual OOB store fails outright.

Mutation-proven: with the N%2 guard removed, the odd-features case corrupts the sentinel band and fails; restored, green. (Both states measured, 20/20 with the fix.)

Allowlist: main's re-sorted file taken wholesale, VT_ROCM_SKINNY inserted once after VT_ROCM_HIPBLASLT.

The one place I checked rather than accepted: fp16/bias/padded-stride arms are deferred with the reasons recorded in .agents/specs/rocm-skinny-gemm.md (the vt::MatmulBT seam has no bias operand; our dispatch precondition is contiguous rows).

[bakon11]

Hey @VikashLoomba — coordinating from the ROCm/gfx1201 (RDNA4, wave32) side so we help here rather than collide.

Heads-up on a real conflict. We have an approved fix in flight (#837) that replaces the single-TLS GetBlas in rocm_matmul_hipblaslt.hip with a dual-slot GetBlas — two BLAS handles so device hops in the 2-GPU MoE prefill stop churning/destroying a single handle (it was deadlocking multi-GPU MoE serving). git merge-tree shows a content conflict in that exact region, since #506 still carries the old single-handle path.

The logic is complementary, though: your skinny-GEMM eligible shapes (M=1..4, N>8/even, K%8, LDS-bound, wave32) return before GetBlas, and every fallback/non-eligible shape still needs the handle. So the clean resolution is keep #837's dual-slot seam and layer your SkinnyGemm enable/arch/dispatch on top — not take either file wholesale. Happy to co-resolve that one file or hand you the seam, whichever you prefer. It just needs to sequence with/after #837.

The bigger offer: your dispatch admits capability major 11 or 12, but the evidence here is gfx1100-only and the notes say the gfx1200 path wasn't ported. We run dual Radeon AI PRO R9700 (gfx1201) and can provide the missing RDNA4 validation — positive candidate witness/dispatch breadcrumb, M={1,4}, real decode N/K, odd-N and N≤8 fallback-guard, K-tail, candidate-vs-BLAS correctness, and a forward/reverse perf A/B. That either earns gfx1201 coverage or tells us to narrow the arch gate to gfx11. Want us to run it once there's an integration head that preserves the dual-slot handle?