"""W4A16 weight-only int4 quantized GEMM, fused unpack+GEMM, B200/SM100.

Scheme: AWQ/GPTQ asymmetric int4, per-group (128 along K) bf16 scales/zeros.
  x:(M,K) bf16, w_q:(K//2,N) uint8 (lo nibble=even-K, hi=odd-K),
  scales/zeros:(K//128,N) bf16, out:(M,N) bf16.
  w[k,n]=(unpack(w_q)[k,n]-zeros[k//128,n])*scales[k//128,n]

M==1 (decode): custom CUDA GEMV (single kernel, cooperative warp reduction,
direct bf16 output). M>1: Triton tiled GEMM with fused dequant.
"""
from __future__ import annotations

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

GROUP_SIZE = 128

# ---------------------------------------------------------------------------
# CUDA GEMV for M == 1
# ---------------------------------------------------------------------------
_CUDA_SRC = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <c10/cuda/CUDAStream.h>

// Vectorized split-K GEMV. wq viewed as uint32 (K2, N/4). Each thread owns CPT
// columns; shared-staged x per group; split-K via grid.y with atomic accum.
// Zero-point and scale are factored out to once-per-group: the hot loop does
// only lo*xe + hi*xo (2 FMA/byte), then y += s*(raw - z*sum_x) per group.
template<int BLK, int SPLITK, int CPT>
__global__ void gemv_kernel(const __nv_bfloat16* __restrict__ x,
                            const uint32_t* __restrict__ wq,
                            const __nv_bfloat16* __restrict__ scales,
                            const __nv_bfloat16* __restrict__ zeros,
                            float* __restrict__ out, int N, int K) {
    const int K2=K/2, N4=N/4, CU=CPT/4;
    int cu0=(blockIdx.x*BLK+threadIdx.x)*CU;
    int col0=cu0*4;
    int gps=(K2/64)/SPLITK, g0=blockIdx.y*gps;
    __shared__ float xs[128];
    float acc[CPT];
    #pragma unroll
    for(int c=0;c<CPT;c++) acc[c]=0.f;
    for(int g=g0; g<g0+gps; g++){
        int kbase=g*128;
        for(int i=threadIdx.x;i<128;i+=BLK) xs[i]=__bfloat162float(x[kbase+i]);
        __syncthreads();
        float s[CPT], z[CPT], ge[CPT];
        #pragma unroll
        for(int c=0;c<CPT;c++){ s[c]=__bfloat162float(scales[g*N+col0+c]); z[c]=__bfloat162float(zeros[g*N+col0+c]); ge[c]=0.f; }
        int kh0=g*64; float xsum=0.f;
        #pragma unroll 8
        for(int kk=0;kk<64;kk++){
            float xe=xs[2*kk], xo=xs[2*kk+1]; xsum+=xe+xo;
            #pragma unroll
            for(int u=0;u<CU;u++){
                uint32_t w=wq[(kh0+kk)*N4+cu0+u];
                #pragma unroll
                for(int bb=0;bb<4;bb++){
                    int c=u*4+bb;
                    uint8_t b=(w>>(bb*8))&0xFF;
                    ge[c]+= (float)(b&0xF)*xe + (float)(b>>4)*xo;
                }
            }
        }
        #pragma unroll
        for(int c=0;c<CPT;c++) acc[c]+= s[c]*(ge[c]-z[c]*xsum);
        __syncthreads();
    }
    #pragma unroll
    for(int c=0;c<CPT;c++) atomicAdd(&out[col0+c],acc[c]);
}

template<int BLK, int CPT>
static void launch(int splitk, dim3 grid, cudaStream_t stream,
                   const __nv_bfloat16* x, const uint32_t* wq,
                   const __nv_bfloat16* s, const __nv_bfloat16* z, float* out, int N, int K){
    #define LK(SK) gemv_kernel<BLK,SK,CPT><<<grid,BLK,0,stream>>>(x,wq,s,z,out,N,K)
    if(splitk==4) LK(4); else if(splitk==8) LK(8); else if(splitk==16) LK(16); else LK(32);
    #undef LK
}

