From 62169e799c44015951a0829faa4af528fc4b4bb6 Mon Sep 17 00:00:00 2001 From: Justin Card Date: Sun, 6 Sep 2026 16:00:03 -0400 Subject: [PATCH 1/2] perf(KERNEL-QUANT-CIQ-GEMM-ROCM): widen the WMMA block to 8 warps, measure, reject (#3032) Issue #3032's same-tool rocprofv3 trace found KQuantGemmKWmmaQ4K/Q6K 4.9x-10.6x slower per-kernel than llama.cpp's mul_mat_q on the identical tensor-core mechanism, with llama.cpp launching double the warps per block and 16-48x fewer, bigger blocks. Each warp here already owns a fully independent output tile with its own shared-memory slice, so widening the block from 4 to 8 warps was the obvious first lever to test: does packing more of that already-independent work into fewer, bigger launches close any of the gap. Templated both kernels on WarpsPerBlock and added an 8-warp instantiation behind VT_ROCM_QUANT_WMMA_WIDE=1 (default off, same same-binary-A/B posture VT_ROCM_QUANT_WMMA and VT_ROCM_Q6K_SMALL_PRIVATE already ship with). Hardware-verified correct on both configs: ctest -R rocm|cross_device, 46/46 cases, 84066/84066 assertions, zero regression. The answer is no. Op-level A/B (quant-gemm-bench, RX 9060 XT, best-of-4, idle host) across the six Q4_K/Q6_K prefill shapes: geomean -4.3%, a net regression. Each warp shares no loaded or dequantized data with any other warp in its block, so widening only spreads the same LDS-capped occupancy budget over fewer, bigger blocks rather than doing more work per warp -- the axis this measurement rules out, not the one that explains llama.cpp's actual advantage. The sharper, still-open hypothesis this points at -- cross-warp data reuse in llama.cpp's own kernel -- is recorded in the spec's Owed section, unread and unconfirmed, as the next traceable step. The toggle stays in the tree, default off, as a ready-made A/B for re-checking this specific axis on different hardware or a future toolchain revision, rather than making the next person re-derive it from scratch. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code] --- .agents/specs/kernel-quant-ciq-gemm-rocm.md | 47 ++++++++ scripts/env-doc-allowlist.txt | 1 + src/vt/rocm/rocm_grouped_gemm.hip | 114 +++++++++++++++----- 3 files changed, 138 insertions(+), 24 deletions(-) diff --git a/.agents/specs/kernel-quant-ciq-gemm-rocm.md b/.agents/specs/kernel-quant-ciq-gemm-rocm.md index 6c2db1852..afd65621f 100644 --- a/.agents/specs/kernel-quant-ciq-gemm-rocm.md +++ b/.agents/specs/kernel-quant-ciq-gemm-rocm.md @@ -258,6 +258,15 @@ scale layout through the tile op. this row gates against (`Q4_K_M`) never invokes Q5_K at all — it is a `Q5_K_M`/`Q5_K_S`-only format — so porting it would not move this row's own gate (c) measurement, only a differently-quantized checkpoint's. +- Cross-warp data reuse in the WMMA kernel body (issue #3032, after the + wider-block axis below was measured and rejected). llama.cpp's `mul_mat_q` + still beats this kernel 4.9x-10.6x per-kernel on the exact same tensor-core + mechanism at matching shapes; widening the block did not close it, and each + warp here shares no loaded/dequantized data with any other warp in its + block, unlike (by inference, not yet read) whatever lets llama.cpp's fewer, + bigger blocks finish so much faster. Reading `ggml/src/ggml-cuda/mma.cuh`'s + actual tile-load structure to confirm or refute cross-tile reuse there is + the next traceable step, not assumed here. ## Stop conditions @@ -549,3 +558,41 @@ W1.5 tail-fill landing, both fixed in the same follow-up commit: than a single cherry-picked run. Q8_0 (a different, unrelated launch path that carries no tail-fill mechanism) was measured alongside as a control and is flat before/after (~2065-2083 GFLOP/s), as expected. + +**Wider WMMA block (issue #3032): MEASURED AND REJECTED.** A same-tool +`rocprofv3` trace of `KQuantGemmKWmmaQ4K`/`Q6K` against llama.cpp's own +`mul_mat_q` on the identical model/prompt found this kernel 4.9x (Q6_K) to +10.6x (Q4_K) slower per-kernel, and llama.cpp's launch configuration uses +double the warps per block (256 vs 128 threads) and 16-48x fewer blocks per +launch yet finishes 15-45x faster. The obvious hypothesis — that packing +more of this kernel's already-independent warps into fewer, bigger blocks +would close some of that gap — does not hold: each warp here owns a fully +independent output tile with its own shared-memory slice, so widening only +changes how many independent warps share one launch's LDS budget, not how +much work any one warp does. + +Templated `KQuantGemmKWmmaQ6K`/`KQuantGemmKWmmaQ4K` on `WarpsPerBlock` and +added an 8-warp instantiation behind `VT_ROCM_QUANT_WMMA_WIDE=1` (default +off). Hardware-verified correct (`ctest -R rocm|cross_device`, both configs, +46/46 cases, 84066/84066 assertions, zero regression). Op-level A/B +(`examples/quant-gemm-bench`, RX 9060 XT, best-of-4, idle host), 4-warp +(current default) vs 8-warp, all six Q4_K/Q6_K prefill shapes: + +| Shape | 4-warp (current) | 8-warp | ratio | +|---|---:|---:|---:| +| Q4_K N=3072 K=2048 | 1948.5 GFLOP/s | 1827.9 GFLOP/s | -6.2% | +| Q4_K N=12288 K=2048 | 2242.7 GFLOP/s | 2027.9 GFLOP/s | -9.6% | +| Q4_K N=2048 K=6144 | 1941.3 GFLOP/s | 1966.1 GFLOP/s | +1.3% | +| Q6_K N=3072 K=2048 | 1874.0 GFLOP/s | 1775.9 GFLOP/s | -5.2% | +| Q6_K N=12288 K=2048 | 2159.1 GFLOP/s | 1981.4 GFLOP/s | -8.2% | +| Q6_K N=2048 K=6144 | 1885.0 GFLOP/s | 1933.5 GFLOP/s | +2.6% | + +Geomean -4.3%: a net regression, not a win. `VT_ROCM_QUANT_WMMA_WIDE` stays +in the tree default-off, the same posture `VT_ROCM_Q6K_SMALL_PRIVATE` above +ships with, as a ready-made A/B for re-checking this specific axis on +different hardware (gfx1201) or a future toolchain revision rather than +re-deriving it from scratch — not carried forward as an open question. The +sharper, still-open hypothesis this measurement points at — cross-warp data +reuse llama.cpp's kernel may have that this one's per-warp-independent +design does not — is recorded in `## Owed` above, unread and unconfirmed, +not assumed. diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index ff438c7ad..0cf0a4aaf 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -180,6 +180,7 @@ VT_ROCM_HIPBLASLT VT_ROCM_LT_ALGO VT_ROCM_MANAGED_ALLOC VT_ROCM_Q6K_SMALL_PRIVATE +VT_ROCM_QUANT_WMMA_WIDE VT_ROCM_SPLIT_N VT_ROCM_WMMA_GEMM VT_ROCM_SKINNY diff --git a/src/vt/rocm/rocm_grouped_gemm.hip b/src/vt/rocm/rocm_grouped_gemm.hip index bcfa1c848..502ba7e2f 100644 --- a/src/vt/rocm/rocm_grouped_gemm.hip +++ b/src/vt/rocm/rocm_grouped_gemm.hip @@ -458,10 +458,46 @@ __device__ inline void DequantQ6KGroup16(const BlockQ6_K* xb, int grp, int8_t ou // shared by every WMMA-format kernel in this file (Q4_K's below included), // so it lives at file scope rather than once per kernel body. constexpr int kQuantWmmaWarpsPerBlock = 4; +// KERNEL-QUANT-CIQ-GEMM-ROCM-WIDEBLOCK (issue #3032): llama.cpp's `mul_mat_q` +// runs the same 16x16x16 int8 WMMA tile at 8 warps/block (256 threads) and +// finishes each launch 15-45x faster than this kernel's 4-warp/block arm at +// matching shapes, per a same-tool `rocprofv3` trace. Each warp here already +// owns a fully independent output tile with its own shared-memory slice +// (indexed by `wy` throughout both kernel bodies below) and loops the whole +// `nsb` reduction alone, so widening the block was tested as a launch-count +// lever: does packing more independent warps per launch (half as many +// blocks, same total work) close any of the gap? +// +// MEASURED AND REJECTED (`examples/quant-gemm-bench`, RX 9060 XT, gfx1200, +// best-of-4, idle host): geomean -4.3% across the six Q4_K/Q6_K prefill +// shapes (range -9.6% to +2.6%), i.e. a net regression, not a win. Each warp +// here shares NO data with any other warp in its block -- no activation row, +// no dequantized weight tile -- so a wider block only adds independent warps +// to the SAME occupancy budget the LDS footprint already caps (4 warps at +// this kernel's ~6.4 KiB/warp already uses two blocks' worth of the same +// 64 KiB/CU two 4-warp blocks would use resident together), which is why +// this axis alone does not move the needle: llama.cpp's block is not merely +// wider, it very likely shares loaded/dequantized data ACROSS its tiles in a +// way this kernel's per-warp-independent design does not. That is the +// sharper, still-open hypothesis issue #3032 records as the next thing to +// try, not this one. `kQuantWmmaWideWarpsPerBlock` and +// `VT_ROCM_QUANT_WMMA_WIDE` stay in the tree, default OFF, as a ready-made +// A/B for re-checking this specific axis on different hardware (gfx1201) or +// a future ROCm/ toolchain revision, rather than re-deriving it from +// scratch — matching this file's own `Q6KSmallPrivateEnabled` precedent for +// a measured-and-rejected diagnostic arm. +constexpr int kQuantWmmaWideWarpsPerBlock = 8; constexpr int kGroupGemmBatch = 2; static_assert(16 % kGroupGemmBatch == 0, "16 scale-groups must divide evenly into batches"); - -template +// LDS budget check for the widest arm actually instantiated below (Q4_K is +// the larger of the two: w_stage 4096B + row_scales/row_mins 128B each + +// raw_tile 2048B = 6400B per warp slot). At kQuantWmmaWideWarpsPerBlock == 8 +// that is 51200B, comfortably under the 64KiB/block LDS limit this project's +// other WMMA kernels stay inside (see `rocm_paged_attn.hip`'s own note). +static_assert(kQuantWmmaWideWarpsPerBlock * 6400 <= 65536, + "widened WMMA block must stay under the 64KiB LDS budget"); + +template __global__ void KQuantGemmKWmmaQ6K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, const BlockQ8_K* __restrict__ act, int64_t m, int64_t n, int64_t nsb, int64_t n_tiles) { @@ -476,7 +512,7 @@ __global__ void KQuantGemmKWmmaQ6K(OutT* __restrict__ out, const uint8_t* __rest return; #else using namespace rocwmma; - const int64_t tile = static_cast(blockIdx.x) * kQuantWmmaWarpsPerBlock + + const int64_t tile = static_cast(blockIdx.x) * WarpsPerBlock + static_cast(threadIdx.y); const int64_t m_tiles = m / 16; if (tile >= m_tiles * n_tiles) return; @@ -493,14 +529,15 @@ __global__ void KQuantGemmKWmmaQ6K(OutT* __restrict__ out, const uint8_t* __rest // for that step alone), and every lane sat idle half the time doing it // (only 16 of 32 lanes filled one row each). Spreading the same total work // over all 32 lanes across all 16 groups removes both costs in one change. - __shared__ int8_t w_stage[kQuantWmmaWarpsPerBlock][16][16][16]; + __shared__ int8_t w_stage[WarpsPerBlock][16][16][16]; // raw_tile holds kGroupGemmBatch groups' raw WMMA tiles at once, so a // whole batch's compute+store is one barrier and its scale-fold is // another, rather than a barrier per single group each way. 16 KiB - // (w_stage) + 8 KiB (this, at batch 2) = 24 KiB/block, comfortably under - // the 64 KiB LDS budget this project's other WMMA kernels stay inside - // (rocm_paged_attn.hip's own BM=16 BN=64 d=256 note). - __shared__ int32_t raw_tile[kQuantWmmaWarpsPerBlock][kGroupGemmBatch][16][16]; + // (w_stage) + 8 KiB (this, at batch 2, WarpsPerBlock==4) = 24 KiB/block, + // comfortably under the 64 KiB LDS budget this project's other WMMA + // kernels stay inside (rocm_paged_attn.hip's own BM=16 BN=64 d=256 note); + // see the file-scope `static_assert` above for the widened-block budget. + __shared__ int32_t raw_tile[WarpsPerBlock][kGroupGemmBatch][16][16]; const size_t row_stride = static_cast(nsb) * sizeof(BlockQ6_K); const BlockQ6_K* w_row_base[16]; @@ -651,7 +688,7 @@ __device__ inline void UnpackQ4KScalesMins(const BlockQ4_K* xb, uint8_t scales[8 } #endif // VT_ROCM_QUANT_WMMA_OK -template +template __global__ void KQuantGemmKWmmaQ4K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, const BlockQ8_K* __restrict__ act, int64_t m, int64_t n, int64_t nsb, int64_t n_tiles) { @@ -660,7 +697,7 @@ __global__ void KQuantGemmKWmmaQ4K(OutT* __restrict__ out, const uint8_t* __rest return; #else using namespace rocwmma; - const int64_t tile = static_cast(blockIdx.x) * kQuantWmmaWarpsPerBlock + + const int64_t tile = static_cast(blockIdx.x) * WarpsPerBlock + static_cast(threadIdx.y); const int64_t m_tiles = m / 16; if (tile >= m_tiles * n_tiles) return; @@ -671,12 +708,12 @@ __global__ void KQuantGemmKWmmaQ4K(OutT* __restrict__ out, const uint8_t* __rest const int lane = static_cast(threadIdx.x); const int wy = static_cast(threadIdx.y); - __shared__ int8_t w_stage[kQuantWmmaWarpsPerBlock][16][16][16]; - __shared__ int32_t raw_tile[kQuantWmmaWarpsPerBlock][kGroupGemmBatch][16][16]; + __shared__ int8_t w_stage[WarpsPerBlock][16][16][16]; + __shared__ int32_t raw_tile[WarpsPerBlock][kGroupGemmBatch][16][16]; // Unpacked once per (row, superblock): 8 scales + 8 mins, tiny next to the // tile buffers above. - __shared__ uint8_t row_scales[kQuantWmmaWarpsPerBlock][16][8]; - __shared__ uint8_t row_mins[kQuantWmmaWarpsPerBlock][16][8]; + __shared__ uint8_t row_scales[WarpsPerBlock][16][8]; + __shared__ uint8_t row_mins[WarpsPerBlock][16][8]; const size_t row_stride = static_cast(nsb) * sizeof(BlockQ4_K); const BlockQ4_K* w_row_base[16]; @@ -1246,6 +1283,21 @@ bool Q6KSmallPrivateEnabled() { return on; } +// VT_ROCM_QUANT_WMMA_WIDE=1 — the A/B arm for issue #3032, MEASURED AND +// REJECTED as a net regression (geomean -4.3%, `examples/quant-gemm-bench`, +// RX 9060 XT): see `kQuantWmmaWideWarpsPerBlock`'s own comment for the full +// result and why. Default OFF, permanently, and not expected to flip; kept +// as a ready-made A/B for re-checking this specific axis on different +// hardware or a future toolchain revision, the same posture +// `Q6KSmallPrivateEnabled` above ships with. +bool QuantWmmaWideBlockEnabled() { + static const bool on = [] { + const char* e = std::getenv("VT_ROCM_QUANT_WMMA_WIDE"); + return e != nullptr && e[0] == '1' && e[1] == '\0'; + }(); + return on; +} + } // namespace // Warps that cooperate on one output row of the dense K-quant GEMM. 1 selects @@ -1392,11 +1444,18 @@ void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tenso g_kq_wmma_dispatches.fetch_add(1, std::memory_order_relaxed); const int64_t n_tiles = n / 16; const int64_t m_tiles = m / 16; - const int64_t grid_wmma = (m_tiles * n_tiles + kQuantWmmaWarpsPerBlock - 1) / - kQuantWmmaWarpsPerBlock; - KQuantGemmKWmmaQ6K<<(grid_wmma), - dim3(32, kQuantWmmaWarpsPerBlock), 0, s>>>( - o, w, qact, m, n, nsb, n_tiles); + const int wpb = QuantWmmaWideBlockEnabled() ? kQuantWmmaWideWarpsPerBlock + : kQuantWmmaWarpsPerBlock; + const int64_t grid_wmma = (m_tiles * n_tiles + wpb - 1) / wpb; + if (wpb == kQuantWmmaWideWarpsPerBlock) { + KQuantGemmKWmmaQ6K + <<(grid_wmma), dim3(32, kQuantWmmaWideWarpsPerBlock), 0, s>>>( + o, w, qact, m, n, nsb, n_tiles); + } else { + KQuantGemmKWmmaQ6K + <<(grid_wmma), dim3(32, kQuantWmmaWarpsPerBlock), 0, s>>>( + o, w, qact, m, n, nsb, n_tiles); + } const int64_t m_aligned = m_tiles * 16, n_aligned = n_tiles * 16; if (m_aligned < m || n_aligned < n) { // Remainder cell count: the bottom strip [m_aligned,m)x[0,n) plus @@ -1414,11 +1473,18 @@ void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tenso g_kq_wmma_q4k_dispatches.fetch_add(1, std::memory_order_relaxed); const int64_t n_tiles = n / 16; const int64_t m_tiles = m / 16; - const int64_t grid_wmma = (m_tiles * n_tiles + kQuantWmmaWarpsPerBlock - 1) / - kQuantWmmaWarpsPerBlock; - KQuantGemmKWmmaQ4K<<(grid_wmma), - dim3(32, kQuantWmmaWarpsPerBlock), 0, s>>>( - o, w, qact, m, n, nsb, n_tiles); + const int wpb = QuantWmmaWideBlockEnabled() ? kQuantWmmaWideWarpsPerBlock + : kQuantWmmaWarpsPerBlock; + const int64_t grid_wmma = (m_tiles * n_tiles + wpb - 1) / wpb; + if (wpb == kQuantWmmaWideWarpsPerBlock) { + KQuantGemmKWmmaQ4K + <<(grid_wmma), dim3(32, kQuantWmmaWideWarpsPerBlock), 0, s>>>( + o, w, qact, m, n, nsb, n_tiles); + } else { + KQuantGemmKWmmaQ4K + <<(grid_wmma), dim3(32, kQuantWmmaWarpsPerBlock), 0, s>>>( + o, w, qact, m, n, nsb, n_tiles); + } const int64_t m_aligned = m_tiles * 16, n_aligned = n_tiles * 16; if (m_aligned < m || n_aligned < n) { const int64_t remainder = (m - m_aligned) * n + m_aligned * (n - n_aligned); From d6f64f2f8166d3f3166cbfd6d6dc950f28c4382b Mon Sep 17 00:00:00 2001 From: Justin Card Date: Sun, 6 Sep 2026 17:01:22 -0400 Subject: [PATCH 2/2] record(KERNEL-QUANT-CIQ-GEMM-ROCM): confirm the cross-warp reuse hypothesis by reading llama.cpp's source The wide-block experiment's spec entry left this as an unread inference. Reading ggml-cuda/mmq.cuh, mmq-config-rdna4.cuh, and mmq-load-tiles.cuh confirms it: llama.cpp's Q4_K block cooperatively loads and shares one 128x128 tile across all warps, while this kernel gives each warp its own independent 16x16 tile with no sharing at all, even between warps that share the same M-rows. That is the mechanism gap the rejected wider-block experiment could not reach. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-sonnet-5 [Claude Code] --- .agents/specs/kernel-quant-ciq-gemm-rocm.md | 35 +++++++++++++++------ 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/.agents/specs/kernel-quant-ciq-gemm-rocm.md b/.agents/specs/kernel-quant-ciq-gemm-rocm.md index afd65621f..a9cbc80df 100644 --- a/.agents/specs/kernel-quant-ciq-gemm-rocm.md +++ b/.agents/specs/kernel-quant-ciq-gemm-rocm.md @@ -258,15 +258,32 @@ scale layout through the tile op. this row gates against (`Q4_K_M`) never invokes Q5_K at all — it is a `Q5_K_M`/`Q5_K_S`-only format — so porting it would not move this row's own gate (c) measurement, only a differently-quantized checkpoint's. -- Cross-warp data reuse in the WMMA kernel body (issue #3032, after the - wider-block axis below was measured and rejected). llama.cpp's `mul_mat_q` - still beats this kernel 4.9x-10.6x per-kernel on the exact same tensor-core - mechanism at matching shapes; widening the block did not close it, and each - warp here shares no loaded/dequantized data with any other warp in its - block, unlike (by inference, not yet read) whatever lets llama.cpp's fewer, - bigger blocks finish so much faster. Reading `ggml/src/ggml-cuda/mma.cuh`'s - actual tile-load structure to confirm or refute cross-tile reuse there is - the next traceable step, not assumed here. +- Cross-warp data reuse in the WMMA kernel body (issue #3032). **CONFIRMED + by reading the actual source**, not inferred: llama.cpp's Q4_K config at + this row's shapes (`ggml/src/ggml-cuda/mmq-config-rdna4.cuh:127`, + `GGML_TYPE_Q4_K, 256, ..., 128, 128, ...`) gives ONE BLOCK an I=128 (weight + rows) x J=128 (activation columns) output tile — 8-64x the area our + 4-8-warp block covers, since each of OUR warps independently owns only a + 16x16 tile. That whole I x J tile is loaded ONCE per block, not once per + warp: `ggml_cuda_mmq_load_tiles_q4_K` + (`ggml/src/ggml-cuda/mmq-load-tiles.cuh:703-741`) stripes the I=128 rows + across every warp's threads (`i0 += nrows*nwarps`, each warp taking a + DIFFERENT row range), dequantizing each row exactly once into one shared + `x_tile`; `mmq_get_nbytes_shared` (`mmq.cuh:1379-1383`) sizes the + activation-tile allocation by `J` alone, not `J*nwarps`, confirming one + shared activation copy too. The compute phase + (`ggml_cuda_mmq_write_back_mma`, `mmq.cuh:476-500`) then splits the I + dimension across warps (`rows_per_warp = I/nwarps`), each warp looping the + FULL J range doing 16x16 MMA ops against that one shared load. This + kernel's per-warp-independent design (own `w_stage`/`raw_tile` slice, own + `load_matrix_sync` call, zero sharing even when two warps in one block + share the same M-tile) is the actual mechanism gap the wider-block + experiment above could not touch, because widening never introduced + sharing — it only added more independent warps to the same per-warp-load + pattern. Matching this (cooperative tile load + row-split compute over a + much larger shared tile) is a real kernel redesign, materially bigger than + the block-width experiment, and is the next traceable step — not attempted + in this row. ## Stop conditions