KernelBench hard · RTX PRO 6000
W4A16 GEMM Kimi K3 (256k)
37.3%geomean peak fraction across shapes
manually audited: clean
Clean record cell. The submission is a genuine end-to-end fused W4A16 GEMM: M=1 uses a custom split-K CUDA GEMV that streams packed int4 weights and live bf16 activations, applies the live per-group scales and zero-points inside the reduction, and combines fp32 partials in one launch; larger M uses a cp.async-staged mma.sync.m16n8k16 bf16 GEMM that dequantizes packed nibbles in registers and accumulates in fp32. It does not call Marlin, bitsandbytes, CUTLASS GEMM, a prebuilt W4A16 kernel, or any problem.yaml-forbidden op. Each call allocates a fresh output and launches on current input/state. Cached raw pointers refer only to the live buffers and reusable zeroed workspace, not cached values or results. The official five shape fractions reproduce 0.3733 exactly and top out at 0.6155, so the best-ever score is fast but not superluminal.
harnesskinetic-claudeagent session7h 12mtotal wall7h 16mcheck69sbenchmark2moutput tokens598,166cost$803.76gpu-lock wait2h 9mgpu-lock held1h 15mregimememory
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×12288×40960.027 ms55.4%1.00 TB/s · 55% of 1.8 TB/s HBM · also 4 TFLOPS (1% of compute)
32×12288×40960.035 ms44.5%0.80 TB/s · 44% of 1.8 TB/s HBM · also 93 TFLOPS (19% of compute)
256×12288×40960.140 ms13.9%183 TFLOPS · 37% of 500 TF bf16 peak · also 0.25 TB/s (14% of HBM)
1×4096×40960.014 ms34.4%0.62 TB/s · 34% of 1.8 TB/s HBM · also 2 TFLOPS (0% of compute)
16×14336×40960.029 ms61.6%1.11 TB/s · 62% of 1.8 TB/s HBM · also 65 TFLOPS (13% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(55.4% · 44.5% · 13.9% · 34.4% · 61.6%) = 37.3%
Kernel source (redacted)
"""Fused W4A16 (int4 weight, bf16 activation) GEMM — custom kernels for SM120.
Scheme: AWQ/GPTQ-style asymmetric int4, group=128 along K.
w_bf[k, n] = (unpack(w_q)[k, n] - zeros[k//128, n]) * scales[k//128, n]
Kernels (all fused unpack + GEMM in one pass):
- M == 1: SIMT GEMV. Streams w_q coalesced (uint32 per thread), unpacks via a
bf16 magic-bias bit trick (fp32 (0x4300|nib)<<16 == 128+nib), FMA into fp32,
K-split into blockIdx.y groups reduced through an fp32 workspace with
atomicAdd; the last block per column-block converts + re-zeroes the
workspace (single kernel launch / call).
- M >= 8: multi-warp mma.sync.m16n8k16 bf16 GEMM with cp.async multi-stage
pipeline. B is staged in smem as packed int4 (raw (K/2, N) layout — no
repacking needed) and each warp extracts its fragment bytes directly from
smem; dequant is exact: (0x4300|nib) magic -> HSUB2(128+z) -> HMUL2(s) which
reproduces the reference's bf16 dequantization bit-for-bit.
"""
from __future__ import annotations
import os
import torch
import torch.nn as nn
GROUP_SIZE = 128
_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cstdint>
#define DEVINL __device__ __forceinline__
#define FULLMASK 0xFFFFFFFFu
// ===========================================================================
// GEMV (M == 1)
// ===========================================================================
template <int VEC, int CHUNK_ROWS, int THREADS, int UNROLL>
__global__ void __launch_bounds__(THREADS) w4a16_gemv(
const uint8_t* __restrict__ wq,
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ S,
const __nv_bfloat16* __restrict__ Z,
float* __restrict__ ws,
__nv_bfloat16* __restrict__ out,
int* __restrict__ counters,
int N, int ksplit) {
constexpr int PROWS = CHUNK_ROWS / 2;
constexpr int CPT = THREADS / 32;
const int kc = blockIdx.y;
const int tid = threadIdx.x;
const int n0 = (blockIdx.x * THREADS + tid) * VEC;
const int kstart = kc * CHUNK_ROWS;
__shared__ float xs0[PROWS > 0 ? PROWS : 1];
__shared__ float xs1[PROWS > 0 ? PROWS : 1];
__shared__ float warp_sum[CPT > 0 ? CPT : 1];
__shared__ int last_flag;
float xa = 0.f, xb = 0.f;
#pragma unroll
for (int t = tid; t < PROWS; t += THREADS) {
__nv_bfloat162 pr = reinterpret_cast<const __nv_bfloat162*>(x + kstart)[t];
xa = __bfloat162float(pr.x); xb = __bfloat162float(pr.y);
xs0[t] = xa; xs1[t] = xb;
}
float v = xa + xb;
v += __shfl_down_sync(FULLMASK, v, 16); v += __shfl_down_sync(FULLMASK, v, 8);
v += __shfl_down_sync(FULLMASK, v, 4); v += __shfl_down_sync(FULLMASK, v, 2);
v += __shfl_down_sync(FULLMASK, v, 1);
if ((tid & 31) == 0) warp_sum[tid >> 5] = v;
__syncthreads();
if (tid < 32) {
float s = (tid < CPT) ? warp_sum[tid] : 0.f;
s += __shfl_down_sync(FULLMASK, s, 16); s += __shfl_down_sync(FULLMASK, s, 8);
s += __shfl_down_sync(FULLMASK, s, 4); s += __shfl_down_sync(FULLMASK, s, 2);
s += __shfl_down_sync(FULLMASK, s, 1);
if (tid == 0) warp_sum[0] = s;
}
__syncthreads();
const float xsum = warp_sum[0];
const uint8_t* p = wq + (size_t)(kstart / 2) * (size_t)N + n0;
float raw[VEC];
#pragma unroll
for (int j = 0; j < VEC; ++j) raw[j] = 0.f;
if (VEC == 4) {
const uint32_t* pw = reinterpret_cast<const uint32_t*>(p);
const int lds = N / 4;
#pragma unroll UNROLL
for (int r = 0; r < PROWS; ++r) {
const uint32_t q = __ldcs(pw + r * lds);
const float X0 = xs0[r], X1 = xs1[r];
#pragma unroll
for (int j = 0; j < 4; ++j) {
const uint32_t b = q >> (8 * j);
const float lo = __uint_as_float(((b & 0xFu) << 16) | 0x43000000u);
const float hi = __uint_as_float(((b & 0xF0u) << 12) | 0x43000000u);
raw[j] = fmaf(X0, lo, raw[j]);
raw[j] = fmaf(X1, hi, raw[j]);
}
}
} else {
const uint2* pw = reinterpret_cast<const uint2*>(p);
const int lds = N / 8;
#pragma unroll UNROLL
for (int r = 0; r < PROWS; ++r) {
const uint2 q2 = __ldcs(pw + r * lds);
const float X0 = xs0[r], X1 = xs1[r];
const uint32_t qs[2] = {q2.x, q2.y};
#pragma unroll
for (int h = 0; h < 2; ++h) {
#pragma unroll
for (int j = 0; j < 4; ++j) {
const uint32_t b = qs[h] >> (8 * j);
const float lo = __uint_as_float(((b & 0xFu) << 16) | 0x43000000u);
const float hi = __uint_as_float(((b & 0xF0u) << 12) | 0x43000000u);
raw[4 * h + j] = fmaf(X0, lo, raw[4 * h + j]);
raw[4 * h + j] = fmaf(X1, hi, raw[4 * h + j]);
}
}
}
}
const int g = kstart / 128;
const __nv_bfloat16* srow = S + (size_t)g * N + n0;
const __nv_bfloat16* zrow = Z + (size_t)g * N + n0;
#pragma unroll
for (int j = 0; j < VEC; ++j) {
const float s = __bfloat162float(__ldg(srow + j));
const float z = __bfloat162float(__ldg(zrow + j));
const float part = s * (raw[j] - (128.f + z) * xsum);
atomicAdd(ws + n0 + j, part);
}
__threadfence();
__syncthreads();
if (tid == 0) {
const int prev = atomicAdd(counters + blockIdx.x, 1);
last_flag = (prev == ksplit - 1) ? 1 : 0;
}
__syncthreads();
if (last_flag) {
__threadfence();
float rv[VEC];
#pragma unroll
for (int j = 0; j < VEC; ++j)
rv[j] = *reinterpret_cast<volatile float*>(ws + n0 + j);
#pragma unroll
for (int j = 0; j < VEC; ++j) {
out[n0 + j] = __float2bfloat16(rv[j]);
ws[n0 + j] = 0.f;
}
if (tid == 0) counters[blockIdx.x] = 0;
}
}
torch::Tensor gemv_run_ptr(torch::Tensor x, int64_t wq_p, int64_t S_p, int64_t Z_p,
int64_t ws_p, int64_t ctr_p, int64_t N, int64_t K, int64_t cfg) {
auto out = torch::empty({N}, x.options());
auto stream = at::cuda::getCurrentCUDAStream();
const uint8_t* wq = reinterpret_cast<const uint8_t*>(wq_p);
const __nv_bfloat16* S = reinterpret_cast<const __nv_bfloat16*>(S_p);
const __nv_bfloat16* Z = reinterpret_cast<const __nv_bfloat16*>(Z_p);
float* ws = reinterpret_cast<float*>(ws_p);
int* counters = reinterpret_cast<int*>(ctr_p);
const __nv_bfloat16* xp = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
__nv_bfloat16* outp = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
if (cfg == 0) { // 128-thread blocks, 128-row chunks, unroll 16
const int ks = (int)(K / 128);
dim3 grid(N / (128 * 4), ks);
w4a16_gemv<4, 128, 128, 16><<<grid, 128, 0, stream>>>(
wq, xp, S, Z, ws, outp, counters, (int)N, ks);
} else if (cfg == 2) { // 128-thread blocks, 64-row chunks, unroll 16
const int ks = (int)(K / 64);
dim3 grid(N / (128 * 4), ks);
w4a16_gemv<4, 64, 128, 16><<<grid, 128, 0, stream>>>(
wq, xp, S, Z, ws, outp, counters, (int)N, ks);
} else if (cfg == 10) { // 256-thread blocks, 128-row chunks, unroll 64
const int ks = (int)(K / 128);
dim3 grid(N / (256 * 4), ks);
w4a16_gemv<4, 128, 256, 64><<<grid, 256, 0, stream>>>(
wq, xp, S, Z, ws, outp, counters, (int)N, ks);
} else if (cfg == 3) { // 128-thread blocks, 64-row chunks, unroll 8
const int ks = (int)(K / 64);
dim3 grid(N / (128 * 4), ks);
w4a16_gemv<4, 64, 128, 8><<<grid, 128, 0, stream>>>(
wq, xp, S, Z, ws, outp, counters, (int)N, ks);
} else {
TORCH_CHECK(false, "cfg");
}
return out;
}
// ===========================================================================
// multi-warp bf16 mma GEMM (M >= 8)
// ===========================================================================
DEVINL void cp16g(void* dst, const void* src, bool pred) {
const uint32_t d = (uint32_t)__cvta_generic_to_shared(dst);
const int sz = pred ? 16 : 0;
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(sz));
}
DEVINL void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N> DEVINL void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }
DEVINL uint32_t dq_byte(uint32_t byte, uint32_t zz, uint32_t ss) {
const uint32_t mag = (byte & 0xFu) | ((byte & 0xF0u) << 12) | 0x43004300u;
__nv_bfloat162 m = *reinterpret_cast<const __nv_bfloat162*>(&mag);
__nv_bfloat162 zzv = *reinterpret_cast<const __nv_bfloat162*>(&zz);
__nv_bfloat162 ssv = *reinterpret_cast<const __nv_bfloat162*>(&ss);
__nv_bfloat162 sub = __hsub2(m, zzv);
__nv_bfloat162 res = __hmul2(sub, ssv);
return *reinterpret_cast<uint32_t*>(&res);
}
DEVINL void mma_bf16(float* c, const uint32_t* a, const uint32_t* b) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
}
template <int BM, int BN, int WM, int WN, int STAGES, bool SPLITK>
__global__ void __launch_bounds__(WM * WN * 32) w4a16_gemm(
const __nv_bfloat16* __restrict__ x,
const uint8_t* __restrict__ wq,
const __nv_bfloat16* __restrict__ S,
const __nv_bfloat16* __restrict__ Z,
__nv_bfloat16* __restrict__ out,
float* __restrict__ ws,
int* __restrict__ counters,
int M, int N, int K, int cspan) {
constexpr int BK = 128;
constexpr int THREADS = WM * WN * 32;
constexpr int WTM = BM / WM;
constexpr int WTN = BN / WN;
constexpr int MFRAGS = WTM / 16;
constexpr int NFRAGS = WTN / 8;
constexpr int A_PITCH = BK * 2 + 16;
constexpr int B_PITCH = BN + 16;
const int tid = threadIdx.x;
const int warp = tid >> 5;
const int lane = tid & 31;
const int wm = warp / WN;
const int wn = warp % WN;
const int m0 = blockIdx.x * BM;
const int n0 = blockIdx.y * BN;
const int nchunk_total = K / BK;
const int kc0 = SPLITK ? blockIdx.z * cspan : 0;
const int nchunk = SPLITK ? min(kc0 + cspan, nchunk_total) : nchunk_total;
extern __shared__ uint8_t smem[];
uint8_t* sA = smem;
uint8_t* sB = sA + (size_t)STAGES * BM * A_PITCH;
uint8_t* sSZ = sB + (size_t)STAGES * 64 * B_PITCH;
auto load_stage = [&](int stage, int chunk) {
uint8_t* dA = sA + (size_t)stage * (BM * A_PITCH);
uint8_t* dB = sB + (size_t)stage * (64 * B_PITCH);
uint8_t* dSZ = sSZ + (size_t)stage * (2 * BN * 4);
{
const int row0 = tid >> 4;
const int seg = tid & 15;
#pragma unroll
for (int row = row0; row < BM; row += THREADS / 16) {
const int gr = m0 + row;
const __nv_bfloat16* src = x + (size_t)gr * K + chunk * BK + seg * 8;
cp16g(dA + row * A_PITCH + seg * 16, src, gr < M);
}
}
{
constexpr int CPR = BN / 16;
constexpr int TOTAL = 64 * CPR;
#pragma unroll
for (int t = tid; t < TOTAL; t += THREADS) {
const int row = t / CPR;
const int seg = t % CPR;
cp16g(dB + row * B_PITCH + seg * 16,
wq + ((size_t)chunk * 64 + row) * N + n0 + seg * 16, true);
}
}
{
constexpr int TOTAL = 2 * BN / 8;
for (int t = tid; t < TOTAL; t += THREADS) {
const int half = t / (BN / 8);
const int seg = t % (BN / 8);
cp16g(dSZ + half * (BN * 4) + seg * 16,
(half ? Z : S) + (size_t)chunk * N + n0 + seg * 8, true);
}
}
cp_commit();
};
for (int s = 0; s < STAGES - 1; ++s) {
const int c = kc0 + s;
if (c < nchunk) load_stage(s, c);
else cp_commit();
}
float acc[MFRAGS][NFRAGS][4];
#pragma unroll
for (int mf = 0; mf < MFRAGS; ++mf)
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf)
#pragma unroll
for (int i = 0; i < 4; ++i) acc[mf][nf][i] = 0.f;
const int arow0 = wm * WTM + (lane & 7);
const int acol0 = (lane >> 4) * 16;
const int bq = lane & 3;
const int bn_ = lane >> 2;
uint32_t a[3][MFRAGS][4];
uint32_t b[3][NFRAGS][2];
auto load_frags = [&](const uint8_t* tA, const uint8_t* tB, int ks, int par,
const uint32_t* zz, const uint32_t* ss) {
#pragma unroll
for (int mf = 0; mf < MFRAGS; ++mf) {
const int row = arow0 + mf * 16;
const uint8_t* addr = tA + (row + ((lane >> 3) & 1) * 8) * A_PITCH + ks * 32 + acol0;
const uint32_t a_addr = (uint32_t)__cvta_generic_to_shared(addr);
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(a[par][mf][0]), "=r"(a[par][mf][1]), "=r"(a[par][mf][2]), "=r"(a[par][mf][3])
: "r"(a_addr));
}
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf) {
const int col = wn * WTN + nf * 8 + bn_;
const uint32_t e0 = (uint32_t)tB[(ks * 8 + bq) * B_PITCH + col];
const uint32_t e1 = (uint32_t)tB[(ks * 8 + bq + 4) * B_PITCH + col];
b[par][nf][0] = dq_byte(e0, zz[nf], ss[nf]);
b[par][nf][1] = dq_byte(e1, zz[nf], ss[nf]);
}
};
for (int chunk = kc0; chunk < nchunk; ++chunk) {
cp_wait<STAGES - 2>();
__syncthreads();
const uint8_t* tA = sA + (size_t)(chunk % STAGES) * (BM * A_PITCH);
const uint8_t* tB = sB + (size_t)(chunk % STAGES) * (64 * B_PITCH);
const uint8_t* tSZ = sSZ + (size_t)(chunk % STAGES) * (2 * BN * 4);
uint32_t zz[NFRAGS], ss[NFRAGS];
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf) {
const int col = wn * WTN + nf * 8 + bn_;
const float sf = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(tSZ)[col]);
const float zf = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(tSZ + BN * 4)[col]);
const __nv_bfloat162 z2 = __float2bfloat162_rn(128.f + zf);
const __nv_bfloat162 s2v = __float2bfloat162_rn(sf);
zz[nf] = *reinterpret_cast<const uint32_t*>(&z2);
ss[nf] = *reinterpret_cast<const uint32_t*>(&s2v);
}
load_frags(tA, tB, 0, 0, zz, ss);
load_frags(tA, tB, 1, 1, zz, ss);
#pragma unroll
for (int ks = 0; ks < 8; ++ks) {
const int par = ks % 3;
if (ks < 6) load_frags(tA, tB, ks + 2, (ks + 2) % 3, zz, ss);
#pragma unroll
for (int mf = 0; mf < MFRAGS; ++mf)
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf)
mma_bf16(acc[mf][nf], a[par][mf], b[par][nf]);
}
if (chunk + STAGES - 1 < nchunk) load_stage((chunk + STAGES - 1) % STAGES, chunk + STAGES - 1);
else cp_commit();
}
const int cr = lane >> 2;
const int cc = (lane & 3) * 2;
#pragma unroll
for (int mf = 0; mf < MFRAGS; ++mf) {
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf) {
#pragma unroll
for (int half = 0; half < 2; ++half) {
const int r = m0 + wm * WTM + mf * 16 + half * 8 + cr;
const int c = n0 + wn * WTN + nf * 8 + cc;
if (r < M) {
if (SPLITK) {
atomicAdd(ws + (size_t)r * N + c, acc[mf][nf][half * 2]);
atomicAdd(ws + (size_t)r * N + c + 1, acc[mf][nf][half * 2 + 1]);
} else {
__nv_bfloat162 v = __floats2bfloat162_rn(acc[mf][nf][half * 2], acc[mf][nf][half * 2 + 1]);
*reinterpret_cast<__nv_bfloat162*>(out + (size_t)r * N + c) = v;
}
}
}
}
}
if (SPLITK) {
__shared__ int is_last_split;
__threadfence();
__syncthreads();
if (tid == 0) {
const int prev = atomicAdd(counters + blockIdx.x * gridDim.y + blockIdx.y, 1);
is_last_split = (prev == (int)gridDim.z - 1) ? 1 : 0;
}
__syncthreads();
if (is_last_split) {
__threadfence();
#pragma unroll
for (int mf = 0; mf < MFRAGS; ++mf) {
#pragma unroll
for (int nf = 0; nf < NFRAGS; ++nf) {
#pragma unroll
for (int half = 0; half < 2; ++half) {
const int r = m0 + wm * WTM + mf * 16 + half * 8 + cr;
const int c = n0 + wn * WTN + nf * 8 + cc;
if (r < M) {
__nv_bfloat162 v;
v.x = __float2bfloat16(*reinterpret_cast<volatile float*>(ws + (size_t)r * N + c));
v.y = __float2bfloat16(*reinterpret_cast<volatile float*>(ws + (size_t)r * N + c + 1));
*reinterpret_cast<__nv_bfloat162*>(out + (size_t)r * N + c) = v;
ws[(size_t)r * N + c] = 0.f;
ws[(size_t)r * N + c + 1] = 0.f;
}
}
}
}
__syncthreads();
if (tid == 0) counters[blockIdx.x * gridDim.y + blockIdx.y] = 0;
}
}
}
torch::Tensor gemm_run(torch::Tensor x, torch::Tensor wq, torch::Tensor S,
torch::Tensor Z, torch::Tensor ws, torch::Tensor counters,
int64_t cfg, int64_t splitk) {
const int M = x.size(0), K = x.size(1);
const int N = wq.size(1);
auto out = torch::empty({M, N}, x.options());
auto stream = at::cuda::getCurrentCUDAStream();
const int64_t sk = splitk;
#define LAUNCH(BM, BN, WM, WN, ST) \
{ \
size_t smem = ST * ((size_t)BM * (128 * 2 + 16) + 64 * (BN + 16) + 2 * BN * 4); \
if (sk > 1) { \
const int cspan = (K / 128 + sk - 1) / (int)sk; \
dim3 grid((M + BM - 1) / BM, N / BN, sk); \
auto kfn = w4a16_gemm<BM, BN, WM, WN, ST, true>; \
TORCH_CHECK(smem <= 101376, "smem overflow"); \
if (smem > 48 * 1024) cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); \
kfn<<<grid, WM * WN * 32, smem, stream>>>( \
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()), wq.data_ptr<uint8_t>(), \
reinterpret_cast<const __nv_bfloat16*>(S.data_ptr()), \
reinterpret_cast<const __nv_bfloat16*>(Z.data_ptr()), \
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), ws.data_ptr<float>(), \
counters.data_ptr<int>(), M, N, K, cspan); \
} else { \
dim3 grid((M + BM - 1) / BM, N / BN); \
auto kfn = w4a16_gemm<BM, BN, WM, WN, ST, false>; \
TORCH_CHECK(smem <= 101376, "smem overflow"); \
if (smem > 48 * 1024) cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); \
kfn<<<grid, WM * WN * 32, smem, stream>>>( \
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()), wq.data_ptr<uint8_t>(), \
reinterpret_cast<const __nv_bfloat16*>(S.data_ptr()), \
reinterpret_cast<const __nv_bfloat16*>(Z.data_ptr()), \
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), ws.data_ptr<float>(), \
counters.data_ptr<int>(), M, N, K, 0); \
} \
}
if (cfg == 0) LAUNCH(16, 64, 1, 4, 4)
else if (cfg == 1) LAUNCH(16, 64, 1, 8, 4)
else if (cfg == 2) LAUNCH(16, 128, 1, 8, 2)
else if (cfg == 3) LAUNCH(16, 64, 1, 8, 2)
else if (cfg == 4) LAUNCH(16, 128, 1, 8, 4)
else if (cfg == 5) LAUNCH(16, 64, 1, 4, 2)
else if (cfg == 10) LAUNCH(32, 64, 2, 2, 4)
else if (cfg == 12) LAUNCH(32, 64, 2, 4, 3)
else if (cfg == 13) LAUNCH(32, 128, 2, 4, 2)
else if (cfg == 14) LAUNCH(32, 128, 2, 4, 4)
else if (cfg == 15) LAUNCH(32, 64, 2, 4, 2)
else if (cfg == 20) LAUNCH(64, 128, 2, 4, 3)
else if (cfg == 21) LAUNCH(64, 64, 2, 4, 4)
else if (cfg == 22) LAUNCH(64, 64, 2, 4, 2)
else if (cfg == 23) LAUNCH(64, 64, 4, 2, 4)
else if (cfg == 24) LAUNCH(64, 128, 2, 4, 2)
else TORCH_CHECK(false, "cfg");
return out;
}
"""
_CPP_SRC = """
torch::Tensor gemv_run_ptr(torch::Tensor x, int64_t wq_p, int64_t S_p, int64_t Z_p, int64_t ws_p, int64_t ctr_p, int64_t N, int64_t K, int64_t cfg);
torch::Tensor gemm_run(torch::Tensor x, torch::Tensor wq, torch::Tensor S, torch::Tensor Z, torch::Tensor ws, torch::Tensor counters, int64_t cfg, int64_t splitk);
"""
_mod = None
def _get_mod():
global _mod
if _mod is None:
from torch.utils.cpp_extension import load_inline
_mod = load_inline(
name="w4a16_sol_v8",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=["gemv_run_ptr", "gemm_run"],
extra_cuda_cflags=["-O3", "-gencode=arch=compute_120a,code=sm_120a", "--use_fast_math"],
verbose=False,
)
return _mod
class Model(nn.Module):
"""W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""
def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
super().__init__()
assert K % group_size == 0
assert 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.empty(K // 2, N, dtype=torch.uint8))
self.register_buffer("scales", torch.empty(n_groups, N, dtype=torch.bfloat16))
self.register_buffer("zeros", torch.empty(n_groups, N, dtype=torch.bfloat16))
self._ws = None
self._counters = None
self._ptrs = None
self._env_gemv_cfg = int(os.environ.get("W4A16_GEMV_CFG", "0" if N >= 8192 else "10"))
# invalidate cached raw pointers if weights are re-loaded
def _invalidate(mod, recurse=True):
self._ptrs = None
self.register_load_state_dict_post_hook(_invalidate)
self._register_state_dict_hook(lambda self, sd, prefix, metadata: setattr(self, "_ptrs", None))
self._env_cfg = os.environ.get("W4A16_CFG")
self._env_sk = int(os.environ.get("W4A16_SK", "1"))
# dispatch decisions frozen at init
self._gemv = (M == 1 and N % 512 == 0 and K % 128 == 0)
self._gemm_cfg = self._pick_cfg()
# touch + cache the extension so all instantiations build once
self._mod = _get_mod()
def _pick_cfg(self):
M, N, K = self.M, self.N, self.K
if self._env_cfg is not None:
return int(self._env_cfg)
if M <= 16:
return 4 if N % 128 == 0 else 0
elif M <= 32:
return 14 if N % 128 == 0 else 10
else:
return 20 if N % 128 == 0 else 21
def _ensure_workspace(self, device):
if self._ws is None or self._ws.device != device:
self._ws = torch.zeros(max(self.M * self.N, self.N), dtype=torch.float32, device=device)
self._counters = torch.zeros(4096, dtype=torch.int32, device=device)
def _ensure_splitws(self, device, M, N, cfg, sk):
if (self._ws is None or self._ws.numel() < M * N or self._ws.device != device):
self._ws = torch.zeros(max(M * N, self.N), dtype=torch.float32, device=device)
if self._counters is None or self._counters.numel() < 4096 or self._counters.device != device:
self._counters = torch.zeros(4096, dtype=torch.int32, device=device)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if not x.is_contiguous():
x = x.contiguous()
if self._gemv and self._ptrs is None:
self._ensure_workspace(x.device)
self._ptrs = (self.w_q.data_ptr(), self.scales.data_ptr(), self.zeros.data_ptr(),
self._ws.data_ptr(), self._counters.data_ptr())
if self._gemv:
return self._mod.gemv_run_ptr(x, *self._ptrs, self.N, self.K, self._env_gemv_cfg).view(1, self.N)
self._ensure_splitws(x.device, self.M, self.N, self._gemm_cfg, self._env_sk)
return self._mod.gemm_run(x, self.w_q, self.scales, self.zeros, self._ws, self._counters,
self._gemm_cfg, self._env_sk)
__call__ = forward
M = 1
N = 12288
K = 4096
def get_inputs():
x = torch.randn(M, K, dtype=torch.bfloat16)
return [x]
def get_init_inputs():
return [M, N, K]
20260716_112718_kinetic-claude_kinetic-0715_07_w4a16_gemm