"""W4A16 weight-only int4 quantized GEMM for B200 (SM100, Blackwell).

Fused int4 unpack + per-group dequant + GEMM in Triton. AWQ/GPTQ-style asymmetric
int4 with explicit bf16 zero-points and per-group bf16 scales, group size 128.

Weight packing: w_q is (K//2, N) uint8; two int4 per byte, low nibble = even-K
row, high nibble = odd-K row.

Dequant (cvt-free via bf16 bitcast magic): zero-points are integers 0..15, so
  raw   = bitcast_bf16(0x4300 | nibble) = 128 + nibble   (exact)
  z_adj = bitcast_bf16(0x4300 | z)      = 128 + z         (exact)
  w_bf  = (raw - z_adj) * scale                            (exact integer diff)
The integer subtraction at magnitude ~135 is exact in bf16, so no fp cvt is
needed for the unpack -- only free bitcasts.
"""
from __future__ import annotations

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

GROUP_SIZE = 128


def _configs():
    cfgs = []
    for bm in (16, 32, 64, 128):
        for bn in (64, 128, 256):
            for nw in (4, 8):
                for ns in (3, 4):
                    cfgs.append(
                        triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": 128},
                                      num_warps=nw, num_stages=ns)
                    )
    return cfgs


@triton.autotune(configs=_configs(), key=["M", "N", "K"])
@triton.jit
def w4a16_gemm_kernel(
    x_ptr, wq_ptr, scales_ptr, zeros_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wkh, stride_wn,
    stride_sg, stride_sn,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_SIZE_K: tl.constexpr,
    GROUP_M: tl.constexpr = 8,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    rk = tl.arange(0, BLOCK_K)
    BK_HALF: tl.constexpr = BLOCK_K // 2
    rkh = tl.arange(0, BK_HALF)

    m_mask = rm < M
    n_mask = rn < N

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    x_rm = x_ptr + rm[:, None] * stride_xm
    w_rn = wq_ptr + rn[None, :] * stride_wn

    for ki in range(0, K, BLOCK_K):
        g_idx = ki // GROUP_SIZE_K

        x = tl.load(x_rm + (ki + rk)[None, :] * stride_xk, mask=m_mask[:, None], other=0.0)
        x2 = tl.reshape(x, (BLOCK_M, BK_HALF, 2))
        x_even, x_odd = tl.split(x2)

        kh = (ki // 2) + rkh
        w_packed = tl.load(w_rn + kh[:, None] * stride_wkh, mask=n_mask[None, :], other=0)
        raw_lo = tl.cast((w_packed & 0xF).to(tl.uint16) | 0x4300, tl.bfloat16, bitcast=True)
        raw_hi = tl.cast(((w_packed >> 4) & 0xF).to(tl.uint16) | 0x4300, tl.bfloat16, bitcast=True)

        z = tl.load(zeros_ptr + g_idx * stride_sg + rn * stride_sn, mask=n_mask, other=0.0)
        s = tl.load(scales_ptr + g_idx * stride_sg + rn * stride_sn, mask=n_mask, other=0.0)
        z_adj = tl.cast(z.to(tl.uint16) | 0x4300, tl.bfloat16, bitcast=True)

        w_lo = (raw_lo - z_adj[None, :]) * s[None, :]
        w_hi = (raw_hi - z_adj[None, :]) * s[None, :]

        acc += tl.dot(x_even, w_lo)
        acc += tl.dot(x_odd, w_hi)

    out_ptrs = out_ptr + rm[:, None] * stride_om + rn[None, :] * stride_on
    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


class Model(nn.Module):
    """W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""

    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        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:
        x = x.contiguous()
        M = x.shape[0]
        N, K = self.N, self.K
        out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
        grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
        w4a16_gemm_kernel[grid](
            x, self.w_q, self.scales, self.zeros, out,
            M, N, K,
            x.stride(0), x.stride(1),
            self.w_q.stride(0), self.w_q.stride(1),
            self.scales.stride(0), self.scales.stride(1),
            out.stride(0), out.stride(1),
            GROUP_SIZE_K=self.group_size,
        )
        return out


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]