cuda: register-stage the plain-matmul WMMA K-loop prefetch (1.1-2.1x)

matmul_rm_wmma_kernel (the C = A(M,K) @ B(K,N) row-major FP16/BF16 GEMM
behind the plain matmul() op) had the same single-buffered K-loop as the
ABT kernel: each slab's global load stalled between the two
__syncthreads with nothing to hide its HBM latency.

Apply the identical register-staged prefetch (rm_prefetch_A/B ->
rm_deposit_A/B): hoist the next slab's global reads into per-thread int4
registers before the WMMA compute so HBM latency overlaps the tensor
cores, then do the fast register->shared deposit at the top of the next
iteration. Shared memory stays single-buffered (occupancy unchanged --
doubling the tiles was measured slower on the ABT twin), and the
mma_sync order is untouched, so FP32 accumulation is identical.

4090 FP16 matmul() (100 iters):
  512^3           0.0166 -> 0.0108 ms  (1.53x)
  1024^3          0.0368 -> 0.0182 ms  (2.02x)
  2048^3          0.1353 -> 0.1205 ms  (1.12x)
  4096^3          1.2062 -> 0.9691 ms  (1.24x)
  64x4096x4096    0.1059 -> 0.0512 ms  (2.07x)  <- short-M linear-like
  128x4096x4096   0.1060 -> 0.0507 ms  (2.09x)
  256x14336x4096  0.4001 -> 0.2733 ms  (1.46x)

Parity green on every consumer: matmul (fp32+fp16), matmul_backward,
matmul_parity (CPU<->GPU), cpu_gpu_parity. Adds tests/bench_matmul_rm.cpp
(brotensor_bench_matmul_rm, a hand-run microbench, not a ctest) as the
measurement harness.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: wlejon

src/cuda/matmul.cu

@@ -195,6 +195,103 @@ constexpr int RM_FRAGS_N = RM_WN / RM_WMMA_N;                // 2
 constexpr int RM_LDA_SMEM = RM_BK + 8;   // 40, for A tile (BM rows, BK cols)
 constexpr int RM_LDB_SMEM = RM_BN + 8;   // 72, for B tile (BK rows, BN cols)
 
