inferencelayer 0.2.8

//! cuBLAS prefill GEMM shim (P0′ of BEATING_VLLM.md) — the NVIDIA tensor-core fast path for the
//! LLM prefill matmuls, reached the ONLY way that works on NVIDIA: cuBLAS `GemmEx`, not wgpu
//! (the spike proved wgpu's cooperative-matrix is unavailable on the NVIDIA Vulkan backend).
//!
//! PORTABILITY BY CONSTRUCTION: this whole module is behind `#[cfg(feature = "cudarc")]`, and
//! `cudarc` dlopens libcuda/libcublas at RUNTIME (`dynamic-loading`, zero link-time deps). A
//! default build never compiles this and is byte-identical to today; Metal/CPU/any-Vulkan hosts
//! are untouched. The engine's job at the call site is: if this path is present AND a CUDA device
//! is reachable, use it for prefill GEMMs; otherwise fall through to the WGSL kernel unchanged.
//!
//! Mirrors `cuda_tower::CudaGlmTower::gemm` 1:1 (f16 in, `CUBLAS_COMPUTE_32F` ⇒ tensor cores with
//! fp32 accumulate — the wgpu `f16a` acceptance class), so an eventual serving integration reuses
//! the arm the vision tower already validates.
use anyhow::{Context as _, Result};
use std::sync::Arc;

use cudarc::cublas::{result as blas, sys as blas_sys, CudaBlas};
use cudarc::driver::{
    CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
    PushKernelArg,
};
use cudarc::nvrtc::{compile_ptx_with_opts, CompileOptions};
use half::f16;

/// Decode-side Q4 GEMV. **This is the win path, not the parity path.** Decode is memory-bound:
/// Qwen3-30B-A3B activates ~3.3B params/token, so our `q4_0` weights are ~2.1 GB/token against
/// vLLM's BF16 ~6.6 GB/token — a ~4x structural bandwidth advantage that only materialises if the
/// kernel actually saturates HBM. The WGSL arm does not on NVIDIA (28.5 tok/s measured); this is
/// the CUDA arm that should.
///
/// `q4_0` layout, matching `weights::quantize_q4_0` / the WGSL kernels: 32 weights per block,
/// one f32 scale + four u32 of packed nibbles (= 20 bytes / 32 weights = 5 bits/weight).
/// Nibble ORDER here mirrors the WGSL `unpack4xU8(w & 0x0F0F0F0F)` / `>>4` pair; exact index
/// agreement is re-gated at wire-up time — this kernel is a BANDWIDTH probe, and byte volume
/// (which decides decode tok/s) is layout-exact regardless.
const Q4_GEMV_KERNEL: &str = r#"
// NVRTC compiles at RUNTIME and ships no cuda_fp16.h, and requiring CUDA *headers* on a serving
// host would undo the point of a dlopen-only fast path. So every half operation below is inline
// PTX — available to NVRTC unconditionally, no include path, no toolkit.
__device__ __forceinline__ float h2f(unsigned short h) {
    // the u16 must be moved into a .f16 register first — cvt.f32.f16 will not take a .u16 operand
    float f;
    asm("{ .reg .f16 t; mov.b16 t, %1; cvt.f32.f16 %0, t; }" : "=f"(f) : "h"(h));
    return f;
}
__device__ __forceinline__ unsigned int pack_h2(float lo, float hi) {
    unsigned int d;
    asm("{ .reg .f16 l, h; cvt.rn.f16.f32 l, %1; cvt.rn.f16.f32 h, %2; mov.b32 %0, {l,h}; }"
        : "=r"(d) : "f"(lo), "f"(hi));
    return d;
}
__device__ __forceinline__ unsigned int hadd2(unsigned int a, unsigned int b) {
    unsigned int d; asm("add.f16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); return d;
}
__device__ __forceinline__ unsigned int hfma2(unsigned int a, unsigned int b, unsigned int c) {
    unsigned int d; asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(c)); return d;
}
__device__ __forceinline__ void unpack_h2(unsigned int a, float* lo, float* hi) {
    asm("{ .reg .f16 l, h; mov.b32 {l,h}, %2; cvt.f32.f16 %0, l; cvt.f32.f16 %1, h; }"
        : "=f"(*lo), "=f"(*hi) : "r"(a));
}

