dashq_kernel.py

"""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}"
        )