import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

HIDDEN = 1024
INTERMEDIATE = 3072
NUM_Q = 16
NUM_KV = 8
HEAD_DIM = 128
NUM_LAYERS = 4
NOISE_CHUNK = 1024

_CPP_SRC = r"""
#include <torch/extension.h>
#include <pybind11/stl.h>

torch::Tensor decode_steps_cuda(torch::Tensor hidden, torch::Tensor noise,
                                std::vector<torch::Tensor> weights,
                                std::vector<torch::Tensor> k_caches,
                                std::vector<torch::Tensor> v_caches,
                                int64_t start_pos);
"""

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

namespace {

constexpr int HIDDEN = 1024;
constexpr int INTER = 3072;
constexpr int NQ = 16;
constexpr int NKV = 8;
constexpr int HD = 128;
constexpr int QDIM = NQ * HD;
constexpr int KVDIM = NKV * HD;
constexpr int QKV_DIM = QDIM + 2 * KVDIM;
constexpr int NLAYERS = 4;
constexpr int WPB = 11;
constexpr int THREADS = 256;
constexpr int ROWS = THREADS / 32;
constexpr int CHUNK = 128;
constexpr float EPS = 1e-6f;
constexpr float ATT_SCALE = 0.08838834764831845f;

using bf16 = __nv_bfloat16;

__device__ __forceinline__ float warp_sum(float v) {
#pragma unroll
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
  return v;
}

__device__ __forceinline__ float warp_max(float v) {
#pragma unroll
  for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o));
  return v;
}

__device__ __forceinline__ float round_bf16(float x) {
  return __bfloat162float(__float2bfloat16(x));
}

template <bool MAX>
__device__ __forceinline__ float block_reduce(float v, float* red) {
  const int lane = threadIdx.x & 31;
  const int nw = blockDim.x >> 5;
  v = MAX ? warp_max(v) : warp_sum(v);
  if (lane == 0) red[threadIdx.x >> 5] = v;
  __syncthreads();
  float r = lane < nw ? red[lane] : (MAX ? -INFINITY : 0.f);
  r = MAX ? warp_max(r) : warp_sum(r);
  __syncthreads();
  return r;
}

__device__ __forceinline__ void unpack8(const uint4 raw, float* f) {
  const __nv_bfloat162* p = reinterpret_cast<const __nv_bfloat162*>(&raw);
#pragma unroll
  for (int j = 0; j < 4; ++j) {
    const float2 t = __bfloat1622float2(p[j]);
    f[2 * j] = t.x;
    f[2 * j + 1] = t.y;
  }
}

__device__ __forceinline__ void unpack4(const uint2 raw, float* f) {
  const float2 a = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&raw.x));
  const float2 b = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&raw.y));
  f[0] = a.x;
  f[1] = a.y;
  f[2] = b.x;
  f[3] = b.y;
}

template <int K>
__device__ __forceinline__ void load_row(const bf16* __restrict__ wrow, int lane, uint4* regs) {
  constexpr int NV = K / 256;
  const uint4* wv = reinterpret_cast<const uint4*>(wrow);
#pragma unroll
  for (int i = 0; i < NV; ++i) regs[i] = wv[lane + 32 * i];
}

template <int K>
__device__ __forceinline__ float dot_regs(const uint4* regs, const float* __restrict__ sx, int lane) {
  constexpr int NV = K / 256;
  float acc = 0.f;
#pragma unroll
  for (int i = 0; i < NV; ++i) {
    float f[8];
    unpack8(regs[i], f);
    const float* xs = sx + 8 * (lane + 32 * i);
#pragma unroll
    for (int j = 0; j < 8; ++j) acc = fmaf(f[j], xs[j], acc);
  }
  return warp_sum(acc);
}

