"""Fused W4A16 (int4 weight, bf16 activation) GEMM — custom kernels for SM120.

Scheme: AWQ/GPTQ-style asymmetric int4, group=128 along K.
  w_bf[k, n] = (unpack(w_q)[k, n] - zeros[k//128, n]) * scales[k//128, n]

Kernels (all fused unpack + GEMM in one pass):
  - M == 1: SIMT GEMV. Streams w_q coalesced (uint32 per thread), unpacks via a
    bf16 magic-bias bit trick (fp32 (0x4300|nib)<<16 == 128+nib), FMA into fp32,
    K-split into blockIdx.y groups reduced through an fp32 workspace with
    atomicAdd; the last block per column-block converts + re-zeroes the
    workspace (single kernel launch / call).
  - M >= 8: multi-warp mma.sync.m16n8k16 bf16 GEMM with cp.async multi-stage
    pipeline. B is staged in smem as packed int4 (raw (K/2, N) layout — no
    repacking needed) and each warp extracts its fragment bytes directly from
    smem; dequant is exact: (0x4300|nib) magic -> HSUB2(128+z) -> HMUL2(s) which
    reproduces the reference's bf16 dequantization bit-for-bit.
"""
from __future__ import annotations

import os
import torch
import torch.nn as nn

GROUP_SIZE = 128

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cstdint>

#define DEVINL __device__ __forceinline__
#define FULLMASK 0xFFFFFFFFu