torch::Tensor gemv(torch::Tensor x, torch::Tensor wq32, torch::Tensor scales,
                   torch::Tensor zeros, int N, int K, int splitk, int cpt, int blk){
    auto out=torch::zeros({1,N}, x.options().dtype(torch::kFloat32));
    auto stream=c10::cuda::getCurrentCUDAStream();
    const auto* xp=(const __nv_bfloat16*)x.data_ptr();
    const auto* wp=(const uint32_t*)wq32.data_ptr();
    const auto* sp=(const __nv_bfloat16*)scales.data_ptr();
    const auto* zp=(const __nv_bfloat16*)zeros.data_ptr();
    float* op=out.data_ptr<float>();
    int cols_per_block = blk*cpt;
    dim3 grid((N + cols_per_block - 1)/cols_per_block, splitk);
    if(blk==64){
        if(cpt==4) launch<64,4>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
        else if(cpt==8) launch<64,8>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
        else launch<64,16>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
    } else {
        if(cpt==4) launch<128,4>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
        else if(cpt==8) launch<128,8>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
        else launch<128,16>(splitk,grid,stream,xp,wp,sp,zp,op,N,K);
    }
    return out.to(torch::kBFloat16);
}
'''
_CUDA_CPP = "torch::Tensor gemv(torch::Tensor x, torch::Tensor wq32, torch::Tensor scales, torch::Tensor zeros, int N, int K, int splitk, int cpt, int blk);"

_cuda_mod = None
def _get_cuda():
    global _cuda_mod
    if _cuda_mod is None:
        _cuda_mod = load_inline(
            name="w4a16_gemv", cpp_sources=_CUDA_CPP, cuda_sources=_CUDA_SRC,
            functions=["gemv"], extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100a"],
            verbose=False,
        )
    return _cuda_mod


# ---------------------------------------------------------------------------
# Triton tiled GEMM for M > 1
# ---------------------------------------------------------------------------
def _cfg(BM, BN, SK, w, s):
    return triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "SPLIT_K": SK}, num_warps=w, num_stages=s)

@triton.autotune(
    configs=[
        # small-M (bandwidth bound): split-K + small tiles
        _cfg(16, 64, 4, 4, 3), _cfg(16, 64, 8, 4, 3), _cfg(16, 128, 4, 4, 3),
        _cfg(16, 64, 16, 4, 3), _cfg(16, 64, 2, 4, 4), _cfg(16, 32, 16, 4, 3),
        _cfg(32, 64, 2, 4, 4), _cfg(32, 64, 4, 4, 3), _cfg(32, 64, 8, 4, 3),
        _cfg(32, 64, 16, 4, 3), _cfg(32, 128, 2, 4, 3), _cfg(32, 64, 1, 4, 3),
        _cfg(32, 32, 16, 4, 3),
        # large-M (compute bound): bigger tiles, no split-K
        _cfg(64, 64, 1, 4, 3), _cfg(128, 64, 1, 4, 3), _cfg(128, 128, 1, 8, 3),
        _cfg(64, 128, 1, 8, 3), _cfg(128, 256, 1, 8, 4), _cfg(256, 64, 1, 8, 3),
        _cfg(128, 64, 1, 8, 4), _cfg(128, 128, 1, 8, 4),
    ],
    key=["M", "N", "K"],
    reset_to_zero=["out_ptr"],
)
@triton.jit
def _gemm_kernel(
    xe_ptr, xo_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
    M, N, K2,
    stride_xm, stride_xk, stride_wk, stride_wn,
    stride_sg, stride_sn, stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
    SPLIT_K: tl.constexpr, GH: tl.constexpr,
):
    pid_n = tl.program_id(0)
    pid_m = tl.program_id(1)
    pid_k = tl.program_id(2)
    n_groups = K2 // GH
    g_per = n_groups // SPLIT_K
    g0 = pid_k * g_per
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_kh = tl.arange(0, GH)
    m_mask = offs_m < M
    n_mask = offs_n < N
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for g in range(g0, g0 + g_per):
        kh = g * GH + offs_kh
        wq = tl.load(wq_ptr + kh[:, None] * stride_wk + offs_n[None, :] * stride_wn,
                     mask=n_mask[None, :], other=0)
        lo = (wq & 0xF).to(tl.bfloat16)
        hi = ((wq >> 4) & 0xF).to(tl.bfloat16)
        z = tl.load(z_ptr + g * stride_sg + offs_n * stride_sn, mask=n_mask, other=0.0).to(tl.bfloat16)
        s = tl.load(s_ptr + g * stride_sg + offs_n * stride_sn, mask=n_mask, other=0.0).to(tl.float32)
        lo = lo - z[None, :]
        hi = hi - z[None, :]
        xe = tl.load(xe_ptr + offs_m[:, None] * stride_xm + kh[None, :] * stride_xk, mask=m_mask[:, None], other=0.0)
        xo = tl.load(xo_ptr + offs_m[:, None] * stride_xm + kh[None, :] * stride_xk, mask=m_mask[:, None], other=0.0)
        acc += s[None, :] * (tl.dot(xe, lo) + tl.dot(xo, hi))
    om = offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
    mask = m_mask[:, None] & n_mask[None, :]
    if SPLIT_K == 1:
        tl.store(out_ptr + om, acc, mask=mask)
    else:
        tl.atomic_add(out_ptr + om, acc, mask=mask)


def _gemm(x, w_q, scales, zeros, N, K):
    M = x.shape[0]
    K2 = K // 2
    xe = x[:, 0::2].contiguous()
    xo = x[:, 1::2].contiguous()
    out = torch.zeros((M, N), dtype=torch.float32, device=x.device)
    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), triton.cdiv(M, meta["BLOCK_M"]), meta["SPLIT_K"])
    _gemm_kernel[grid](
        xe, xo, w_q, scales, zeros, out,
        M, N, K2,
        xe.stride(0), xe.stride(1), w_q.stride(0), w_q.stride(1),
        scales.stride(0), scales.stride(1), out.stride(0), out.stride(1),
        GH=GROUP_SIZE // 2,
    )
    return out.to(torch.bfloat16)


def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    lo = w_q[0::2].to(torch.uint8) & 0xF
    hi = w_q[1::2].to(torch.uint8) & 0xF
    return (lo | (hi << 4)).contiguous()


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

        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
        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).view(K, N)
        scales_2d = scales.squeeze(1).to(torch.bfloat16)
        zeros_2d = zeros.squeeze(1).to(torch.bfloat16)
        w_packed = _pack_int4(w_q)

        self.register_buffer("w_q", w_packed)
        self.register_buffer("scales", scales_2d)
        self.register_buffer("zeros", zeros_2d)
        self._wq32 = None
        self._gemv_cfg = None

    def _gemv(self, x):
        mod = _get_cuda()
        if self._wq32 is None:
            self._wq32 = self.w_q.view(torch.int32).contiguous()
        xf = x.reshape(-1).contiguous()
        if self._gemv_cfg is None:
            # tiny one-shot autotune over (splitk, cpt, blk)
            best, bcfg = 1e9, (16, 4, 128)
            torch.cuda.synchronize()
            for blk in (64, 128):
                for cpt in (4, 8):
                    for sk in (4, 8, 16, 32):
                        try:
                            for _ in range(3):
                                mod.gemv(xf, self._wq32, self.scales, self.zeros, self.N, self.K, sk, cpt, blk)
                            torch.cuda.synchronize()
                            s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
                            s.record()
                            for _ in range(5):
                                mod.gemv(xf, self._wq32, self.scales, self.zeros, self.N, self.K, sk, cpt, blk)
                            e.record(); torch.cuda.synchronize()
                            t = s.elapsed_time(e)
                            if t < best:
                                best, bcfg = t, (sk, cpt, blk)
                        except Exception:
                            pass
            self._gemv_cfg = bcfg
        sk, cpt, blk = self._gemv_cfg
        return mod.gemv(xf, self._wq32, self.scales, self.zeros, self.N, self.K, sk, cpt, blk)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M = x.shape[0]
        if M == 1:
            try:
                return self._gemv(x)
            except Exception:
                pass
        return _gemm(x, self.w_q, self.scales, self.zeros, self.N, self.K)


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]