"""Fused W4A16 weight-only int4 quantized GEMM for H100 (SM90 Hopper).

AWQ/GPTQ-style scheme:
  x:      (M, K)               bf16
  w_q:    (K // 2, N)          uint8   -- two int4 weights packed per byte (low nibble = even-K, high = odd-K)
  scales: (K // group, N)      bf16
  zeros:  (K // group, N)      bf16    -- asymmetric (stored already as float zero-point)
  out:    (M, N)                bf16

Approach:
  - M == 1 (decode): custom CUDA GEMV kernel via load_inline, maximizing HBM bandwidth
  - M > 1 (prefill/batched): Triton tiled GEMM kernel with fused dequant using tl.dot tensor cores
"""
from __future__ import annotations

import os

import torch
import torch.nn as nn
import triton
import triton.language as tl
from torch.utils.cpp_extension import load_inline

OP_TYPE = "gemm_w4a16"
SUPPORTED_PRECISIONS = ["int4_bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

GROUP_SIZE = 128

# ---------------------------------------------------------------------------
# CUDA GEMV kernel for M == 1 (decode, bandwidth-bound)
# ---------------------------------------------------------------------------

_gemv_cuda_source = r"""
#include <cuda_bf16.h>
#include <torch/extension.h>

#define GROUP_SIZE 128
#define HALF_GROUP 64
#define NUM_GROUPS 32

__global__ void w4a16_gemv_kernel(
    const __nv_bfloat16* __restrict__ x,       // (1, K) bf16
    const unsigned char* __restrict__ wq,      // (K/2, N) uint8 packed int4
    const __nv_bfloat16* __restrict__ scales,  // (K/128, N) bf16
    const __nv_bfloat16* __restrict__ zeros,   // (K/128, N) bf16
    __nv_bfloat16* __restrict__ out,           // (1, N) bf16
    int N
) {
    const int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;

    float acc = 0.0f;

    #pragma unroll
    for (int g = 0; g < NUM_GROUPS; ++g) {
        const int g_offset = g * N;
        float sc = __bfloat162float(scales[g_offset + n]);
        float zf = __bfloat162float(zeros[g_offset + n]);

        const int wq_base = g * HALF_GROUP * N + n;
        const int k_base = g * GROUP_SIZE;

        #pragma unroll 4
        for (int kk = 0; kk < HALF_GROUP; ++kk) {
            int k_even = k_base + 2 * kk;

            float x_even = __bfloat162float(x[k_even]);
            float x_odd  = __bfloat162float(x[k_even + 1]);

            unsigned char w_byte = __ldg(&wq[wq_base + kk * N]);

            float w_even_f = (float)(w_byte & 0x0F) - zf;
            float w_odd_f  = (float)((w_byte >> 4) & 0x0F) - zf;

            acc += x_even * w_even_f * sc;
            acc += x_odd  * w_odd_f  * sc;
        }
    }

    out[n] = __float2bfloat16(acc);
}

// Host-side launch wrapper (compiled with nvcc so <<<>>> works)
torch::Tensor gemv_forward(
    torch::Tensor x,
    torch::Tensor wq,
    torch::Tensor scales,
    torch::Tensor zeros
) {
    int N = wq.size(1);
    auto out = torch::empty({1, N}, x.options().dtype(torch::kBFloat16));

    const int BLOCK_N = 256;
    const int grid = (N + BLOCK_N - 1) / BLOCK_N;

    w4a16_gemv_kernel<<<grid, BLOCK_N>>>(
        reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
        reinterpret_cast<const unsigned char*>(wq.data_ptr()),
        reinterpret_cast<const __nv_bfloat16*>(scales.data_ptr()),
        reinterpret_cast<const __nv_bfloat16*>(zeros.data_ptr()),
        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
        N
    );

    return out;
}
"""

_gemv_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor gemv_forward(
    torch::Tensor x,
    torch::Tensor wq,
    torch::Tensor scales,
    torch::Tensor zeros
);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("gemv_forward", &gemv_forward, "W4A16 GEMV forward");
}
"""

# Build the CUDA extension (cached after first build)
_gemv_module = None

def _get_gemv_module():
    global _gemv_module
    if _gemv_module is not None:
        return _gemv_module

    # Write a marker file so the cache distinguishes different CUDA versions
    extra_cuda_cflags = [
        "-O3",
        "--use_fast_math",
        "-Xptxas", "-O3",
        "-Xptxas", "-maxrregcount=64",
        "--extended-lambda",
    ]

    _gemv_module = load_inline(
        name="w4a16_gemv_cuda",
        cpp_sources=[_gemv_cpp_source],
        cuda_sources=[_gemv_cuda_source],
        extra_cuda_cflags=extra_cuda_cflags,
        verbose=False,
    )
    return _gemv_module


# ---------------------------------------------------------------------------
# Triton GEMM kernel for M > 1 (prefill / batched)
# ---------------------------------------------------------------------------

@triton.jit
def _w4a16_gemm_kernel(
    x_ptr,
    wq_ptr,
    scales_ptr,
    zeros_ptr,
    out_ptr,
    M,
    N,
    stride_x_m,
    stride_x_k,
    stride_wq_kh,
    stride_wq_n,
    stride_sc_g,
    stride_sc_n,
    stride_z_g,
    stride_z_n,
    stride_out_m,
    stride_out_n,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP_M: tl.constexpr,
    GROUP: tl.constexpr,
    HALF_GROUP: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
):
    """Tiled GEMM with fused int4 dequant for W4A16 (M > 1).

    For each group of GROUP=128 K positions:
      1. Load scales & zeros.
      2. Load packed uint8 weights, extract lo/hi nibbles.
      3. Dequantize both halves.
      4. Load activations for even and odd K positions separately.
      5. Two tl.dot calls (even K and odd K), accumulated in fp32.

    This avoids interleaving the unpacked K dimension into a single tensor,
    which is awkward in Triton.  Each tl.dot is
        (BLOCK_M, HALF_GROUP) @ (HALF_GROUP, BLOCK_N).
    """
    pid = tl.program_id(0)

    # L2-friendly swizzling (standard Triton GEMM pattern)
    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

    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

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

    for g in range(NUM_GROUPS):
        # --- scale & zero for this group ---
        sc = tl.load(scales_ptr + g * stride_sc_g + offs_n[None, :] * stride_sc_n,
                     mask=n_mask[None, :], other=0.0)
        z = tl.load(zeros_ptr + g * stride_z_g + offs_n[None, :] * stride_z_n,
                    mask=n_mask[None, :], other=0.0)
        sc_f32 = sc.to(tl.float32)
        z_f32 = z.to(tl.float32)

        # --- packed weights: (HALF_GROUP, BLOCK_N) uint8 ---
        offs_packed_k = g * HALF_GROUP + tl.arange(0, HALF_GROUP)
        wq_ptrs = wq_ptr + offs_packed_k[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
        w_packed = tl.load(wq_ptrs, mask=n_mask[None, :], other=0)

        # --- unpack nibbles ---
        w_lo = (w_packed & 0xF).to(tl.float32)
        w_hi = ((w_packed >> 4) & 0xF).to(tl.float32)

        # --- dequantize ---
        w_lo_deq = (w_lo - z_f32) * sc_f32
        w_hi_deq = (w_hi - z_f32) * sc_f32

        # --- activations for even K positions ---
        k_even_start = g * GROUP
        offs_k_even = k_even_start + tl.arange(0, HALF_GROUP) * 2
        a_ptrs_even = x_ptr + offs_m[:, None] * stride_x_m + offs_k_even[None, :] * stride_x_k
        a_even = tl.load(a_ptrs_even, mask=m_mask[:, None], other=0.0)

        # --- activations for odd K positions ---
        offs_k_odd = k_even_start + tl.arange(0, HALF_GROUP) * 2 + 1
        a_ptrs_odd = x_ptr + offs_m[:, None] * stride_x_m + offs_k_odd[None, :] * stride_x_k
        a_odd = tl.load(a_ptrs_odd, mask=m_mask[:, None], other=0.0)

        # --- tensor-core matmul (two halves, avoiding K-interleaving) ---
        acc += tl.dot(a_even, w_lo_deq.to(tl.bfloat16))
        acc += tl.dot(a_odd, w_hi_deq.to(tl.bfloat16))

    # --- store output ---
    out_ptrs = out_ptr + offs_m[:, None] * stride_out_m + offs_n[None, :] * stride_out_n
    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


# ---------------------------------------------------------------------------
# Packing utility
# ---------------------------------------------------------------------------

def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    """Pack (K, N) uint8 in [0,15] into (K//2, N) uint8.

    Even rows go in the low nibble, odd rows in the high nibble.
    """
    K, N = w_q.shape
    assert K % 2 == 0
    lo = w_q[0::2].to(torch.uint8) & 0xF
    hi = w_q[1::2].to(torch.uint8) & 0xF
    return (lo | (hi << 4)).contiguous()


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

M = 1
N = 12288
K = 4096


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

    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K % group_size == 0, "K must be divisible by group_size"
        assert K % 2 == 0, "K must be even (int4 packing)"
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size

        # Synthetic quant: reproducible random weight, per-group asym quant
        torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
        w_full = torch.randn(K, N, dtype=torch.float32) * 0.02

        w_g = w_full.view(n_groups, group_size, N)
        w_min = w_g.min(dim=1, keepdim=True).values   # (n_groups, 1, N)
        w_max = w_g.max(dim=1, keepdim=True).values
        scales = (w_max - w_min).clamp_min(1e-8) / 15.0
        zeros = (-w_min / scales).round().clamp(0, 15)

        w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8)
        w_q = w_q.view(K, N)

        scales_2d = scales.squeeze(1).to(torch.bfloat16)   # (n_groups, N)
        zeros_2d = zeros.squeeze(1).to(torch.bfloat16)
        w_packed = _pack_int4(w_q)                          # (K//2, N)

        self.register_buffer("w_q", w_packed)     # uint8
        self.register_buffer("scales", scales_2d) # bf16
        self.register_buffer("zeros", zeros_2d)   # bf16

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.to(torch.bfloat16)
        M_val, N_val = x.shape[0], self.N

        if M_val == 1:
            # --- Decode: CUDA GEMV (memory-bandwidth-optimized) ---
            mod = _get_gemv_module()
            return mod.gemv_forward(
                x.contiguous(),
                self.w_q.contiguous(),
                self.scales.contiguous(),
                self.zeros.contiguous(),
            )

        # --- Prefill / batched: Triton tiled GEMM (tensor cores via tl.dot) ---
        num_groups = self.K // self.group_size
        half_group = self.group_size // 2

        BLOCK_M = min(M_val, 64)
        BLOCK_N = 128
        GROUP_M = 8
        grid = (triton.cdiv(M_val, BLOCK_M) * triton.cdiv(N_val, BLOCK_N),)
        out = torch.empty(M_val, N_val, dtype=torch.bfloat16, device=x.device)
        _w4a16_gemm_kernel[grid](
            x,
            self.w_q,
            self.scales,
            self.zeros,
            out,
            M_val,
            N_val,
            x.stride(0),
            x.stride(1) if x.ndim > 1 and x.stride(1) else 1,
            self.w_q.stride(0),
            self.w_q.stride(1),
            self.scales.stride(0),
            self.scales.stride(1),
            self.zeros.stride(0),
            self.zeros.stride(1),
            out.stride(0),
            out.stride(1),
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
            GROUP_M=GROUP_M,
            GROUP=self.group_size,
            HALF_GROUP=half_group,
            NUM_GROUPS=num_groups,
        )
        return out


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


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