"""Fused W4A16 weight-only quantized GEMM. CUDA split-K GEMV for decode (M=1); Triton fused dequant+GEMM for prefill (M>1). Scheme matches reference.py: x: (M, K) bf16 w_q: (K//2, N) uint8 -- low nibble = even-K, high nibble = odd-K scales: (K//128, N) bf16 zeros: (K//128, N) bf16 out: (M, N) bf16 w_bf[k,n] = bf16((unpack(w_q)[k,n] - zeros[k//128,n]) * scales[k//128,n]) """ from __future__ import annotations from pathlib import Path import torch import torch.nn as nn import triton import triton.language as tl from torch.utils.cpp_extension import load_inline OP_TYPE = "gemm_w4a16" SUPPORTED_PRECISIONS = ["int4_bf16"] HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"] GROUP_SIZE = 128 # --------------------------------------------------------------------------- # CUDA decode (M=1): split-K high-occupancy GEMV # --------------------------------------------------------------------------- _CUDA_SRC = (Path(__file__).parent / "w4a16_cuda.cu").read_text() _CPP_SRC = ( "torch::Tensor w4a16_forward(torch::Tensor x, torch::Tensor w_q, " "torch::Tensor scales, torch::Tensor zeros);" ) _mod = None def _cuda(): global _mod if _mod is None: _mod = load_inline( name="w4a16_gemm_final", cpp_sources=_CPP_SRC, cuda_sources=_CUDA_SRC, functions=["w4a16_forward"], extra_cuda_cflags=[ "-O3", "--use_fast_math", "-lineinfo", "-gencode=arch=compute_120a,code=sm_120a", ], verbose=False, ) return _mod # --------------------------------------------------------------------------- # Triton prefill (M>1): fused unpack + dequant + tl.dot # --------------------------------------------------------------------------- @triton.autotune( configs=[ triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 16, "BLOCK_N": 256, "BLOCK_K": 64}, num_stages=2, num_warps=4), triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 32, "BLOCK_N": 256, "BLOCK_K": 64}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 64}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 128}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=2, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=3, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64}, num_stages=2, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128}, num_stages=2, num_warps=4), triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 128}, num_stages=2, num_warps=4), ], key=["M", "N", "K"], ) @triton.jit def _w4a16_gemm_kernel( x_ptr, w_ptr, s_ptr, z_ptr, out_ptr, M, N, K, stride_xm, stride_xk, stride_wk, stride_wn, stride_sg, stride_sn, stride_zg, stride_zn, stride_om, stride_on, GROUP_SIZE: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid = tl.program_id(0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) GROUP_M: tl.constexpr = 8 num_pid_in_group = GROUP_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_M group_size_m = min(num_pid_m - first_pid_m, GROUP_M) pid_m = first_pid_m + (pid % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for kt in range(tl.cdiv(K, BLOCK_K)): k_base = kt * BLOCK_K k_offs = k_base + offs_k x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk x = tl.load( x_ptrs, mask=(offs_m[:, None] < M) & (k_offs[None, :] < K), other=0.0, ) pk = k_offs // 2 w_ptrs = w_ptr + pk[:, None] * stride_wk + offs_n[None, :] * stride_wn w_mask = (k_offs[:, None] < K) & (offs_n[None, :] < N) w_packed = tl.load(w_ptrs, mask=w_mask, other=0) is_odd = (k_offs % 2)[:, None] w_i4 = tl.where(is_odd == 0, w_packed & 0xF, (w_packed >> 4) & 0xF).to(tl.float32) g_offs = k_offs // GROUP_SIZE s_ptrs = s_ptr + g_offs[:, None] * stride_sg + offs_n[None, :] * stride_sn z_ptrs = z_ptr + g_offs[:, None] * stride_zg + offs_n[None, :] * stride_zn scales = tl.load(s_ptrs, mask=w_mask, other=0.0).to(tl.float32) zeros = tl.load(z_ptrs, mask=w_mask, other=0.0).to(tl.float32) # bf16-round dequant matches reference numerics under large activations w_bf = ((w_i4 - zeros) * scales).to(tl.bfloat16) acc += tl.dot(x, w_bf) out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on tl.store( out_ptrs, acc.to(tl.bfloat16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), ) def _triton_gemm(x, w_q, scales, zeros, group_size): M, K = x.shape N = w_q.shape[1] out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) grid = lambda meta: ( triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["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), GROUP_SIZE=group_size, ) return out def w4a16_gemm(x, w_q, scales, zeros, group_size=GROUP_SIZE): x = x.contiguous() w_q = w_q.contiguous() scales = scales.contiguous() zeros = zeros.contiguous() if x.shape[0] == 1: return _cuda().w4a16_forward(x, w_q, scales, zeros) return _triton_gemm(x, w_q, scales, zeros, group_size) 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.ones(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: return w4a16_gemm(x, self.w_q, self.scales, self.zeros, self.group_size) M, N, K = 1, 12288, 4096 def get_inputs(): return [torch.randn(M, K, dtype=torch.bfloat16)] def get_init_inputs(): return [M, N, K]