From fc26771ce555c99a15ebb70906d473ad246ec5bf Mon Sep 17 00:00:00 2001 From: Roto Le Date: Fri, 11 Sep 2026 16:33:11 -0700 Subject: [PATCH] Rock: cap K-copy vectorization at kPerBlock (fix int4 M=1 GEMV on WMMA) 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. --- mlir/lib/Dialect/Rock/utility/loweringUtils.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/Rock/utility/loweringUtils.cpp b/mlir/lib/Dialect/Rock/utility/loweringUtils.cpp index 54ace7f945a6..7571519002b3 100644 --- a/mlir/lib/Dialect/Rock/utility/loweringUtils.cpp +++ b/mlir/lib/Dialect/Rock/utility/loweringUtils.cpp @@ -1142,7 +1142,13 @@ computeCopyPerThread(Type elementType, int64_t copyPerThread, int64_t kPerBlock, copyKPerThread = copyPerThread / copyDPerThread; dim = GemmDimension::MorN; } else { - copyKPerThread = math_util::gcd(maxVlen, copyPerThread); + // Cap the K-vectorization at kPerBlock: with narrow types (e.g. int4, + // maxVlen=32) the greedy gcd(maxVlen, copyPerThread) can exceed kPerBlock + // (e.g. 16), which makes a thread copy more K elements than the block tile + // holds -> "incoherent tuning parameters". Capping keeps the copy coherent + // (still divides copyPerThread and maxVlen; thread tiling is unchanged since + // copyKPerThread*copyDPerThread == copyPerThread). + copyKPerThread = math_util::gcd(math_util::gcd(maxVlen, copyPerThread), kPerBlock); copyDPerThread = copyPerThread / copyKPerThread; dim = GemmDimension::K; }