extern "C" __global__ void gemv_q4(
    const float* __restrict__ scales,
    const uint4* __restrict__ quants,
    const float* __restrict__ x,
    float*       __restrict__ y,
    const int nblk)
{
    const int row = blockIdx.x;
    const long base = (long)row * (long)nblk;
    float acc = 0.f;
    for (int b = threadIdx.x; b < nblk; b += blockDim.x) {
        const uint4 q = quants[base + b];
        const float s = scales[base + b];
        const float* xp = x + ((long)b << 5);
        unsigned int w[4] = { q.x, q.y, q.z, q.w };
        float sub = 0.f;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const unsigned int v = w[i];
            // low nibble of each byte, then high nibble — the WGSL q4_lo/q4_hi pair.
            #pragma unroll
            for (int byte = 0; byte < 4; ++byte) {
                const int lo = (int)((v >> (8 * byte)) & 0xFu) - 8;
                const int hi = (int)((v >> (8 * byte + 4)) & 0xFu) - 8;
                sub = fmaf((float)lo, xp[i * 8 + byte], sub);
                sub = fmaf((float)hi, xp[i * 8 + 4 + byte], sub);
            }
        }
        acc = fmaf(s, sub, acc);
    }
    __shared__ float red[256];
    red[threadIdx.x] = acc;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if ((int)threadIdx.x < off) red[threadIdx.x] += red[threadIdx.x + off];
        __syncthreads();
    }
    if (threadIdx.x == 0) y[row] = red[0];
}

// v2: ONE WARP PER ROW (8 rows per 256-thread block). Reduction is __shfl_down only — no shared
// memory, no __syncthreads — and 8x more rows in flight per block, which is what keeps enough
// loads outstanding to saturate HBM on a GEMV this thin.
extern "C" __global__ void gemv_q4_warp(
    const float* __restrict__ scales,
    const uint4* __restrict__ quants,
    const float* __restrict__ x,
    float*       __restrict__ y,
    const int nblk,
    const int rows)
{
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int row  = blockIdx.x * (blockDim.x >> 5) + warp;
    if (row >= rows) return;
    const long base = (long)row * (long)nblk;
    float acc = 0.f;
    for (int b = lane; b < nblk; b += 32) {
        const uint4 q = quants[base + b];
        const float s = scales[base + b];
        const float* xp = x + ((long)b << 5);
        unsigned int w[4] = { q.x, q.y, q.z, q.w };
        float sub = 0.f;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const unsigned int v = w[i];
            #pragma unroll
            for (int byte = 0; byte < 4; ++byte) {
                const int lo = (int)((v >> (8 * byte)) & 0xFu) - 8;
                const int hi = (int)((v >> (8 * byte + 4)) & 0xFu) - 8;
                sub = fmaf((float)lo, xp[i * 8 + byte], sub);
                sub = fmaf((float)hi, xp[i * 8 + 4 + byte], sub);
            }
        }
        acc = fmaf(s, sub, acc);
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0) y[row] = acc;
}

// v3: the one that actually streams HBM. v1/v2 read x straight from global, where lane L wants
// x[(lane + 32k)*32 .. +32] — a 128-BYTE STRIDE ACROSS LANES, so every one of the 32 scalar loads
// fans out into 32 separate cache lines. Measured cost: 124 GB/s, 6% of peak.
//
// Fix, both halves:
//   (a) stage the x tile in SHARED once per block, loaded coalesced, reused by all 8 rows;
//   (b) pad the per-q-block stride 32 -> 33 floats so lane L's read of element j lands in bank
//       (L + j) % 32 — distinct across the warp instead of all 32 lanes hitting one bank.
#define TILEB 64            /* q-blocks per K-tile: 64*32 = 2048 activations */
#define PAD   33            /* padded floats per q-block (32 + 1) */
extern "C" __global__ void gemv_q4_sh(
    const float* __restrict__ scales,
    const uint4* __restrict__ quants,
    const float* __restrict__ x,
    float*       __restrict__ y,
    const int nblk,
    const int rows)
{
    __shared__ float xs[TILEB * PAD];
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int row  = blockIdx.x * 8 + warp;      /* 8 warps = 8 rows per block */
    const long base = (long)row * (long)nblk;
    float acc = 0.f;

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        // coalesced global read, scattered into the padded shared layout
        for (int i = threadIdx.x; i < TILEB * 32; i += blockDim.x) {
            const int b = i >> 5, j = i & 31;
            const int gi = (t0 << 5) + i;
            xs[b * PAD + j] = ((t0 + b) < nblk) ? x[gi] : 0.f;
        }
        __syncthreads();
        if (row < rows) {
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                const uint4 q = quants[base + t0 + b];
                const float s = scales[base + t0 + b];
                const float* xp = xs + b * PAD;
                unsigned int w[4] = { q.x, q.y, q.z, q.w };
                float sub = 0.f;
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const unsigned int v = w[i];
                    #pragma unroll
                    for (int byte = 0; byte < 4; ++byte) {
                        const int lo = (int)((v >> (8 * byte)) & 0xFu) - 8;
                        const int hi = (int)((v >> (8 * byte + 4)) & 0xFu) - 8;
                        sub = fmaf((float)lo, xp[i * 8 + byte], sub);
                        sub = fmaf((float)hi, xp[i * 8 + 4 + byte], sub);
                    }
                }
                acc = fmaf(s, sub, acc);
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0 && row < rows) y[row] = acc;
}

