perf(KERNEL-QUANT-CIQ-GEMM-ROCM): widen the WMMA block to 8 warps, measure, reject

.agents/specs/kernel-quant-ciq-gemm-rocm.md

@@ -258,6 +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). **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
 
@@ -549,3 +575,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.

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

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 <typename OutT>
+// 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 <typename OutT, int WarpsPerBlock>
 __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<int64_t>(blockIdx.x) * kQuantWmmaWarpsPerBlock +
+  const int64_t tile = static_cast<int64_t>(blockIdx.x) * WarpsPerBlock +
                        static_cast<int64_t>(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<size_t>(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 <typename OutT>
+template <typename OutT, int WarpsPerBlock>
 __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<int64_t>(blockIdx.x) * kQuantWmmaWarpsPerBlock +
+  const int64_t tile = static_cast<int64_t>(blockIdx.x) * WarpsPerBlock +
                        static_cast<int64_t>(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<int>(threadIdx.x);
   const int wy = static_cast<int>(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<size_t>(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<OutT><<<static_cast<unsigned>(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<OutT, kQuantWmmaWideWarpsPerBlock>
+              <<<static_cast<unsigned>(grid_wmma), dim3(32, kQuantWmmaWideWarpsPerBlock), 0, s>>>(
+                  o, w, qact, m, n, nsb, n_tiles);
+        } else {
+          KQuantGemmKWmmaQ6K<OutT, kQuantWmmaWarpsPerBlock>
+              <<<static_cast<unsigned>(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<OutT><<<static_cast<unsigned>(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<OutT, kQuantWmmaWideWarpsPerBlock>
+              <<<static_cast<unsigned>(grid_wmma), dim3(32, kQuantWmmaWideWarpsPerBlock), 0, s>>>(
+                  o, w, qact, m, n, nsb, n_tiles);
+        } else {
+          KQuantGemmKWmmaQ4K<OutT, kQuantWmmaWarpsPerBlock>
+              <<<static_cast<unsigned>(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);