Skip to content

Rock: cap K-copy vectorization at kPerBlock (fix int4 M=1 GEMV on WMMA) - #2463

Draft
rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/rock-cap-kcopy-at-kperblock
Draft

rlegithub wants to merge 1 commit into
ROCm:developfrom
rlegithub:pr/rock-cap-kcopy-at-kperblock

Conversation

@rlegithub

Copy link
Copy Markdown

Summary

In computeCopyPerThread (the kpack != 1 / K-vectorization branch), the LDS-copy K-vectorization width is maxVlen = 128 / elem_bits. For int4 that is 32, so the greedy copyKPerThread = gcd(maxVlen, copyPerThread) can exceed the tile's kPerBlock (e.g. 16 for the M=1 decode GEMV configs). A thread then tries to copy 32 K-elements from a 16-deep block tile → "gemmA per thread copy smaller than per block copy, incoherent tuning parameters" → rock.gridwise_gemm_accel fails to legalize, so the int4 M=1 GEMV has no valid kernel.

Fix — cap the K-copy vectorization at kPerBlock:

copyKPerThread = gcd(gcd(maxVlen, copyPerThread), kPerBlock);

copyKPerThread still divides both copyPerThread and maxVlen, and the thread tiling is unchanged (copyK*copyD == copyPerThread, so (kPerBlock/copyK)*(dPerBlock/copyD) == blockSize still holds). FP16 (maxVlen = 8) and int4 prefill (configs with kPerBlock = 32) are unaffected; only the narrow-type case where maxVlen > kPerBlock changes.

This makes the in-kernel int4 M=1 decode GEMV compile on gfx1151 (RDNA3.5, WMMA), unblocking int4 (W4A16) LLM decode.

Test plan

  • CI
  • Verified on gfx1151 (Strix Halo): int4 M=1 decode GEMV compiles in-kernel; int4 LLM decodes coherently. FP16 / prefill paths byte-for-byte unchanged.

computeCopyPerThread greedily set copyKPerThread = gcd(maxVlen, copyPerThread).
For narrow types (int4: maxVlen = 128/4 = 32) this can exceed kPerBlock (e.g.
16), so a thread would copy more K elements than the block tile holds, tripping
"gemmA per thread copy smaller than per block copy, incoherent tuning parameters"
and failing to legalize rock.gridwise_gemm_accel. This is exactly the int4 decode
GEMV (M=1, kPerBlock=16) case on gfx1151/RDNA3.5 WMMA; f16 (maxVlen=8) and int4
prefill (kPerBlock=32 configs) were unaffected.

Cap the K-vectorization at kPerBlock: copyKPerThread =
gcd(gcd(maxVlen, copyPerThread), kPerBlock). Still divides copyPerThread and
maxVlen, and thread tiling is unchanged (copyKPerThread*copyDPerThread ==
copyPerThread). Unblocks in-kernel int4 decode; f16/prefill paths unchanged.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant