W4A16 GEMM Claude Opus 5

"""W4A16 (AWQ/GPTQ-style asymmetric int4, group=128) GEMM for SM120 Blackwell.

Fused unpack + GEMM in a single kernel launch.  Two code paths:

  * ``gemv_kernel`` (M <= 4): SIMT, fp32 accumulators, 128-bit weight loads,
    nibble unpack via one ``lop3`` per bf16x2 (0x4300 | v == 128.0f + v exactly),
    split-K with an in-kernel "last block reduces" fixup so the whole op stays a
    single launch.
  * ``gemm_kernel`` (M >= 5): bf16 tensor cores, ``mma.m16n8k16``.  Weights are
    repacked once at first use into per-lane B-fragment order, so a warp pulls
    512 contiguous bytes per (128k x 8n) tile and the four ``lop3`` outputs land
    directly in mma operand registers -- no shared memory and no shuffles for
    the weights at all.  The dequant is two ``fma.rn.bf16x2`` per packed pair and
    is bit-exact against the reference: ``(128+v) - (128+z)`` is exact in bf16
    (both sides are integers in the same binade, so the subtraction cannot round)
    and the following multiply by ``s`` reproduces ``bf16((v-z)*s)`` exactly, so
    only fp32 accumulation order differs from ``x @ dequant(w)``.

Neither path ever materialises a dequantized bf16 weight copy.

Where it stands, and what the remaining gap actually is.  Measured peak_fraction
per shape against a same-harness ceiling (scratch/bwprobe.py: a pure streaming
read of the same byte count, through the same L2 flush and the same ~4.2us
cuda-event floor, so the two are directly comparable):

    shape          0       1       2       3       4
    achieved  0.5491  0.4237  0.1497  0.3449  0.5787
    ceiling   0.6736  0.6843  0.7154  0.5049  0.7039

0 and 4 are within 1.2x of a flat stream and are done.  2 is the outlier, and
it is not the memory system: it is the only shape where the mma pipe is the
cost.  The pipe-throughput ceiling for it was originally read as 52us from
scratch/mmaprobe.py, which is wrong in a way worth recording, because the probe
runs the tensor pipe alone at 2.56 GHz while the real kernel lights every pipe
at once and clocks down to 1.93 GHz.  Correcting for that, and for
``sm__pipe_tensor_cycles_active`` being summed over sub-partitions rather than
wall-clock, the honest numbers are 4.20 active pipe-cycles per mma in the probe
against 4.82 here: the mma stream already runs at 87% of the probe's per-cycle
efficiency, and shape 2's floor is ~69us, not 52us.

What the other 33% of the pipe's idle cycles are waiting on, by ablation
(scratch/ab.py, each -D removes one piece of work and keeps the rest):

    weight loads (KB_NOWL / KB_WSAME)   1.21x / 1.19x
    dequant      (KB_NODQ)              1.08x
    ldmatrix     (KB_NOLDS)             1.07x

The weight-load fifth is L2 bandwidth, not DRAM and not latency: M=256 over
BM=32 is 8 row-tiles per column-tile, each re-reading the same 256 KB of
weights, so 25 MB of unique weights become 201 MB of L2 requests.  DRAM traffic
is already at the 35 MB theoretical minimum, which is why KB_WSAME (same
instruction count, one shared tile) buys almost exactly what KB_NOWL (no loads
at all) does.  Killing that amplification needs the weight k-slice staged in
shared memory and shared across row-tiles, i.e. a CUTLASS-style multistage
pipeline.  Every cheaper approximation of it measured worse: BM=64/128 (cfgs
30, 35) and WM>1 (cfgs 55-69) all pay for the wider tile in accumulator
registers, and registers are the binding constraint here -- 126 regs x 256
threads is 32256 per block, so 65536/32256 = 2.03 makes two blocks per SM the
hard ceiling and any accumulator growth drops it to one.
"""
from __future__ import annotations

import os

import torch
import torch.nn as nn

GROUP_SIZE = 128

os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")

_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdint.h>

#define DEVI __device__ __forceinline__

// ---------------------------------------------------------------------------
// int4 -> bf16 bit tricks
// ---------------------------------------------------------------------------
// bf16 with exponent field 0x86 (2^7) and mantissa m has value 128 + m.  Four
// mantissa bits are enough for a nibble, so 0x4300 | v is exactly (128.0f + v)
// in bf16.  One lop3 turns two nibbles of a 32-bit word into a packed bf16x2:
// low half from bits [shift+3:shift], high half from bits [shift+19:shift+16].
DEVI uint32_t nib2bf16x2(uint32_t w, int shift) {
    uint32_t r;
    asm("lop3.b32 %0, %1, %2, %3, 0xea;"
        : "=r"(r)
        : "r"(w >> shift), "n"(0x000f000fu), "n"(0x43004300u));
    return r;
}

DEVI float bf16lo_to_f32(uint32_t h) { return __int_as_float(h << 16); }
DEVI float bf16hi_to_f32(uint32_t h) { return __int_as_float(h & 0xffff0000u); }

DEVI uint32_t hfma2(uint32_t a, uint32_t b, uint32_t c) {
    uint32_t d;
    asm("fma.rn.bf16x2 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(c));
    return d;
}

DEVI void mma16816(float &d0, float &d1, float &d2, float &d3,
                   uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                   uint32_t b0, uint32_t b1) {
    asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// 16-byte global -> shared copy that does not pass through registers.  The
// src-size operand is the zero-fill trick: 0 writes 16 bytes of zeros without
// touching src, which is exactly the out-of-range-row padding the gemm needs.
DEVI void cpasync16(void* dst, const void* src, bool valid) {
    const uint32_t d = static_cast<uint32_t>(__cvta_generic_to_shared(dst));
    const int sz = valid ? 16 : 0;
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
                 :: "r"(d), "l"(src), "r"(sz) : "memory");
}

DEVI void cpasync_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); }
DEVI void cpasync_wait() { asm volatile("cp.async.wait_group 0;" ::: "memory"); }

// Whole mma A fragment (four 8x8 bf16 tiles) in one instruction instead of four
// LDS.  ldmatrix's per-matrix distribution -- lane l holds row l/4, columns
// 2*(l%4) and +1 -- is exactly the m16n8k16 A layout, and the four matrices come
// back in {rows 0-7 x cols 0-7, rows 8-15 x cols 0-7, rows 0-7 x cols 8-15,
// rows 8-15 x cols 8-15} order, which is exactly a0..a3.
DEVI void ldmatrix_x4(uint32_t& d0, uint32_t& d1, uint32_t& d2, uint32_t& d3,
                      const void* p) {
    const uint32_t a = static_cast<uint32_t>(__cvta_generic_to_shared(p));
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
                 : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) : "r"(a));
}

// ---------------------------------------------------------------------------
// GEMV: out[m, n] = sum_k x[m, k] * (unpack(wq)[k, n] - z[k/128, n]) * s[k/128, n]
//
// Grid  : (N / BN, SPLIT)
// Block : THREADS threads, a multiple of BN; each thread owns 16 consecutive
//         columns and loads 128 bits (16 columns x 2 k) per step.
//
// THREADS > BN is how this kernel buys occupancy.  ncu on shape 3 puts it at
// 43.6% of peak warps with a long-scoreboard stall of 17.6 -- pure load-latency
// starvation -- and the occupancy is set entirely by grid size, since total
// threads is (N / BN) * split * THREADS and THREADS used to be pinned to BN.
// Raising split instead would work too, but every extra split writes and re-reads
// another N floats of partials, and at shape 3's split of 32 that overhead is
// already 12% of the byte budget.  Widening the block splits k across more warps
// *inside* the block, where the existing red[] reduction already sums them, so
// the occupancy arrives free of extra traffic.
// ---------------------------------------------------------------------------
// KB_VMINB: same carveout lever as the gemm's MINB_OF (see there).  red[] is
// THREADS*17 floats, so a 128-thread block already wants 8.7 KB and lands in the
// 16 KB partition -- two resident blocks where the SM has room for far more.
// This kernel is small enough that the registers are there, so the request is
// free until the cap bites.  Measured curve at THREADS=128, medians of 9
// interleaved rounds, shape 0 / shape 3 peak_fraction:
//     1: 0.469 0.304   4: 0.466 0.310   6: 0.572 0.357   <-- best
//     8: 0.502 0.348  10: 0.373 0.185  12: 0.307 0.137
// 6 is +22% / +17% over no request.  The cliff after it is the register cap:
// 65536/(6*128) = 85 registers, and each further block drops that by ~14 until
// the kernel spills.  ncu shows none of this -- it normalises the carveout away.
#ifndef KB_VMINB
#define KB_VMINB 6
#endif

template <int M, int BN, int THREADS>
__global__ __launch_bounds__(THREADS, KB_VMINB) void gemv_kernel(
    const __nv_bfloat16* __restrict__ x,
    const uint8_t* __restrict__ wq,
    const __nv_bfloat16* __restrict__ scales,
    const __nv_bfloat16* __restrict__ zeros,
    __nv_bfloat16* __restrict__ out,
    float* __restrict__ partials,
    int* __restrict__ counters,
    int N, int K, int kh_per_split, int split) {
    constexpr int TPR = BN / 16;          // threads per packed-row
    constexpr int RPI = THREADS / TPR;    // packed rows consumed per iteration
    constexpr int ITERS_PER_GROUP = 64 / RPI;
    static_assert(RPI <= 64 && 64 % RPI == 0, "bad tile");
    static_assert(THREADS % BN == 0 && (BN & (BN - 1)) == 0, "THREADS must be a multiple of a power-of-two BN");

    const int tid = threadIdx.x;
    const int cg = tid % TPR;             // which 16-column slice
    const int ro = tid / TPR;             // row offset within an iteration
    const int n_tile = blockIdx.x;
    const int sp = blockIdx.y;
    const int n0 = n_tile * BN + cg * 16;

    const int kh_begin = sp * kh_per_split;
    const int kh_end = kh_begin + kh_per_split;

    // stride 17 (odd) makes both the strided store and the strided load
    // conflict-free / near conflict-free across the 32 smem banks.
    __shared__ float red[THREADS * 17];

    float outacc[M][16];
#pragma unroll
    for (int m = 0; m < M; ++m)
#pragma unroll
        for (int j = 0; j < 16; ++j) outacc[m][j] = 0.f;

    const uint32_t* wq32 = reinterpret_cast<const uint32_t*>(wq);

    for (int kh = kh_begin; kh < kh_end; kh += 64) {
        const int g = (kh * 2) / 128;
        float acc[M][16];
#pragma unroll
        for (int m = 0; m < M; ++m)
#pragma unroll
            for (int j = 0; j < 16; ++j) acc[m][j] = 0.f;
        float xsum[M];
#pragma unroll
        for (int m = 0; m < M; ++m) xsum[m] = 0.f;

#pragma unroll
        for (int it = 0; it < ITERS_PER_GROUP; ++it) {
            const int r = kh + it * RPI + ro;
            // 4 x 32-bit = 16 columns x (even k, odd k)
            uint4 wv = *reinterpret_cast<const uint4*>(wq32 + (size_t)r * (N / 4) + n0 / 4);
            const uint32_t words[4] = {wv.x, wv.y, wv.z, wv.w};

            float xe[M], xo[M];
#pragma unroll
            for (int m = 0; m < M; ++m) {
                uint32_t xp = *reinterpret_cast<const uint32_t*>(x + (size_t)m * K + 2 * r);
                xe[m] = bf16lo_to_f32(xp);
                xo[m] = bf16hi_to_f32(xp);
                xsum[m] += xe[m] + xo[m];
            }

#pragma unroll
            for (int w = 0; w < 4; ++w) {
                const uint32_t W = words[w];
                // h0={c0.lo, c2.lo} h1={c0.hi, c2.hi} h2={c1.lo, c3.lo} h3={c1.hi, c3.hi}
                // (byte b of W is column 4w+b; low nibble = even k, high = odd k)
                uint32_t h0 = nib2bf16x2(W, 0);
                uint32_t h1 = nib2bf16x2(W, 4);
                uint32_t h2 = nib2bf16x2(W, 8);
                uint32_t h3 = nib2bf16x2(W, 12);
                float v0e = bf16lo_to_f32(h0), v2e = bf16hi_to_f32(h0);
                float v0o = bf16lo_to_f32(h1), v2o = bf16hi_to_f32(h1);
                float v1e = bf16lo_to_f32(h2), v3e = bf16hi_to_f32(h2);
                float v1o = bf16lo_to_f32(h3), v3o = bf16hi_to_f32(h3);
#pragma unroll
                for (int m = 0; m < M; ++m) {
                    acc[m][4 * w + 0] += xe[m] * v0e + xo[m] * v0o;
                    acc[m][4 * w + 1] += xe[m] * v1e + xo[m] * v1o;
                    acc[m][4 * w + 2] += xe[m] * v2e + xo[m] * v2o;
                    acc[m][4 * w + 3] += xe[m] * v3e + xo[m] * v3o;
                }
            }
        }

        // group tail: out += s * (acc - (128 + z) * xsum).  16 columns of bf16
        // scales/zeros = 32 bytes each, fetched as two 128-bit loads.
        const uint4* svp = reinterpret_cast<const uint4*>(scales + (size_t)g * N + n0);
        const uint4* zvp = reinterpret_cast<const uint4*>(zeros + (size_t)g * N + n0);
        uint4 sa = svp[0], sb = svp[1];
        uint4 za = zvp[0], zb = zvp[1];
        const uint32_t sw[8] = {sa.x, sa.y, sa.z, sa.w, sb.x, sb.y, sb.z, sb.w};
        const uint32_t zw[8] = {za.x, za.y, za.z, za.w, zb.x, zb.y, zb.z, zb.w};
#pragma unroll
        for (int j2 = 0; j2 < 8; ++j2) {
            float s_lo = bf16lo_to_f32(sw[j2]), s_hi = bf16hi_to_f32(sw[j2]);
            float z_lo = bf16lo_to_f32(zw[j2]) + 128.0f;
            float z_hi = bf16hi_to_f32(zw[j2]) + 128.0f;
#pragma unroll
            for (int m = 0; m < M; ++m) {
                outacc[m][2 * j2 + 0] += s_lo * (acc[m][2 * j2 + 0] - z_lo * xsum[m]);
                outacc[m][2 * j2 + 1] += s_hi * (acc[m][2 * j2 + 1] - z_hi * xsum[m]);
            }
        }
    }

    // ---- reduce the RPI partial rows inside the block --------------------
    float col[M];
#pragma unroll
    for (int m = 0; m < M; ++m) {
        if (m) __syncthreads();
#pragma unroll
        for (int j = 0; j < 16; ++j) red[tid * 17 + j] = outacc[m][j];
        __syncthreads();
        float t = 0.f;
        // Wrapped, so the THREADS - BN threads that own no output column still
        // read in range; they just recompute a column someone else stores.
        const int lane_cg = (tid & (BN - 1)) >> 4;   // 0..TPR-1
        const int lane_j = tid % 16;
#pragma unroll 8
        for (int r = 0; r < RPI; ++r) t += red[((r * TPR) + lane_cg) * 17 + lane_j];
        col[m] = t;
    }

    const int myn = n_tile * BN + tid;   // only tid < BN owns an output column
    if (split == 1) {
        if (tid < BN) {
#pragma unroll
            for (int m = 0; m < M; ++m) out[(size_t)m * N + myn] = __float2bfloat16(col[m]);
        }
        return;
    }

    // Single-launch split-K: publish this split's partial, then whichever block
    // arrives last for this column tile does the cross-split reduction.  Loads
    // bypass L1 (volatile) because the producers ran on other SMs.
    volatile float* vp = partials;
    if (tid < BN) {
#pragma unroll
        for (int m = 0; m < M; ++m)
            vp[((size_t)sp * M + m) * N + myn] = col[m];
    }
    __threadfence();
    __shared__ bool last;
    if (tid == 0) {
        int old = atomicAdd(&counters[n_tile], 1);
        last = (old == split - 1);
        if (last) counters[n_tile] = 0;
    }
    __syncthreads();
    if (!last || tid >= BN) return;
#pragma unroll
    for (int m = 0; m < M; ++m) {
        float t = 0.f;
        for (int s = 0; s < split; ++s) t += vp[((size_t)s * M + m) * N + myn];
        out[(size_t)m * N + myn] = __float2bfloat16(t);
    }
}

// ---------------------------------------------------------------------------
// GEMM (tensor cores).
//
// Repacked weight layout, built once on the host:
//     wr[(kc * (N/8) + nt) * 512 + t * 16 + (4*w + bi)]
// where kc = k/128, nt = n/8, t = 4*c + r is the lane id (c = column within the
// 8-wide n tile = the mma B-fragment's groupID, r = lane%4 = the k-pair index),
// w = 0..3 and bi = 0..3.  A lane's 16 bytes are therefore one uint4 load, and
// for word w the four nib2bf16x2 outputs at shifts 0/4/8/12 are exactly
//     b0(kstep 2w), b1(kstep 2w), b0(kstep 2w+1), b1(kstep 2w+1)
// of mma.m16n8k16 -- no permutation, no smem staging for the weights.
//
// Grid  : (N / BN, ceil(M / BM))
// ---------------------------------------------------------------------------
// WDG = weight prefetch depth in super-chunks.  Bytes in flight per warp are
// NT*WDG*512, and the small-M shapes are latency-bound on exactly that number;
// the large-M shape is register-bound instead, so the depth is per-config.
// ---- ablation hooks: numerically wrong on purpose, used only to attribute the
// gemm's stall time to one stage.  KB_NODQ drops the unpack+scale ALU chain,
// KB_NOLDS drops the shared-memory A-fragment reads, KB_NOWL drops the weight
// global loads.  None is defined in a normal build, where all three expand to
// exactly the code they replace.
#ifdef KB_NOLDS
#define LOADA(mi, KK)                                                                \
    a[mi][0] = a[mi][1] = a[mi][2] = a[mi][3] = 0x3f803f80u | ((KK) & 1);
#elif defined(KB_NOLDM)
// four separate LDS, kept only so scratch/ab.py can measure what ldmatrix buys
#define LOADA(mi, KK)                                                                \
    {                                                                                \
        const int rb_ = xofs32 + (((wk * BM) + m0 + (mi) * 16 + gid) * XW            \
                                 + (KK) * 8 + tig);                                  \
        a[mi][0] = xsm32[rb_];                                                       \
        a[mi][1] = xsm32[rb_ + 8 * XW];                                              \
        a[mi][2] = xsm32[rb_ + 4];                                                   \
        a[mi][3] = xsm32[rb_ + 8 * XW + 4];                                          \
    }
#else
#define LOADA(mi, KK)                                                                \
    ldmatrix_x4(a[mi][0], a[mi][1], a[mi][2], a[mi][3],                              \
                &xsmb[xofsb + ((wk * BM) + m0 + (mi) * 16 + lrow) * XSTRIDE          \
                      + (KK) * 16 + lkoff]);
#endif

#ifdef KB_NOWL
#define WLOAD(dst, off) do { uint32_t v_ = 0x12345678u ^ (uint32_t)(off);             \
    dst[0] = v_; dst[1] = v_; dst[2] = v_; dst[3] = v_; } while (0)
#elif defined(KB_WSAME)
// same instruction count and same load latency as the real thing, but every
// block reads the one 512 B tile, so L2 traffic collapses to nothing.  Separates
// "L2 bandwidth bound" from "weight-load latency bound".
#define WLOAD(dst, off) do { uint4 t_ = wrv[(off) & 31];                               \
    dst[0] = t_.x; dst[1] = t_.y; dst[2] = t_.z; dst[3] = t_.w; } while (0)
#else
#define WLOAD(dst, off) do { uint4 t_ = wrv[off];                                      \
    dst[0] = t_.x; dst[1] = t_.y; dst[2] = t_.z; dst[3] = t_.w; } while (0)
#endif

#ifdef KB_NODQ
#define DEQUANT(W_, S0, S1, ti)                                                      \
    const uint32_t b0 = (W_) ^ sp[ti];                                               \
    const uint32_t b1 = (W_) ^ np[ti];
#else
#define DEQUANT(W_, S0, S1, ti)                                                      \
    const uint32_t d0 = hfma2(nib2bf16x2(W_, S0), ONE2, np[ti]);                     \
    const uint32_t d1 = hfma2(nib2bf16x2(W_, S1), ONE2, np[ti]);                     \
    const uint32_t b0 = hfma2(d0, sp[ti], ZERO2);                                    \
    const uint32_t b1 = hfma2(d1, sp[ti], ZERO2);
#endif

// XSM picks how the x slab reaches shared memory, and the right answer is not
// the same for every tile.  XSM=0 stages through registers with two barriers per
// super-chunk; XSM=1 uses a double-buffered cp.async with one.  cp.async is worth
// +1.7% on the M=16 tile, which is grid-starved (~112 of 188 SMs get a block) so
// in-block overlap is the only overlap it has, and -12% on the M=32 tile, where
// the register path's load-before-the-barrier already keeps one global load per
// warp in flight and the cp.async wait just serialises it.  Measured with
// scratch/ab.py, which interleaves variants in one warmed process; single
// benchmark.py runs cannot resolve this on a box whose SM clock floats 2-3 GHz.
// The second __launch_bounds__ argument is not about registers here, it is about
// the shared-memory carveout.  The driver sizes the L1/shared split from the
// .minnctapersm directive, so a tile wanting 17.4 KB of shared and not asking
// for anything gets the 32 KB partition -- one resident block per SM, 8 warps,
// 16.5% occupancy, and ncu blaming 36% of the stall cycles on
// math_pipe_throttle purely for want of warps to hide the tensor pipe.  Asking
// for two makes the driver reserve enough for two blocks: +20% on the M=256
// tile and +4% on M=32, from SASS that cuobjdump confirms is byte-identical.
//
// Only for tiles of at most 512 threads.  Two 1024-thread blocks exceed the
// 1536-thread limit, so the request cannot be honoured there and would only cap
// registers at 32 and spill.  Three is worse everywhere: 65536/(3*256) leaves
// 85 registers against the 126 this tile needs, and the spill costs more than
// the occupancy buys (measured 0.837x on M=256).
#ifdef KB_MINB
#define MINB_OF(THR_) KB_MINB
#else
#define MINB_OF(THR_) ((THR_) <= 512 ? 2 : 1)
#endif

template <int BM, int BN, int WM, int WN, int WK, int WDG, int XSM>
__global__ __launch_bounds__(32 * WM * WN * WK, MINB_OF(32 * WM * WN * WK))
void gemm_kernel(
    const __nv_bfloat16* __restrict__ x,
    const uint8_t* __restrict__ wr,
    const __nv_bfloat16* __restrict__ sc,
    const __nv_bfloat16* __restrict__ zr,
    __nv_bfloat16* __restrict__ out,
    int M, int N, int K) {
    constexpr int WARPS = WM * WN * WK;
    constexpr int THREADS = 32 * WARPS;
    constexpr int MT = BM / (WM * 16);         // m-tiles per warp
    constexpr int NT = BN / (WN * 8);          // n-tiles per warp
    constexpr int ACCN = MT * NT * 4;
    constexpr int XSTRIDE = 136;               // bf16 elements; 272 B, 16 B aligned
    constexpr int XW = XSTRIDE / 2;            // uint32 words per staged x row
    constexpr int XVEC = BM * 16 * WK;         // uint4 loads per staged x tile
    constexpr int XWORDS = WK * BM * XW;
    // Two x buffers when they fit in the 48 KB static limit: the copy for
    // super-chunk su+1 then overlaps the whole mma block for su and the loop
    // needs one barrier instead of two.  The serial single-buffer cp.async
    // schedule measured worse than register staging everywhere, so a tile too
    // tall to double-buffer just falls back to registers rather than to it.
    constexpr int XASYNC = (XSM == 1 && 2 * XWORDS * 4 <= 48 * 1024) ? 1 : 0;
    constexpr int XBUF = XASYNC ? 2 : 1;
    constexpr int RWORDS = (WK - 1) * (WM * WN) * 32 * ACCN;
    constexpr int SWORDS = XBUF * XWORDS > RWORDS ? XBUF * XWORDS : RWORDS;
    static_assert(BM == WM * MT * 16, "bad BM");
    static_assert(BN == WN * NT * 8, "bad BN");
    static_assert(XVEC % THREADS == 0, "bad x staging");

    __shared__ __align__(16) uint32_t smem[SWORDS];
    __nv_bfloat16* xsm = reinterpret_cast<__nv_bfloat16*>(smem);
    const __nv_bfloat16* xsmb = xsm;

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int wk = warp / (WM * WN);
    const int wmn = warp % (WM * WN);
    const int wm = wmn / WN;
    const int wn = wmn % WN;
    const int gid = lane >> 2;
    const int tig = lane & 3;
    // ldmatrix.x4 address contract: lanes 0-15 hand it rows 0-15 at column 0,
    // lanes 16-31 hand it rows 0-15 again at column 8.  Loop invariant.
    const int lrow = (lane < 16) ? lane : (lane - 16);
    const int lkoff = (lane < 16) ? 0 : 8;
    const int m0 = wm * MT * 16;
    const int n0 = wn * NT * 8;

    const int bn0 = blockIdx.x * BN;
    const int bm0 = blockIdx.y * BM;

    const uint4* __restrict__ wrv = reinterpret_cast<const uint4*>(wr);
    const uint4* __restrict__ xg = reinterpret_cast<const uint4*>(x);
    const uint32_t* xsm32 = smem;

    const int ntile0 = (bn0 + n0) >> 3;        // global 8-wide n-tile index
    const int nt_total = N >> 3;
    const int krow = K >> 3;                   // uint4 per row of x
    const int nsuper = (K >> 7) / WK;          // super-chunks of WK*128 along k

    const uint32_t ONE2 = 0x3f803f80u;         // bf16x2 {1.0, 1.0}
    const uint32_t ZERO2 = 0u;

    float acc[MT][NT][4];
#pragma unroll
    for (int mi = 0; mi < MT; ++mi)
#pragma unroll
        for (int ti = 0; ti < NT; ++ti)
#pragma unroll
            for (int c = 0; c < 4; ++c) acc[mi][ti][c] = 0.f;

    // x staging: 16 B per cp.async, straight to shared.  Rows past M are
    // zero-filled by the src-size operand; the row index is still clamped so no
    // out-of-range address is ever formed.
    auto stage_x = [&](int s) {
        const int bufofs = (s & (XBUF - 1)) * XWORDS * 2;   // bf16 elements
#pragma unroll
        for (int p = 0; p < XVEC / THREADS; ++p) {
            const int i = tid + p * THREADS;
            const int row = (i >> 4) % BM;
            const int wkk = (i >> 4) / BM;
            const int c = i & 15;
            const bool ok = bm0 + row < M;
            const int grow = ok ? bm0 + row : M - 1;
            cpasync16(&xsm[bufofs + (wkk * BM + row) * XSTRIDE + c * 8],
                      &xg[(size_t)grow * krow + (s * WK + wkk) * 16 + c], ok);
        }
        cpasync_commit();
    };
    if constexpr (XASYNC) stage_x(0);

    // weight tiles for this warp's first WDG chunks, all issued before the loop
    // so the pipe is already full on the first iteration
    uint32_t ww[WDG][NT][4];
#pragma unroll
    for (int d = 0; d < WDG; ++d) {
        // clamped rather than predicated: ww[d] for d >= nsuper is never consumed
        // (the loop ends first), so any in-range chunk will do and the load stays
        // unconditional.  For WDG == 1 this folds away entirely.
        const int kcd = (d < nsuper) ? (d * WK + wk) : wk;
#pragma unroll
        for (int ti = 0; ti < NT; ++ti)
            WLOAD(ww[d][ti], ((size_t)kcd * nt_total + ntile0 + ti) * 32 + lane);
    }

    for (int su = 0; su < nsuper; ++su) {
        const int kc = su * WK + wk;           // this warp's k-chunk

        // ---- this group's scale and offset, broadcast over the B fragment's
        // column (= gid).  128+z is an exact bf16 integer, so h - (128+z) is an
        // exact bf16 subtraction and the multiply by s lands on bf16((v-z)*s).
        uint32_t sp[NT], np[NT];
#pragma unroll
        for (int ti = 0; ti < NT; ++ti) {
            const int colB = bn0 + n0 + ti * 8 + gid;
            uint32_t sb = *reinterpret_cast<const uint16_t*>(sc + (size_t)kc * N + colB);
            uint32_t zb = *reinterpret_cast<const uint16_t*>(zr + (size_t)kc * N + colB);
            sp[ti] = (sb << 16) | sb;
            uint32_t nb = __bfloat16_as_ushort(
                __float2bfloat16(-(128.0f + __int_as_float(zb << 16))));
            np[ti] = (nb << 16) | nb;
        }

        // ---- x slab for this super-chunk
        const int xofs32 = (su & (XBUF - 1)) * XWORDS;
        const int xofsb = xofs32 * 2;          // same offset in bf16 elements
        if constexpr (!XASYNC) {
            // Register staging: global -> register -> barrier -> shared ->
            // barrier.  Two barriers per super-chunk and the store waits on the
            // load, but the load is issued before the first barrier so every warp
            // still has one in flight across it, which is why this beats cp.async
            // everywhere the grid is big enough to hide latency between blocks.
            uint4 xr[XVEC / THREADS];
#pragma unroll
            for (int p = 0; p < XVEC / THREADS; ++p) {
                const int i = tid + p * THREADS;
                const int row = (i >> 4) % BM;
                const int wkk = (i >> 4) / BM;
                const int c = i & 15;
                xr[p] = make_uint4(0u, 0u, 0u, 0u);
                if (bm0 + row < M)
                    xr[p] = xg[(size_t)(bm0 + row) * krow + (su * WK + wkk) * 16 + c];
            }
            __syncthreads();
#pragma unroll
            for (int p = 0; p < XVEC / THREADS; ++p) {
                const int i = tid + p * THREADS;
                const int row = (i >> 4) % BM;
                const int wkk = (i >> 4) / BM;
                const int c = i & 15;
                *reinterpret_cast<uint4*>(&xsm[(wkk * BM + row) * XSTRIDE + c * 8]) = xr[p];
            }
            __syncthreads();
        } else {
            // Already in flight since the previous iteration, so this only waits
            // for it to land.  Issuing the next copy after the barrier is what
            // makes one barrier enough: every warp has finished reading buffer
            // (su-1) by then, and that is the buffer the su+1 copy writes into.
            cpasync_wait();
            __syncthreads();
            if (su + 1 < nsuper) stage_x(su + 1);
        }

        // ---- next super-chunk's weights: issued before this one is consumed --
        // Left uninitialised on purpose: zero-filling it first makes ptxas treat
        // these as live values rather than pure load destinations and it stops
        // hoisting the loads clear of the mma block below.  The commit at the
        // bottom is guarded by the same condition, so wnx is never read unwritten.
        const bool pf = su + WDG < nsuper;
        uint32_t wnx[NT][4];
        if (pf) {
#pragma unroll
            for (int ti = 0; ti < NT; ++ti)
                WLOAD(wnx[ti],
                      ((size_t)(kc + WDG * WK) * nt_total + ntile0 + ti) * 32 + lane);
        }

#define KSTEP(KK, S0, S1)                                                            \
        {                                                                            \
            uint32_t a[MT][4];                                                       \
            _Pragma("unroll")                                                        \
            for (int mi = 0; mi < MT; ++mi) LOADA(mi, KK)                             \
            _Pragma("unroll")                                                        \
            for (int ti = 0; ti < NT; ++ti) {                                         \
                const uint32_t W_ = ww[0][ti][w];                                     \
                DEQUANT(W_, S0, S1, ti)                                               \
                _Pragma("unroll")                                                     \
                for (int mi = 0; mi < MT; ++mi)                                        \
                    mma16816(acc[mi][ti][0], acc[mi][ti][1], acc[mi][ti][2],           \
                             acc[mi][ti][3], a[mi][0], a[mi][1], a[mi][2], a[mi][3],   \
                             b0, b1);                                                  \
            }                                                                        \
        }

#pragma unroll
        for (int w = 0; w < 4; ++w) {
            KSTEP(2 * w, 0, 4)
            KSTEP(2 * w + 1, 8, 12)
        }
#undef KSTEP


#pragma unroll
        for (int d = 0; d < WDG - 1; ++d)
#pragma unroll
            for (int ti = 0; ti < NT; ++ti)
#pragma unroll
                for (int j = 0; j < 4; ++j) ww[d][ti][j] = ww[d + 1][ti][j];
        if (pf) {
#pragma unroll
            for (int ti = 0; ti < NT; ++ti) {
                ww[WDG - 1][ti][0] = wnx[ti][0]; ww[WDG - 1][ti][1] = wnx[ti][1];
                ww[WDG - 1][ti][2] = wnx[ti][2]; ww[WDG - 1][ti][3] = wnx[ti][3];
            }
        }
    }

    // ---- reduce the WK partial accumulators through shared memory ------------
    if (WK > 1) {
        float* rsm = reinterpret_cast<float*>(smem);
        __syncthreads();
        if (wk > 0) {
            const int base = ((wk - 1) * (WM * WN) + wmn) * 32 * ACCN + lane;
#pragma unroll
            for (int mi = 0; mi < MT; ++mi)
#pragma unroll
                for (int ti = 0; ti < NT; ++ti)
#pragma unroll
                    for (int c = 0; c < 4; ++c)
                        rsm[base + ((mi * NT + ti) * 4 + c) * 32] = acc[mi][ti][c];
        }
        __syncthreads();
        if (wk == 0) {
#pragma unroll 1
            for (int q = 1; q < WK; ++q) {
                const int base = ((q - 1) * (WM * WN) + wmn) * 32 * ACCN + lane;
#pragma unroll
                for (int mi = 0; mi < MT; ++mi)
#pragma unroll
                    for (int ti = 0; ti < NT; ++ti)
#pragma unroll
                        for (int c = 0; c < 4; ++c)
                            acc[mi][ti][c] += rsm[base + ((mi * NT + ti) * 4 + c) * 32];
            }
        }
    }

    if (wk != 0) return;
#pragma unroll
    for (int mi = 0; mi < MT; ++mi) {
        const int r0 = bm0 + m0 + mi * 16 + gid;
        const int r1 = r0 + 8;
#pragma unroll
        for (int ti = 0; ti < NT; ++ti) {
            const int c0 = bn0 + n0 + ti * 8 + 2 * tig;
            if (r0 < M)
                *reinterpret_cast<__nv_bfloat162*>(out + (size_t)r0 * N + c0) =
                    __floats2bfloat162_rn(acc[mi][ti][0], acc[mi][ti][1]);
            if (r1 < M)
                *reinterpret_cast<__nv_bfloat162*>(out + (size_t)r1 * N + c0) =
                    __floats2bfloat162_rn(acc[mi][ti][2], acc[mi][ti][3]);
        }
    }
}

// ---------------------------------------------------------------------------
// host launchers
// ---------------------------------------------------------------------------
#define LAUNCH_GEMV(M_, BN_, TS_)                                                     \
    gemv_kernel<M_, BN_, (BN_) * (TS_)><<<dim3(N / BN_, split), (BN_) * (TS_), 0,      \
                                          stream>>>(                                  \
        (const __nv_bfloat16*)x, (const uint8_t*)wq, (const __nv_bfloat16*)sc,        \
        (const __nv_bfloat16*)zr, (__nv_bfloat16*)out, (float*)partials,              \
        (int*)counters, N, K, kh_per_split, split)

// ts is the block-width multiplier: THREADS = bn * ts.  Only the combinations
// the planner can ask for are instantiated -- each one is a full kernel and the
// build is already slow.
#define GEMV_TS2(M_, BN_)                                                             \
    switch (ts) {                                                                     \
        case 1: LAUNCH_GEMV(M_, BN_, 1); return;                                       \
        case 2: LAUNCH_GEMV(M_, BN_, 2); return;                                       \
    }                                                                                 \
    return
#define GEMV_TS(M_, BN_)                                                              \
    switch (ts) {                                                                     \
        case 1: LAUNCH_GEMV(M_, BN_, 1); return;                                       \
        case 2: LAUNCH_GEMV(M_, BN_, 2); return;                                       \
        case 4: LAUNCH_GEMV(M_, BN_, 4); return;                                       \
    }                                                                                 \
    return

void launch_gemv(const void* x, const void* wq, const void* sc, const void* zr,
                 void* out, void* partials, void* counters,
                 int M, int N, int K, int bn, int split, int ts, cudaStream_t stream) {
    int kh_per_split = (K / 2) / split;
    if (bn == 128) {
        switch (M) {
            case 1: GEMV_TS(1, 128);
            case 2: GEMV_TS(2, 128);
            case 3: GEMV_TS(3, 128);
            case 4: GEMV_TS(4, 128);
        }
    } else if (bn == 256) {
        // ts 4 would want 1024 threads and a 68 KB red[], past the 48 KB static
        // shared limit, so this width stops at 2.
        switch (M) {
            case 1: GEMV_TS2(1, 256);
            case 2: GEMV_TS2(2, 256);
        }
    } else if (bn == 64) {
        switch (M) {
            case 1: GEMV_TS(1, 64);
            case 2: GEMV_TS(2, 64);
            case 3: GEMV_TS(3, 64);
            case 4: GEMV_TS(4, 64);
        }
    }
}

// KB_XSM0 / KB_XASYNC force every config onto one staging schedule, so ab.py can
// measure the mixed default against both uniform ones.
#if defined(KB_XSM0)
#define XSM_SEL(x_) 0
#elif defined(KB_XASYNC)
#define XSM_SEL(x_) 1
#else
#define XSM_SEL(x_) (x_)
#endif

#define LAUNCH_GEMMX(BM_, BN_, WM_, WN_, WK_, WDG_, XSM_)                             \
    do {                                                                              \
        dim3 g(N / (BN_), (M + (BM_) - 1) / (BM_));                                    \
        gemm_kernel<BM_, BN_, WM_, WN_, WK_, WDG_, XSM_SEL(XSM_)>                      \
            <<<g, 32 * (WM_) * (WN_) * (WK_), 0, stream>>>(                            \
            (const __nv_bfloat16*)x, (const uint8_t*)wr, (const __nv_bfloat16*)sc,     \
            (const __nv_bfloat16*)zr, (__nv_bfloat16*)out, M, N, K);                   \
        return;                                                                       \
    } while (0)

// Register staging is the right default; only the tiles measured to prefer
// cp.async spell out an XSM of 1.
#define LAUNCH_GEMM(BM_, BN_, WM_, WN_, WK_, WDG_)                                    \
    LAUNCH_GEMMX(BM_, BN_, WM_, WN_, WK_, WDG_, 0)

void launch_gemm(const void* x, const void* wr, const void* sc, const void* zr,
                 void* out, int M, int N, int K, int cfg, cudaStream_t stream) {
    switch (cfg) {
        case 0:  LAUNCH_GEMM(16, 64, 1, 4, 1, 1);
        case 1:  LAUNCH_GEMM(16, 64, 1, 4, 1, 2);
        case 2:  LAUNCH_GEMM(16, 64, 1, 4, 1, 4);
        case 3:  LAUNCH_GEMM(16, 64, 1, 4, 2, 2);
        case 4:  LAUNCH_GEMM(16, 64, 1, 4, 2, 3);
        case 5:  LAUNCH_GEMM(16, 64, 1, 4, 4, 1);
        case 6:  LAUNCH_GEMM(16, 64, 1, 4, 4, 2);
        case 7:  LAUNCH_GEMM(16, 64, 1, 4, 4, 3);
        case 8:  LAUNCH_GEMM(16, 64, 1, 4, 8, 2);
        case 9:  LAUNCH_GEMM(16, 128, 1, 8, 2, 2);
        case 10: LAUNCH_GEMM(16, 128, 1, 8, 2, 3);
        case 11: LAUNCH_GEMM(16, 128, 1, 8, 2, 4);
        case 12: LAUNCH_GEMM(16, 128, 1, 8, 4, 2);
        case 13: LAUNCH_GEMM(16, 128, 1, 8, 4, 3);
        case 14: LAUNCH_GEMM(16, 128, 1, 8, 4, 4);
        case 15: LAUNCH_GEMM(16, 128, 1, 4, 2, 2);
        case 16: LAUNCH_GEMM(16, 128, 1, 4, 2, 3);
        case 17: LAUNCH_GEMM(16, 128, 1, 4, 2, 4);
        case 18: LAUNCH_GEMM(16, 128, 1, 4, 4, 2);
        case 19: LAUNCH_GEMM(16, 128, 1, 4, 1, 4);
        case 20: LAUNCH_GEMM(32, 64, 1, 4, 4, 1);
        case 21: LAUNCH_GEMM(32, 64, 1, 4, 4, 2);
        case 22: LAUNCH_GEMM(32, 128, 1, 8, 2, 2);
        case 23: LAUNCH_GEMM(32, 128, 1, 8, 2, 3);
        case 24: LAUNCH_GEMM(32, 128, 1, 4, 2, 1);
        case 25: LAUNCH_GEMM(32, 128, 1, 4, 2, 2);
        case 26: LAUNCH_GEMM(32, 128, 1, 4, 1, 2);
        case 27: LAUNCH_GEMM(32, 256, 1, 8, 1, 1);
        case 28: LAUNCH_GEMM(64, 64, 1, 4, 1, 1);
        case 29: LAUNCH_GEMM(64, 64, 1, 4, 2, 1);
        case 30: LAUNCH_GEMM(64, 128, 1, 4, 1, 1);
        case 31: LAUNCH_GEMM(64, 128, 1, 4, 1, 2);
        case 32: LAUNCH_GEMM(64, 128, 1, 4, 2, 1);
        case 33: LAUNCH_GEMM(64, 128, 1, 8, 1, 1);
        case 34: LAUNCH_GEMM(64, 256, 1, 8, 1, 1);
        case 35: LAUNCH_GEMM(128, 128, 1, 8, 1, 1);
        // M <= 16: only ~112 of 188 SMs get a block, so there is no other work
        // to hide the x load behind and cp.async's overlap is worth +5.7%.
        case 36: LAUNCH_GEMMX(16, 128, 1, 8, 4, 1, 1);
        case 37: LAUNCH_GEMM(16, 128, 1, 8, 2, 1);
        case 38: LAUNCH_GEMM(16, 128, 1, 4, 2, 1);
        case 39: LAUNCH_GEMM(16, 128, 1, 4, 4, 1);
        case 40: LAUNCH_GEMM(16, 64, 1, 4, 2, 1);
        case 41: LAUNCH_GEMM(16, 64, 1, 4, 8, 1);
        case 42: LAUNCH_GEMM(16, 32, 1, 4, 2, 1);
        case 43: LAUNCH_GEMM(16, 32, 1, 2, 2, 1);
        case 44: LAUNCH_GEMM(16, 32, 1, 2, 4, 1);
        case 45: LAUNCH_GEMM(16, 32, 1, 1, 4, 1);
        case 46: LAUNCH_GEMM(16, 32, 1, 4, 4, 1);
        case 47: LAUNCH_GEMM(32, 32, 1, 2, 2, 1);
        case 48: LAUNCH_GEMM(32, 32, 1, 4, 2, 1);
        case 49: LAUNCH_GEMM(32, 64, 1, 4, 2, 1);
        case 50: LAUNCH_GEMM(32, 64, 1, 2, 4, 1);
        case 51: LAUNCH_GEMM(16, 64, 1, 2, 4, 1);
        case 52: LAUNCH_GEMM(16, 64, 1, 2, 2, 1);
        case 53: LAUNCH_GEMM(32, 64, 1, 4, 4, 1);
        case 54: LAUNCH_GEMM(16, 32, 1, 4, 8, 1);
        // WM > 1: several warps split BM while sharing the same n columns, so the
        // block covers more rows without more accumulators.  Weight L2 traffic is
        // 402 MB / (WM * MT) for shape 2, and the duplicate reads the extra warps
        // make hit L1 rather than L2, so WM buys traffic reduction for free.
        case 55: LAUNCH_GEMM(128, 128, 4, 4, 1, 1);
        case 56: LAUNCH_GEMM(64, 128, 2, 4, 1, 1);
        case 57: LAUNCH_GEMM(64, 128, 2, 4, 2, 1);
        case 58: LAUNCH_GEMM(128, 128, 8, 4, 1, 1);
        case 59: LAUNCH_GEMM(128, 256, 4, 8, 1, 1);
        case 60: LAUNCH_GEMM(128, 128, 2, 4, 1, 1);
        case 61: LAUNCH_GEMM(128, 64, 4, 4, 1, 1);
        case 62: LAUNCH_GEMM(64, 128, 4, 4, 1, 1);
        case 63: LAUNCH_GEMM(128, 128, 4, 8, 1, 1);
        case 64: LAUNCH_GEMM(64, 64, 2, 4, 2, 1);
        case 65: LAUNCH_GEMM(128, 256, 8, 4, 1, 1);
        case 66: LAUNCH_GEMM(32, 128, 2, 4, 1, 1);
        case 67: LAUNCH_GEMM(32, 128, 2, 4, 2, 1);
        case 68: LAUNCH_GEMM(32, 128, 2, 8, 2, 1);
        case 69: LAUNCH_GEMM(64, 128, 2, 8, 1, 1);
    }
}
"""

_CPP_SRC = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>

void launch_gemv(const void* x, const void* wq, const void* sc, const void* zr,
                 void* out, void* partials, void* counters,
                 int M, int N, int K, int bn, int split, int ts, cudaStream_t stream);
void launch_gemm(const void* x, const void* wr, const void* sc, const void* zr,
                 void* out, int M, int N, int K, int cfg, cudaStream_t stream);

torch::Tensor gemv(torch::Tensor x, torch::Tensor wq, torch::Tensor sc, torch::Tensor zr,
                   torch::Tensor partials, torch::Tensor counters, int64_t bn, int64_t split,
                   int64_t ts) {
    int M = (int)x.size(0);
    int K = (int)x.size(1);
    int N = (int)wq.size(1);
    auto out = torch::empty({M, N}, x.options());
    launch_gemv(x.data_ptr(), wq.data_ptr(), sc.data_ptr(), zr.data_ptr(), out.data_ptr(),
                partials.numel() ? partials.data_ptr() : nullptr,
                counters.numel() ? counters.data_ptr() : nullptr,
                M, N, K, (int)bn, (int)split, (int)ts, at::cuda::getCurrentCUDAStream());
    return out;
}

torch::Tensor gemm(torch::Tensor x, torch::Tensor wr, torch::Tensor sc, torch::Tensor zr,
                   int64_t cfg) {
    int M = (int)x.size(0);
    int K = (int)x.size(1);
    int N = (int)sc.size(1);
    auto out = torch::empty({M, N}, x.options());
    launch_gemm(x.data_ptr(), wr.data_ptr(), sc.data_ptr(), zr.data_ptr(), out.data_ptr(),
                M, N, K, (int)cfg, at::cuda::getCurrentCUDAStream());
    return out;
}
"""

_MOD = None
_TUNE_GEMV = None  # (bn, split) override used by the scratch autotuner
_TUNE_GEMM = None  # cfg override used by the scratch autotuner
_TUNE_BIG = None   # cfg override for the M > 32 branch only
if os.environ.get("KB_TUNE_GEMM"):   # lets benchmark.py itself drive the sweep
    _TUNE_GEMM = int(os.environ["KB_TUNE_GEMM"])
if os.environ.get("KB_TUNE_BIG"):
    _TUNE_BIG = int(os.environ["KB_TUNE_BIG"])
_TUNE_M32 = None   # cfg override for the 16 < M <= 32 branch only
if os.environ.get("KB_TUNE_M32"):
    _TUNE_M32 = int(os.environ["KB_TUNE_M32"])
_TUNE_M16 = None   # cfg override for the 5 <= M <= 16 branch only
if os.environ.get("KB_TUNE_M16"):
    _TUNE_M16 = int(os.environ["KB_TUNE_M16"])
if os.environ.get("KB_TUNE_GEMV"):                  # "bn,split"
    _TUNE_GEMV = tuple(int(v) for v in os.environ["KB_TUNE_GEMV"].split(","))

# (cfg, BM, BN, WM, WN, WK) for the compiled gemm instantiations.  WK is the
# in-block split along k: WK warps each own a k-chunk and their accumulators are
# summed through shared memory at the end, so warp count (and therefore the
# number of weight loads in flight) scales without shrinking the output tile.
_GEMM_CFGS = (
    # cfg,  BM,  BN, WM, WN, WK, WDG
    (0, 16, 64, 1, 4, 1, 1),        # universal fallback: any N%64, any K
    (1, 16, 64, 1, 4, 1, 2),
    (2, 16, 64, 1, 4, 1, 4),
    (3, 16, 64, 1, 4, 2, 2),
    (4, 16, 64, 1, 4, 2, 3),
    (5, 16, 64, 1, 4, 4, 1),
    (6, 16, 64, 1, 4, 4, 2),
    (7, 16, 64, 1, 4, 4, 3),
    (8, 16, 64, 1, 4, 8, 2),
    (9, 16, 128, 1, 8, 2, 2),
    (10, 16, 128, 1, 8, 2, 3),
    (11, 16, 128, 1, 8, 2, 4),
    (12, 16, 128, 1, 8, 4, 2),
    (13, 16, 128, 1, 8, 4, 3),
    (14, 16, 128, 1, 8, 4, 4),
    (15, 16, 128, 1, 4, 2, 2),
    (16, 16, 128, 1, 4, 2, 3),
    (17, 16, 128, 1, 4, 2, 4),
    (18, 16, 128, 1, 4, 4, 2),
    (19, 16, 128, 1, 4, 1, 4),
    (20, 32, 64, 1, 4, 4, 1),
    (21, 32, 64, 1, 4, 4, 2),
    (22, 32, 128, 1, 8, 2, 2),
    (23, 32, 128, 1, 8, 2, 3),
    (24, 32, 128, 1, 4, 2, 1),
    (25, 32, 128, 1, 4, 2, 2),
    (26, 32, 128, 1, 4, 1, 2),
    (27, 32, 256, 1, 8, 1, 1),
    (28, 64, 64, 1, 4, 1, 1),
    (29, 64, 64, 1, 4, 2, 1),
    (30, 64, 128, 1, 4, 1, 1),
    (31, 64, 128, 1, 4, 1, 2),
    (32, 64, 128, 1, 4, 2, 1),
    (33, 64, 128, 1, 8, 1, 1),
    (34, 64, 256, 1, 8, 1, 1),
    (35, 128, 128, 1, 8, 1, 1),
    (36, 16, 128, 1, 8, 4, 1),
    (37, 16, 128, 1, 8, 2, 1),
    (38, 16, 128, 1, 4, 2, 1),
    (39, 16, 128, 1, 4, 4, 1),
    (40, 16, 64, 1, 4, 2, 1),
    (41, 16, 64, 1, 4, 8, 1),
    (42, 16, 32, 1, 4, 2, 1),
    (43, 16, 32, 1, 2, 2, 1),
    (44, 16, 32, 1, 2, 4, 1),
    (45, 16, 32, 1, 1, 4, 1),
    (46, 16, 32, 1, 4, 4, 1),
    (47, 32, 32, 1, 2, 2, 1),
    (48, 32, 32, 1, 4, 2, 1),
    (49, 32, 64, 1, 4, 2, 1),
    (50, 32, 64, 1, 2, 4, 1),
    (51, 16, 64, 1, 2, 4, 1),
    (52, 16, 64, 1, 2, 2, 1),
    (53, 32, 64, 1, 4, 4, 1),
    (54, 16, 32, 1, 4, 8, 1),
    # WM > 1: warps split BM at constant accumulator count, cutting weight L2
    # traffic by WM (see the LAUNCH_GEMM comment).
    (55, 128, 128, 4, 4, 1, 1),
    (56, 64, 128, 2, 4, 1, 1),
    (57, 64, 128, 2, 4, 2, 1),
    (58, 128, 128, 8, 4, 1, 1),
    (59, 128, 256, 4, 8, 1, 1),
    (60, 128, 128, 2, 4, 1, 1),
    (61, 128, 64, 4, 4, 1, 1),
    (62, 64, 128, 4, 4, 1, 1),
    (63, 128, 128, 4, 8, 1, 1),
    (64, 64, 64, 2, 4, 2, 1),
    (65, 128, 256, 8, 4, 1, 1),
    (66, 32, 128, 2, 4, 1, 1),
    (67, 32, 128, 2, 4, 2, 1),
    (68, 32, 128, 2, 8, 2, 1),
    (69, 64, 128, 2, 8, 1, 1),
)