// v4: v3 + the two things production already implies.
//   (a) f16 SCALES — the engine's WGSL GEMV binds `array<f16>`; an f32 scale is 4 of every 20
//       bytes, i.e. 20%% of decode traffic spent on scales. f16 -> 18 B per 32 weights.
//   (b) 2-way unroll over q-blocks, so two uint4 loads are in flight per lane instead of one
//       (memory-level parallelism is what a GEMV this thin actually starves on).
extern "C" __global__ void gemv_q4_v4(
    const unsigned short* __restrict__ scales,
    const uint4*  __restrict__ quants,
    const float*  __restrict__ x,
    float*        __restrict__ y,
    const int nblk,
    const int rows)
{
    __shared__ float xs[TILEB * PAD];
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int row  = blockIdx.x * 8 + warp;
    const long base = (long)row * (long)nblk;
    float acc = 0.f;

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        for (int i = threadIdx.x; i < TILEB * 32; i += blockDim.x) {
            const int b = i >> 5, j = i & 31;
            xs[b * PAD + j] = ((t0 + b) < nblk) ? x[(t0 << 5) + i] : 0.f;
        }
        __syncthreads();
        if (row < rows) {
            // two q-blocks per lane per step: b and b+32
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 64) {
                const int b2 = b + 32;
                const bool has2 = (b2 < TILEB) && ((t0 + b2) < nblk);
                const uint4 qa = quants[base + t0 + b];
                const uint4 qb = has2 ? quants[base + t0 + b2] : make_uint4(0u, 0u, 0u, 0u);
                const float sa = h2f(scales[base + t0 + b]);
                const float sb = has2 ? h2f(scales[base + t0 + b2]) : 0.f;
                const float* xa = xs + b * PAD;
                const float* xb = xs + (has2 ? b2 : b) * PAD;
                unsigned int wa[4] = { qa.x, qa.y, qa.z, qa.w };
                unsigned int wb[4] = { qb.x, qb.y, qb.z, qb.w };
                float suba = 0.f, subb = 0.f;
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    #pragma unroll
                    for (int byte = 0; byte < 4; ++byte) {
                        const int la = (int)((wa[i] >> (8 * byte)) & 0xFu) - 8;
                        const int ha = (int)((wa[i] >> (8 * byte + 4)) & 0xFu) - 8;
                        suba = fmaf((float)la, xa[i * 8 + byte], suba);
                        suba = fmaf((float)ha, xa[i * 8 + 4 + byte], suba);
                        const int lb = (int)((wb[i] >> (8 * byte)) & 0xFu) - 8;
                        const int hb = (int)((wb[i] >> (8 * byte + 4)) & 0xFu) - 8;
                        subb = fmaf((float)lb, xb[i * 8 + byte], subb);
                        subb = fmaf((float)hb, xb[i * 8 + 4 + byte], subb);
                    }
                }
                acc = fmaf(sa, suba, acc);
                if (has2) acc = fmaf(sb, subb, acc);
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0 && row < rows) y[row] = acc;
}