// ===========================================================================
// GEMV (M == 1)
// ===========================================================================
template <int VEC, int CHUNK_ROWS, int THREADS, int UNROLL>
__global__ void __launch_bounds__(THREADS) w4a16_gemv(
    const uint8_t* __restrict__ wq,
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ S,
    const __nv_bfloat16* __restrict__ Z,
    float* __restrict__ ws,
    __nv_bfloat16* __restrict__ out,
    int* __restrict__ counters,
    int N, int ksplit) {
  constexpr int PROWS = CHUNK_ROWS / 2;
  constexpr int CPT = THREADS / 32;
  const int kc = blockIdx.y;
  const int tid = threadIdx.x;
  const int n0 = (blockIdx.x * THREADS + tid) * VEC;
  const int kstart = kc * CHUNK_ROWS;

  __shared__ float xs0[PROWS > 0 ? PROWS : 1];
  __shared__ float xs1[PROWS > 0 ? PROWS : 1];
  __shared__ float warp_sum[CPT > 0 ? CPT : 1];
  __shared__ int last_flag;

  float xa = 0.f, xb = 0.f;
#pragma unroll
  for (int t = tid; t < PROWS; t += THREADS) {
    __nv_bfloat162 pr = reinterpret_cast<const __nv_bfloat162*>(x + kstart)[t];
    xa = __bfloat162float(pr.x); xb = __bfloat162float(pr.y);
    xs0[t] = xa; xs1[t] = xb;
  }
  float v = xa + xb;
  v += __shfl_down_sync(FULLMASK, v, 16); v += __shfl_down_sync(FULLMASK, v, 8);
  v += __shfl_down_sync(FULLMASK, v, 4); v += __shfl_down_sync(FULLMASK, v, 2);
  v += __shfl_down_sync(FULLMASK, v, 1);
  if ((tid & 31) == 0) warp_sum[tid >> 5] = v;
  __syncthreads();
  if (tid < 32) {
    float s = (tid < CPT) ? warp_sum[tid] : 0.f;
    s += __shfl_down_sync(FULLMASK, s, 16); s += __shfl_down_sync(FULLMASK, s, 8);
    s += __shfl_down_sync(FULLMASK, s, 4); s += __shfl_down_sync(FULLMASK, s, 2);
    s += __shfl_down_sync(FULLMASK, s, 1);
    if (tid == 0) warp_sum[0] = s;
  }
  __syncthreads();
  const float xsum = warp_sum[0];

  const uint8_t* p = wq + (size_t)(kstart / 2) * (size_t)N + n0;
  float raw[VEC];
#pragma unroll
  for (int j = 0; j < VEC; ++j) raw[j] = 0.f;

  if (VEC == 4) {
    const uint32_t* pw = reinterpret_cast<const uint32_t*>(p);
    const int lds = N / 4;
#pragma unroll UNROLL
    for (int r = 0; r < PROWS; ++r) {
      const uint32_t q = __ldcs(pw + r * lds);
      const float X0 = xs0[r], X1 = xs1[r];
#pragma unroll
      for (int j = 0; j < 4; ++j) {
        const uint32_t b = q >> (8 * j);
        const float lo = __uint_as_float(((b & 0xFu) << 16) | 0x43000000u);
        const float hi = __uint_as_float(((b & 0xF0u) << 12) | 0x43000000u);
        raw[j] = fmaf(X0, lo, raw[j]);
        raw[j] = fmaf(X1, hi, raw[j]);
      }
    }
  } else {
    const uint2* pw = reinterpret_cast<const uint2*>(p);
    const int lds = N / 8;
#pragma unroll UNROLL
    for (int r = 0; r < PROWS; ++r) {
      const uint2 q2 = __ldcs(pw + r * lds);
      const float X0 = xs0[r], X1 = xs1[r];
      const uint32_t qs[2] = {q2.x, q2.y};
#pragma unroll
      for (int h = 0; h < 2; ++h) {
#pragma unroll
        for (int j = 0; j < 4; ++j) {
          const uint32_t b = qs[h] >> (8 * j);
          const float lo = __uint_as_float(((b & 0xFu) << 16) | 0x43000000u);
          const float hi = __uint_as_float(((b & 0xF0u) << 12) | 0x43000000u);
          raw[4 * h + j] = fmaf(X0, lo, raw[4 * h + j]);
          raw[4 * h + j] = fmaf(X1, hi, raw[4 * h + j]);
        }
      }
    }
  }

  const int g = kstart / 128;
  const __nv_bfloat16* srow = S + (size_t)g * N + n0;
  const __nv_bfloat16* zrow = Z + (size_t)g * N + n0;
#pragma unroll
  for (int j = 0; j < VEC; ++j) {
    const float s = __bfloat162float(__ldg(srow + j));
    const float z = __bfloat162float(__ldg(zrow + j));
    const float part = s * (raw[j] - (128.f + z) * xsum);
    atomicAdd(ws + n0 + j, part);
  }

  __threadfence();
  __syncthreads();
  if (tid == 0) {
    const int prev = atomicAdd(counters + blockIdx.x, 1);
    last_flag = (prev == ksplit - 1) ? 1 : 0;
  }
  __syncthreads();
  if (last_flag) {
    __threadfence();
    float rv[VEC];
#pragma unroll
    for (int j = 0; j < VEC; ++j)
      rv[j] = *reinterpret_cast<volatile float*>(ws + n0 + j);
#pragma unroll
    for (int j = 0; j < VEC; ++j) {
      out[n0 + j] = __float2bfloat16(rv[j]);
      ws[n0 + j] = 0.f;
    }
    if (tid == 0) counters[blockIdx.x] = 0;
  }
}