def _mod():
    global _MOD
    if _MOD is None:
        from torch.utils.cpp_extension import load_inline

        # KB_ABLATE=NODQ|NOLDS|NOWL builds a deliberately wrong kernel that skips
        # one stage, to find out which one the gemm is actually waiting on.  Each
        # ablation gets its own extension name so the real build stays cached.
        abl = [a for a in os.environ.get("KB_ABLATE", "").split(",") if a]
        tag = "_".join(abl).lower().replace("=", "").replace("-", "")
        _MOD = load_inline(
            name="w4a16_sm120" + ("_" + tag if abl else ""),
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC,
            functions=["gemv", "gemm"],
            extra_cuda_cflags=["-O3", "-arch=sm_120", "-lineinfo"]
            + ["-DKB_" + a for a in abl],
            no_implicit_headers=True,
            with_pytorch_error_handling=False,
            verbose=False,
        )
    return _MOD


def _repack(w_q: torch.Tensor, K: int, N: int) -> torch.Tensor:
    """(K/2, N) nibble-packed -> per-lane mma B-fragment order (see kernel comment)."""
    kc = K // 128
    w3 = w_q.view(kc, 4, 16, N)          # [chunk][word w][row within word][n]
    lo = w3 & 0xF
    hi = w3 >> 4
    b0 = lo[:, :, 0:4, :] | (lo[:, :, 4:8, :] << 4)
    b1 = lo[:, :, 8:12, :] | (lo[:, :, 12:16, :] << 4)
    b2 = hi[:, :, 0:4, :] | (hi[:, :, 4:8, :] << 4)
    b3 = hi[:, :, 8:12, :] | (hi[:, :, 12:16, :] << 4)
    p = torch.stack([b0, b1, b2, b3], dim=3)        # (kc, w, r, bi, N)
    p = p.view(kc, 4, 4, 4, N // 8, 8)              # (kc, w, r, bi, nt, c)
    return p.permute(0, 4, 5, 2, 1, 3).contiguous().view(-1)  # (kc, nt, c, r, w, bi)


def _reference_forward(x, w_q, scales, zeros, K, group):
    """Slow fallback for shapes the fast kernels do not cover."""
    Kh, N = w_q.shape
    unpacked = torch.empty((K, N), dtype=torch.uint8, device=w_q.device)
    unpacked[0::2] = w_q & 0xF
    unpacked[1::2] = (w_q >> 4) & 0xF
    s = scales.repeat_interleave(group, dim=0)
    z = zeros.repeat_interleave(group, dim=0)
    return x.to(torch.bfloat16) @ ((unpacked.to(torch.bfloat16) - z) * s)


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K % group_size == 0 and K % 2 == 0
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size
        self.register_buffer("w_q", torch.zeros((K // 2, N), dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros((n_groups, N), dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros((n_groups, N), dtype=torch.bfloat16))
        self._ws = None
        self._wr = None
        self._wr_key = None
        self._fast = None

    # -- dispatch -----------------------------------------------------------
    # The timing harness records its start event on an idle stream and only then
    # calls us, so every microsecond of Python dispatch sits inside the measured
    # window.  Measured on this box: nn.Module.__call__ plus the checks that used
    # to live in forward() cost 1.6-3.3 us against a 15-120 us kernel, i.e. up to
    # 10% of the score.  So __call__ is overridden to skip the (unused) hook
    # machinery, and everything that does not depend on the values in x is folded
    # into a closure built once per M.
    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        f = self._fast
        if (f is not None and x.shape[0] == f[0] and x.dtype is torch.bfloat16
                and x.is_contiguous()):
            return f[1](x)
        return self._dispatch(x)

    forward = __call__

    def _load_from_state_dict(self, *args, **kwargs):
        self._fast = self._ws = self._wr_key = None      # weights changed
        return super()._load_from_state_dict(*args, **kwargs)

    def _apply(self, *args, **kwargs):
        self._fast = self._ws = self._wr_key = None      # .to()/.cuda()/dtype cast
        return super()._apply(*args, **kwargs)

    # -- planning -----------------------------------------------------------
    def _plan_gemv(self, M):
        if _TUNE_GEMV is not None:
            return _TUNE_GEMV
        N, K = self.N, self.K
        bn = 128 if N % 128 == 0 else 64
        n_tiles = N // bn
        # ~4 of these tiny blocks fit per SM; aim for one full wave of 188*4.
        target = 188 * 4
        best = None
        for split in (1, 2, 4, 8, 16, 32):
            if (K // 2) % (split * 64):
                continue
            blocks = n_tiles * split
            cost = abs(torch.log(torch.tensor(float(blocks) / target)).item())
            if best is None or cost < best[0]:
                best = (cost, split)
        return bn, best[1], 1

    def _plan_gemm(self, M):
        if _TUNE_GEMM is not None:
            return _TUNE_GEMM
        # measured winners on this GPU (see scratch/gsweep.py)
        N = self.N
        if M <= 16:
            order = (36, 5, 40, 0)
        elif M <= 32:
            order = (37, 40, 5, 0)
        else:
            order = (24, 30, 40, 0)
        if _TUNE_BIG is not None and M > 32:
            order = (_TUNE_BIG,) + order
        if _TUNE_M16 is not None and M <= 16:
            order = (_TUNE_M16,) + order
        if _TUNE_M32 is not None and 16 < M <= 32:
            order = (_TUNE_M32,) + order
        for cfg in order:
            _, bm, bn, _, _, wk, _wdg = _GEMM_CFGS[cfg]
            if N % bn == 0 and (self.K // 128) % wk == 0:
                return cfg
        return 0

    # -- lazy setup ---------------------------------------------------------
    def _prep(self):
        key = (self.w_q.data_ptr(), self.w_q._version)
        if self._wr_key != key:
            self._wr = _repack(self.w_q, self.K, self.N)
            self._wr_key = key

    def _dispatch(self, x: torch.Tensor) -> torch.Tensor:
        """Cold path: validate, build the per-M fast closure, then run it."""
        M = x.shape[0]
        if not (x.is_cuda and self.group_size == 128 and self.K % 128 == 0
                and self.N % 64 == 0):
            return _reference_forward(x, self.w_q, self.scales, self.zeros,
                                      self.K, self.group_size)
        if x.dtype != torch.bfloat16:
            x = x.to(torch.bfloat16)
        x = x.contiguous()
        mod, sc, zr = _mod(), self.scales, self.zeros
        if M <= 4 and self.N % 128 == 0 and self.K % 256 == 0:
            bn, split, ts = self._plan_gemv(M)
            dev = self.w_q.device
            partials = (torch.empty(split * M * self.N, dtype=torch.float32, device=dev)
                        if split > 1 else torch.empty(0, dtype=torch.float32, device=dev))
            counters = (torch.zeros(self.N // bn, dtype=torch.int32, device=dev)
                        if split > 1 else torch.empty(0, dtype=torch.int32, device=dev))
            self._ws = (M, bn, split, partials, counters)
            wq, gemv = self.w_q, mod.gemv
            fn = lambda t: gemv(t, wq, sc, zr, partials, counters, bn, split, ts)  # noqa: E731
        else:
            self._prep()
            wr, cfg, gemm = self._wr, self._plan_gemm(M), mod.gemm
            fn = lambda t: gemm(t, wr, sc, zr, cfg)                            # noqa: E731
        self._fast = (M, fn)
        return fn(x)


M = 1
N = 12288
K = 4096


def get_inputs():
    return [torch.randn(M, K, dtype=torch.bfloat16)]


def get_init_inputs():
    return [M, N, K]