// v5: v4 + AWQ/FasterTransformer-style dequant. The 32 int->float `cvt` per q-block run on a
// quarter-rate pipe and are the prime suspect for the last ~25%% of peak. This builds halves
// DIRECTLY with bit ops: OR a nibble into the mantissa of 1024.0h (0x6400) so the value arrives
// as (1024 + q) in half, then one __hfma2 folds out the +1024 and the q4_0 -8 zero point.
// Activations are staged as half2 as well, halving shared traffic and letting one __hfma2 do
// two MACs.
extern "C" __global__ void gemv_q4_v5(
    const unsigned short* __restrict__ scales,
    const uint4*  __restrict__ quants,
    const float*  __restrict__ x,
    float*        __restrict__ y,
    const int nblk,
    const int rows)
{
    __shared__ unsigned int xh[TILEB * 17];   /* 16 half2 per q-block + 1 pad slot */
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int row  = blockIdx.x * 8 + warp;
    const long base = (long)row * (long)nblk;
    float acc = 0.f;
    const unsigned int bias = 0xE408E408u;   /* half2(-1032.0) = -(1024 + 8), both lanes */

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        // Stage activations as half2, PERMUTED to match the 2-op nibble extraction below: masking
        // a u32 with 0x000F000F yields the nibbles at bits[0:4] and bits[16:20] together, i.e.
        // elements (8i + r) and (8i + r + 4). Pair them here so the inner loop needs no shuffling.
        for (int i = threadIdx.x; i < TILEB * 16; i += blockDim.x) {
            const int b = i >> 4, j = i & 15;
            const int u = j >> 2, r = j & 3;              /* u32 index, pair within it */
            const int e = (b << 5) + (u << 3) + r;        /* element (8u + r) of q-block b */
            const int gi = (t0 << 5) + e;
            xh[b * 17 + j] = ((t0 + b) < nblk) ? pack_h2(x[gi], x[gi + 4]) : 0u;
        }
        __syncthreads();
        if (row < rows) {
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                const uint4 q = quants[base + t0 + b];
                const float s = h2f(scales[base + t0 + b]);
                const unsigned int* xp = xh + b * 17;
                unsigned int w[4] = { q.x, q.y, q.z, q.w };
                unsigned int sum = 0u;
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const unsigned int v = w[i];
                    // Two ops per PAIR: OR the nibble into the mantissa of 1024.0h (0x6400) so the
                    // value lands as (1024 + q) in half, then one add folds out 1024 and the -8
                    // zero point. No int->float cvt anywhere.
                    #pragma unroll
                    for (int r = 0; r < 4; ++r) {
                        const unsigned int m = ((v >> (4 * r)) & 0x000F000Fu) | 0x64006400u;
                        sum = hfma2(hadd2(m, bias), xp[i * 4 + r], sum);
                    }
                }
                float lo, hi;
                unpack_h2(sum, &lo, &hi);
                acc = fmaf(s, lo + hi, acc);
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0 && row < rows) y[row] = acc;
}
"#;

/// A cuBLAS handle bound to CUDA device 0. Holds nothing model-specific — it is a reusable GEMM
/// engine for the prefill matmuls (QKV, MLP gate/up/down, LM head).
pub struct CublasPrefill {
    _ctx: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    blas: CudaBlas,
}

impl CublasPrefill {
    /// Fails cleanly (no panic) if no CUDA device / cuBLAS is reachable — the caller treats an
    /// `Err` as "CUDA fast path unavailable, use WGSL".
    pub fn new() -> Result<Self> {
        let ctx = CudaContext::new(0).context("cuda device 0 (no NVIDIA GPU / driver?)")?;
        let stream = ctx.default_stream();
        let blas = CudaBlas::new(stream.clone()).context("cublas handle")?;
        Ok(Self { _ctx: ctx, stream, blas })
    }

    /// Upload a host f16 slice to the device (weights/activations are prepared once, off the hot path).
    pub fn upload_f16(&self, host: &[f16]) -> Result<CudaSlice<f16>> {
        Ok(self.stream.memcpy_stod(host)?)
    }

    pub fn alloc_out(&self, len: usize) -> Result<CudaSlice<f32>> {
        Ok(self.stream.alloc_zeros::<f32>(len)?)
    }

    /// `y[m, n] = x[m, k] · Wᵀ`, W row-major `[n, k]`, both operands f16, output f32 — one
    /// `cublasGemmEx` on the tensor cores. Identical call shape to `cuda_tower::gemm`.
    pub fn gemm(&self, m: usize, n: usize, k: usize,
                w: &CudaSlice<f16>, x: &CudaSlice<f16>, y: &mut CudaSlice<f32>) -> Result<()> {
        let alpha: f32 = 1.0;
        let beta: f32 = 0.0;
        let (pw, _gw) = w.device_ptr(&self.stream);
        let (px, _gx) = x.device_ptr(&self.stream);
        let (py, _gy) = y.device_ptr_mut(&self.stream);
        unsafe {
            blas::gemm_ex(
                *self.blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                n as i32, m as i32, k as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16F, k as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16F, k as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, n as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            )
            .map_err(|e| anyhow::anyhow!("gemm_ex: {e:?}"))?;
        }
        Ok(())
    }

