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

Fused int4-unpack + GEMM. Custom paths:
  - M == 1:  fp32-magic int4 GEMV (split-K, one kernel).
  - M >= 2:  TMA-staged mma.sync kernel with repacked weights, split-K combine.
  - Fallback: Triton fused-dequant kernel for odd shapes.
"""
from __future__ import annotations

import torch
import torch.nn as nn

OP_TYPE = "gemm_w4a16"
SUPPORTED_PRECISIONS = ["int4_bf16"]
HARDWARE_REQUIRED = ["B200"]

GROUP_SIZE = 128

# ----------------------------------------------------------------------------
# Extension sources (GEMV + MMA built once at import)
# ----------------------------------------------------------------------------
_CPP_DECL = r"""
#include <torch/extension.h>
torch::Tensor repack(torch::Tensor wq);
torch::Tensor gemv_w4(torch::Tensor x, torch::Tensor wq, torch::Tensor sc, torch::Tensor zr,
                      torch::Tensor wksp, torch::Tensor counters, long ksplit, long gs, long wpg, long unroll);
torch::Tensor w4_mma(torch::Tensor x, torch::Tensor wr, torch::Tensor sc, torch::Tensor zr,
                     torch::Tensor wksp, torch::Tensor counters, long variant, long splitk, long dbg);