torch::Tensor gemv_run_ptr(torch::Tensor x, int64_t wq_p, int64_t S_p, int64_t Z_p,
                           int64_t ws_p, int64_t ctr_p, int64_t N, int64_t K, int64_t cfg) {
  auto out = torch::empty({N}, x.options());
  auto stream = at::cuda::getCurrentCUDAStream();
  const uint8_t* wq = reinterpret_cast<const uint8_t*>(wq_p);
  const __nv_bfloat16* S = reinterpret_cast<const __nv_bfloat16*>(S_p);
  const __nv_bfloat16* Z = reinterpret_cast<const __nv_bfloat16*>(Z_p);
  float* ws = reinterpret_cast<float*>(ws_p);
  int* counters = reinterpret_cast<int*>(ctr_p);
  const __nv_bfloat16* xp = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
  __nv_bfloat16* outp = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
  if (cfg == 0) {  // 128-thread blocks, 128-row chunks, unroll 16
    const int ks = (int)(K / 128);
    dim3 grid(N / (128 * 4), ks);
    w4a16_gemv<4, 128, 128, 16><<<grid, 128, 0, stream>>>(
        wq, xp, S, Z, ws, outp, counters, (int)N, ks);
  } else if (cfg == 2) {  // 128-thread blocks, 64-row chunks, unroll 16
    const int ks = (int)(K / 64);
    dim3 grid(N / (128 * 4), ks);
    w4a16_gemv<4, 64, 128, 16><<<grid, 128, 0, stream>>>(
        wq, xp, S, Z, ws, outp, counters, (int)N, ks);
  } else if (cfg == 10) {  // 256-thread blocks, 128-row chunks, unroll 64
    const int ks = (int)(K / 128);
    dim3 grid(N / (256 * 4), ks);
    w4a16_gemv<4, 128, 256, 64><<<grid, 256, 0, stream>>>(
        wq, xp, S, Z, ws, outp, counters, (int)N, ks);
  } else if (cfg == 3) {  // 128-thread blocks, 64-row chunks, unroll 8
    const int ks = (int)(K / 64);
    dim3 grid(N / (128 * 4), ks);
    w4a16_gemv<4, 64, 128, 8><<<grid, 128, 0, stream>>>(
        wq, xp, S, Z, ws, outp, counters, (int)N, ks);
  } else {
    TORCH_CHECK(false, "cfg");
  }
  return out;
}

// ===========================================================================
// multi-warp bf16 mma GEMM (M >= 8)
// ===========================================================================
DEVINL void cp16g(void* dst, const void* src, bool pred) {
  const uint32_t d = (uint32_t)__cvta_generic_to_shared(dst);
  const int sz = pred ? 16 : 0;
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(d), "l"(src), "r"(sz));
}
DEVINL void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N> DEVINL void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }

DEVINL uint32_t dq_byte(uint32_t byte, uint32_t zz, uint32_t ss) {
  const uint32_t mag = (byte & 0xFu) | ((byte & 0xF0u) << 12) | 0x43004300u;
  __nv_bfloat162 m = *reinterpret_cast<const __nv_bfloat162*>(&mag);
  __nv_bfloat162 zzv = *reinterpret_cast<const __nv_bfloat162*>(&zz);
  __nv_bfloat162 ssv = *reinterpret_cast<const __nv_bfloat162*>(&ss);
  __nv_bfloat162 sub = __hsub2(m, zzv);
  __nv_bfloat162 res = __hmul2(sub, ssv);
  return *reinterpret_cast<uint32_t*>(&res);
}

DEVINL void mma_bf16(float* c, const uint32_t* a, const uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
      : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
}