template <int K>
__device__ void rms_prologue(const float* __restrict__ x, const bf16* __restrict__ w, float* sx, float* red) {
  float ss = 0.f;
  for (int i = threadIdx.x; i < K; i += THREADS) {
    const float v = x[i];
    sx[i] = v;
    ss += v * v;
  }
  ss = block_reduce<false>(ss, red);
  const float rs = rsqrtf(ss / K + EPS);
  for (int i = threadIdx.x; i < K; i += THREADS) sx[i] = sx[i] * rs * __bfloat162float(w[i]);
  __syncthreads();
}

__global__ void mix_kernel(const bf16* __restrict__ noise, float* __restrict__ h) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  h[i] = round_bf16(0.5f * __bfloat162float(noise[i]) + 0.5f * h[i]);
}

__global__ void __launch_bounds__(THREADS) qkv_kernel(
    const float* __restrict__ x, const bf16* __restrict__ ln, const bf16* __restrict__ wq,
    const bf16* __restrict__ wk, const bf16* __restrict__ wv, float* __restrict__ qkv) {
  __shared__ float sx[HIDDEN];
  __shared__ float red[32];
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int r = blockIdx.x * ROWS + warp;
  const bf16* w;
  int rr;
  if (r < QDIM) {
    w = wq;
    rr = r;
  } else if (r < QDIM + KVDIM) {
    w = wk;
    rr = r - QDIM;
  } else {
    w = wv;
    rr = r - QDIM - KVDIM;
  }
  uint4 wr[HIDDEN / 256];
  load_row<HIDDEN>(w + (size_t)rr * HIDDEN, lane, wr);
  rms_prologue<HIDDEN>(x, ln, sx, red);
  const float acc = dot_regs<HIDDEN>(wr, sx, lane);
  if (lane == 0) qkv[r] = acc;
}

__global__ void qk_norm_rope_kernel(float* __restrict__ qkv, const bf16* __restrict__ qn,
                                    const bf16* __restrict__ kn, bf16* __restrict__ kc,
                                    bf16* __restrict__ vc, int max_seq, int pos) {
  const int head = blockIdx.x;
  const int lane = threadIdx.x;
  if (head >= NQ + NKV) {
    const int kvh = head - NQ - NKV;
    const float* src = qkv + QDIM + KVDIM + kvh * HD;
    bf16* dst = vc + ((size_t)kvh * max_seq + pos) * HD;
#pragma unroll
    for (int j = 0; j < 4; ++j) dst[lane + 32 * j] = __float2bfloat16(src[lane + 32 * j]);
    return;
  }
  const bool is_q = head < NQ;
  float* x = qkv + (is_q ? head * HD : QDIM + (head - NQ) * HD);
  const bf16* nw = is_q ? qn : kn;
  float v[4];
  float ss = 0.f;
#pragma unroll
  for (int j = 0; j < 4; ++j) {
    v[j] = x[lane + 32 * j];
    ss += v[j] * v[j];
  }
  ss = warp_sum(ss);
  const float rs = rsqrtf(ss / HD + EPS);
#pragma unroll
  for (int j = 0; j < 4; ++j) v[j] = v[j] * rs * __bfloat162float(nw[lane + 32 * j]);
#pragma unroll
  for (int p = 0; p < 2; ++p) {
    const int i = lane + 32 * p;
    const float inv = 1.0f / powf(10000.0f, (float)i / 64.0f);
    float s, c;
    sincosf((float)pos * inv, &s, &c);
    const float x1 = v[p];
    const float x2 = v[p + 2];
    v[p] = x1 * c - x2 * s;
    v[p + 2] = x1 * s + x2 * c;
  }
  if (is_q) {
#pragma unroll
    for (int j = 0; j < 4; ++j) x[lane + 32 * j] = v[j];
  } else {
    bf16* dst = kc + ((size_t)(head - NQ) * max_seq + pos) * HD;
#pragma unroll
    for (int j = 0; j < 4; ++j) dst[lane + 32 * j] = __float2bfloat16(v[j]);
  }
}