"""

_GEMV_SRC = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>

#define MAGICV 0x46000000u

template <unsigned IMM>
__device__ __forceinline__ unsigned lop3v(unsigned a, unsigned b, unsigned c) {
    unsigned r;
    asm("lop3.b32 %0, %1, %2, %3, %4;" : "=r"(r) : "r"(a), "r"(b), "r"(c), "n"(IMM));
    return r;
}
__device__ __forceinline__ uint4 ld_ef(const void* p, unsigned long long pol) {
    uint4 v;
    asm volatile("ld.global.nc.L2::cache_hint.v4.u32 {%0,%1,%2,%3}, [%4], %5;"
                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p), "l"(pol));
    return v;
}
struct alignas(32) b32x { unsigned long long a, b, c, d; };
__device__ __forceinline__ b32x ld_sz(const void* p) {
    b32x v;
    asm volatile("ld.global.nc.v4.u64 {%0,%1,%2,%3}, [%4];"
                 : "=l"(v.a), "=l"(v.b), "=l"(v.c), "=l"(v.d) : "l"(p));
    return v;
}
__device__ __forceinline__ void unpack_bf16_16(b32x v, float* f) {
    unsigned long long vs[4] = {v.a, v.b, v.c, v.d};
    #pragma unroll
    for (int q = 0; q < 4; q++) {
        unsigned lo = (unsigned)(vs[q] & 0xFFFFFFFFull);
        unsigned hi = (unsigned)(vs[q] >> 32);
        f[4*q+0] = __uint_as_float(lo << 16);
        f[4*q+1] = __uint_as_float(lo & 0xFFFF0000u);
        f[4*q+2] = __uint_as_float(hi << 16);
        f[4*q+3] = __uint_as_float(hi & 0xFFFF0000u);
    }
}

__device__ __constant__ float PAE_TAB[4] = {0x1p10f, 0x1p2f,  0x1p-6f, 0x1p-6f};

template <int GS, int WPG, int UNROLL>
__global__ void __launch_bounds__(32 * GS * WPG) gemv_w4_kernel(
    const unsigned char* __restrict__ wq,
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ sc,
    const __nv_bfloat16* __restrict__ zr,
    __nv_bfloat16* __restrict__ out,
    float* __restrict__ wksp,
    int* __restrict__ counters,
    int N, int KH, int ksplit)
{
    constexpr int WY = GS * WPG;
    const int lane = threadIdx.x;
    const int wy = threadIdx.y;
    const int tile = blockIdx.y;
    const int slice = blockIdx.x;
    const int col0 = tile * 512 + lane * 16;

    const int g_local = wy / WPG;
    const int r_off = wy % WPG;
    const int g0 = slice * GS + g_local;
    const int row0 = g0 * 64 + r_off;

    float acc[16];
    #pragma unroll
    for (int i = 0; i < 16; i++) acc[i] = 0.f;

    unsigned long long pol;
    asm("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol));

    float A[16], C[16], Dz[16];
    {
        b32x sv = ld_sz(sc + (size_t)g0 * N + col0);
        b32x zv = ld_sz(zr + (size_t)g0 * N + col0);
        float s[16], z[16];
        unpack_bf16_16(sv, s);
        unpack_bf16_16(zv, z);
        #pragma unroll
        for (int i = 0; i < 16; i++) {
            int b = i & 3;
            A[i] = s[i] * PAE_TAB[b];
            C[i] = -8192.f * A[i];
            Dz[i] = -s[i] * z[i];
        }
    }
    float gsum = 0.f;
    const __nv_bfloat16* xp = x + (size_t)row0 * 2;
    const unsigned char* wp = wq + (size_t)row0 * N + col0;
    constexpr int ROWS = 64 / WPG;
    #pragma unroll
    for (int rb = 0; rb < ROWS; rb += UNROLL) {
        uint4 v[UNROLL];
        float xe[UNROLL], xo[UNROLL], xo16[UNROLL];
        #pragma unroll
        for (int u = 0; u < UNROLL; u++) {
            v[u] = ld_ef(wp + (size_t)(rb + u) * WPG * N, pol);
            unsigned xv = *(const unsigned*)(xp + (size_t)(rb + u) * WPG * 2);
            xo[u] = __uint_as_float(xv & 0xFFFF0000u);
            xe[u] = __uint_as_float(xv << 16);
            xo16[u] = xo[u] * 0.0625f;
            gsum = gsum + xe[u] + xo[u];
        }
        #pragma unroll
        for (int u = 0; u < UNROLL; u++) {
            unsigned ws[4] = {v[u].x, v[u].y, v[u].z, v[u].w};
            #pragma unroll
            for (int j = 0; j < 4; j++) {
                unsigned wu = ws[j];
                {
                    unsigned lo = lop3v<0xEA>(wu, 0x0000000Fu, MAGICV);
                    float te = fmaf(A[4*j+0], __uint_as_float(lo), C[4*j+0]);
                    acc[4*j+0] = fmaf(xe[u], te, acc[4*j+0]);
                    unsigned hi = lop3v<0xEA>(wu, 0x000000F0u, MAGICV);
                    float to = fmaf(A[4*j+0], __uint_as_float(hi), C[4*j+0]);
                    acc[4*j+0] = fmaf(xo16[u], to, acc[4*j+0]);
                }
                {
                    unsigned lo = lop3v<0xEA>(wu, 0x00000F00u, MAGICV);
                    float te = fmaf(A[4*j+1], __uint_as_float(lo), C[4*j+1]);
                    acc[4*j+1] = fmaf(xe[u], te, acc[4*j+1]);
                    unsigned hi = lop3v<0xEA>(wu, 0x0000F000u, MAGICV);
                    float to = fmaf(A[4*j+1], __uint_as_float(hi), C[4*j+1]);
                    acc[4*j+1] = fmaf(xo16[u], to, acc[4*j+1]);
                }
                {
                    unsigned lo = lop3v<0xEA>(wu, 0x000F0000u, MAGICV);
                    float te = fmaf(A[4*j+2], __uint_as_float(lo), C[4*j+2]);
                    acc[4*j+2] = fmaf(xe[u], te, acc[4*j+2]);
                    unsigned hi = lop3v<0xEA>(wu >> 4, 0x000F0000u, MAGICV);
                    float to = fmaf(A[4*j+2], __uint_as_float(hi), C[4*j+2]);
                    acc[4*j+2] = fmaf(xo[u], to, acc[4*j+2]);
                }
                {
                    unsigned lo = lop3v<0xEA>(wu >> 8, 0x000F0000u, MAGICV);
                    float te = fmaf(A[4*j+3], __uint_as_float(lo), C[4*j+3]);
                    acc[4*j+3] = fmaf(xe[u], te, acc[4*j+3]);
                    unsigned hi = lop3v<0xEA>(wu >> 12, 0x000F0000u, MAGICV);
                    float to = fmaf(A[4*j+3], __uint_as_float(hi), C[4*j+3]);
                    acc[4*j+3] = fmaf(xo[u], to, acc[4*j+3]);
                }
            }
        }
    }

    #pragma unroll
    for (int i = 0; i < 16; i++) acc[i] = fmaf(Dz[i], gsum, acc[i]);

    __shared__ float reduce[WY][512];
    __shared__ int done_flag;
    if (WY > 1) {
        #pragma unroll
        for (int i = 0; i < 16; i++) reduce[wy][i * 32 + lane] = acc[i];
        __syncthreads();
        if (wy == 0) {
            #pragma unroll
            for (int i = 0; i < 16; i++) {
                float t = 0.f;
                #pragma unroll
                for (int yy = 0; yy < WY; yy++) t += reduce[yy][i * 32 + lane];
                acc[i] = t;
            }
        }
        __syncthreads();
    }
    if (wy == 0) {
        float* wp_out = wksp + tile * 512 + lane * 16;
        #pragma unroll
        for (int i = 0; i < 16; i++) atomicAdd(wp_out + i, acc[i]);
    }
    __threadfence();
    __syncthreads();
    if (threadIdx.x == 0 && threadIdx.y == 0) {
        int old = atomicAdd(&counters[tile], 1);
        done_flag = (old == ksplit - 1) ? 1 : 0;
    }
    __threadfence();
    __syncthreads();
    if (done_flag && wy == 0) {
        float su[16];
        const float* rp = wksp + tile * 512 + lane * 16;
        #pragma unroll
        for (int i = 0; i < 16; i++) su[i] = rp[i];
        __nv_bfloat16* op = out + tile * 512 + lane * 16;
        #pragma unroll
        for (int i = 0; i < 16; i++) op[i] = __float2bfloat16(su[i]);
        float* zp = wksp + tile * 512 + lane * 16;
        #pragma unroll
        for (int i = 0; i < 16; i++) zp[i] = 0.f;
        if (lane == 0) counters[tile] = 0;
    }
}

#define DISPATCH_GEMV(GS_, WPG_, UNROLL_)                                        \
    gemv_w4_kernel<GS_, WPG_, UNROLL_><<<grid, dim3(32, (GS_) * (WPG_)), 0, stream>>>( \
        (const unsigned char*)wq.data_ptr(), (const __nv_bfloat16*)x.data_ptr(), \
        (const __nv_bfloat16*)sc.data_ptr(), (const __nv_bfloat16*)zr.data_ptr(), \
        (__nv_bfloat16*)out.data_ptr(), wksp.data_ptr<float>(),                  \
        counters.data_ptr<int>(), N, KH, (int)ksplit)

torch::Tensor gemv_w4(torch::Tensor x, torch::Tensor wq, torch::Tensor sc, torch::Tensor zr,
                      torch::Tensor wksp, torch::Tensor counters, long ksplit, long gs, long wpg, long unroll) {
    int N = wq.size(1);
    int KH = wq.size(0);
    auto out = at::empty({N}, x.options());
    dim3 grid(ksplit, N / 512);
    auto stream = at::cuda::getCurrentCUDAStream();
    TORCH_CHECK(32 / ksplit == gs, "ksplit*gs must be 32");
    if (gs == 2 && wpg == 2 && unroll == 8) { DISPATCH_GEMV(2, 2, 8); }
    else if (gs == 1 && wpg == 4 && unroll == 8) { DISPATCH_GEMV(1, 4, 8); }
    else if (gs == 1 && wpg == 4 && unroll == 4) { DISPATCH_GEMV(1, 4, 4); }
    else if (gs == 1 && wpg == 2 && unroll == 8) { DISPATCH_GEMV(1, 2, 8); }
    else if (gs == 2 && wpg == 4 && unroll == 8) { DISPATCH_GEMV(2, 4, 8); }
    else if (gs == 4 && wpg == 2 && unroll == 8) { DISPATCH_GEMV(4, 2, 8); }
    else if (gs == 4 && wpg == 4 && unroll == 8) { DISPATCH_GEMV(4, 4, 8); }
    else if (gs == 2 && wpg == 2 && unroll == 4) { DISPATCH_GEMV(2, 2, 4); }
    else if (gs == 1 && wpg == 1 && unroll == 8) { DISPATCH_GEMV(1, 1, 8); }
    else { TORCH_CHECK(false, "no such variant"); }
    return out;
}
'''

_MMA_SRC = r'''
//
// W4A16 fused mma kernel (bf16, SM100). v6
//
// Same repack/dequant/fragment scheme as v5 (see w4mma_kernel history).
// v6 changes:
//   - two-chunk software pipeline: LDS + dequant for chunks (2c, 2c+1) issued
//     before the mma cluster of chunk 2c to overlap HMMA issue with unpack.
//   - split-K epilogue: fp32 partials with plain vector stores (no atomic RMW
//     storm), one fence, per-tile semaphore, last block reads SPLITK partial
//     slabs, sums, converts to bf16, re-zeroes them for the next call.
//   - globaltimer STAMP hooks for profiling.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>

// ---------------------------------------------------------------------------
// repack kernel: one thread per (n32-tile, k-chunk, lane) -> uint2
// ---------------------------------------------------------------------------
__global__ void repack_w4(const unsigned char* __restrict__ wq,
                          uint2* __restrict__ wr,
                          int N, int KH) {
    int n32tile = blockIdx.x;
    int chunk = blockIdx.y;     // 256
    int lane = threadIdx.x;
    int n128 = n32tile / 4;
    int st = n32tile % 4;
    int col_base = n128 * 128 + st * 32;
    int s = chunk / 8, ch = chunk % 8;
    int br0 = chunk * 8;
    unsigned q[2] = {0u, 0u};
    #pragma unroll
    for (int j = 0; j < 2; j++) {
        #pragma unroll
        for (int half = 0; half < 2; half++) {
            int t = 2 * j + half;
            int col = col_base + 8 * t + lane / 4;
            int rl = br0 + (lane % 4);
            int rh = br0 + (lane % 4) + 4;
            unsigned bl = wq[(size_t)rl * N + col];
            unsigned bh = wq[(size_t)rh * N + col];
            if (half == 0) {
                q[j] |= (bl & 0xF);
                q[j] |= (bl & 0xF0);
                q[j] |= (bh & 0xF) << 8;
                q[j] |= (bh & 0xF0) << 12;
            } else {
                q[j] |= (bl & 0xF) << 12;
                q[j] |= (bl & 0xF0) << 16;
                q[j] |= (bh & 0xF) << 24;
                q[j] |= (bh & 0xF0) << 24;
            }
        }
    }
    size_t u2 = ((((size_t)n128 * 32 + s) * 4 + st) * 8 + ch) * 32 + lane;
    wr[u2] = make_uint2(q[0], q[1]);
}

torch::Tensor repack(torch::Tensor wq) {
    int KH = wq.size(0);
    int N = wq.size(1);
    TORCH_CHECK(N % 128 == 0 && KH == 2048, "shape");
    auto wr = at::zeros({(long)N * KH / 8 * 2}, wq.options().dtype(at::kInt));
    repack_w4<<<dim3((N / 128) * 4, 256), 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        (const unsigned char*)wq.data_ptr(), (uint2*)wr.data_ptr(), N, KH);
    return wr;
}

// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------
template <unsigned IMM>
__device__ __forceinline__ unsigned lop3(unsigned a, unsigned b, unsigned c) {
    unsigned r;
    asm("lop3.b32 %0, %1, %2, %3, %4;" : "=r"(r) : "r"(a), "r"(b), "r"(c), "n"(IMM));
    return r;
}
__device__ __forceinline__ void mma_bf16(float c[4], const unsigned a[4], const unsigned b[2]) {
    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]));
}
__device__ __forceinline__ unsigned pack_bf16x2(float a_hi, float b_lo) {
    unsigned r;
    asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(a_hi), "f"(b_lo));
    return r;
}
__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) {
    uint32_t s = static_cast<uint32_t>(__cvta_generic_to_shared(bar));
    asm volatile("mbarrier.init.shared.b64 [%0], %1;" :: "r"(s), "r"(count));
}
__device__ __forceinline__ void mbar_expect(uint64_t* bar, uint32_t bytes) {
    uint32_t b = static_cast<uint32_t>(__cvta_generic_to_shared(bar));
    asm volatile(
        "{\n\t.reg .b64 t1;\n\tmbarrier.arrive.expect_tx.release.cta.shared::cta.b64 t1, [%0], %1;\n\t}"
        :: "r"(b), "r"(bytes));
}
__device__ __forceinline__ void tma_bulk(void* smem, const void* glob, uint32_t bytes, uint64_t* bar, unsigned long long pol) {
    uint32_t s = static_cast<uint32_t>(__cvta_generic_to_shared(smem));
    uint32_t b = static_cast<uint32_t>(__cvta_generic_to_shared(bar));
    asm volatile(
        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;"
        :: "r"(s), "l"(glob), "r"(bytes), "r"(b), "l"(pol));
}
__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t parity) {
    uint32_t b = static_cast<uint32_t>(__cvta_generic_to_shared(bar));
    asm volatile(
        "{\n\t.reg .pred p;\n\t"
        "WAIT_%=:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1;\n\t"
        "@!p bra WAIT_%=;\n\t}"
        :: "r"(b), "r"(parity));
}

#define MAGIC 0x46000000u
__device__ __constant__ float PAE_TA[4] = {0x1p10f, 0x1p6f, 0x1p2f, 0x1p-6f};  // tile A {0,4,8,16}
__device__ __constant__ float PAE_TB[4] = {0x1p10f, 0x1p2f, 0x1p-2f, 0x1p-6f}; // tile B {0,8,12,16}

template <int TILE>
__device__ __forceinline__ void dequant_tile(unsigned q, unsigned out[2],
                                             const float A[4], const float C[4], float tz) {
    unsigned v0, v1, v2, v3;
    if (TILE == 0) {
        v0 = lop3<0xEA>(q, 0x0000000Fu, MAGIC);
        v1 = lop3<0xEA>(q, 0x000000F0u, MAGIC);
        v2 = lop3<0xEA>(q, 0x00000F00u, MAGIC);
        v3 = lop3<0xEA>(q, 0x000F0000u, MAGIC);
    } else {
        unsigned qq = q >> 12;
        v0 = lop3<0xEA>(qq, 0x0000000Fu, MAGIC);
        v1 = lop3<0xEA>(qq, 0x00000F00u, MAGIC);
        v2 = lop3<0xEA>(qq, 0x0000F000u, MAGIC);
        v3 = lop3<0xEA>(qq, 0x000F0000u, MAGIC);
    }
    float t0 = fmaf(A[0], __uint_as_float(v0), C[0]) + tz;
    float t1 = fmaf(A[1], __uint_as_float(v1), C[1]) + tz;
    float t2 = fmaf(A[2], __uint_as_float(v2), C[2]) + tz;
    float t3 = fmaf(A[3], __uint_as_float(v3), C[3]) + tz;
    out[0] = pack_bf16x2(t1, t0);
    out[1] = pack_bf16x2(t3, t2);
}

// ---------------------------------------------------------------------------
// main kernel
// ---------------------------------------------------------------------------
template <int MT, int SPLITK, int STAGES, int DBGM=0>
__global__ void __launch_bounds__(128) w4_mma_kernel(
    const uint2* __restrict__ wr,
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ sc,
    const __nv_bfloat16* __restrict__ zr,
    __nv_bfloat16* __restrict__ out,
    float* __restrict__ wksp,
    int* __restrict__ counters,
    int M, int N, int K,
    unsigned long long* __restrict__ tbuf)
{
    constexpr int CHUNKS = 256 / SPLITK;
    constexpr int N_STAGES_TOTAL = CHUNKS / 8;
    constexpr int SZ_OFF = STAGES * 8192;
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int n128 = blockIdx.y;
    const int slice = blockIdx.x;
    const int mb = blockIdx.z;
    const int chunk_base = slice * CHUNKS;
    const int row_base = mb * 16 * MT;
    const int tid = threadIdx.x;

    extern __shared__ __align__(16) unsigned char smem[];
    __shared__ __align__(8) uint64_t bars[STAGES];
    __shared__ int done_flag;

    auto stamp = [&]() { unsigned long long t; asm volatile("mov.u64 %0, %globaltimer;" : "=l"(t)); return t; };
    #define STAMP(phase) if (tbuf && tid == 0) tbuf[((size_t)(mb * gridDim.y + n128) * gridDim.x + slice) * 8 + (phase)] = stamp()

    float c[MT][4][4];
    #pragma unroll
    for (int i = 0; i < MT; i++)
        #pragma unroll
        for (int j = 0; j < 4; j++)
            #pragma unroll
            for (int u = 0; u < 4; u++) c[i][j][u] = 0.f;

    unsigned long long pol;
    asm("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol));
    if (tid == 0) {
        #pragma unroll
        for (int s = 0; s < STAGES; s++) mbar_init(&bars[s], 1);
    }
    __syncthreads();

    auto fetch = [&](int s) {
        if (s >= N_STAGES_TOTAL) return;
        int slot = s % STAGES;
        int group = chunk_base / 8 + s;
        if (tid == 0) {
            mbar_expect(&bars[slot], 8192 + 512);
            size_t u2_global = (((size_t)n128 * 32 + group) * 4 * 8 * 32);
            tma_bulk(&smem[slot * 8192], (const unsigned char*)wr + u2_global * 8, 8192, &bars[slot], pol);
            tma_bulk(&smem[SZ_OFF + slot * 512], sc + (size_t)group * N + n128 * 128, 256, &bars[slot], pol);
            tma_bulk(&smem[SZ_OFF + slot * 512 + 256], zr + (size_t)group * N + n128 * 128, 256, &bars[slot], pol);
        }
    };

    for (int s = 0; s < STAGES - 1; s++) fetch(s);

    float A[4][4], C[4][4], tz[4];
    // x slice pointers (via L1, prefetched at start of each group-epoch)
    STAMP(0);

    for (int s = 0; s < N_STAGES_TOTAL; s++) {
        int slot = s % STAGES;
        mbar_wait(&bars[slot], (s / STAGES) & 1);
        if (s == 0) STAMP(1);
        {
            const __nv_bfloat16* shS = (const __nv_bfloat16*)&smem[SZ_OFF + slot * 512];
            const __nv_bfloat16* shZ = (const __nv_bfloat16*)&smem[SZ_OFF + slot * 512 + 256];
            #pragma unroll
            for (int t = 0; t < 4; t++) {
                int col = warp * 32 + 8 * t + lane / 4;
                float sf = __bfloat162float(shS[col]);
                float zf = __bfloat162float(shZ[col]);
                const float* pae = (t % 2 == 0) ? PAE_TA : PAE_TB;
                #pragma unroll
                for (int u = 0; u < 4; u++) {
                    A[t][u] = sf * pae[u];
                    C[t][u] = -8192.f * A[t][u];
                }
                tz[t] = -sf * zf;
            }
        }
        // A-frag x loads via L1 (global); prefetch next group while computing
        __nv_bfloat16* __restrict__ xg = (__nv_bfloat16*)x;
        // two-chunk software pipeline: (load+dequant of ch pair) then mma cluster
        #pragma unroll
        for (int ch2 = 0; ch2 < 4; ch2++) {
            if (s == 0) STAMP(2 + ch2);
            unsigned a[2][MT][4];
            #pragma unroll
            for (int cc = 0; cc < 2; cc++) {
                int ch = ch2 * 2 + cc;
                long kk = ((long)chunk_base + s * 8 + ch) * 16 + 2 * (lane % 4);
                #pragma unroll
                for (int i = 0; i < MT; i++) {
                    if (DBGM == 2 || DBGM == 3) { a[cc][i][0]=0u; a[cc][i][1]=1u; a[cc][i][2]=2u; a[cc][i][3]=3u; } else {
                    int r0g = row_base + i * 16 + lane / 4;
                    int r1g = r0g + 8;
                    bool p0 = r0g < M, p1 = r1g < M;
                    const __nv_bfloat16* x0 = xg + (size_t)r0g * K + kk;
                    const __nv_bfloat16* x1 = xg + (size_t)r1g * K + kk;
                    a[cc][i][0] = p0 ? __ldg((const unsigned*)x0) : 0u;
                    a[cc][i][1] = p1 ? __ldg((const unsigned*)x1) : 0u;
                    a[cc][i][2] = p0 ? __ldg((const unsigned*)(x0 + 8)) : 0u;
                    a[cc][i][3] = p1 ? __ldg((const unsigned*)(x1 + 8)) : 0u;
                    }
                }
            }
            uint2 qA = ((uint2*)&smem[slot * 8192])[(warp * 8 + ch2 * 2) * 32 + lane];
            uint2 qB = ((uint2*)&smem[slot * 8192])[(warp * 8 + ch2 * 2 + 1) * 32 + lane];
            unsigned bA[4][4], bB[4][4];
            #pragma unroll
            for (int j = 0; j < 2; j++) {
                unsigned qa = j == 0 ? qA.x : qA.y;
                unsigned qb = j == 0 ? qB.x : qB.y;
                if (DBGM == 1 || DBGM == 3) {
                    bA[2*j+0][0] = qa; bA[2*j+0][1] = qa >> 8; bA[2*j+1][0] = qa; bA[2*j+1][1] = qa >> 8;
                    bB[2*j+0][0] = qb; bB[2*j+0][1] = qb >> 8; bB[2*j+1][0] = qb; bB[2*j+1][1] = qb >> 8;
                } else {
                    dequant_tile<0>(qa, (unsigned*)&bA[2*j+0], A[2*j+0], C[2*j+0], tz[2*j+0]);
                    dequant_tile<1>(qa, (unsigned*)&bA[2*j+1], A[2*j+1], C[2*j+1], tz[2*j+1]);
                    dequant_tile<0>(qb, (unsigned*)&bB[2*j+0], A[2*j+0], C[2*j+0], tz[2*j+0]);
                    dequant_tile<1>(qb, (unsigned*)&bB[2*j+1], A[2*j+1], C[2*j+1], tz[2*j+1]);
                }
            }
            #pragma unroll
            for (int t = 0; t < 4; t++) {
                #pragma unroll
                for (int i = 0; i < MT; i++) {
                    mma_bf16(c[i][t], a[0][i], (unsigned*)&bA[t]);
                }
            }
            #pragma unroll
            for (int t = 0; t < 4; t++) {
                #pragma unroll
                for (int i = 0; i < MT; i++) {
                    mma_bf16(c[i][t], a[1][i], (unsigned*)&bB[t]);
                }
            }
        }
        if (s == 0) STAMP(3);
        __syncthreads();
        fetch(s + STAGES - 1);
    }
    STAMP(4);

    if (SPLITK == 1) {
        #pragma unroll
        for (int i = 0; i < MT; i++) {
            #pragma unroll
            for (int t = 0; t < 4; t++) {
                int col0 = n128 * 128 + warp * 32 + 8 * t + 2 * (lane % 4);
                int r0 = row_base + i * 16 + lane / 4;
                unsigned v01 = pack_bf16x2(c[i][t][1], c[i][t][0]);
                unsigned v23 = pack_bf16x2(c[i][t][3], c[i][t][2]);
                if (r0 < M) *(unsigned*)(out + (size_t)r0 * N + col0) = v01;
                if (r0 + 8 < M) *(unsigned*)(out + (size_t)(r0 + 8) * N + col0) = v23;
            }
        }
    } else {
        // fp32 partials to per-slice slabs [slice][row][col], plain vector stores
        #pragma unroll
        for (int i = 0; i < MT; i++) {
            #pragma unroll
            for (int t = 0; t < 4; t++) {
                int col0 = n128 * 128 + warp * 32 + 8 * t + 2 * (lane % 4);
                int r0 = row_base + i * 16 + lane / 4;
                if (r0 < M) {
                    float2 v01 = make_float2(c[i][t][0], c[i][t][1]);
                    *(float2*)&wksp[(((size_t)slice * M + r0) * N) + col0] = v01;
                }
                if (r0 + 8 < M) {
                    float2 v23 = make_float2(c[i][t][2], c[i][t][3]);
                    *(float2*)&wksp[(((size_t)slice * M + r0 + 8) * N) + col0] = v23;
                }
            }
        }
        __threadfence();
        __syncthreads();
        if (tid == 0) {
            int idx = mb * (N / 128) + n128;
            int old = atomicAdd(&counters[idx], 1);
            done_flag = (old == SPLITK - 1) ? 1 : 0;
        }
        __syncthreads();
        if (done_flag) {
            __threadfence();
            int idx = mb * (N / 128) + n128;
            if (tid == 0) counters[idx] = 0;
            int rows = min(16 * MT, M - row_base);
            for (int r = 0; r < rows; r++) {
                int cr = row_base + r;
                // 128 threads: each thread does one column (scalar; region = 128 cols)
                int cc = n128 * 128 + tid;
                float sum = 0.f;
                #pragma unroll
                for (int sl = 0; sl < SPLITK; sl++) sum += wksp[(((size_t)sl * M + cr) * N) + cc];
                out[(size_t)cr * N + cc] = __float2bfloat16(sum);
                #pragma unroll
                for (int sl = 0; sl < SPLITK; sl++) wksp[(((size_t)sl * M + cr) * N) + cc] = 0.f;
            }
        }
    }
    STAMP(5);
    #undef STAMP
}

// ---------------------------------------------------------------------------
// host dispatch
// ---------------------------------------------------------------------------
static int smem_bytes(int stages, int mt) {
    return stages * (8192 + 512);
}

static unsigned long long* tbuf_ptr = nullptr;
void set_tbuf(uint64_t* p) { tbuf_ptr = (unsigned long long*)p; }
template <int MT, int SPLITK, int STAGES, int DBGM_>
static void launch_v4(dim3 grid, cudaStream_t stream, torch::Tensor& x, torch::Tensor& wr,
                      torch::Tensor& sc, torch::Tensor& zr, torch::Tensor& out,
                      torch::Tensor& wksp, torch::Tensor& counters, int M, int N, int K) {
    int bytes = smem_bytes(STAGES, MT);
    static int configured = 0;
    if (bytes > 48 * 1024 && !configured) {
        cudaError_t aerr = cudaFuncSetAttribute(w4_mma_kernel<MT, SPLITK, STAGES, DBGM_>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
        TORCH_CHECK(aerr == cudaSuccess, "attr set failed: ", cudaGetErrorString(aerr), " (bytes=", bytes, ")");
        configured = 1;
    }
    w4_mma_kernel<MT, SPLITK, STAGES, DBGM_><<<grid, 128, bytes, stream>>>(
        (const uint2*)wr.data_ptr(), (const __nv_bfloat16*)x.data_ptr(),
        (const __nv_bfloat16*)sc.data_ptr(), (const __nv_bfloat16*)zr.data_ptr(),
        (__nv_bfloat16*)out.data_ptr(), wksp.data_ptr<float>(), counters.data_ptr<int>(), M, N, K,
        tbuf_ptr);
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "kernel launch failed: ", cudaGetErrorString(err), " (bytes=", bytes, ")");
}

torch::Tensor w4_mma_set_tbuf(torch::Tensor t) { set_tbuf((uint64_t*)t.data_ptr()); return t; }

torch::Tensor w4_mma(torch::Tensor x, torch::Tensor wr, torch::Tensor sc, torch::Tensor zr,
                     torch::Tensor wksp, torch::Tensor counters,
                     long variant, long splitk, long dbg) {
    int M = x.size(0);
    int K = x.size(1);
    int N = sc.size(1);
    auto out = at::empty({M, N}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    #define LAUNCH(MT_, SPLITK_, STAGES_) \
        { dim3 grid(SPLITK_, N / 128, (M + 16 * MT_ - 1) / (16 * MT_)); \
          launch_v4<MT_, SPLITK_, STAGES_, 0>(grid, stream, x, wr, sc, zr, out, wksp, counters, M, N, K); }
    if (dbg == 1) { launch_v4<1, 4, 4, 1>(dim3(4, N/128, 1), stream, x, wr, sc, zr, out, wksp, counters, M, N, K); return out; }
    if (dbg == 2) { launch_v4<1, 4, 4, 2>(dim3(4, N/128, 1), stream, x, wr, sc, zr, out, wksp, counters, M, N, K); return out; }
    if (dbg == 3) { launch_v4<1, 4, 4, 3>(dim3(4, N/128, 1), stream, x, wr, sc, zr, out, wksp, counters, M, N, K); return out; }
    if (variant == 1 && splitk == 4) { LAUNCH(1, 4, 4); }
    else if (variant == 1 && splitk == 8) { LAUNCH(1, 8, 4); }
    else if (variant == 1 && splitk == 16) { LAUNCH(1, 16, 3); }
    else if (variant == 1 && splitk == 2) { LAUNCH(1, 2, 4); }
    else if (variant == 1 && splitk == 1) { LAUNCH(1, 1, 4); }
    else if (variant == 2 && splitk == 4) { LAUNCH(2, 4, 4); }
    else if (variant == 2 && splitk == 2) { LAUNCH(2, 2, 4); }
    else if (variant == 2 && splitk == 1) { LAUNCH(2, 1, 3); }
    else if (variant == 2 && splitk == 8) { LAUNCH(2, 8, 3); }
    else if (variant == 4 && splitk == 1) { LAUNCH(4, 1, 3); }
    else if (variant == 4 && splitk == 2) { LAUNCH(4, 2, 3); }
    else { TORCH_CHECK(false, "no such variant"); }
    return out;
    #undef LAUNCH
}
'''


_EXT = None


def _build_ext():
    global _EXT
    if _EXT is None:
        from torch.utils.cpp_extension import load_inline
        _EXT = load_inline(
            name="w4a16_ext_v6",
            cpp_sources=[_CPP_DECL],
            cuda_sources=[_GEMV_SRC, _MMA_SRC],
            functions=["repack", "gemv_w4", "w4_mma"],
            extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_100,code=sm_100"],
            verbose=False,
        )
    return _EXT


# ----------------------------------------------------------------------------
# Triton fused-dequant fallback (odd shapes)
# ----------------------------------------------------------------------------
import triton
import triton.language as tl


@triton.jit
def _w4_gemm_kernel(
    x_ptr, w_ptr, s_ptr, z_ptr, o_ptr,
    M, N, K, stride_xm, stride_wn, stride_om,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    m_mask = offs_m < M
    offs_ke = tl.arange(0, BK // 2)
    offs_pk = tl.arange(0, BK // 2)
    for kb in range(0, K // BK):
        x_even = tl.load(x_ptr + offs_m[:, None] * stride_xm + (kb * BK + 2 * offs_ke)[None, :],
                         mask=m_mask[:, None], other=0.0)
        x_odd = tl.load(x_ptr + offs_m[:, None] * stride_xm + (kb * BK + 2 * offs_ke + 1)[None, :],
                        mask=m_mask[:, None], other=0.0)
        wpt = tl.load(w_ptr + (kb * BK // 2 + offs_pk)[:, None] * stride_wn + offs_n[None, :])
        lo = (wpt & 0xF).to(tl.bfloat16)
        hi = (wpt >> 4).to(tl.bfloat16)
        g = (kb * BK) // 128
        s = tl.load(s_ptr + g * N + offs_n)
        z = tl.load(z_ptr + g * N + offs_n)
        lo = (lo - z) * s
        hi = (hi - z) * s
        acc = tl.dot(x_even, lo, acc)
        acc = tl.dot(x_odd, hi, acc)
    o_ptrs = o_ptr + offs_m[:, None] * stride_om + offs_n[None, :]
    tl.store(o_ptrs, acc.to(tl.bfloat16), mask=m_mask[:, None])


def _triton_w4_gemm(x, wq, sc, zr):
    M, K = x.shape
    N = wq.shape[1]
    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    BM = 16 if M <= 16 else (32 if M <= 32 else 64)
    BN, BK, warps, stages = 128, 128, 4, 4
    grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
    _w4_gemm_kernel[grid](x, wq, sc, zr, out, M, N, K, x.stride(0), wq.stride(0), out.stride(0),
                          BM=BM, BN=BN, BK=BK, num_warps=warps, num_stages=stages)
    return out


# ----------------------------------------------------------------------------
# Model
# ----------------------------------------------------------------------------
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
        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))
        self._wr = None
        self._g_wksp = None
        self._g_counters = None
        self._m_wksp = None
        self._m_counters = None
        self._device = None

    # -- weights repack hook --------------------------------------------------
    def load_state_dict(self, *args, **kwargs):
        ret = super().load_state_dict(*args, **kwargs)
        self._invalidate_repack()
        return ret

    def _invalidate_repack(self):
        self._wr = None  # rebuilt lazily on next forward

    def _plan(self, x):
        """Kernel dispatch config per (M, N). Returns (path, ...)."""
        M, N, K = x.shape[0], self.N, self.K
        if K % 2 or self.N % 128 or (self.K // self.group_size) != 32:
            return ("triton",)
        if M == 1 and N % 512 == 0:
            # gemv path: (ksplit, gs, wpg, unroll)
            if N >= 12288:
                return ("gemv", (16, 2, 2, 8))
            return ("gemv", (32, 1, 4, 8))
        if N % 128 == 0:
            if M <= 16:
                return ("mma", 1, 4 if N >= 12288 else 8)
            if M <= 32:
                return ("mma", 2, 4 if N >= 12288 else 8)
            return ("mma", 4, 1)
        return ("triton",)

    def _ensure(self, x):
        dev = x.device
        ext = _build_ext()
        if self._device != dev:
            self._device = dev
            self._wr = None
            self._g_wksp = None
            self._g_counters = None
            self._m_wksp = None
            self._m_counters = None
        if self._wr is None and self.N % 128 == 0 and self.w_q.is_cuda:
            self._wr = ext.repack(self.w_q)
        return ext

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._run(x)

    # thin __call__ to skip nn.Module hook machinery in the hot path
    __call__ = None  # replaced below

    def _run(self, x: torch.Tensor) -> torch.Tensor:
        ext = self._ensure(x)
        M = x.shape[0]
        N, K = self.N, self.K
        plan = self._plan(x)
        kind = plan[0]
        if kind == "gemv":
            ksplit, gs, wpg, unroll = plan[1]
            if self._g_wksp is None:
                self._g_wksp = torch.zeros(N, dtype=torch.float32, device=x.device)
                self._g_counters = torch.zeros(N // 512, dtype=torch.int32, device=x.device)
            return ext.gemv_w4(x.view(-1), self.w_q, self.scales, self.zeros,
                               self._g_wksp, self._g_counters, ksplit, gs, wpg, unroll).view(1, N)
        if kind == "mma":
            mt, ks = plan[1], plan[2]
            mblocks = (M + 16 * mt - 1) // (16 * mt)
            if self._m_wksp is None:
                self._m_wksp = torch.zeros(32 * M, N, dtype=torch.float32, device=x.device) if ks > 1 else torch.zeros(1, N, dtype=torch.float32, device=x.device)
                self._m_counters = torch.zeros(max(1, mblocks) * (N // 128), dtype=torch.int32, device=x.device)
            return ext.w4_mma(x, self._wr, self.scales, self.zeros,
                              self._m_wksp, self._m_counters, mt, ks, 0)
        return _triton_w4_gemm(x, self.w_q, self.scales, self.zeros)


def _model_call(self, x):
    return self._run(x)


Model.__call__ = _model_call

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]