template <int BM, int BN, int WM, int WN, int STAGES, bool SPLITK>
__global__ void __launch_bounds__(WM * WN * 32) w4a16_gemm(
    const __nv_bfloat16* __restrict__ x,
    const uint8_t* __restrict__ wq,
    const __nv_bfloat16* __restrict__ S,
    const __nv_bfloat16* __restrict__ Z,
    __nv_bfloat16* __restrict__ out,
    float* __restrict__ ws,
    int* __restrict__ counters,
    int M, int N, int K, int cspan) {
  constexpr int BK = 128;
  constexpr int THREADS = WM * WN * 32;
  constexpr int WTM = BM / WM;
  constexpr int WTN = BN / WN;
  constexpr int MFRAGS = WTM / 16;
  constexpr int NFRAGS = WTN / 8;
  constexpr int A_PITCH = BK * 2 + 16;
  constexpr int B_PITCH = BN + 16;

  const int tid = threadIdx.x;
  const int warp = tid >> 5;
  const int lane = tid & 31;
  const int wm = warp / WN;
  const int wn = warp % WN;

  const int m0 = blockIdx.x * BM;
  const int n0 = blockIdx.y * BN;
  const int nchunk_total = K / BK;
  const int kc0 = SPLITK ? blockIdx.z * cspan : 0;
  const int nchunk = SPLITK ? min(kc0 + cspan, nchunk_total) : nchunk_total;

  extern __shared__ uint8_t smem[];
  uint8_t* sA = smem;
  uint8_t* sB = sA + (size_t)STAGES * BM * A_PITCH;
  uint8_t* sSZ = sB + (size_t)STAGES * 64 * B_PITCH;

  auto load_stage = [&](int stage, int chunk) {
    uint8_t* dA = sA + (size_t)stage * (BM * A_PITCH);
    uint8_t* dB = sB + (size_t)stage * (64 * B_PITCH);
    uint8_t* dSZ = sSZ + (size_t)stage * (2 * BN * 4);
    {
      const int row0 = tid >> 4;
      const int seg = tid & 15;
#pragma unroll
      for (int row = row0; row < BM; row += THREADS / 16) {
        const int gr = m0 + row;
        const __nv_bfloat16* src = x + (size_t)gr * K + chunk * BK + seg * 8;
        cp16g(dA + row * A_PITCH + seg * 16, src, gr < M);
      }
    }
    {
      constexpr int CPR = BN / 16;
      constexpr int TOTAL = 64 * CPR;
#pragma unroll
      for (int t = tid; t < TOTAL; t += THREADS) {
        const int row = t / CPR;
        const int seg = t % CPR;
        cp16g(dB + row * B_PITCH + seg * 16,
              wq + ((size_t)chunk * 64 + row) * N + n0 + seg * 16, true);
      }
    }
    {
      constexpr int TOTAL = 2 * BN / 8;
      for (int t = tid; t < TOTAL; t += THREADS) {
        const int half = t / (BN / 8);
        const int seg = t % (BN / 8);
        cp16g(dSZ + half * (BN * 4) + seg * 16,
              (half ? Z : S) + (size_t)chunk * N + n0 + seg * 8, true);
      }
    }
    cp_commit();
  };

  for (int s = 0; s < STAGES - 1; ++s) {
    const int c = kc0 + s;
    if (c < nchunk) load_stage(s, c);
    else cp_commit();
  }

  float acc[MFRAGS][NFRAGS][4];
#pragma unroll
  for (int mf = 0; mf < MFRAGS; ++mf)
#pragma unroll
    for (int nf = 0; nf < NFRAGS; ++nf)
#pragma unroll
      for (int i = 0; i < 4; ++i) acc[mf][nf][i] = 0.f;

  const int arow0 = wm * WTM + (lane & 7);
  const int acol0 = (lane >> 4) * 16;
  const int bq = lane & 3;
  const int bn_ = lane >> 2;

  uint32_t a[3][MFRAGS][4];
  uint32_t b[3][NFRAGS][2];

  auto load_frags = [&](const uint8_t* tA, const uint8_t* tB, int ks, int par,
                        const uint32_t* zz, const uint32_t* ss) {
#pragma unroll
    for (int mf = 0; mf < MFRAGS; ++mf) {
      const int row = arow0 + mf * 16;
      const uint8_t* addr = tA + (row + ((lane >> 3) & 1) * 8) * A_PITCH + ks * 32 + acol0;
      const uint32_t a_addr = (uint32_t)__cvta_generic_to_shared(addr);
      asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                   : "=r"(a[par][mf][0]), "=r"(a[par][mf][1]), "=r"(a[par][mf][2]), "=r"(a[par][mf][3])
                   : "r"(a_addr));
    }
#pragma unroll
    for (int nf = 0; nf < NFRAGS; ++nf) {
      const int col = wn * WTN + nf * 8 + bn_;
      const uint32_t e0 = (uint32_t)tB[(ks * 8 + bq) * B_PITCH + col];
      const uint32_t e1 = (uint32_t)tB[(ks * 8 + bq + 4) * B_PITCH + col];
      b[par][nf][0] = dq_byte(e0, zz[nf], ss[nf]);
      b[par][nf][1] = dq_byte(e1, zz[nf], ss[nf]);
    }
  };

  for (int chunk = kc0; chunk < nchunk; ++chunk) {
    cp_wait<STAGES - 2>();
    __syncthreads();
    const uint8_t* tA = sA + (size_t)(chunk % STAGES) * (BM * A_PITCH);
    const uint8_t* tB = sB + (size_t)(chunk % STAGES) * (64 * B_PITCH);
    const uint8_t* tSZ = sSZ + (size_t)(chunk % STAGES) * (2 * BN * 4);

    uint32_t zz[NFRAGS], ss[NFRAGS];
#pragma unroll
    for (int nf = 0; nf < NFRAGS; ++nf) {
      const int col = wn * WTN + nf * 8 + bn_;
      const float sf = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(tSZ)[col]);
      const float zf = __bfloat162float(reinterpret_cast<const __nv_bfloat16*>(tSZ + BN * 4)[col]);
      const __nv_bfloat162 z2 = __float2bfloat162_rn(128.f + zf);
      const __nv_bfloat162 s2v = __float2bfloat162_rn(sf);
      zz[nf] = *reinterpret_cast<const uint32_t*>(&z2);
      ss[nf] = *reinterpret_cast<const uint32_t*>(&s2v);
    }

    load_frags(tA, tB, 0, 0, zz, ss);
    load_frags(tA, tB, 1, 1, zz, ss);
