"""W4A16 weight-only quantized GEMM — fused unpack + GEMM (Triton). Scheme (AWQ/GPTQ-style asymmetric int4, group_size=128 along K): x: (M, K) bf16 w_q: (K // 2, N) uint8 (low nibble = even-K row, high nibble = odd-K row) scales: (K // 128, N) bf16 zeros: (K // 128, N) bf16 out[m, n] = sum_k x[m, k] * (unpack(w_q)[k, n] - zeros[k // 128, n]) * scales[k // 128, n] Two kernels: * M == 1 decode: split-K GEMV. Each program owns a (BLOCK_N x groups) slab, accumulates s * (sum_k q*x - z * sum_k x) per group in fp32, writes fp32 partials; a tiny second kernel reduces partials and casts to bf16. * M >= 2: tensor-core GEMM. Per 128-wide K group, the packed tile (64, BN) is dequantized into an even-K and an odd-K bf16 tile (no row interleave needed: two tl.dot calls against x's even / odd columns). """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl GROUP_SIZE = 128 # --------------------------------------------------------------------------- # GEMV (M == 1) # --------------------------------------------------------------------------- @triton.jit def _w4a16_gemv_kernel( x_ptr, # (K,) bf16 wq_ptr, # (K//2, N) uint8 s_ptr, # (K//128, N) bf16 z_ptr, # (K//128, N) bf16 part_ptr, # (SPLITS, N) fp32 partials N, GROUPS_PER_SPLIT: tl.constexpr, BLOCK_N: tl.constexpr, ): pid_n = tl.program_id(0) pid_k = tl.program_id(1) cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) acc = tl.zeros((BLOCK_N,), dtype=tl.float32) for g in range(GROUPS_PER_SPLIT): gid = pid_k * GROUPS_PER_SPLIT + g # 64 packed rows == 128 unpacked K rows == exactly one quant group rows = gid * 64 + tl.arange(0, 64) b = tl.load(wq_ptr + rows[:, None] * N + cols[None, :]) # (64, BN) u8 xg = tl.load(x_ptr + gid * 128 + tl.arange(0, 128)).to(tl.float32) xp = tl.reshape(xg, (64, 2)) xe, xo = tl.split(xp) # even-K, odd-K activations (64,) lo = (b & 0xF).to(tl.float32) hi = (b >> 4).to(tl.float32) qsum = tl.sum(lo * xe[:, None] + hi * xo[:, None], axis=0) # (BN,) xsum = tl.sum(xg, axis=0) s = tl.load(s_ptr + gid * N + cols).to(tl.float32) z = tl.load(z_ptr + gid * N + cols).to(tl.float32) acc += s * (qsum - z * xsum) tl.store(part_ptr + pid_k * N + cols, acc) @triton.jit def _reduce_kernel( part_ptr, # (SPLITS, N) fp32 out_ptr, # (N,) bf16 N, SPLITS: tl.constexpr, BLOCK_N: tl.constexpr, ): pid = tl.program_id(0) cols = pid * BLOCK_N + tl.arange(0, BLOCK_N) acc = tl.zeros((BLOCK_N,), dtype=tl.float32) for s in range(SPLITS): acc += tl.load(part_ptr + s * N + cols) tl.store(out_ptr + cols, acc.to(tl.bfloat16)) # --------------------------------------------------------------------------- # GEMM (M >= 2) # --------------------------------------------------------------------------- @triton.jit def _w4a16_gemm_kernel( x_ptr, # (M, K) bf16 wq_ptr, # (K//2, N) uint8 s_ptr, # (K//128, N) bf16 z_ptr, # (K//128, N) bf16 out_ptr, # (M, N) bf16 (zero-initialised when SPLIT_K > 1) M, N, K, GROUPS_PER_SPLIT: tl.constexpr, SPLIT_K: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): pid_n = tl.program_id(0) pid_m = tl.program_id(1) pid_k = tl.program_id(2) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) m_mask = offs_m < M acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for g in range(GROUPS_PER_SPLIT): gid = pid_k * GROUPS_PER_SPLIT + g # x group tile: (BM, 128) -> even / odd K halves (BM, 64) xk = gid * 128 + tl.arange(0, 128) xt = tl.load( x_ptr + offs_m[:, None] * K + xk[None, :], mask=m_mask[:, None], other=0.0, ) xp = tl.reshape(xt, (BLOCK_M, 64, 2)) xe, xo = tl.split(xp) rows = gid * 64 + tl.arange(0, 64) b = tl.load(wq_ptr + rows[:, None] * N + offs_n[None, :]) # (64, BN) u8 s = tl.load(s_ptr + gid * N + offs_n).to(tl.float32)[None, :] z = tl.load(z_ptr + gid * N + offs_n).to(tl.float32)[None, :] w_lo = (((b & 0xF).to(tl.float32) - z) * s).to(tl.bfloat16) w_hi = (((b >> 4).to(tl.float32) - z) * s).to(tl.bfloat16) acc = tl.dot(xe, w_lo, acc) acc = tl.dot(xo, w_hi, acc) out_offs = offs_m[:, None] * N + offs_n[None, :] if SPLIT_K == 1: tl.store(out_ptr + out_offs, acc.to(tl.bfloat16), mask=m_mask[:, None]) else: # out_ptr is a zeroed fp32 workspace here; cast kernel finishes the job tl.atomic_add(out_ptr + out_offs, acc, mask=m_mask[:, None]) @triton.jit def _cast_kernel(ws_ptr, out_ptr, n_elem, BLOCK: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) v = tl.load(ws_ptr + offs, mask=offs < n_elem, other=0.0) tl.store(out_ptr + offs, v.to(tl.bfloat16), mask=offs < n_elem) # --------------------------------------------------------------------------- # Host-side model # --------------------------------------------------------------------------- def _pack_int4(w_q: torch.Tensor) -> torch.Tensor: K, N = w_q.shape lo = w_q[0::2].to(torch.uint8) & 0xF hi = w_q[1::2].to(torch.uint8) & 0xF return (lo | (hi << 4)).contiguous() 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 torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K)) w_full = torch.randn(K, N, dtype=torch.float32) * 0.02 w_g = w_full.view(n_groups, group_size, N) w_min = w_g.min(dim=1, keepdim=True).values w_max = w_g.max(dim=1, keepdim=True).values scales = (w_max - w_min).clamp_min(1e-8) / 15.0 zeros = (-w_min / scales).round().clamp(0, 15) w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N) self.register_buffer("w_q", _pack_int4(w_q)) self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16)) self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16)) def forward(self, x: torch.Tensor) -> torch.Tensor: M = x.shape[0] N, K = self.N, self.K n_groups = K // self.group_size if M == 1: SPLITS = 16 BLOCK_N = 256 gps = n_groups // SPLITS part = torch.empty((SPLITS, N), dtype=torch.float32, device=x.device) out = torch.empty((1, N), dtype=torch.bfloat16, device=x.device) _w4a16_gemv_kernel[(N // BLOCK_N, SPLITS)]( x, self.w_q, self.scales, self.zeros, part, N, GROUPS_PER_SPLIT=gps, BLOCK_N=BLOCK_N, num_warps=4, ) _reduce_kernel[(N // 1024,)]( part, out, N, SPLITS=SPLITS, BLOCK_N=1024, num_warps=4, ) return out # tensor-core path if M <= 16: BM, BN, SPLIT_K, warps = 16, 128, 4, 4 elif M <= 32: BM, BN, SPLIT_K, warps = 32, 128, 4, 4 else: BM, BN, SPLIT_K, warps = 128, 128, 1, 8 gps = n_groups // SPLIT_K out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) grid = (N // BN, triton.cdiv(M, BM), SPLIT_K) if SPLIT_K > 1: ws = torch.zeros((M, N), dtype=torch.float32, device=x.device) _w4a16_gemm_kernel[grid]( x, self.w_q, self.scales, self.zeros, ws, M, N, K, GROUPS_PER_SPLIT=gps, SPLIT_K=SPLIT_K, BLOCK_M=BM, BLOCK_N=BN, num_warps=warps, ) n_elem = M * N _cast_kernel[(triton.cdiv(n_elem, 4096),)](ws, out, n_elem, BLOCK=4096) else: _w4a16_gemm_kernel[grid]( x, self.w_q, self.scales, self.zeros, out, M, N, K, GROUPS_PER_SPLIT=gps, SPLIT_K=SPLIT_K, BLOCK_M=BM, BLOCK_N=BN, num_warps=warps, ) return out 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]