    pub fn sync(&self) -> Result<()> {
        self.stream.synchronize()?;
        Ok(())
    }

    /// Compile + bind the `q4_0` decode GEMV (NVRTC, at startup — the only JIT in this module).
    pub fn q4_gemv(&self) -> Result<CudaQ4Gemv> {
        // NVRTC defaults to compute_52, where the f16x2 ops v5 needs do not exist (sm_53+).
        // compute_70 is the floor that covers every GPU this engine targets (V100 and up); the
        // driver JITs the PTX to the actual device on load.
        let ptx = compile_ptx_with_opts(
            Q4_GEMV_KERNEL,
            CompileOptions {
                arch: Some("compute_70"),
                ..Default::default()
            },
        )
        .map_err(|e| anyhow::anyhow!("nvrtc: {e:?}"))?;
        let module = self._ctx.load_module(ptx).context("load ptx")?;
        Ok(CudaQ4Gemv {
            stream: self.stream.clone(),
            f_block: module.load_function("gemv_q4").context("gemv_q4")?,
            f_warp: module.load_function("gemv_q4_warp").context("gemv_q4_warp")?,
            f_sh: module.load_function("gemv_q4_sh").context("gemv_q4_sh")?,
            f_v4: module.load_function("gemv_q4_v4").context("gemv_q4_v4")?,
            f_v5: module.load_function("gemv_q4_v5").context("gemv_q4_v5")?,
        })
    }

    /// Best-of-`reps`-rounds TFLOP/s at (m,n,k). Synthetic operands; the number is the tensor-core
    /// ceiling this shim delivers — compare against the WGSL/shader-ALU `gemm-coop-bench`.
    pub fn bench_tflops(&self, m: usize, n: usize, k: usize, reps: usize) -> Result<f64> {
        let w = self.upload_f16(&vec![f16::from_f32(0.02); n * k])?;
        let x = self.upload_f16(&vec![f16::from_f32(0.02); m * k])?;
        let mut y = self.alloc_out(m * n)?;
        for _ in 0..3 { self.gemm(m, n, k, &w, &x, &mut y)?; } // warm
        self.sync()?;
        let mut best = f64::MAX;
        for _ in 0..5 {
            let t0 = std::time::Instant::now();
            for _ in 0..reps { self.gemm(m, n, k, &w, &x, &mut y)?; }
            self.sync()?;
            best = best.min(t0.elapsed().as_secs_f64());
        }
        Ok(2.0 * m as f64 * n as f64 * k as f64 * reps as f64 / best / 1e12)
    }
}

/// The `q4_0` decode GEMV bound to a stream — decode's whole cost is streaming weights, so this
/// kernel's sustained GB/s IS the decode tok/s ceiling.
pub struct CudaQ4Gemv {
    stream: Arc<CudaStream>,
    f_block: CudaFunction,
    f_warp: CudaFunction,
    f_sh: CudaFunction,
    f_v4: CudaFunction,
    f_v5: CudaFunction,
}

/// Which `q4_0` GEMV to launch. v1/v2 read activations from global with a 128 B inter-lane stride
/// (measured 124 GB/s = 6% of H100 peak); v3 stages them in padded shared memory.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum GemvKernel {
    /// One 256-thread block per row, shared-memory tree reduction.
    BlockPerRow,
    /// One warp per row (8 rows/block), `__shfl_down` reduction.
    WarpPerRow,
    /// v3: warp-per-row + x staged in shared, stride padded 32→33 (bank-conflict free).
    SharedX,
    /// v4: v3 + f16 scales (production format, 18 B/block) + 2-way q-block unroll.
    F16ScalesUnroll,
    /// v5: v4 + bit-trick nibble→half dequant and `__hfma2` (no int→float `cvt`).
    Half2BitTrick,
}