#pragma unroll
    for (int ks = 0; ks < 8; ++ks) {
      const int par = ks % 3;
      if (ks < 6) load_frags(tA, tB, ks + 2, (ks + 2) % 3, zz, ss);
#pragma unroll
      for (int mf = 0; mf < MFRAGS; ++mf)
#pragma unroll
        for (int nf = 0; nf < NFRAGS; ++nf)
          mma_bf16(acc[mf][nf], a[par][mf], b[par][nf]);
    }
    if (chunk + STAGES - 1 < nchunk) load_stage((chunk + STAGES - 1) % STAGES, chunk + STAGES - 1);
    else cp_commit();
  }

  const int cr = lane >> 2;
  const int cc = (lane & 3) * 2;
#pragma unroll
  for (int mf = 0; mf < MFRAGS; ++mf) {
#pragma unroll
    for (int nf = 0; nf < NFRAGS; ++nf) {
#pragma unroll
      for (int half = 0; half < 2; ++half) {
        const int r = m0 + wm * WTM + mf * 16 + half * 8 + cr;
        const int c = n0 + wn * WTN + nf * 8 + cc;
        if (r < M) {
          if (SPLITK) {
            atomicAdd(ws + (size_t)r * N + c, acc[mf][nf][half * 2]);
            atomicAdd(ws + (size_t)r * N + c + 1, acc[mf][nf][half * 2 + 1]);
          } else {
            __nv_bfloat162 v = __floats2bfloat162_rn(acc[mf][nf][half * 2], acc[mf][nf][half * 2 + 1]);
            *reinterpret_cast<__nv_bfloat162*>(out + (size_t)r * N + c) = v;
          }
        }
      }
    }
  }

  if (SPLITK) {
    __shared__ int is_last_split;
    __threadfence();
    __syncthreads();
    if (tid == 0) {
      const int prev = atomicAdd(counters + blockIdx.x * gridDim.y + blockIdx.y, 1);
      is_last_split = (prev == (int)gridDim.z - 1) ? 1 : 0;
    }
    __syncthreads();
    if (is_last_split) {
      __threadfence();
#pragma unroll
      for (int mf = 0; mf < MFRAGS; ++mf) {
#pragma unroll
        for (int nf = 0; nf < NFRAGS; ++nf) {
#pragma unroll
          for (int half = 0; half < 2; ++half) {
            const int r = m0 + wm * WTM + mf * 16 + half * 8 + cr;
            const int c = n0 + wn * WTN + nf * 8 + cc;
            if (r < M) {
              __nv_bfloat162 v;
              v.x = __float2bfloat16(*reinterpret_cast<volatile float*>(ws + (size_t)r * N + c));
              v.y = __float2bfloat16(*reinterpret_cast<volatile float*>(ws + (size_t)r * N + c + 1));
              *reinterpret_cast<__nv_bfloat162*>(out + (size_t)r * N + c) = v;
              ws[(size_t)r * N + c] = 0.f;
              ws[(size_t)r * N + c + 1] = 0.f;
            }
          }
        }
      }
      __syncthreads();
      if (tid == 0) counters[blockIdx.x * gridDim.y + blockIdx.y] = 0;
    }
  }
}