__global__ void __launch_bounds__(THREADS) attn_split_kernel(
    const float* __restrict__ qkv, const bf16* __restrict__ kc, const bf16* __restrict__ vc,
    int max_seq, int pos, int nch, float* __restrict__ pm, float* __restrict__ pl,
    float* __restrict__ po) {
  __shared__ float sm[ROWS][2];
  __shared__ float sl[ROWS][2];
  __shared__ float so[ROWS][2][HD];
  const int kvh = blockIdx.x;
  const int chunk = blockIdx.y;
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int c0 = chunk * CHUNK;
  const int c1 = min(c0 + CHUNK, pos + 1);
  const float4 qa = *reinterpret_cast<const float4*>(qkv + (2 * kvh) * HD + 4 * lane);
  const float4 qb = *reinterpret_cast<const float4*>(qkv + (2 * kvh + 1) * HD + 4 * lane);
  const size_t head_off = (size_t)kvh * max_seq * HD + 4 * lane;
  const bf16* kp = kc + head_off;
  const bf16* vp = vc + head_off;
  float ma = -INFINITY, mb = -INFINITY, la = 0.f, lb = 0.f;
  float oa[4] = {0.f, 0.f, 0.f, 0.f};
  float ob[4] = {0.f, 0.f, 0.f, 0.f};
  constexpr int U = 4;
  for (int base = c0 + warp * U; base < c1; base += ROWS * U) {
    uint2 kr[U], vr[U];
#pragma unroll
    for (int j = 0; j < U; ++j) {
      if (base + j < c1) {
        kr[j] = *reinterpret_cast<const uint2*>(kp + (size_t)(base + j) * HD);
        vr[j] = *reinterpret_cast<const uint2*>(vp + (size_t)(base + j) * HD);
      }
    }
#pragma unroll
    for (int j = 0; j < U; ++j) {
      if (base + j < c1) {
        float k[4], v[4];
        unpack4(kr[j], k);
        unpack4(vr[j], v);
        const float da = warp_sum(qa.x * k[0] + qa.y * k[1] + qa.z * k[2] + qa.w * k[3]) * ATT_SCALE;
        const float db = warp_sum(qb.x * k[0] + qb.y * k[1] + qb.z * k[2] + qb.w * k[3]) * ATT_SCALE;
        const float mna = fmaxf(ma, da);
        const float ca = expf(ma - mna);
        const float pa = expf(da - mna);
        la = la * ca + pa;
#pragma unroll
        for (int d4 = 0; d4 < 4; ++d4) oa[d4] = oa[d4] * ca + pa * v[d4];
        ma = mna;
        const float mnb = fmaxf(mb, db);
        const float cb = expf(mb - mnb);
        const float pb = expf(db - mnb);
        lb = lb * cb + pb;
#pragma unroll
        for (int d4 = 0; d4 < 4; ++d4) ob[d4] = ob[d4] * cb + pb * v[d4];
        mb = mnb;
      }
    }
  }
#pragma unroll
  for (int j = 0; j < 4; ++j) {
    so[warp][0][4 * lane + j] = oa[j];
    so[warp][1][4 * lane + j] = ob[j];
  }
  if (lane == 0) {
    sm[warp][0] = ma;
    sm[warp][1] = mb;
    sl[warp][0] = la;
    sl[warp][1] = lb;
  }
  __syncthreads();
  const int h = threadIdx.x >> 7;
  const int d = threadIdx.x & (HD - 1);
  float M = -INFINITY;
#pragma unroll
  for (int w = 0; w < ROWS; ++w) M = fmaxf(M, sm[w][h]);
  float L = 0.f;
  float O = 0.f;
#pragma unroll
  for (int w = 0; w < ROWS; ++w) {
    const float f = expf(sm[w][h] - M);
    L += sl[w][h] * f;
    O += so[w][h][d] * f;
  }
  const int qh = 2 * kvh + h;
  if (d == 0) {
    pm[qh * nch + chunk] = M;
    pl[qh * nch + chunk] = L;
  }
  po[((size_t)qh * nch + chunk) * HD + d] = O;
}

