"""Triton kernels for group-wise asymmetric integer weights. Weights are stored K-major: W_q has shape (K // elements_per_word, N) with values packed along K, and scale/zero have shape (K // group_size, N). Three kernels are selected by the number of input rows M: M == 1 GEMV 2 <= M <= 32 fused dequantize-GEMM with split-K M > 32 fused dequantize-GEMM, accumulator kept in registers Supported bit widths are 1, 2, 3, 4 and 8. 3-bit is stored as a 2-bit plane plus a 1-bit plane, so it occupies exactly 3 bits per weight. """ from __future__ import annotations from typing import Optional import torch import torch.nn as nn try: import triton import triton.language as tl TRITON_AVAILABLE = True except Exception: # triton is optional TRITON_AVAILABLE = False SUPPORTED_NBITS = (1, 2, 3, 4, 8) _GEMM_CONFIG_CACHE = {} if TRITON_AVAILABLE: @triton.jit def _dashq_gemv_kernel( x_ptr, w_ptr, lo_ptr, s_ptr, z_ptr, y_ptr, N, K, NBITS: tl.constexpr, EPS: tl.constexpr, GS: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): pid_n = tl.program_id(0) pid_k = tl.program_id(1) * 2 offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_N), BLOCK_N) # 2 * BLOCK_K == GS, so a program covers exactly one scale group. k_m = (pid_k * BLOCK_K) // GS scales = tl.load(s_ptr + k_m * N + offs_n).to(tl.float32) zeros = tl.load(z_ptr + k_m * N + offs_n).to(tl.float32) acc = tl.zeros((BLOCK_N,), dtype=tl.float32) offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) for _ in tl.static_range(2): a = tl.load(x_ptr + offs_k, eviction_policy="evict_last").to(tl.float32) if NBITS == 3: hw = tl.load( w_ptr + (offs_k // 16)[:, None] * N + offs_n[None, :], eviction_policy="evict_first", ) lw = tl.load( lo_ptr + (offs_k // 32)[:, None] * N + offs_n[None, :], eviction_policy="evict_first", ) q = (((hw >> (((offs_k % 16) * 2)[:, None])) & 3) << 1) | ( (lw >> ((offs_k % 32)[:, None])) & 1 ) else: wv = tl.load( w_ptr + (offs_k // EPS)[:, None] * N + offs_n[None, :], eviction_policy="evict_first", ) q = (wv >> (((offs_k % EPS) * NBITS)[:, None])) & ((1 << NBITS) - 1) b = (q.to(tl.float32) - zeros[None, :]) * scales[None, :] acc += tl.sum(a[:, None] * b, axis=0) offs_k += BLOCK_K tl.atomic_add(y_ptr + offs_n, acc, sem="relaxed") @triton.jit def _dashq_gemm_kernel( x_ptr, w_ptr, lo_ptr, s_ptr, z_ptr, y_ptr, M, N, K, NBITS: tl.constexpr, EPS: tl.constexpr, GS: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SPLIT_K: tl.constexpr, OUT_DTYPE: tl.constexpr, ): """y[M, N] = x[M, K] @ dequantize(w)[K, N] BLOCK_K divides the group size, so a K-tile lies inside one group and the scale/zero load is a single (1, BLOCK_N) vector. """ pid_m = tl.program_id(0) pid_n = 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) mask_m = offs_m < M mask_n = offs_n < N acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for t in range(pid_k, tl.cdiv(K, BLOCK_K), SPLIT_K): k0 = t * BLOCK_K offs_k = k0 + tl.arange(0, BLOCK_K) mask_k = offs_k < K x = tl.load(x_ptr + offs_m[:, None] * K + offs_k[None, :], mask=mask_m[:, None] & mask_k[None, :], other=0.0) if NBITS == 3: hw = tl.load(w_ptr + (offs_k // 16)[:, None] * N + offs_n[None, :], mask=mask_k[:, None] & mask_n[None, :], other=0) lw = tl.load(lo_ptr + (offs_k // 32)[:, None] * N + offs_n[None, :], mask=mask_k[:, None] & mask_n[None, :], other=0) q = (((hw >> (((offs_k % 16) * 2)[:, None])) & 3) << 1) | ( (lw >> ((offs_k % 32)[:, None])) & 1) else: wv = tl.load(w_ptr + (offs_k // EPS)[:, None] * N + offs_n[None, :], mask=mask_k[:, None] & mask_n[None, :], other=0) q = (wv >> (((offs_k % EPS) * NBITS)[:, None])) & ((1 << NBITS) - 1) g = k0 // GS s = tl.load(s_ptr + g * N + offs_n, mask=mask_n, other=0.0).to(tl.float32) z = tl.load(z_ptr + g * N + offs_n, mask=mask_n, other=0.0).to(tl.float32) w = (q.to(tl.float32) - z[None, :]) * s[None, :] acc += tl.dot(x, w.to(x.dtype), out_dtype=tl.float32) out = acc.to(OUT_DTYPE) y_ptrs = y_ptr + offs_m[:, None] * N + offs_n[None, :] if SPLIT_K == 1: tl.store(y_ptrs, out, mask=mask_m[:, None] & mask_n[None, :]) else: tl.atomic_add(y_ptrs, out, mask=mask_m[:, None] & mask_n[None, :], sem="relaxed") def _pack_kmajor(q_kn: torch.Tensor, bits: int) -> torch.Tensor: """(K, N) codes -> (K // eps, N) int32, value k in word k // eps.""" K, N = q_kn.shape eps = 32 // bits v = q_kn.to(torch.int32).reshape(K // eps, eps, N) words = torch.zeros(K // eps, N, dtype=torch.int32, device=q_kn.device) for j in range(eps): words |= v[:, j, :] << (bits * j) return words def _unpack_kmajor(words: torch.Tensor, bits: int, K: int) -> torch.Tensor: eps = 32 // bits WK, N = words.shape shifts = (torch.arange(eps, device=words.device, dtype=torch.int32) * bits).view(1, eps, 1) q = (words.view(WK, 1, N) >> shifts) & ((1 << bits) - 1) return q.reshape(WK * eps, N)[:K] class TritonQuantLinear(nn.Module): """Linear layer over group-wise asymmetric integer weights.""" def __init__( self, W_int: torch.Tensor, # (out_features, in_features) integer codes scale: torch.Tensor, # (out_features, num_groups) zero: torch.Tensor, # (out_features, num_groups) nbits: int, group_size: int, bias: Optional[torch.Tensor] = None, out_dtype: torch.dtype = torch.float16, block_n: int = 128, num_warps: int = 1, ) -> None: super().__init__() if not TRITON_AVAILABLE: raise RuntimeError("Triton is not available.") if nbits not in SUPPORTED_NBITS: raise ValueError(f"Unsupported nbits: {nbits}") out_features, in_features = W_int.shape if in_features % group_size != 0: raise ValueError("in_features must be divisible by group_size.") if group_size % 2 != 0: raise ValueError("group_size must be even.") self.out_features = out_features self.in_features = in_features self.nbits = int(nbits) self.group_size = int(group_size) self.out_dtype = out_dtype self.block_n = int(block_n) self.num_warps = int(num_warps) self.block_k = self.group_size // 2 q_kn = W_int.t().contiguous().to(torch.uint8) if nbits == 3: self.register_buffer("W_q", _pack_kmajor(q_kn >> 1, 2)) self.register_buffer("W_lo", _pack_kmajor(q_kn & 1, 1)) self.eps = 16 else: self.register_buffer("W_q", _pack_kmajor(q_kn, nbits)) self.register_buffer("W_lo", torch.zeros(1, dtype=torch.int32, device=q_kn.device)) self.eps = 32 // nbits del q_kn self.register_buffer("scale", scale.t().contiguous().to(out_dtype)) self.register_buffer("zero", zero.t().contiguous().to(out_dtype)) if bias is not None: self.register_buffer("bias", bias.detach().clone().to(out_dtype)) else: self.bias = None # The GEMV accumulates with atomics, so it starts from the bias. acc_init = torch.zeros(out_features, dtype=torch.float32, device=self.W_q.device) if bias is not None: acc_init.copy_(self.bias.float()) self.register_buffer("_acc_init", acc_init) self.register_buffer("_acc", acc_init.clone()) self._grid = ( (out_features + self.block_n - 1) // self.block_n, in_features // self.group_size, ) def dequantize_weight(self, dtype: torch.dtype) -> torch.Tensor: """Returns W^T with shape (in_features, out_features).""" if self.nbits == 3: q = (_unpack_kmajor(self.W_q, 2, self.in_features).to(torch.int32) << 1) | ( _unpack_kmajor(self.W_lo, 1, self.in_features).to(torch.int32) ) else: q = _unpack_kmajor(self.W_q, self.nbits, self.in_features) s = self.scale.repeat_interleave(self.group_size, dim=0).to(dtype) z = self.zero.repeat_interleave(self.group_size, dim=0).to(dtype) return (q.to(dtype) - z) * s # (BLOCK_M, BLOCK_N, SPLIT_K, num_warps, num_stages), largest tile first; # the first entry that fits in shared memory is cached per shape. _SMALL_M_CONFIGS = ((16, 64, 8, 4, 2), (16, 64, 4, 4, 1)) _LARGE_M_CONFIGS = ((128, 128, 1, 8, 4), (128, 128, 1, 8, 3), (128, 64, 1, 4, 3), (64, 64, 1, 4, 2)) def _gemm(self, x2d: torch.Tensor) -> torch.Tensor: M = x2d.shape[0] N, K, gs = self.out_features, self.in_features, self.group_size block_k = min(gs, 32) configs = self._SMALL_M_CONFIGS if M <= 32 else self._LARGE_M_CONFIGS cache_key = (M <= 32, N, K, gs, self.nbits) if cache_key in _GEMM_CONFIG_CACHE: configs = (_GEMM_CONFIG_CACHE[cache_key],) tl_dtype = tl.float16 if self.out_dtype == torch.float16 else tl.bfloat16 last_err = None for cfg in configs: block_m, block_n, split_k, warps, stages = cfg split_k = min(split_k, max(1, K // block_k)) alloc = torch.empty if split_k == 1 else torch.zeros y = alloc(M, N, dtype=self.out_dtype, device=x2d.device) grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n), split_k) try: _dashq_gemm_kernel[grid]( x2d, self.W_q, self.W_lo, self.scale, self.zero, y, M, N, K, self.nbits, self.eps, gs, block_m, block_n, block_k, split_k, tl_dtype, num_warps=warps, num_stages=stages, ) except triton.runtime.errors.OutOfResources as exc: last_err = exc continue _GEMM_CONFIG_CACHE[cache_key] = cfg return y raise last_err def forward(self, x: torch.Tensor) -> torch.Tensor: shape = x.shape tokens = x.numel() // shape[-1] if tokens == 1 and x.is_cuda: self._acc.copy_(self._acc_init) _dashq_gemv_kernel[self._grid]( x.reshape(-1), self.W_q, self.W_lo, self.scale, self.zero, self._acc, self.out_features, self.in_features, self.nbits, self.eps, self.group_size, self.block_n, self.block_k, num_warps=self.num_warps, ) return self._acc.to(x.dtype).reshape(*shape[:-1], self.out_features) x2d = x.reshape(tokens, -1) if x.is_cuda and TRITON_AVAILABLE: out = self._gemm(x2d) else: out = x2d @ self.dequantize_weight(x.dtype) if self.bias is not None: out = out + self.bias.to(out.dtype) return out.to(x.dtype).reshape(*shape[:-1], self.out_features) def extra_repr(self) -> str: return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"nbits={self.nbits}, group_size={self.group_size}" )