+// One thread's share of a K-slab, as int4s (8 x 16-bit each).
+constexpr int RM_A_SLAB_REGS = (RM_BM * RM_BK / 8) / RM_THREADS_PER_CTA;  // 2
+constexpr int RM_B_SLAB_REGS = (RM_BK * RM_BN / 8) / RM_THREADS_PER_CTA;  // 2
+
+// Register-staged software prefetch — same latency-hiding transform as
+// fp16_matmul.cu's ABT kernel. The single-buffered loop stalled on each slab's
+// global load between the two __syncthreads; here the next slab's GLOBAL reads
+// are issued (into registers) *before* the WMMA compute so HBM latency overlaps
+// the tensor-core math, and the fast register->shared deposit happens at the top
+// of the next iteration. Shared memory stays single-buffered (occupancy is
+// unchanged) and the mma_sync order is identical, so results are bit-for-bit the
+// same. A tile is (BM,BK) accessed along K (identical to ABT's A); B tile is
+// (BK,BN) accessed along N — the natural (K,N) row-major layout, no transpose.
+template <typename T>
+__device__ __forceinline__ void rm_prefetch_A(
+    int4 (&a_reg)[RM_A_SLAB_REGS], const T* __restrict__ A,
+    int M, int K, int k0, int block_m, int tid) {
+    using TR = rm_traits<T>;
+    constexpr int kEPL = 8;
+    #pragma unroll
+    for (int li = 0; li < RM_A_SLAB_REGS; ++li) {
+        const int lin = tid + li * RM_THREADS_PER_CTA;
+        const int row = lin / (RM_BK / kEPL);
+        const int gcol = (lin % (RM_BK / kEPL)) * kEPL;
+        const int grow = block_m + row;
+        const int gk   = k0 + gcol;
+        T tmp[kEPL];
+        if (grow < M && gk + kEPL <= K) {
+            *reinterpret_cast<int4*>(tmp) =
+                *reinterpret_cast<const int4*>(&A[grow * K + gk]);
+        } else {
+            #pragma unroll
+            for (int q = 0; q < kEPL; ++q) {
+                const int gk_q = gk + q;
+                tmp[q] = (grow < M && gk_q < K) ? A[grow * K + gk_q]
+                                                : TR::from_f32(0.0f);
+            }
+        }
+        a_reg[li] = *reinterpret_cast<int4*>(tmp);
+    }
+}
+
+template <typename T>
+__device__ __forceinline__ void rm_prefetch_B(
+    int4 (&b_reg)[RM_B_SLAB_REGS], const T* __restrict__ B,
+    int N, int K, int k0, int block_n, int tid) {
+    using TR = rm_traits<T>;
+    constexpr int kEPL = 8;
+    #pragma unroll
+    for (int li = 0; li < RM_B_SLAB_REGS; ++li) {
+        const int lin = tid + li * RM_THREADS_PER_CTA;
+        const int row = lin / (RM_BN / kEPL);   // K-tile row
+        const int gcol = (lin % (RM_BN / kEPL)) * kEPL;  // N-tile col
+        const int gk = k0 + row;
+        const int gn = block_n + gcol;
+        T tmp[kEPL];
+        if (gk < K && gn + kEPL <= N) {
+            *reinterpret_cast<int4*>(tmp) =
+                *reinterpret_cast<const int4*>(&B[size_t(gk) * N + gn]);
+        } else {
+            #pragma unroll
+            for (int q = 0; q < kEPL; ++q) {
+                const int gn_q = gn + q;
+                tmp[q] = (gk < K && gn_q < N) ? B[size_t(gk) * N + gn_q]
+                                              : TR::from_f32(0.0f);
+            }
+        }
+        b_reg[li] = *reinterpret_cast<int4*>(tmp);
+    }
+}
+
+template <typename T>
+__device__ __forceinline__ void rm_deposit_A(
+    const int4 (&a_reg)[RM_A_SLAB_REGS], T As[RM_BM][RM_LDA_SMEM], int tid) {
+    constexpr int kEPL = 8;
+    #pragma unroll
+    for (int li = 0; li < RM_A_SLAB_REGS; ++li) {
+        const int lin = tid + li * RM_THREADS_PER_CTA;
+        const int row = lin / (RM_BK / kEPL);
+        const int gcol = (lin % (RM_BK / kEPL)) * kEPL;
+        *reinterpret_cast<int4*>(&As[row][gcol]) = a_reg[li];
+    }
+}
+
+template <typename T>
+__device__ __forceinline__ void rm_deposit_B(
+    const int4 (&b_reg)[RM_B_SLAB_REGS], T Bs[RM_BK][RM_LDB_SMEM], int tid) {
+    constexpr int kEPL = 8;
+    #pragma unroll
+    for (int li = 0; li < RM_B_SLAB_REGS; ++li) {
+        const int lin = tid + li * RM_THREADS_PER_CTA;
+        const int row = lin / (RM_BN / kEPL);
+        const int gcol = (lin % (RM_BN / kEPL)) * kEPL;
+        *reinterpret_cast<int4*>(&Bs[row][gcol]) = b_reg[li];
+    }
+}
+
 template <typename T>
 __launch_bounds__(RM_THREADS_PER_CTA)
 __global__ void matmul_rm_wmma_kernel(const T* __restrict__ A,
@@ -223,72 +320,28 @@ __global__ void matmul_rm_wmma_kernel(const T* __restrict__ A,
         }
     }
 
