dashq_kernel.py

"""Triton weight-only GEMV backend for DASH-Q packed checkpoints.

Group-wise asymmetric integer weights (the format DASH-Q emits) are stored
K-major so that a decode-time GEMV reads each packed word exactly once with
fully coalesced loads:

    W_q : (K // elements_per_word, N) int32   (packed along K, N contiguous)
    s,z : (K // group_size, N)                (one group per program)

Supported bit widths: 2, 3, 4, 8 (and 1). 3-bit uses two bit-planes -- a
2-bit plane plus a 1-bit plane -- which is exactly 3 bits per weight and is
not covered by existing kernel libraries.

Batched inputs (prefill) fall back to an unpack-and-matmul path that uses the
same K-major buffers, so the original torch buffers can be released.
"""
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:  # pragma: no cover - triton is an optional dependency
    TRITON_AVAILABLE = False

SUPPORTED_NBITS = (1, 2, 3, 4, 8)


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)

        # one scale/zero group per program (2 * BLOCK_K == GS)
        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")


def _pack_kmajor(q_kn: torch.Tensor, bits: int) -> torch.Tensor:
    """(K, N) uint8 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):
    """Decode-optimized replacement for a DASH-Q PackedQuantizedLinear."""

    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 for the Triton backend: {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

        # accumulator is seeded with the bias, so the kernel never adds it
        # (each K-split program contributes once via atomic_add)
        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 as (in_features, out_features), matching the K-major layout."""
        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

    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)

        w_t = self.dequantize_weight(x.dtype)
        out = x.reshape(tokens, -1) @ w_t
        if self.bias is not None:
            out = out + self.bias.to(x.dtype)
        return out.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}, backend=triton"
        )