__global__ void __launch_bounds__(HD) attn_combine_kernel(
    const float* __restrict__ pm, const float* __restrict__ pl, const float* __restrict__ po,
    int nch, float* __restrict__ out) {
  extern __shared__ float sw[];
  __shared__ float red[32];
  const int qh = blockIdx.x;
  const int t = threadIdx.x;
  const float* m = pm + (size_t)qh * nch;
  const float* l = pl + (size_t)qh * nch;
  float mx = -INFINITY;
  for (int c = t; c < nch; c += HD) mx = fmaxf(mx, m[c]);
  mx = block_reduce<true>(mx, red);
  for (int c = t; c < nch; c += HD) sw[c] = expf(m[c] - mx);
  __syncthreads();
  float ls = 0.f;
  for (int c = t; c < nch; c += HD) ls += l[c] * sw[c];
  const float L = block_reduce<false>(ls, red);
  const float* o = po + (size_t)qh * nch * HD + t;
  float acc = 0.f;
  for (int c = 0; c < nch; ++c) acc += sw[c] * o[(size_t)c * HD];
  out[qh * HD + t] = acc / L;
}

__global__ void __launch_bounds__(THREADS) oproj_kernel(
    const float* __restrict__ attn, const bf16* __restrict__ wo, const float* __restrict__ resid,
    float* __restrict__ out) {
  __shared__ float sx[QDIM];
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int r = blockIdx.x * ROWS + warp;
  uint4 wr[QDIM / 256];
  load_row<QDIM>(wo + (size_t)r * QDIM, lane, wr);
  for (int i = threadIdx.x; i < QDIM; i += THREADS) sx[i] = attn[i];
  __syncthreads();
  const float acc = dot_regs<QDIM>(wr, sx, lane);
  if (lane == 0) out[r] = resid[r] + acc;
}

__global__ void __launch_bounds__(THREADS) gateup_kernel(
    const float* __restrict__ x, const bf16* __restrict__ ln, const bf16* __restrict__ wg,
    const bf16* __restrict__ wu, float* __restrict__ act) {
  __shared__ float sx[HIDDEN];
  __shared__ float red[32];
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int r = blockIdx.x * ROWS + warp;
  uint4 wgr[HIDDEN / 256], wur[HIDDEN / 256];
  load_row<HIDDEN>(wg + (size_t)r * HIDDEN, lane, wgr);
  load_row<HIDDEN>(wu + (size_t)r * HIDDEN, lane, wur);
  rms_prologue<HIDDEN>(x, ln, sx, red);
  const float g = dot_regs<HIDDEN>(wgr, sx, lane);
  const float u = dot_regs<HIDDEN>(wur, sx, lane);
  if (lane == 0) act[r] = g / (1.f + expf(-g)) * u;
}

__global__ void __launch_bounds__(THREADS) down_kernel(
    const float* __restrict__ act, const bf16* __restrict__ wd, const float* __restrict__ resid,
    float* __restrict__ out) {
  __shared__ float sx[INTER];
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int r = blockIdx.x * ROWS + warp;
  uint4 wr[INTER / 256];
  load_row<INTER>(wd + (size_t)r * INTER, lane, wr);
  for (int i = threadIdx.x; i < INTER; i += THREADS) sx[i] = act[i];
  __syncthreads();
  const float acc = dot_regs<INTER>(wr, sx, lane);
  if (lane == 0) out[r] = round_bf16(resid[r] + acc);
}

}