-    for (int k0 = 0; k0 < K; k0 += RM_BK) {
-        // ---- Load A tile (RM_BM x RM_BK): natural row-major (M,K) access ----
-        {
-            constexpr int kElemsPerLoad = 8;  // int4 = 8 x 16-bit
-            constexpr int kTotalElems   = RM_BM * RM_BK;
-            constexpr int kLoadsTotal   = kTotalElems / kElemsPerLoad;
-            constexpr int kLoadsPerThr  = kLoadsTotal / RM_THREADS_PER_CTA;
-
-            #pragma unroll
-            for (int li = 0; li < kLoadsPerThr; ++li) {
-                const int lin = tid + li * RM_THREADS_PER_CTA;
-                const int row = lin / (RM_BK / kElemsPerLoad);
-                const int col_grp = lin % (RM_BK / kElemsPerLoad);
-                const int gcol = col_grp * kElemsPerLoad;
-                const int grow = block_m + row;
-                const int gk   = k0 + gcol;
-
-                T tmp[kElemsPerLoad];
-                if (grow < M && gk + kElemsPerLoad <= K) {
-                    const int4* src = reinterpret_cast<const int4*>(&A[grow * K + gk]);
-                    *reinterpret_cast<int4*>(tmp) = *src;
-                } else {
-                    #pragma unroll
-                    for (int q = 0; q < kElemsPerLoad; ++q) {
-                        const int gk_q = gk + q;
-                        tmp[q] = (grow < M && gk_q < K) ? A[grow * K + gk_q] : TR::from_f32(0.0f);
-                    }
-                }
-                *reinterpret_cast<int4*>(&As[row][gcol]) = *reinterpret_cast<int4*>(tmp);
-            }
-        }
+    // K loop with register-staged prefetch (see rm_prefetch_* above): the next
+    // slab's global reads are hoisted ahead of the WMMA math so HBM latency
+    // overlaps the tensor cores. Shared memory is single-buffered; the mma_sync
+    // order is unchanged, so FP32 accumulation is bit-for-bit identical.
+    int4 a_reg[RM_A_SLAB_REGS];
+    int4 b_reg[RM_B_SLAB_REGS];
+    rm_prefetch_A<T>(a_reg, A, M, K, 0, block_m, tid);
+    rm_prefetch_B<T>(b_reg, B, N, K, 0, block_n, tid);
 
-        // ---- Load B tile (RM_BK x RM_BN): natural row-major (K,N) access,
-        // no transpose — this is the key difference from the ABT kernel. ----
-        {
-            constexpr int kElemsPerLoad = 8;
-            constexpr int kTotalElems   = RM_BK * RM_BN;
-            constexpr int kLoadsTotal   = kTotalElems / kElemsPerLoad;
-            constexpr int kLoadsPerThr  = kLoadsTotal / RM_THREADS_PER_CTA;
+    for (int k0 = 0; k0 < K; k0 += RM_BK) {
+        // Deposit the slab prefetched last iteration, then publish it.
+        rm_deposit_A<T>(a_reg, As, tid);
+        rm_deposit_B<T>(b_reg, Bs, tid);
+        __syncthreads();
 
-            #pragma unroll
-            for (int li = 0; li < kLoadsPerThr; ++li) {
-                const int lin = tid + li * RM_THREADS_PER_CTA;
-                const int row = lin / (RM_BN / kElemsPerLoad);   // K-tile row
-                const int col_grp = lin % (RM_BN / kElemsPerLoad);
-                const int gcol = col_grp * kElemsPerLoad;        // N-tile col
-                const int gk   = k0 + row;
-                const int gn   = block_n + gcol;
-
-                T tmp[kElemsPerLoad];
-                if (gk < K && gn + kElemsPerLoad <= N) {
-                    const int4* src = reinterpret_cast<const int4*>(&B[size_t(gk) * N + gn]);
-                    *reinterpret_cast<int4*>(tmp) = *src;
-                } else {
-                    #pragma unroll
-                    for (int q = 0; q < kElemsPerLoad; ++q) {
-                        const int gn_q = gn + q;
-                        tmp[q] = (gk < K && gn_q < N) ? B[size_t(gk) * N + gn_q] : TR::from_f32(0.0f);
-                    }
-                }
-                *reinterpret_cast<int4*>(&Bs[row][gcol]) = *reinterpret_cast<int4*>(tmp);
-            }
+        // Issue the next slab's global loads BEFORE compute so they overlap.
+        const int k_next = k0 + RM_BK;
+        if (k_next < K) {
+            rm_prefetch_A<T>(a_reg, A, M, K, k_next, block_m, tid);
+            rm_prefetch_B<T>(b_reg, B, N, K, k_next, block_n, tid);
         }
 
-        __syncthreads();
-
         // ---- Compute on shared mem tiles ----
         // A frag: row_major from As, leading dim RM_LDA_SMEM, sub-tile at
         //         (warp_m*WM+i*WMMA_M, kk).

tests/CMakeLists.txt

@@ -126,6 +126,8 @@ if(BROTENSOR_HAS_GPU)
     brotensor_add_test(brotensor_test_int8_attention test_int8_attention.cpp)
     if(BROTENSOR_WITH_CUDA)
         brotensor_add_test(brotensor_test_streams test_streams.cpp)
+        # Hand-run microbench for the plain FP16 matmul WMMA path (not a ctest).
+        brotensor_add_executable(brotensor_bench_matmul_rm bench_matmul_rm.cpp)
     endif()
 
     # CPU↔GPU parity (GPU-only).

tests/bench_matmul_rm.cpp

@@ -0,0 +1,69 @@
+// Microbench for the plain FP16 matmul WMMA path (matmul_rm_wmma_kernel):
+//   C(M,N) = A(M,K) @ B(K,N), row-major, FP16 storage / FP32 accumulate.
+// Times a few representative shapes on CUDA. Not a ctest — built as a target
+// and run by hand to compare the register-prefetch kernel vs baseline.
+
+#include <brotensor/ops.h>
+#include <brotensor/runtime.h>
+#include <brotensor/tensor.h>
+
+#include <chrono>
+#include <cstdint>
+#include <cstdio>
+#include <random>
+#include <vector>
+
+using brotensor::Device;
+using brotensor::Tensor;
+
+static std::vector<uint16_t> to_fp16(const std::vector<float>& v) {
+    std::vector<uint16_t> o(v.size());
+    for (size_t i = 0; i < v.size(); ++i)
+        o[i] = brotensor::fp32_to_fp16_bits(v[i]);
+    return o;
+}
+
+static void bench(int M, int N, int K, int iters) {
+    std::mt19937 rng(0xC0DEu);
+    std::uniform_real_distribution<float> dist(-0.5f, 0.5f);
+    std::vector<float> A(size_t(M) * K), B(size_t(K) * N);
+    for (auto& v : A) v = dist(rng);
+    for (auto& v : B) v = dist(rng);
+    auto Ah = to_fp16(A), Bh = to_fp16(B);
+    Tensor Ag = Tensor::from_host_fp16_on(Device::CUDA, Ah.data(), M, K);
+    Tensor Bg = Tensor::from_host_fp16_on(Device::CUDA, Bh.data(), K, N);
+
+    Tensor Cg;
+    brotensor::matmul(Ag, Bg, Cg);  // warm-up
+    brotensor::sync_all();
+
+    using clk = std::chrono::steady_clock;
+    const auto t0 = clk::now();
+    for (int i = 0; i < iters; ++i) {
+        Tensor tmp;
+        brotensor::matmul(Ag, Bg, tmp);
+    }
+    brotensor::sync_all();
+    const double ms =
+        std::chrono::duration<double, std::milli>(clk::now() - t0).count();
+    const double gflop = 2.0 * M * N * K * 1e-9;
+    std::printf("  M=%-6d N=%-6d K=%-6d : %8.4f ms/iter  %8.1f GFLOP/s\n",
+                M, N, K, ms / iters, gflop / (ms / iters * 1e-3));
+}
+
+int main() {
+    brotensor::init();
+    std::setvbuf(stdout, nullptr, _IONBF, 0);
+    std::printf("[bench] matmul_rm FP16 WMMA path\n");
+    const int it = 100;
+    // Square GEMMs across scales.
+    bench(512, 512, 512, it);
+    bench(1024, 1024, 1024, it);
+    bench(2048, 2048, 2048, it);
+    bench(4096, 4096, 4096, it);
+    // Skewed shapes (short M, wide K/N — transformer-linear-like).
+    bench(64, 4096, 4096, it);
+    bench(128, 4096, 4096, it);
+    bench(256, 14336, 4096, it);
+    return 0;
+}