"""W4A16 weight-only quantized GEMM (Triton, fused unpack + matmul). Scheme: AWQ/GPTQ-style asymmetric int4 with explicit zero-points and per-group bf16 scales. Inputs: x: (M, K) bf16 w_q: (K//2, N) uint8 -- two int4 weights packed per byte -- (low nibble = even-K row, high = odd) scales: (K//group, N) bf16 zeros: (K//group, N) bf16 -- stored as float zero-point out: (M, N) bf16 Dequant (per group along K): w_bf[k, n] = (w_q_unpacked[k, n] - zeros[k // group, n]) * scales[k // group, n] We never call any prebuilt W4A16 kernel (marlin/bnb/etc.) - the unpack+matmul is fused in one Triton kernel. """ from __future__ import annotations import torch import triton import triton.language as tl GROUP_SIZE = 128 # --------------------------------------------------------------------------- # Triton kernel # --------------------------------------------------------------------------- # Tile sizes: # BLOCK_M -- rows of the output per program (handles M=1 with masking) # BLOCK_N -- columns of the output per program # BLOCK_K -- K reduction block. We set this == GROUP_SIZE so each iteration # is one quant group, which lets us load scales/zeros once and # dequant a whole tile. # # Inner dot: x_even/w_bf_lo and x_odd/w_bf_hi are each (M, 64) @ (64, N), so # the inner reduction is 64 (K-axis is the *post-split* K, i.e. BLOCK_K/2). # Hopper WGMMA needs K in {16, 32, 64, ...}, so 64 is fine. # --------------------------------------------------------------------------- @triton.jit def _w4a16_gemm_kernel( x_ptr, w_ptr, scales_ptr, zeros_ptr, out_ptr, M, N, K, stride_xm, stride_xk, stride_wk, stride_wn, stride_sk, stride_sn, stride_zk, stride_zn, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_SIZE: tl.constexpr, ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) HALF_K: tl.constexpr = BLOCK_K // 2 n_groups: tl.constexpr = 0 # set at runtime via K acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # Pre-build stride-scaled base pointers. x_row_base = x_ptr + offs_m[:, None] * stride_xm w_base = w_ptr + offs_n[None, :] * stride_wn n_groups = K // BLOCK_K for g in range(0, n_groups): k_start = g * BLOCK_K kp_start = g * HALF_K # ---- x tile (BLOCK_M, BLOCK_K) bf16, fully contiguous in K ---- x = tl.load( x_row_base + (k_start + tl.arange(0, BLOCK_K))[None, :] * stride_xk, mask=offs_m[:, None] < M, other=0.0, ) # ---- packed weights (HALF_K, BLOCK_N) uint8 ---- w_packed = tl.load( w_base + (kp_start + tl.arange(0, HALF_K))[:, None] * stride_wk ) # ---- unpack + dequant ---- w_lo = (w_packed & 0xF).to(tl.bfloat16) w_hi = (w_packed >> 4).to(tl.bfloat16) scales = tl.load(scales_ptr + g * stride_sk + offs_n * stride_sn) zeros_ = tl.load(zeros_ptr + g * stride_zk + offs_n * stride_zn) w_bf_lo = (w_lo - zeros_[None, :]) * scales[None, :] w_bf_hi = (w_hi - zeros_[None, :]) * scales[None, :] # ---- split x into even/odd K rows via reshape+split ---- # x of shape (M, BLOCK_K=128) -> reshape to (M, 64, 2) # x_reshape[m, i, c] = x[m, 2i + c] (c is the parity) # tl.split peels off the size-2 minor dim. x_r = tl.reshape(x, (BLOCK_M, HALF_K, 2)) x_even, x_odd = tl.split(x_r) acc = tl.dot(x_even, w_bf_lo, acc) acc = tl.dot(x_odd, w_bf_hi, acc) out = acc.to(tl.bfloat16) tl.store( out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on, out, mask=offs_m[:, None] < M, ) # --------------------------------------------------------------------------- # Python launch wrapper # --------------------------------------------------------------------------- def _pick_block_m(M: int) -> int: # Tile size choices for the M dim. Triton recompiles per BLOCK_M, so we # only include a small handful of sizes. For tiny M (decode), pad to 16 # so the WGMMA / mma instructions can use their M=16 m-tile. if M <= 16: return 16 if M <= 32: return 32 if M <= 64: return 64 return 128 def w4a16_gemm(x: torch.Tensor, w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor) -> torch.Tensor: """Fused W4A16 GEMM: y = x @ dequant(w_q, scales, zeros).""" M, K = x.shape _, N = w_q.shape out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) BLOCK_M = _pick_block_m(M) BLOCK_N = 128 BLOCK_K = GROUP_SIZE # == 128 grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _w4a16_gemm_kernel[grid]( x, w_q, scales, zeros, out, M, N, K, x.stride(0), x.stride(1), w_q.stride(0), w_q.stride(1), scales.stride(0), scales.stride(1), zeros.stride(0), zeros.stride(1), out.stride(0), out.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SIZE=GROUP_SIZE, num_warps=4, num_stages=3, ) return out # --------------------------------------------------------------------------- # Model + get_inputs / get_init_inputs (must match reference.py interface). # The buffers are registered so the harness's state_dict load (strict=True) # works after the reference model has been built and quantized. # --------------------------------------------------------------------------- class Model(torch.nn.Module): 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 # Buffers are placeholders; the real values are loaded by the harness # from the reference model's state_dict. Names must match reference.py. 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)) def forward(self, x: torch.Tensor) -> torch.Tensor: return w4a16_gemm(x, self.w_q, self.scales, self.zeros) 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]