W4A16 GEMM MiniMax-M3

"""W4A16 weight-only quantized GEMM (Triton, fused unpack + matmul).

Scheme: AWQ/GPTQ-style asymmetric int4 with explicit zero-points and per-group
bf16 scales. Inputs:
  x:      (M, K)               bf16
  w_q:    (K//2, N)            uint8   -- two int4 weights packed per byte
                                        -- (low nibble = even-K row, high = odd)
  scales: (K//group, N)        bf16
  zeros:  (K//group, N)        bf16    -- stored as float zero-point
  out:    (M, N)               bf16

Dequant (per group along K):
  w_bf[k, n] = (w_q_unpacked[k, n] - zeros[k // group, n]) * scales[k // group, n]

We never call any prebuilt W4A16 kernel (marlin/bnb/etc.) - the unpack+matmul
is fused in one Triton kernel.
"""
from __future__ import annotations

import torch
import triton
import triton.language as tl

GROUP_SIZE = 128


# ---------------------------------------------------------------------------
# Triton kernel
# ---------------------------------------------------------------------------
# Tile sizes:
#   BLOCK_M  -- rows of the output per program (handles M=1 with masking)
#   BLOCK_N  -- columns of the output per program
#   BLOCK_K  -- K reduction block. We set this == GROUP_SIZE so each iteration
#               is one quant group, which lets us load scales/zeros once and
#               dequant a whole tile.
#
# Inner dot: x_even/w_bf_lo and x_odd/w_bf_hi are each (M, 64) @ (64, N), so
# the inner reduction is 64 (K-axis is the *post-split* K, i.e. BLOCK_K/2).
# Hopper WGMMA needs K in {16, 32, 64, ...}, so 64 is fine.
# ---------------------------------------------------------------------------
@triton.jit
def _w4a16_gemm_kernel(
    x_ptr, w_ptr, scales_ptr, zeros_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wk, stride_wn,
    stride_sk, stride_sn,
    stride_zk, stride_zn,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_SIZE: 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)

    HALF_K: tl.constexpr = BLOCK_K // 2
    n_groups: tl.constexpr = 0  # set at runtime via K

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Pre-build stride-scaled base pointers.
    x_row_base = x_ptr + offs_m[:, None] * stride_xm
    w_base = w_ptr + offs_n[None, :] * stride_wn

    n_groups = K // BLOCK_K
    for g in range(0, n_groups):
        k_start = g * BLOCK_K
        kp_start = g * HALF_K

        # ---- x tile (BLOCK_M, BLOCK_K) bf16, fully contiguous in K ----
        x = tl.load(
            x_row_base + (k_start + tl.arange(0, BLOCK_K))[None, :] * stride_xk,
            mask=offs_m[:, None] < M,
            other=0.0,
        )

        # ---- packed weights (HALF_K, BLOCK_N) uint8 ----
        w_packed = tl.load(
            w_base + (kp_start + tl.arange(0, HALF_K))[:, None] * stride_wk
        )

        # ---- unpack + dequant ----
        w_lo = (w_packed & 0xF).to(tl.bfloat16)
        w_hi = (w_packed >> 4).to(tl.bfloat16)

        scales = tl.load(scales_ptr + g * stride_sk + offs_n * stride_sn)
        zeros_ = tl.load(zeros_ptr + g * stride_zk + offs_n * stride_zn)

        w_bf_lo = (w_lo - zeros_[None, :]) * scales[None, :]
        w_bf_hi = (w_hi - zeros_[None, :]) * scales[None, :]

        # ---- split x into even/odd K rows via reshape+split ----
        # x of shape (M, BLOCK_K=128) -> reshape to (M, 64, 2)
        #   x_reshape[m, i, c] = x[m, 2i + c]  (c is the parity)
        # tl.split peels off the size-2 minor dim.
        x_r = tl.reshape(x, (BLOCK_M, HALF_K, 2))
        x_even, x_odd = tl.split(x_r)

        acc = tl.dot(x_even, w_bf_lo, acc)
        acc = tl.dot(x_odd, w_bf_hi, acc)

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


# ---------------------------------------------------------------------------
# Python launch wrapper
# ---------------------------------------------------------------------------
def _pick_block_m(M: int) -> int:
    # Tile size choices for the M dim.  Triton recompiles per BLOCK_M, so we
    # only include a small handful of sizes.  For tiny M (decode), pad to 16
    # so the WGMMA / mma instructions can use their M=16 m-tile.
    if M <= 16:
        return 16
    if M <= 32:
        return 32
    if M <= 64:
        return 64
    return 128


def w4a16_gemm(x: torch.Tensor, w_q: torch.Tensor,
               scales: torch.Tensor, zeros: torch.Tensor) -> torch.Tensor:
    """Fused W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""
    M, K = x.shape
    _, N = w_q.shape

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

    BLOCK_M = _pick_block_m(M)
    BLOCK_N = 128
    BLOCK_K = GROUP_SIZE  # == 128

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, 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),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_SIZE=GROUP_SIZE,
        num_warps=4,
        num_stages=3,
    )
    return out


# ---------------------------------------------------------------------------
# Model + get_inputs / get_init_inputs (must match reference.py interface).
# The buffers are registered so the harness's state_dict load (strict=True)
# works after the reference model has been built and quantized.
# ---------------------------------------------------------------------------
class Model(torch.nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K % group_size == 0
        assert K % 2 == 0
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size
        # Buffers are placeholders; the real values are loaded by the harness
        # from the reference model's state_dict. Names must match reference.py.
        self.register_buffer("w_q", torch.empty(K // 2, N, dtype=torch.uint8))
        self.register_buffer("scales", torch.empty(n_groups, N, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.empty(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)


M = 1
N = 12288
K = 4096


def get_inputs():
    x = torch.randn(M, K, dtype=torch.bfloat16)
    return [x]


def get_init_inputs():
    return [M, N, K]