"""W4A16 weight-only quantized GEMM (AWQ/GPTQ-style asym int4, group=128). Fused unpack + GEMM in Triton. w_q is (K//2, N) uint8: low nibble = even-K row, high nibble = odd-K row. We unpack + dequant inside the GEMM main loop and never materialize the bf16 weight. Each loop iteration processes exactly one quant group (128 contiguous K = 64 packed rows), so scale/zero are loaded once per group as (BLOCK_N,) row vectors and broadcast. Activations are loaded contiguously (BLOCK_M, 128) then split into even/odd K halves with tl.split, matching the low/high nibbles. Two K=64 sub-dots accumulate into fp32. Dispatch: * M == 1 : split-K kernel (occupancy-limited decode) -> fp32 scratch -> convert * M > 1 : direct GEMM kernel, bf16 store """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl GROUP_SIZE = 128 # --------------------------------------------------------------------------- # Direct GEMM kernel (M > 1): bf16 store, BLOCK_K fixed to one group (128). # --------------------------------------------------------------------------- def _direct_configs(): cfgs = [] for bm, bn, w, s in [ (16, 128, 4, 4), (16, 256, 8, 4), (16, 64, 4, 4), (32, 128, 4, 4), (32, 256, 8, 4), (32, 128, 8, 4), (64, 128, 4, 4), (64, 256, 8, 4), (64, 128, 8, 4), (128, 128, 8, 4), (128, 256, 8, 4), (64, 64, 4, 4), ]: cfgs.append(triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_warps=w, num_stages=s)) return cfgs @triton.autotune(configs=_direct_configs(), key=["M", "N", "K"]) @triton.jit def _w4a16_direct( x_ptr, wq_ptr, scl_ptr, zero_ptr, out_ptr, M, N, K, stride_xm, stride_xk, stride_wq_kh, stride_wq_n, stride_scl_g, stride_scl_n, stride_zero_g, stride_zero_n, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: 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) m_mask = offs_m < M n_mask = offs_n < N BLOCK_K2: tl.constexpr = 64 # 64 packed rows = 128 K = one group offs_k2 = tl.arange(0, BLOCK_K2) offs_k = tl.arange(0, 128) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) n_groups = K // 128 x_base = x_ptr + offs_m[:, None] * stride_xm wq_base = wq_ptr + offs_n[None, :] * stride_wq_n sc_off = offs_n * stride_scl_n zr_off = offs_n * stride_zero_n for g in range(0, n_groups): kh = g * BLOCK_K2 + offs_k2 packed = tl.load(wq_base + kh[:, None] * stride_wq_kh, mask=n_mask[None, :], other=0) lo = (packed & 0xF).to(tl.float32) hi = ((packed >> 4) & 0xF).to(tl.float32) scl = tl.load(scl_ptr + g * stride_scl_g + sc_off, mask=n_mask, other=0.0).to(tl.float32) zero = tl.load(zero_ptr + g * stride_zero_g + zr_off, mask=n_mask, other=0.0).to(tl.float32) w_lo = ((lo - zero[None, :]) * scl[None, :]).to(tl.bfloat16) w_hi = ((hi - zero[None, :]) * scl[None, :]).to(tl.bfloat16) k0 = g * 128 x_tile = tl.load(x_base + (k0 + offs_k)[None, :] * stride_xk, mask=m_mask[:, None], other=0.0) x_e, x_o = tl.split(tl.reshape(x_tile, (BLOCK_M, BLOCK_K2, 2))) acc += tl.dot(x_e, w_lo, out_dtype=tl.float32) acc += tl.dot(x_o, w_hi, out_dtype=tl.float32) out = acc.to(tl.bfloat16) tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on, out, mask=m_mask[:, None] & n_mask[None, :]) # --------------------------------------------------------------------------- # Split-K kernel (M == 1): fp32 atomic accumulation into scratch. # --------------------------------------------------------------------------- def _splitk_configs(): cfgs = [] for bn, sp, w, s in [ (64, 1, 4, 4), (128, 1, 4, 4), (256, 1, 8, 4), (64, 2, 4, 4), (128, 2, 4, 4), (256, 2, 8, 4), (64, 4, 4, 4), (128, 4, 4, 4), (256, 4, 8, 4), (64, 8, 4, 4), (128, 8, 4, 4), (64, 4, 2, 3), (64, 8, 2, 3), ]: cfgs.append(triton.Config({"BLOCK_N": bn, "SPLIT_K": sp}, num_warps=w, num_stages=s)) return cfgs @triton.autotune(configs=_splitk_configs(), key=["M", "N", "K"]) @triton.jit def _w4a16_splitk( x_ptr, wq_ptr, scl_ptr, zero_ptr, out_ptr, M, N, K, stride_xm, stride_xk, stride_wq_kh, stride_wq_n, stride_scl_g, stride_scl_n, stride_zero_g, stride_zero_n, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr, ): pid_n = tl.program_id(0) pid_k = tl.program_id(1) offs_m = tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) m_mask = offs_m < M n_mask = offs_n < N BLOCK_K2: tl.constexpr = 64 offs_k2 = tl.arange(0, BLOCK_K2) offs_k = tl.arange(0, 128) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) n_groups = K // 128 groups_per_split = (n_groups + SPLIT_K - 1) // SPLIT_K g_start = pid_k * groups_per_split g_end = min(g_start + groups_per_split, n_groups) x_base = x_ptr + offs_m[:, None] * stride_xm wq_base = wq_ptr + offs_n[None, :] * stride_wq_n sc_off = offs_n * stride_scl_n zr_off = offs_n * stride_zero_n for g in range(g_start, g_end): kh = g * BLOCK_K2 + offs_k2 packed = tl.load(wq_base + kh[:, None] * stride_wq_kh, mask=n_mask[None, :], other=0) lo = (packed & 0xF).to(tl.float32) hi = ((packed >> 4) & 0xF).to(tl.float32) scl = tl.load(scl_ptr + g * stride_scl_g + sc_off, mask=n_mask, other=0.0).to(tl.float32) zero = tl.load(zero_ptr + g * stride_zero_g + zr_off, mask=n_mask, other=0.0).to(tl.float32) w_lo = ((lo - zero[None, :]) * scl[None, :]).to(tl.bfloat16) w_hi = ((hi - zero[None, :]) * scl[None, :]).to(tl.bfloat16) k0 = g * 128 x_tile = tl.load(x_base + (k0 + offs_k)[None, :] * stride_xk, mask=m_mask[:, None], other=0.0) x_e, x_o = tl.split(tl.reshape(x_tile, (BLOCK_M, BLOCK_K2, 2))) acc += tl.dot(x_e, w_lo, out_dtype=tl.float32) acc += tl.dot(x_o, w_hi, out_dtype=tl.float32) out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on full_mask = m_mask[:, None] & n_mask[None, :] if SPLIT_K == 1: tl.store(out_ptrs, acc, mask=full_mask) else: tl.atomic_add(out_ptrs, acc, mask=full_mask) def _w4a16_gemm(x, wq, scales, zeros, M, N, K): if M == 1: scratch = torch.zeros((M, N), dtype=torch.float32, device=x.device) def grid(meta): return (triton.cdiv(N, meta["BLOCK_N"]), meta["SPLIT_K"]) _w4a16_splitk[grid]( x, wq, scales, zeros, scratch, M, N, K, x.stride(0), x.stride(1), wq.stride(0), wq.stride(1), scales.stride(0), scales.stride(1), zeros.stride(0), zeros.stride(1), scratch.stride(0), scratch.stride(1), BLOCK_M=16, ) return scratch.to(torch.bfloat16) out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) def grid(meta): return (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"])) _w4a16_direct[grid]( x, wq, scales, zeros, out, M, N, K, x.stride(0), x.stride(1), wq.stride(0), wq.stride(1), scales.stride(0), scales.stride(1), zeros.stride(0), zeros.stride(1), out.stride(0), out.stride(1), ) return out 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)) def forward(self, x: torch.Tensor) -> torch.Tensor: M, K = x.shape N = self.N x = x.contiguous() return _w4a16_gemm(x, self.w_q, self.scales, self.zeros, M, N, K)