torch::Tensor decode_steps_cuda(torch::Tensor hidden, torch::Tensor noise,
                                std::vector<torch::Tensor> weights,
                                std::vector<torch::Tensor> k_caches,
                                std::vector<torch::Tensor> v_caches,
                                int64_t start_pos) {
  const int n_steps = static_cast<int>(noise.size(0));
  const int pos0 = static_cast<int>(start_pos);
  const int max_seq = static_cast<int>(k_caches[0].size(1));
  const int nch_max = (pos0 + n_steps + CHUNK - 1) / CHUNK;
  const auto f32 = hidden.options().dtype(torch::kFloat32);
  torch::Tensor hbuf = torch::empty({2 * HIDDEN}, f32);
  torch::Tensor qkv = torch::empty({QKV_DIM}, f32);
  torch::Tensor attn = torch::empty({QDIM}, f32);
  torch::Tensor res1 = torch::empty({HIDDEN}, f32);
  torch::Tensor act = torch::empty({INTER}, f32);
  torch::Tensor pm = torch::empty({NQ * nch_max}, f32);
  torch::Tensor pl = torch::empty({NQ * nch_max}, f32);
  torch::Tensor po = torch::empty({NQ * nch_max * HD}, f32);
  hbuf.narrow(0, 0, HIDDEN).copy_(hidden);
  float* hb = hbuf.data_ptr<float>();
  float* qkv_p = qkv.data_ptr<float>();
  float* attn_p = attn.data_ptr<float>();
  float* res1_p = res1.data_ptr<float>();
  float* act_p = act.data_ptr<float>();
  float* pm_p = pm.data_ptr<float>();
  float* pl_p = pl.data_ptr<float>();
  float* po_p = po.data_ptr<float>();
  std::vector<const bf16*> wp(NLAYERS * WPB);
  for (int i = 0; i < NLAYERS * WPB; ++i) wp[i] = reinterpret_cast<const bf16*>(weights[i].data_ptr());
  std::vector<bf16*> kp(NLAYERS), vp(NLAYERS);
  for (int b = 0; b < NLAYERS; ++b) {
    kp[b] = reinterpret_cast<bf16*>(k_caches[b].data_ptr());
    vp[b] = reinterpret_cast<bf16*>(v_caches[b].data_ptr());
  }
  const bf16* nz = reinterpret_cast<const bf16*>(noise.data_ptr());
  cudaStream_t s = at::cuda::getCurrentCUDAStream();
  for (int st = 0; st < n_steps; ++st) {
    const int pos = pos0 + st;
    const int nch = (pos + CHUNK) / CHUNK;
    mix_kernel<<<HIDDEN / THREADS, THREADS, 0, s>>>(nz + (size_t)st * HIDDEN, hb);
    for (int b = 0; b < NLAYERS; ++b) {
      const bf16* const* ws = wp.data() + b * WPB;
      const float* in = hb + (b & 1) * HIDDEN;
      float* out = hb + ((b + 1) & 1) * HIDDEN;
      qkv_kernel<<<QKV_DIM / ROWS, THREADS, 0, s>>>(in, ws[0], ws[1], ws[2], ws[3], qkv_p);
      qk_norm_rope_kernel<<<NQ + 2 * NKV, 32, 0, s>>>(qkv_p, ws[4], ws[5], kp[b], vp[b], max_seq, pos);
      attn_split_kernel<<<dim3(NKV, nch), THREADS, 0, s>>>(qkv_p, kp[b], vp[b], max_seq, pos, nch, pm_p, pl_p, po_p);
      attn_combine_kernel<<<NQ, HD, nch * sizeof(float), s>>>(pm_p, pl_p, po_p, nch, attn_p);
      oproj_kernel<<<HIDDEN / ROWS, THREADS, 0, s>>>(attn_p, ws[6], in, res1_p);
      gateup_kernel<<<INTER / ROWS, THREADS, 0, s>>>(res1_p, ws[7], ws[8], ws[9], act_p);
      down_kernel<<<HIDDEN / ROWS, THREADS, 0, s>>>(act_p, ws[10], res1_p, out);
    }
  }
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return hbuf.narrow(0, 0, HIDDEN).to(torch::kBFloat16);
}
"""

_ext = load_inline(
    name="megaqwen_decode_v1",
    cpp_sources=[_CPP_SRC],
    cuda_sources=[_CUDA_SRC],
    functions=["decode_steps_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "-lineinfo", "-gencode=arch=compute_120,code=sm_120"],
    verbose=False,
)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        H, I, D = HIDDEN, INTERMEDIATE, HEAD_DIM
        self.input_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
        self.q_proj = nn.Parameter(torch.empty(NUM_Q * D, H, dtype=torch.bfloat16))
        self.k_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
        self.v_proj = nn.Parameter(torch.empty(NUM_KV * D, H, dtype=torch.bfloat16))
        self.q_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
        self.k_norm = nn.Parameter(torch.ones(D, dtype=torch.bfloat16))
        self.o_proj = nn.Parameter(torch.empty(H, NUM_Q * D, dtype=torch.bfloat16))
        self.post_ln = nn.Parameter(torch.ones(H, dtype=torch.bfloat16))
        self.gate_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
        self.up_proj = nn.Parameter(torch.empty(I, H, dtype=torch.bfloat16))
        self.down_proj = nn.Parameter(torch.empty(H, I, dtype=torch.bfloat16))
        for p in (self.q_proj, self.k_proj, self.v_proj, self.o_proj,
                  self.gate_proj, self.up_proj, self.down_proj):
            nn.init.normal_(p, std=0.02)


class Model(nn.Module):
    def __init__(self, num_layers: int = NUM_LAYERS, max_seq: int = 131072):
        super().__init__()
        self.num_layers = num_layers
        self.max_seq = max_seq
        self.blocks = nn.ModuleList([Block() for _ in range(num_layers)])


def _weights(model):
    out = []
    for blk in model.blocks:
        out += [blk.input_ln, blk.q_proj, blk.k_proj, blk.v_proj, blk.q_norm, blk.k_norm,
                blk.o_proj, blk.post_ln, blk.gate_proj, blk.up_proj, blk.down_proj]
    return out


def _caches(model, device):
    shape = (NUM_KV, model.max_seq, HEAD_DIM)
    k = [torch.empty(shape, device=device, dtype=torch.bfloat16) for _ in range(model.num_layers)]
    v = [torch.empty(shape, device=device, dtype=torch.bfloat16) for _ in range(model.num_layers)]
    return k, v


def _advance(model, h, k_caches, v_caches, start, n_steps, noise_seed):
    gen = torch.Generator(device="cpu")
    gen.manual_seed(noise_seed)
    weights = _weights(model)
    device = h.device
    for done in range(0, n_steps, NOISE_CHUNK):
        m = min(NOISE_CHUNK, n_steps - done)
        noise = torch.stack(
            [torch.randn(HIDDEN, generator=gen, dtype=torch.bfloat16) for _ in range(m)]
        ).to(device)
        h = _ext.decode_steps_cuda(h, noise, weights, k_caches, v_caches, start + done)
    return h


@torch.no_grad()
def prefill(model, ctx_len, seed):
    device = model.blocks[0].q_proj.device
    assert ctx_len <= model.max_seq
    gen = torch.Generator(device="cpu")
    gen.manual_seed(seed)
    h = torch.randn(HIDDEN, generator=gen, dtype=torch.bfloat16).to(device)
    k_caches, v_caches = _caches(model, device)
    h = _advance(model, h, k_caches, v_caches, 0, ctx_len, seed + 1)
    return h, k_caches, v_caches


@torch.no_grad()
def decode_steps(model, hidden, k_caches, v_caches, start_pos, n_steps, seed):
    h = _advance(model, hidden, k_caches, v_caches, start_pos, n_steps, seed + 2)
    return h, k_caches, v_caches


@torch.no_grad()
def run(ctx_len, decode_steps, seed, model=None):
    if model is None:
        model = Model(NUM_LAYERS, max(ctx_len + decode_steps, 512)).cuda()
    elif model.max_seq < ctx_len + decode_steps:
        raise ValueError(
            f"model.max_seq={model.max_seq} < ctx_len+decode_steps={ctx_len + decode_steps}"
        )
    h, k_caches, v_caches = prefill(model, ctx_len, seed)
    h = _advance(model, h, k_caches, v_caches, ctx_len, decode_steps, seed + 2)
    return {"last_hidden": h}