KernelBench hard · RTX PRO 6000
W4A16 GEMM GPT-5.6 Sol
19.8%geomean peak fraction across shapes
manually audited: clean
Genuine fused W4A16 Triton implementation. Every active path reads the live bf16 activation, packed uint8 int4 weights, bf16 scales, and bf16 zero-points; unpacks the low/even and high/odd nibbles inside the kernel; performs bf16 tensor-core dots into fp32; and returns bf16. M=1 decode uses 16-way split-K with a newly zeroed fp32 accumulation buffer and a final authored bf16 cast kernel, while M>1 uses one direct fused GEMM kernel. There is no dequantized- weight materialization, library GEMM, output/input cache, CUDA graph, value- dependent branch, checker/seed sniff, or template mutation. Official correctness passed all five shapes, three seeds, and nominal/small/large activation cases. Independent arbitrary-state, packing, mutation, cancellation, and repeatability tests also passed. The canonical 0.1980 score is supported by the solution-only benchmark and independently recomputes from its five reported fractions.
harnesscodexagent session37mtotal wall38mcheck15sbenchmark3soutput tokens36,050gpu-lock wait0sgpu-lock held18sregimememory
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×12288×40960.045 ms32.7%0.59 TB/s · 33% of 1.8 TB/s HBM · also 2 TFLOPS (0% of compute)
32×12288×40960.077 ms20.1%0.36 TB/s · 20% of 1.8 TB/s HBM · also 42 TFLOPS (8% of compute)
256×12288×40960.204 ms9.6%126 TFLOPS · 25% of 500 TF bf16 peak · also 0.17 TB/s (10% of HBM)
1×4096×40960.025 ms19.5%0.35 TB/s · 20% of 1.8 TB/s HBM · also 1 TFLOPS (0% of compute)
16×14336×40960.071 ms24.8%0.45 TB/s · 25% of 1.8 TB/s HBM · also 26 TFLOPS (5% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(32.7% · 20.1% · 9.6% · 19.5% · 24.8%) = 19.8%
Kernel source (redacted)
"""Fused W4A16 GEMM for SM120.
The packed weight is consumed directly by the Triton kernel. In particular,
there is deliberately no persistent/cached dequantized weight: every byte is
unpacked next to the tensor-core dot which consumes it.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import triton
import triton.language as tl
GROUP_SIZE = 128
@triton.jit
def _w4a16_gemv_kernel(
x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr,
):
pid_n = tl.program_id(0)
rn = pid_n * BN + tl.arange(0, BN)
rp = tl.arange(0, 64)
acc = tl.zeros((BN,), tl.float32)
for group in tl.static_range(0, K // 128):
packed = tl.load(
q_ptr + (group * 64 + rp[:, None]) * N + rn[None, :]
)
qlo = (packed & 15).to(tl.bfloat16)
qhi = (packed >> 4).to(tl.bfloat16)
s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
wlo = ((qlo - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
whi = ((qhi - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
xlo = tl.load(x_ptr + group * 128 + 2 * rp)
xhi = tl.load(x_ptr + group * 128 + 2 * rp + 1)
products = (
xlo[:, None].to(tl.float32) * wlo.to(tl.float32)
+ xhi[:, None].to(tl.float32) * whi.to(tl.float32)
)
acc += tl.sum(products, axis=0)
tl.store(out_ptr + rn, acc.to(tl.bfloat16))
@triton.jit
def _w4a16_direct_kernel(
x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
rm = pid_m * BM + tl.arange(0, BM)
rn = pid_n * BN + tl.arange(0, BN)
rk = tl.arange(0, BK)
acc = tl.zeros((BM, BN), tl.float32)
for k0 in tl.static_range(0, K, BK):
k = k0 + rk
x = tl.load(
x_ptr + rm[:, None] * K + k[None, :],
mask=rm[:, None] < M,
other=0.0,
)
packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
shift = (k[:, None] & 1) * 4
q = ((packed >> shift) & 15).to(tl.bfloat16)
group = k0 // 128
s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
acc += tl.dot(x, w)
tl.store(
out_ptr + rm[:, None] * N + rn[None, :],
acc.to(tl.bfloat16),
mask=rm[:, None] < M,
)
@triton.jit
def _w4a16_splitk_kernel(
x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
N: tl.constexpr, K: tl.constexpr, BN: tl.constexpr,
SPLIT_K: tl.constexpr,
):
split = tl.program_id(0)
pid_n = tl.program_id(1)
rm = tl.arange(0, 16)
rn = pid_n * BN + tl.arange(0, BN)
rk = tl.arange(0, 128)
acc = tl.zeros((16, BN), tl.float32)
groups_per_split: tl.constexpr = (K // 128) // SPLIT_K
for gi in tl.static_range(0, groups_per_split):
group = split * groups_per_split + gi
k = group * 128 + rk
x = tl.load(x_ptr + rm[:, None] * K + k[None, :], mask=rm[:, None] < 1,
other=0.0)
packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
q = ((packed >> ((k[:, None] & 1) * 4)) & 15).to(tl.bfloat16)
s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
acc += tl.dot(x, w)
tl.atomic_add(out_ptr + rm[:, None] * N + rn[None, :], acc,
mask=rm[:, None] < 1)
@triton.jit
def _fp32_to_bf16_kernel(src_ptr, dst_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
tl.store(dst_ptr + offsets, tl.load(src_ptr + offsets, mask=offsets < N),
mask=offsets < N)
@triton.jit
def _w4a16_grouped_kernel(
x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
BM: tl.constexpr, BN: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
rm = pid_m * BM + tl.arange(0, BM)
rn = pid_n * BN + tl.arange(0, BN)
rk = tl.arange(0, 128)
acc = tl.zeros((BM, BN), tl.float32)
for group in tl.static_range(0, K // 128):
k = group * 128 + rk
x = tl.load(
x_ptr + rm[:, None] * K + k[None, :],
mask=rm[:, None] < M,
other=0.0,
)
packed = tl.load(q_ptr + (k[:, None] // 2) * N + rn[None, :])
q = ((packed >> ((k[:, None] & 1) * 4)) & 15).to(tl.bfloat16)
s = tl.load(scales_ptr + group * N + rn).to(tl.float32)
z = tl.load(zeros_ptr + group * N + rn).to(tl.float32)
qdot = tl.dot(x, q)
xsum = tl.sum(x.to(tl.float32), axis=1)
acc += (qdot - xsum[:, None] * z[None, :]) * s[None, :]
tl.store(
out_ptr + rm[:, None] * N + rn[None, :],
acc.to(tl.bfloat16),
mask=rm[:, None] < M,
)
@triton.jit
def _w4a16_interleave_kernel(
x_ptr, q_ptr, scales_ptr, zeros_ptr, out_ptr,
M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
BM: tl.constexpr, BN: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
rm = pid_m * BM + tl.arange(0, BM)
rn = pid_n * BN + tl.arange(0, BN)
rp = tl.arange(0, 64)
acc = tl.zeros((BM, BN), tl.float32)
for group in tl.static_range(0, K // 128):
packed = tl.load(
q_ptr + (group * 64 + rp[:, None]) * N + rn[None, :]
)
qlo = (packed & 15).to(tl.bfloat16)
qhi = (packed >> 4).to(tl.bfloat16)
# interleave operates on the last dimension. Two logical transposes
# therefore turn [64, BN] packed rows into [128, BN] K-major nibbles
# while retaining a single global load for each source byte.
q = tl.trans(tl.interleave(tl.trans(qlo), tl.trans(qhi)))
s = tl.load(scales_ptr + group * N + rn).to(tl.bfloat16)
z = tl.load(zeros_ptr + group * N + rn).to(tl.bfloat16)
w = ((q - z[None, :]).to(tl.bfloat16) * s[None, :]).to(tl.bfloat16)
x = tl.load(
x_ptr + rm[:, None] * K + group * 128 + tl.arange(0, 128)[None, :],
mask=rm[:, None] < M,
other=0.0,
)
acc += tl.dot(x, w)
tl.store(
out_ptr + rm[:, None] * N + rn[None, :],
acc.to(tl.bfloat16),
mask=rm[:, None] < M,
)
def _launch(x: torch.Tensor, q: torch.Tensor, scales: torch.Tensor,
zeros: torch.Tensor, M: int, N: int, K: int) -> torch.Tensor:
if M <= 1:
partial = torch.zeros((M, N), dtype=torch.float32, device=x.device)
if N == 4096:
bn, warps, stages = 64, 8, 2
else:
bn, warps, stages = 128, 16, 4
split_k = 16
_w4a16_splitk_kernel[(split_k, triton.cdiv(N, bn))](
x, q, scales, zeros, partial,
N=N, K=K, BN=bn, SPLIT_K=split_k,
num_warps=warps, num_stages=stages,
)
out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
_fp32_to_bf16_kernel[(triton.cdiv(N, 256),)](
partial, out, N=N, BLOCK=256, num_warps=4,
)
return out
out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
if M <= 16:
bm, bn, bk, warps, stages = 16, 64, 128, 8, 2
elif M <= 32:
bm, bn, bk, warps, stages = 32, 64, 128, 8, 2
else:
bm, bn, bk, warps, stages = 32, 64, 128, 8, 2
grid = (triton.cdiv(M, bm), triton.cdiv(N, bn))
_w4a16_direct_kernel[grid](
x, q, scales, zeros, out,
M=M, N=N, K=K, BM=bm, BN=bn, BK=bk,
num_warps=warps, num_stages=stages,
)
return out
class Model(nn.Module):
def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
super().__init__()
assert group_size == GROUP_SIZE
self.M, self.N, self.K = M, N, K
self.group_size = group_size
self.register_buffer("w_q", torch.empty((K // 2, N), dtype=torch.uint8))
self.register_buffer("scales", torch.empty((K // group_size, N), dtype=torch.bfloat16))
self.register_buffer("zeros", torch.empty((K // group_size, N), dtype=torch.bfloat16))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return _launch(x, self.w_q, self.scales, self.zeros,
self.M, self.N, self.K)
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]
20260709_184850_codex_gpt-5.6-sol_07_w4a16_gemm