W4A16 GEMM Claude Opus 4.8

"""W4A16 weight-only quantized GEMM (AWQ/GPTQ-style asym int4, group=128).

Fused unpack + GEMM in Triton. w_q is (K//2, N) uint8: low nibble = even-K row,
high nibble = odd-K row. We unpack + dequant inside the GEMM main loop and never
materialize the bf16 weight.

Each loop iteration processes exactly one quant group (128 contiguous K = 64
packed rows), so scale/zero are loaded once per group as (BLOCK_N,) row vectors
and broadcast. Activations are loaded contiguously (BLOCK_M, 128) then split into
even/odd K halves with tl.split, matching the low/high nibbles. Two K=64 sub-dots
accumulate into fp32.

Dispatch:
  * M == 1  : split-K kernel (occupancy-limited decode) -> fp32 scratch -> convert
  * M  > 1  : direct GEMM kernel, bf16 store
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl

GROUP_SIZE = 128


# ---------------------------------------------------------------------------
# Direct GEMM kernel (M > 1): bf16 store, BLOCK_K fixed to one group (128).
# ---------------------------------------------------------------------------
def _direct_configs():
    cfgs = []
    for bm, bn, w, s in [
        (16, 128, 4, 4), (16, 256, 8, 4), (16, 64, 4, 4),
        (32, 128, 4, 4), (32, 256, 8, 4), (32, 128, 8, 4),
        (64, 128, 4, 4), (64, 256, 8, 4), (64, 128, 8, 4),
        (128, 128, 8, 4), (128, 256, 8, 4), (64, 64, 4, 4),
    ]:
        cfgs.append(triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_warps=w, num_stages=s))
    return cfgs


@triton.autotune(configs=_direct_configs(), key=["M", "N", "K"])
@triton.jit
def _w4a16_direct(
    x_ptr, wq_ptr, scl_ptr, zero_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wq_kh, stride_wq_n,
    stride_scl_g, stride_scl_n,
    stride_zero_g, stride_zero_n,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    m_mask = offs_m < M
    n_mask = offs_n < N

    BLOCK_K2: tl.constexpr = 64  # 64 packed rows = 128 K = one group
    offs_k2 = tl.arange(0, BLOCK_K2)
    offs_k = tl.arange(0, 128)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    n_groups = K // 128

    x_base = x_ptr + offs_m[:, None] * stride_xm
    wq_base = wq_ptr + offs_n[None, :] * stride_wq_n
    sc_off = offs_n * stride_scl_n
    zr_off = offs_n * stride_zero_n

    for g in range(0, n_groups):
        kh = g * BLOCK_K2 + offs_k2
        packed = tl.load(wq_base + kh[:, None] * stride_wq_kh, mask=n_mask[None, :], other=0)
        lo = (packed & 0xF).to(tl.float32)
        hi = ((packed >> 4) & 0xF).to(tl.float32)

        scl = tl.load(scl_ptr + g * stride_scl_g + sc_off, mask=n_mask, other=0.0).to(tl.float32)
        zero = tl.load(zero_ptr + g * stride_zero_g + zr_off, mask=n_mask, other=0.0).to(tl.float32)
        w_lo = ((lo - zero[None, :]) * scl[None, :]).to(tl.bfloat16)
        w_hi = ((hi - zero[None, :]) * scl[None, :]).to(tl.bfloat16)

        k0 = g * 128
        x_tile = tl.load(x_base + (k0 + offs_k)[None, :] * stride_xk, mask=m_mask[:, None], other=0.0)
        x_e, x_o = tl.split(tl.reshape(x_tile, (BLOCK_M, BLOCK_K2, 2)))

        acc += tl.dot(x_e, w_lo, out_dtype=tl.float32)
        acc += tl.dot(x_o, w_hi, out_dtype=tl.float32)

    out = acc.to(tl.bfloat16)
    tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on,
             out, mask=m_mask[:, None] & n_mask[None, :])


# ---------------------------------------------------------------------------
# Split-K kernel (M == 1): fp32 atomic accumulation into scratch.
# ---------------------------------------------------------------------------
def _splitk_configs():
    cfgs = []
    for bn, sp, w, s in [
        (64, 1, 4, 4), (128, 1, 4, 4), (256, 1, 8, 4),
        (64, 2, 4, 4), (128, 2, 4, 4), (256, 2, 8, 4),
        (64, 4, 4, 4), (128, 4, 4, 4), (256, 4, 8, 4),
        (64, 8, 4, 4), (128, 8, 4, 4), (64, 4, 2, 3), (64, 8, 2, 3),
    ]:
        cfgs.append(triton.Config({"BLOCK_N": bn, "SPLIT_K": sp}, num_warps=w, num_stages=s))
    return cfgs


@triton.autotune(configs=_splitk_configs(), key=["M", "N", "K"])
@triton.jit
def _w4a16_splitk(
    x_ptr, wq_ptr, scl_ptr, zero_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wq_kh, stride_wq_n,
    stride_scl_g, stride_scl_n,
    stride_zero_g, stride_zero_n,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
):
    pid_n = tl.program_id(0)
    pid_k = tl.program_id(1)

    offs_m = tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    m_mask = offs_m < M
    n_mask = offs_n < N

    BLOCK_K2: tl.constexpr = 64
    offs_k2 = tl.arange(0, BLOCK_K2)
    offs_k = tl.arange(0, 128)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    n_groups = K // 128
    groups_per_split = (n_groups + SPLIT_K - 1) // SPLIT_K
    g_start = pid_k * groups_per_split
    g_end = min(g_start + groups_per_split, n_groups)

    x_base = x_ptr + offs_m[:, None] * stride_xm
    wq_base = wq_ptr + offs_n[None, :] * stride_wq_n
    sc_off = offs_n * stride_scl_n
    zr_off = offs_n * stride_zero_n

    for g in range(g_start, g_end):
        kh = g * BLOCK_K2 + offs_k2
        packed = tl.load(wq_base + kh[:, None] * stride_wq_kh, mask=n_mask[None, :], other=0)
        lo = (packed & 0xF).to(tl.float32)
        hi = ((packed >> 4) & 0xF).to(tl.float32)

        scl = tl.load(scl_ptr + g * stride_scl_g + sc_off, mask=n_mask, other=0.0).to(tl.float32)
        zero = tl.load(zero_ptr + g * stride_zero_g + zr_off, mask=n_mask, other=0.0).to(tl.float32)
        w_lo = ((lo - zero[None, :]) * scl[None, :]).to(tl.bfloat16)
        w_hi = ((hi - zero[None, :]) * scl[None, :]).to(tl.bfloat16)

        k0 = g * 128
        x_tile = tl.load(x_base + (k0 + offs_k)[None, :] * stride_xk, mask=m_mask[:, None], other=0.0)
        x_e, x_o = tl.split(tl.reshape(x_tile, (BLOCK_M, BLOCK_K2, 2)))

        acc += tl.dot(x_e, w_lo, out_dtype=tl.float32)
        acc += tl.dot(x_o, w_hi, out_dtype=tl.float32)

    out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
    full_mask = m_mask[:, None] & n_mask[None, :]
    if SPLIT_K == 1:
        tl.store(out_ptrs, acc, mask=full_mask)
    else:
        tl.atomic_add(out_ptrs, acc, mask=full_mask)


def _w4a16_gemm(x, wq, scales, zeros, M, N, K):
    if M == 1:
        scratch = torch.zeros((M, N), dtype=torch.float32, device=x.device)

        def grid(meta):
            return (triton.cdiv(N, meta["BLOCK_N"]), meta["SPLIT_K"])

        _w4a16_splitk[grid](
            x, wq, scales, zeros, scratch,
            M, N, K,
            x.stride(0), x.stride(1),
            wq.stride(0), wq.stride(1),
            scales.stride(0), scales.stride(1),
            zeros.stride(0), zeros.stride(1),
            scratch.stride(0), scratch.stride(1),
            BLOCK_M=16,
        )
        return scratch.to(torch.bfloat16)

    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)

    def grid(meta):
        return (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"]))

    _w4a16_direct[grid](
        x, wq, scales, zeros, out,
        M, N, K,
        x.stride(0), x.stride(1),
        wq.stride(0), wq.stride(1),
        scales.stride(0), scales.stride(1),
        zeros.stride(0), zeros.stride(1),
        out.stride(0), out.stride(1),
    )
    return out


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.zeros((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:
        M, K = x.shape
        N = self.N
        x = x.contiguous()
        return _w4a16_gemm(x, self.w_q, self.scales, self.zeros, M, N, K)