"""Fused W4A16 weight-only quantized GEMM for RTX PRO 6000 (SM120, CC 12.0). Fuses int4 unpack + per-group dequant + bf16 GEMM in a single Triton tiled pass. 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 = odd-K row scales: (K // 128, N) bf16 zeros: (K // 128, N) bf16 out: (M, N) bf16 w_bf[k,n] = (unpack(w_q)[k,n] - zeros[k//128,n]) * scales[k//128,n] Tile strategy: Tile M and N in blocks, then walk K in chunks of BK (a multiple of GROUP_SIZE=128). Each 128-row group maps to 64 packed-int4 rows and shares one (row of) scale/zero, so the dequant factorizes per group without per-element indexing. Within a group the even/odd K rows map to the low/high nibble of the same packed byte, so we do two dots per group (even-half rows, odd-half rows) instead of a materialized interleave. Accumulate in fp32, store bf16. Decode (M=1) is memory-bound and needs a wide grid over N to fill all SMs, so the autotune set includes small tiles; prefill (larger M) amortizes the x read over a bigger output tile, so it includes large tiles. """ from __future__ import annotations import os os.environ.setdefault("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS", "1") import torch import triton import triton.language as tl GROUP_SIZE = 128 @triton.autotune( configs=[ # compute-bound prefill: big accumulator, amortize x reads over N triton.Config({"BM": 256, "BN": 128, "BK": 128}, num_stages=3, num_warps=8), triton.Config({"BM": 128, "BN": 128, "BK": 128}, num_stages=4, num_warps=8), triton.Config({"BM": 128, "BN": 128, "BK": 256}, num_stages=3, num_warps=8), triton.Config({"BM": 128, "BN": 64, "BK": 128}, num_stages=4, num_warps=4), triton.Config({"BM": 64, "BN": 128, "BK": 128}, num_stages=4, num_warps=4), triton.Config({"BM": 64, "BN": 64, "BK": 128}, num_stages=4, num_warps=4), # memory-bound decode: small tile -> wide grid -> fill all SMs & hide latency triton.Config({"BM": 32, "BN": 32, "BK": 128}, num_stages=4, num_warps=4), triton.Config({"BM": 32, "BN": 64, "BK": 128}, num_stages=4, num_warps=4), triton.Config({"BM": 64, "BN": 32, "BK": 128}, num_stages=4, num_warps=4), ], key=["M", "N", "K"], ) @triton.jit def _w4a16_kernel( x_ptr, wq_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, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BM + tl.arange(0, BM) offs_n = pid_n * BN + tl.arange(0, BN) x_base = x_ptr + offs_m[:, None] * stride_xm # (BM, 1) out_base = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on acc = tl.zeros((BM, BN), dtype=tl.float32) # One quant group = 128 K-rows = 64 packed int4 rows. 64 is a literal => the # tl.arange(0, 64) calls below are constexpr. HALF: tl.constexpr = 64 n_groups_per_tile = BK // 128 for kk in range(0, K, BK): for g in range(n_groups_per_tile): # K rows covered by this group: [gs, gs+128) gs = kk + g * 128 # packed-weight row base for this group: gs maps to wq row gs//2 wq_base = wq_ptr + (gs // 2) * stride_wk + offs_n[None, :] * stride_wn # (1, BN) # scale/zero for this group, broadcast across the group's 128 K-rows & BN cols s_ptr = scales_ptr + (gs // 128) * stride_sk + offs_n[None, :] * stride_sn z_ptr = zeros_ptr + (gs // 128) * stride_zk + offs_n[None, :] * stride_zn scale = tl.load(s_ptr, mask=offs_n[None, :] < N, other=0.0).to(tl.float32) # (1, BN) zero = tl.load(z_ptr, mask=offs_n[None, :] < N, other=0.0).to(tl.float32) # (1, BN) # Load packed int4 weights for this group: (64, BN) -> two nibbles each. # Half-width 64 is a literal constant => the aranges are constexpr. wq = tl.load( wq_base + tl.arange(0, HALF)[:, None] * stride_wk, mask=(gs // 2 + tl.arange(0, HALF)[:, None]) < (K // 2), other=0, ) # (64, BN) uint8 lo = (wq & 0xF).to(tl.float32) # even-K nibble hi = ((wq >> 4) & 0xF).to(tl.float32) # odd-K nibble # Dequant weight halves (broadcast scale/zero over the group) w_lo = (lo - zero) * scale # (64, BN) w_hi = (hi - zero) * scale # (64, BN) # x columns for even rows (gs, gs+2, ...) and odd rows (gs+1, gs+3, ...) even_offs = gs + tl.arange(0, HALF) * 2 # length 64 odd_offs = gs + tl.arange(0, HALF) * 2 + 1 # length 64 x_even = tl.load( x_base + even_offs[None, :] * stride_xk, mask=(offs_m[:, None] < M) & (even_offs[None, :] < K), other=0.0, ) # (BM, 64) bf16 x_odd = tl.load( x_base + odd_offs[None, :] * stride_xk, mask=(offs_m[:, None] < M) & (odd_offs[None, :] < K), other=0.0, ) acc += tl.dot(x_even.to(tl.bfloat16), w_lo.to(tl.bfloat16)) acc += tl.dot(x_odd.to(tl.bfloat16), w_hi.to(tl.bfloat16)) tl.store( out_base, acc.to(tl.bfloat16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), ) 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 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", self._pack_int4(w_q)) self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16)) self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16)) @staticmethod 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() def forward(self, x: torch.Tensor) -> torch.Tensor: M, K = x.shape N = self.N out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) def grid(meta): return (triton.cdiv(M, meta["BM"]), triton.cdiv(N, meta["BN"])) _w4a16_kernel[grid]( x, self.w_q, self.scales, self.zeros, out, M, N, K, x.stride(0), x.stride(1), self.w_q.stride(0), self.w_q.stride(1), self.scales.stride(0), self.scales.stride(1), self.zeros.stride(0), self.zeros.stride(1), out.stride(0), out.stride(1), ) 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]