"""W4A16 weight-only quantized GEMM — baseline (fused kernels to follow).

AWQ/GPTQ-style asymmetric int4, group_size=128 along K.
  w_bf[k, n] = (unpack(w_q)[k, n] - zeros[k // 128, n]) * scales[k // 128, n]
  y = x @ w_bf   (x bf16, y bf16)
"""
from __future__ import annotations

import torch
import torch.nn as nn

GROUP_SIZE = 128


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
        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.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:
        K = self.K
        w = self.w_q
        w_unpacked = torch.empty((K, self.N), dtype=torch.uint8, device=w.device)
        w_unpacked[0::2] = w & 0xF
        w_unpacked[1::2] = (w >> 4) & 0xF
        g = self.group_size
        w_g = w_unpacked.view(K // g, g, self.N).to(torch.bfloat16)
        w_bf = (w_g - self.zeros.unsqueeze(1)) * self.scales.unsqueeze(1)
        w_bf = w_bf.view(K, self.N)
        return x.to(torch.bfloat16) @ w_bf


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]