torch::Tensor gemm_run(torch::Tensor x, torch::Tensor wq, torch::Tensor S,
                       torch::Tensor Z, torch::Tensor ws, torch::Tensor counters,
                       int64_t cfg, int64_t splitk) {
  const int M = x.size(0), K = x.size(1);
  const int N = wq.size(1);
  auto out = torch::empty({M, N}, x.options());
  auto stream = at::cuda::getCurrentCUDAStream();
  const int64_t sk = splitk;

#define LAUNCH(BM, BN, WM, WN, ST) \
  { \
    size_t smem = ST * ((size_t)BM * (128 * 2 + 16) + 64 * (BN + 16) + 2 * BN * 4); \
    if (sk > 1) { \
      const int cspan = (K / 128 + sk - 1) / (int)sk; \
      dim3 grid((M + BM - 1) / BM, N / BN, sk); \
      auto kfn = w4a16_gemm<BM, BN, WM, WN, ST, true>; \
      TORCH_CHECK(smem <= 101376, "smem overflow"); \
      if (smem > 48 * 1024) cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); \
      kfn<<<grid, WM * WN * 32, smem, stream>>>( \
          reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()), wq.data_ptr<uint8_t>(), \
          reinterpret_cast<const __nv_bfloat16*>(S.data_ptr()), \
          reinterpret_cast<const __nv_bfloat16*>(Z.data_ptr()), \
          reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), ws.data_ptr<float>(), \
          counters.data_ptr<int>(), M, N, K, cspan); \
    } else { \
      dim3 grid((M + BM - 1) / BM, N / BN); \
      auto kfn = w4a16_gemm<BM, BN, WM, WN, ST, false>; \
      TORCH_CHECK(smem <= 101376, "smem overflow"); \
      if (smem > 48 * 1024) cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); \
      kfn<<<grid, WM * WN * 32, smem, stream>>>( \
          reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()), wq.data_ptr<uint8_t>(), \
          reinterpret_cast<const __nv_bfloat16*>(S.data_ptr()), \
          reinterpret_cast<const __nv_bfloat16*>(Z.data_ptr()), \
          reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), ws.data_ptr<float>(), \
          counters.data_ptr<int>(), M, N, K, 0); \
    } \
  }

  if (cfg == 0) LAUNCH(16, 64, 1, 4, 4)
  else if (cfg == 1) LAUNCH(16, 64, 1, 8, 4)
  else if (cfg == 2) LAUNCH(16, 128, 1, 8, 2)
  else if (cfg == 3) LAUNCH(16, 64, 1, 8, 2)
  else if (cfg == 4) LAUNCH(16, 128, 1, 8, 4)
  else if (cfg == 5) LAUNCH(16, 64, 1, 4, 2)
  else if (cfg == 10) LAUNCH(32, 64, 2, 2, 4)
  else if (cfg == 12) LAUNCH(32, 64, 2, 4, 3)
  else if (cfg == 13) LAUNCH(32, 128, 2, 4, 2)
  else if (cfg == 14) LAUNCH(32, 128, 2, 4, 4)
  else if (cfg == 15) LAUNCH(32, 64, 2, 4, 2)
  else if (cfg == 20) LAUNCH(64, 128, 2, 4, 3)
  else if (cfg == 21) LAUNCH(64, 64, 2, 4, 4)
  else if (cfg == 22) LAUNCH(64, 64, 2, 4, 2)
  else if (cfg == 23) LAUNCH(64, 64, 4, 2, 4)
  else if (cfg == 24) LAUNCH(64, 128, 2, 4, 2)
  else TORCH_CHECK(false, "cfg");
  return out;
}
"""

_CPP_SRC = """
torch::Tensor gemv_run_ptr(torch::Tensor x, int64_t wq_p, int64_t S_p, int64_t Z_p, int64_t ws_p, int64_t ctr_p, int64_t N, int64_t K, int64_t cfg);
torch::Tensor gemm_run(torch::Tensor x, torch::Tensor wq, torch::Tensor S, torch::Tensor Z, torch::Tensor ws, torch::Tensor counters, int64_t cfg, int64_t splitk);
"""

_mod = None


def _get_mod():
    global _mod
    if _mod is None:
        from torch.utils.cpp_extension import load_inline

        _mod = load_inline(
            name="w4a16_sol_v8",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC,
            functions=["gemv_run_ptr", "gemm_run"],
            extra_cuda_cflags=["-O3", "-gencode=arch=compute_120a,code=sm_120a", "--use_fast_math"],
            verbose=False,
        )
    return _mod


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__()
        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
        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))
        self._ws = None
        self._counters = None
        self._ptrs = None
        self._env_gemv_cfg = int(os.environ.get("W4A16_GEMV_CFG", "0" if N >= 8192 else "10"))
        # invalidate cached raw pointers if weights are re-loaded
        def _invalidate(mod, recurse=True):
            self._ptrs = None
        self.register_load_state_dict_post_hook(_invalidate)
        self._register_state_dict_hook(lambda self, sd, prefix, metadata: setattr(self, "_ptrs", None))
        self._env_cfg = os.environ.get("W4A16_CFG")
        self._env_sk = int(os.environ.get("W4A16_SK", "1"))
        # dispatch decisions frozen at init
        self._gemv = (M == 1 and N % 512 == 0 and K % 128 == 0)
        self._gemm_cfg = self._pick_cfg()
        # touch + cache the extension so all instantiations build once
        self._mod = _get_mod()

    def _pick_cfg(self):
        M, N, K = self.M, self.N, self.K
        if self._env_cfg is not None:
            return int(self._env_cfg)
        if M <= 16:
            return 4 if N % 128 == 0 else 0
        elif M <= 32:
            return 14 if N % 128 == 0 else 10
        else:
            return 20 if N % 128 == 0 else 21

    def _ensure_workspace(self, device):
        if self._ws is None or self._ws.device != device:
            self._ws = torch.zeros(max(self.M * self.N, self.N), dtype=torch.float32, device=device)
            self._counters = torch.zeros(4096, dtype=torch.int32, device=device)

    def _ensure_splitws(self, device, M, N, cfg, sk):
        if (self._ws is None or self._ws.numel() < M * N or self._ws.device != device):
            self._ws = torch.zeros(max(M * N, self.N), dtype=torch.float32, device=device)
        if self._counters is None or self._counters.numel() < 4096 or self._counters.device != device:
            self._counters = torch.zeros(4096, dtype=torch.int32, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        if self._gemv and self._ptrs is None:
            self._ensure_workspace(x.device)
            self._ptrs = (self.w_q.data_ptr(), self.scales.data_ptr(), self.zeros.data_ptr(),
                          self._ws.data_ptr(), self._counters.data_ptr())
        if self._gemv:
            return self._mod.gemv_run_ptr(x, *self._ptrs, self.N, self.K, self._env_gemv_cfg).view(1, self.N)
        self._ensure_splitws(x.device, self.M, self.N, self._gemm_cfg, self._env_sk)
        return self._mod.gemm_run(x, self.w_q, self.scales, self.zeros, self._ws, self._counters,
                                  self._gemm_cfg, self._env_sk)

    __call__ = forward


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]