impl GemvKernel {
    /// q4_0 weight bytes per 32-weight block: 16 B quants + the scale (f32 for v1–v3, f16 after).
    fn bytes_per_block(self) -> f64 {
        match self {
            GemvKernel::BlockPerRow | GemvKernel::WarpPerRow | GemvKernel::SharedX => 20.0,
            _ => 18.0,
        }
    }
    fn uses_f16_scales(self) -> bool {
        self.bytes_per_block() < 20.0
    }
}

impl CudaQ4Gemv {
    /// Sustained weight-traffic GB/s at `rows x cols` (best of 5 rounds of `reps` launches).
    /// Sizes must exceed L2 (50 MB on H100) or this measures cache, not HBM.
    /// `warp_per_row` picks the v2 (shuffle-reduce, 8 rows/block) kernel over the v1 block-per-row.
    pub fn bench_gbps(
        &self,
        rows: usize,
        cols: usize,
        reps: usize,
        kernel: GemvKernel,
    ) -> Result<f64> {
        anyhow::ensure!(cols % 32 == 0, "cols must be a multiple of 32 (q4_0 block)");
        let nblk = cols / 32;
        // v1–v3 bind f32 scales, v4/v5 the production f16. Allocate the one in play.
        let scales = self.stream.alloc_zeros::<f32>(if kernel.uses_f16_scales() { 1 } else { rows * nblk })?;
        let scales16 = self
            .stream
            .alloc_zeros::<f16>(if kernel.uses_f16_scales() { rows * nblk } else { 1 })?;
        let quants = self.stream.alloc_zeros::<u32>(rows * nblk * 4)?;
        let x = self.stream.alloc_zeros::<f32>(cols)?;
        let mut y = self.stream.alloc_zeros::<f32>(rows)?;
        let cfg = LaunchConfig {
            // v1: one 256-thread block per row. v2/v3: 8 rows (8 warps) per block.
            grid_dim: (
                match kernel {
                    GemvKernel::BlockPerRow => rows as u32,
                    _ => rows.div_ceil(8) as u32,
                },
                1,
                1,
            ),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0, // v3's xs[] is statically sized in the kernel
        };
        let nb = nblk as i32;
        let nrows = rows as i32;
        let mut launch = || -> Result<()> {
            unsafe {
                match kernel {
                    GemvKernel::BlockPerRow => self
                        .stream
                        .launch_builder(&self.f_block)
                        .arg(&scales)
                        .arg(&quants)
                        .arg(&x)
                        .arg(&mut y)
                        .arg(&nb)
                        .launch(cfg)?,
                    GemvKernel::WarpPerRow => self
                        .stream
                        .launch_builder(&self.f_warp)
                        .arg(&scales)
                        .arg(&quants)
                        .arg(&x)
                        .arg(&mut y)
                        .arg(&nb)
                        .arg(&nrows)
                        .launch(cfg)?,
                    GemvKernel::SharedX => self
                        .stream
                        .launch_builder(&self.f_sh)
                        .arg(&scales)
                        .arg(&quants)
                        .arg(&x)
                        .arg(&mut y)
                        .arg(&nb)
                        .arg(&nrows)
                        .launch(cfg)?,
                    GemvKernel::F16ScalesUnroll => self
                        .stream
                        .launch_builder(&self.f_v4)
                        .arg(&scales16)
                        .arg(&quants)
                        .arg(&x)
                        .arg(&mut y)
                        .arg(&nb)
                        .arg(&nrows)
                        .launch(cfg)?,
                    GemvKernel::Half2BitTrick => self
                        .stream
                        .launch_builder(&self.f_v5)
                        .arg(&scales16)
                        .arg(&quants)
                        .arg(&x)
                        .arg(&mut y)
                        .arg(&nb)
                        .arg(&nrows)
                        .launch(cfg)?,
                };
            }
            Ok(())
        };
        for _ in 0..3 {
            launch()?;
        }
        self.stream.synchronize()?;
        let mut best = f64::MAX;
        for _ in 0..5 {
            let t0 = std::time::Instant::now();
            for _ in 0..reps {
                launch()?;
            }
            self.stream.synchronize()?;
            best = best.min(t0.elapsed().as_secs_f64());
        }
        // Exact q4_0 weight volume: 16 B quants + the scale (4 B f32, or 2 B f16 in v4/v5).
        let bytes = rows as f64 * nblk as f64 * kernel.bytes_per_block() * reps as f64;
        Ok(bytes / best / 1e9)
    }
}