from __future__ import annotations

import math
from typing import List

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


_CPP = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime_api.h>
#include <vector>
constexpr int NW = 147;
constexpr int NS = 25;
constexpr int NT = 29;
struct Args {
  void* w[NW];
  void* s[NS];
  void* t[NT];
  int ctx;
};
extern "C" cudaError_t launch_coop(const Args*, cudaStream_t);
void launch_step(std::vector<torch::Tensor> w,
                 std::vector<torch::Tensor> s,
                 std::vector<torch::Tensor> t,
                 int64_t context) {
  TORCH_CHECK(w.size() == NW && s.size() == NS && t.size() == NT,
              "internal pointer table size mismatch");
  Args a{};
  for (int i = 0; i < NW; ++i) a.w[i] = w[i].data_ptr();
  for (int i = 0; i < NS; ++i) a.s[i] = s[i].data_ptr();
  for (int i = 0; i < NT; ++i) a.t[i] = t[i].data_ptr();
  a.ctx = (int)context;
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const cudaError_t err = launch_coop(&a, stream);
  TORCH_CHECK(err == cudaSuccess, "cooperative launch failed: ", cudaGetErrorString(err));
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &launch_step);
}
'''

_CUDA = r'''
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <vector>
#include <cmath>

namespace cg = cooperative_groups;
using bf = __nv_bfloat16;
constexpr int NW = 147;
constexpr int NS = 25;
constexpr int NT = 29;
constexpr int MAXOUT = 6144;
constexpr int PST = 8 * MAXOUT;
constexpr int BLOCKS = 188;
constexpr int THREADS = 128;

struct Args {
  void* w[NW];
  void* s[NS];
  void* t[NT];
  int ctx;
};

__device__ __forceinline__ float bfload(const bf* p, int i) {
  return __bfloat162float(p[i]);
}
__device__ __forceinline__ void bfstore(bf* p, int i, float x) {
  p[i] = __float2bfloat16_rn(x);
}
__device__ __forceinline__ float silu(float x) {
  return x / (1.0f + __expf(-x));
}
__device__ __forceinline__ float softplus(float x) {
  return fmaxf(x, 0.0f) + log1pf(__expf(-fabsf(x)));
}
__device__ __forceinline__ float qweight(const uint8_t* wq, const bf* scales,
                                          const bf* zeros, int K, int N,
                                          int k, int n) {
  const uint8_t b = wq[(k >> 1) * N + n];
  const int q = (k & 1) ? (b >> 4) : (b & 15);
  const int gi = (k >> 7) * N + n;
  const float z = bfload(zeros, gi);
  const float sc = bfload(scales, gi);
  const bf d = __float2bfloat16_rn((float)q - z);
  const bf v = __float2bfloat16_rn(__bfloat162float(d) * sc);
  return __bfloat162float(v);
}
__device__ __forceinline__ float xload(const void* x, int i, bool fp32) {
  return fp32 ? ((const float*)x)[i] : bfload((const bf*)x, i);
}
__device__ __forceinline__ float gemv_one(const uint8_t* wq, const bf* scales,
                                           const bf* zeros, const void* x,
                                           int K, int N, int n, int split,
                                           int split_id, bool xfp32) {
  const int chunk = K / split;
  const int ks = split_id * chunk;
  const int ke = ks + chunk;
  float acc = 0.0f;
  for (int g = ks >> 7; g < (ke >> 7); ++g) {
    const int r0 = g * 64;
    const int r1 = r0 + 64;
    const float z = bfload(zeros, g * N + n);
    const float sc = bfload(scales, g * N + n);
    for (int r = r0; r < r1; ++r) {
      const uint8_t b = wq[r * N + n];
      const bf d0 = __float2bfloat16_rn((float)(b & 15) - z);
      const bf d1 = __float2bfloat16_rn((float)(b >> 4) - z);
      const float w0 = __bfloat162float(__float2bfloat16_rn(__bfloat162float(d0) * sc));
      const float w1 = __bfloat162float(__float2bfloat16_rn(__bfloat162float(d1) * sc));
      acc = fmaf(xload(x, 2 * r, xfp32), w0, acc);
      acc = fmaf(xload(x, 2 * r + 1, xfp32), w1, acc);
    }
  }
  return acc;
}

__device__ __forceinline__ float warp_sum(float x) {
  for (int d = 16; d; d >>= 1) x += __shfl_down_sync(0xffffffff, x, d);
  return x;
}
__device__ __forceinline__ float block_sum(float x, float* sh) {
  const int lane = threadIdx.x & 31;
  const int warp = threadIdx.x >> 5;
  x = warp_sum(x);
  if (lane == 0) sh[warp] = x;
  __syncthreads();
  if (warp == 0) {
    x = lane < 4 ? sh[lane] : 0.0f;
    x = warp_sum(x);
    if (lane == 0) sh[0] = x;
  }
  __syncthreads();
  return sh[0];
}
__device__ __forceinline__ float block_sum_root(float x, float* sh) {
  const int lane = threadIdx.x & 31;
  const int warp = threadIdx.x >> 5;
  x = warp_sum(x);
  if (lane == 0) sh[warp] = x;
  __syncthreads();
  if (warp == 0) {
    x = lane < 4 ? sh[lane] : 0.0f;
    x = warp_sum(x);
    if (lane == 0) sh[0] = x;
  }
  return threadIdx.x == 0 ? sh[0] : 0.0f;
}
__device__ __forceinline__ float block_max(float x, float* sh) {
  const int lane = threadIdx.x & 31;
  const int warp = threadIdx.x >> 5;
  for (int d = 16; d; d >>= 1) x = fmaxf(x, __shfl_down_sync(0xffffffff, x, d));
  if (lane == 0) sh[warp] = x;
  __syncthreads();
  if (warp == 0) {
    x = lane < 4 ? sh[lane] : -1.0e30f;
    for (int d = 16; d; d >>= 1) x = fmaxf(x, __shfl_down_sync(0xffffffff, x, d));
    if (lane == 0) sh[0] = x;
  }
  __syncthreads();
  return sh[0];
}

__device__ void rms(const bf* x, const bf* weight, bf* y, float* sh) {
  const int lane = threadIdx.x & 31;
  float ss = 0.0f;
  for (int i = lane; i < 2304; i += 32) {
    const float v = bfload(x, i);
    ss = fmaf(v, v, ss);
  }
  ss = warp_sum(ss);
  ss = __shfl_sync(0xffffffff, ss, 0);
  const float inv = rsqrtf(ss * (1.0f / 2304.0f) + 1.0e-6f);
  for (int i = threadIdx.x; i < 2304; i += THREADS)
    bfstore(y, i, bfload(x, i) * inv * bfload(weight, i));
}

__device__ void reduce_gemv(const float* p, int N, int tile, int splits, bf* out) {
  const int n0 = tile * 256 + threadIdx.x;
  const int n1 = n0 + 128;
  if (n0 < N) {
    float v = 0.0f;
    for (int s = 0; s < splits; ++s) v += p[s * N + n0];
    bfstore(out, n0, v);
  }
  if (n1 < N) {
    float v = 0.0f;
    for (int s = 0; s < splits; ++s) v += p[s * N + n1];
    bfstore(out, n1, v);
  }
}
__device__ void write_partial(const Args& a, int wbase, const void* x,
                              int K, int N, int tile, int splits, int split_id,
                              bool xfp32, float* partial, int pslot,
                              int expert, bool expert_bank) {
  const int n0 = tile * 256 + threadIdx.x;
  const int n1 = n0 + 128;
  const int eszq = (K / 2) * N;
  const int eszg = (K / 128) * N;
  const uint8_t* wq = (const uint8_t*)a.w[wbase];
  const bf* sc = (const bf*)a.w[wbase + 1];
  const bf* ze = (const bf*)a.w[wbase + 2];
  if (expert_bank) {
    wq += (size_t)expert * eszq;
    sc += (size_t)expert * eszg;
    ze += (size_t)expert * eszg;
  }
  const int po = pslot * PST + split_id * N;
  if (n0 < N) partial[po + n0] = gemv_one(wq, sc, ze, x, K, N, n0, splits, split_id, xfp32);
  if (n1 < N) partial[po + n1] = gemv_one(wq, sc, ze, x, K, N, n1, splits, split_id, xfp32);
}

__device__ void router_one(const Args& a, int wb, const bf* x, bf* logits,
                           int e, float* sh) {
  const bf* w = (const bf*)a.w[wb];
  float sum = 0.0f;
  for (int k = threadIdx.x; k < 2304; k += THREADS)
    sum = fmaf(bfload(x, k), bfload(w, e * 2304 + k), sum);
  sum = block_sum_root(sum, sh);
  if (threadIdx.x == 0) bfstore(logits, e, sum);
}

__device__ void beta_one(const Args& a, int wb, const bf* x, bf* beta,
                         int e, float* sh) {
  const bf* w = (const bf*)a.w[wb];
  float sum = 0.0f;
  for (int k = threadIdx.x; k < 2304; k += THREADS)
    sum = fmaf(bfload(x, k), bfload(w, e * 2304 + k), sum);
  sum = block_sum_root(sum, sh);
  if (threadIdx.x == 0) bfstore(beta, e, sum);
}

__device__ void conv_one(const Args& a, int wbase, const bf* val,
                         const bf* old, bf* next, bf* out, int which, int tile) {
  const bf* cw = (const bf*)a.w[wbase];
  for (int z = 0; z < 2; ++z) {
    const int c = tile * 256 + threadIdx.x + z * 128;
    if (c >= 4096) continue;
    const float z0 = bfload(old, 0 * 4096 + c);
    const float z1 = bfload(old, 1 * 4096 + c);
    const float z2 = bfload(old, 2 * 4096 + c);
    const float z3 = bfload(val, c);
    float acc = z0 * bfload(cw, which * 4096 * 4 + c * 4 + 0);
    acc = fmaf(z1, bfload(cw, which * 4096 * 4 + c * 4 + 1), acc);
    acc = fmaf(z2, bfload(cw, which * 4096 * 4 + c * 4 + 2), acc);
    acc = fmaf(z3, bfload(cw, which * 4096 * 4 + c * 4 + 3), acc);
    bfstore(out, c, silu(acc));
    bfstore(next, c, z1);
    bfstore(next, 4096 + c, z2);
    bfstore(next, 8192 + c, z3);
  }
}

__device__ void kda_state(const Args& a, int sb, const bf* q, const bf* k,
                          const bf* v, const bf* g, const bf* beta,
                          float* S, bf* out, int head) {
  const int vi = threadIdx.x;
  const float b = 1.0f / (1.0f + __expf(-bfload(beta, head)));
  float pred = 0.0f;
  for (int ki = 0; ki < 128; ++ki) {
    const int si = (head * 128 + ki) * 128 + vi;
    const float decay = __expf(-softplus(bfload(g, head * 128 + ki)));
    pred = fmaf(S[si] * decay, bfload(k, head * 128 + ki), pred);
  }
  const float delta = bfload(v, head * 128 + vi) - pred;
  float ov = 0.0f;
  for (int ki = 0; ki < 128; ++ki) {
    const int si = (head * 128 + ki) * 128 + vi;
    const float kval = bfload(k, head * 128 + ki);
    const float decay = __expf(-softplus(bfload(g, head * 128 + ki)));
    const float sv = S[si] * decay + b * kval * delta;
    S[si] = sv;
    ov = fmaf(sv, bfload(q, head * 128 + ki) * 0.08838834764831845f, ov);
  }
  bfstore(out, head * 128 + vi, ov);
}

__device__ void choose_top8(const bf* logits, int32_t* ids, float* weights) {
  bool used[64];
  float selected[8];
  for (int i = 0; i < 64; ++i) used[i] = false;
  float mx = -1.0e30f;
  for (int j = 0; j < 8; ++j) {
    int bi = 0;
    float bv = -1.0e30f;
    for (int i = 0; i < 64; ++i) {
      const float v = bfload(logits, i);
      if (!used[i] && v > bv) { bv = v; bi = i; }
    }
    used[bi] = true;
    ids[j] = bi;
    selected[j] = bv;
    mx = fmaxf(mx, bv);
  }
  float den = 0.0f;
  for (int j = 0; j < 8; ++j) den += __expf(selected[j] - mx);
  for (int j = 0; j < 8; ++j) weights[j] = __expf(selected[j] - mx) / den * 2.446f;
}

__device__ void expert_gate_up(const Args& a, int gb, int ub, const bf* x,
                               const int32_t* ids, float* gp, float* up,
                               int slot, int tile, int splits, int split) {
  const int expert = slot == 8 ? 0 : ids[slot];
  const int n0 = tile * 256 + threadIdx.x;
  const int n1 = n0 + 128;
  const int eszq = (2304 / 2) * 1024;
  const int eszg = (2304 / 128) * 1024;
  const uint8_t* gw = (const uint8_t*)a.w[gb] + (size_t)expert * eszq;
  const bf* gs = (const bf*)a.w[gb + 1] + (size_t)expert * eszg;
  const bf* gz = (const bf*)a.w[gb + 2] + (size_t)expert * eszg;
  const uint8_t* uw = (const uint8_t*)a.w[ub] + (size_t)expert * eszq;
  const bf* us = (const bf*)a.w[ub + 1] + (size_t)expert * eszg;
  const bf* uz = (const bf*)a.w[ub + 2] + (size_t)expert * eszg;
  const int po = (slot * splits + split) * 1024;
  if (n0 < 1024) {
    gp[po + n0] = gemv_one(gw, gs, gz, x, 2304, 1024, n0, splits, split, false);
    up[po + n0] = gemv_one(uw, us, uz, x, 2304, 1024, n0, splits, split, false);
  }
  if (n1 < 1024) {
    gp[po + n1] = gemv_one(gw, gs, gz, x, 2304, 1024, n1, splits, split, false);
    up[po + n1] = gemv_one(uw, us, uz, x, 2304, 1024, n1, splits, split, false);
  }
}

__device__ void expert_down(const Args& a, int db, const float* x,
                            const int32_t* ids, float* p, int slot,
                            int tile, int splits, int split) {
  const int expert = slot == 8 ? 0 : ids[slot];
  const int n0 = tile * 256 + threadIdx.x;
  const int n1 = n0 + 128;
  const int eszq = (1024 / 2) * 2304;
  const int eszg = (1024 / 128) * 2304;
  const uint8_t* wq = (const uint8_t*)a.w[db] + (size_t)expert * eszq;
  const bf* sc = (const bf*)a.w[db + 1] + (size_t)expert * eszg;
  const bf* ze = (const bf*)a.w[db + 2] + (size_t)expert * eszg;
  const int po = (slot * splits + split) * 2304;
  if (n0 < 2304) p[po + n0] = gemv_one(wq, sc, ze, x, 1024, 2304, n0, splits, split, true);
  if (n1 < 2304) p[po + n1] = gemv_one(wq, sc, ze, x, 1024, 2304, n1, splits, split, true);
}

__device__ void rope_make(const Args& a, const bf* qfull, const bf* kv,
                          bf* qrot, bf* krot) {
  const int pos = a.ctx;
  for (int z = threadIdx.x; z < 32 * 32; z += THREADS) {
    const int h = z / 32;
    const int i = z & 31;
    const float inv = expf(-logf(10000.0f) * ((2.0f * i) / 64.0f));
    const float ang = (float)pos * inv;
    const float c = cosf(ang), s = sinf(ang);
    const float qe = bfload(qfull, h * 192 + 128 + 2 * i);
    const float qo = bfload(qfull, h * 192 + 129 + 2 * i);
    bfstore(qrot, h * 64 + 2 * i, qe * c - qo * s);
    bfstore(qrot, h * 64 + 2 * i + 1, qe * s + qo * c);
  }
  if (threadIdx.x < 32) {
    const int i = threadIdx.x;
    const float inv = expf(-logf(10000.0f) * ((2.0f * i) / 64.0f));
    const float ang = (float)pos * inv;
    const float c = cosf(ang), s = sinf(ang);
    const float ke = bfload(kv, 512 + 2 * i);
    const float ko = bfload(kv, 513 + 2 * i);
    bfstore(krot, 2 * i, ke * c - ko * s);
    bfstore(krot, 2 * i + 1, ke * s + ko * c);
  }
}

__device__ void rope_cache(const Args& a, const bf* kv, const bf* krot) {
  const int oldN = a.ctx * 512;
  bf* cn = (bf*)a.s[23];
  const bf* co = (const bf*)a.s[21];
  if (cn != co) {
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < oldN; i += gridDim.x * THREADS) cn[i] = co[i];
  }
  if (blockIdx.x == 0) {
    for (int i = threadIdx.x; i < 512; i += THREADS) cn[oldN + i] = kv[i];
  }
  bf* rn = (bf*)a.s[24];
  const bf* ro = (const bf*)a.s[22];
  if (rn != ro) {
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < a.ctx * 64; i += gridDim.x * THREADS) rn[i] = ro[i];
  }
  if (blockIdx.x < 64 && threadIdx.x == 0)
    rn[a.ctx * 64 + blockIdx.x] = krot[blockIdx.x];
}

__device__ void absorb_q(const Args& a, int wb, const bf* q, float* qlat,
                         int block) {
  if (block >= 160) return;
  const int h = block / 5;
  const int split = block % 5;
  const int begin = split * 104;
  const int end = min(begin + 104, 512);
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const uint8_t* wq = (const uint8_t*)a.w[wb];
  const bf* sc = (const bf*)a.w[wb + 1];
  const bf* ze = (const bf*)a.w[wb + 2];
  for (int kb = begin; kb < end; kb += 4) {
    const int k = kb + warp;
    if (k >= end) continue;
    float v = 0.0f;
    for (int j = lane; j < 128; j += 32) {
      const int n = h * 256 + j;
      const float w = qweight(wq, sc, ze, 512, 8192, k, n);
      v = fmaf(w, bfload(q, h * 192 + j), v);
    }
    v = warp_sum(v);
    if (lane == 0) qlat[h * 512 + k] = v;
  }
}

__device__ void mla_score(const Args& a, const bf* qrot, const bf* kcache,
                          const bf* ccache, const float* qlat, float* scores,
                          int block, int parts) {
  if (block >= 32 * parts) return;
  const int h = block % 32;
  const int part = block / 32;
  const int L = a.ctx + 1;
  const int chunk = (L + parts - 1) / parts;
  const int begin = part * chunk;
  const int end = min(begin + chunk, L);
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  for (int tb = begin; tb < end; tb += 4) {
    const int t = tb + warp;
    if (t >= end) continue;
    float dot = 0.0f;
    for (int k = 0; k < 512; k += 32) {
      const float x = bfload(ccache, t * 512 + k + lane);
      dot = fmaf(x, qlat[h * 512 + k + lane], dot);
    }
    dot = warp_sum(dot);
    float rope = 0.0f;
    for (int j = lane; j < 64; j += 32)
      rope = fmaf(bfload(qrot, h * 64 + j), bfload(kcache, t * 64 + j), rope);
    rope = warp_sum(rope);
    if (lane == 0) scores[h * L + t] = (dot + rope) * 0.07216878364870322f;
  }
}

__device__ void mla_softmax(const float* scores, float* probs, int L,
                            int head, float* sh) {
  float mx = -1.0e30f;
  for (int t = threadIdx.x; t < L; t += THREADS)
    mx = fmaxf(mx, scores[head * L + t]);
  mx = block_max(mx, sh);
  float den = 0.0f;
  for (int t = threadIdx.x; t < L; t += THREADS)
    den += __expf(scores[head * L + t] - mx);
  den = block_sum(den, sh);
  for (int t = threadIdx.x; t < L; t += THREADS)
    probs[head * L + t] = __expf(scores[head * L + t] - mx) / den;
  __syncthreads();
}

__device__ void latent_reduce_part(const bf* cache, const float* probs,
                                   float* partial, int L, int h,
                                   int part, int parts) {
  float acc[4] = {0.f, 0.f, 0.f, 0.f};
  const int chunk = (L + parts - 1) / parts;
  const int begin = part * chunk;
  const int end = min(begin + chunk, L);
  for (int t = begin; t < end; ++t) {
    const float p = probs[h * L + t];
    const int k0 = threadIdx.x;
    const int k1 = k0 + 128;
    const int k2 = k0 + 256;
    const int k3 = k0 + 384;
    acc[0] = fmaf(p, bfload(cache, t * 512 + k0), acc[0]);
    acc[1] = fmaf(p, bfload(cache, t * 512 + k1), acc[1]);
    acc[2] = fmaf(p, bfload(cache, t * 512 + k2), acc[2]);
    acc[3] = fmaf(p, bfload(cache, t * 512 + k3), acc[3]);
  }
  const int base = (h * parts + part) * 512 + threadIdx.x;
  partial[base] = acc[0];
  partial[base + 128] = acc[1];
  partial[base + 256] = acc[2];
  partial[base + 384] = acc[3];
}

__device__ void latent_merge(const float* partial, float* latent,
                             int h, int parts) {
  for (int z = 0; z < 4; ++z) {
    const int k = threadIdx.x + z * 128;
    float sum = 0.0f;
    for (int p = 0; p < parts; ++p)
      sum += partial[(h * parts + p) * 512 + k];
    latent[h * 512 + k] = sum;
  }
}

__device__ void v_project(const Args& a, int wb, const float* latent,
                          bf* out, int h) {
  const uint8_t* wq = (const uint8_t*)a.w[wb];
  const bf* sc = (const bf*)a.w[wb + 1];
  const bf* ze = (const bf*)a.w[wb + 2];
  const int n = threadIdx.x;
  float acc = 0.0f;
  const int col = h * 256 + 128 + n;
  for (int k = 0; k < 512; ++k)
    acc = fmaf(latent[h * 512 + k], qweight(wq, sc, ze, 512, 8192, k, col), acc);
  bfstore(out, h * 128 + n, acc);
}

__global__ void decode_kernel(Args a) {
  cg::grid_group grid = cg::this_grid();
  __shared__ float red[128];
  const int bid = blockIdx.x;
  const int tid = threadIdx.x;
  const int mla_parts = gridDim.x >= 320 ? 10 : 5;
  bf* norm = (bf*)a.t[0];
  bf* hres = (bf*)a.t[1];
  bf* hA = (bf*)a.t[2];
  bf* hB = (bf*)a.t[26];
  bf* q = (bf*)a.t[3];
  bf* kbuf = (bf*)a.t[4];
  bf* vbuf = (bf*)a.t[5];
  bf* gbuf = (bf*)a.t[6];
  bf* beta = (bf*)a.t[7];
  bf* attin = (bf*)a.t[8];
  bf* attout = (bf*)a.t[9];
  float* pp = (float*)a.t[10];
  float* pp2 = (float*)a.t[11];
  bf* logits = (bf*)a.t[12];
  int32_t* ids = (int32_t*)a.t[13];
  float* routew = (float*)a.t[14];
  float* gatep = (float*)a.t[15];
  float* upp = (float*)a.t[16];
  float* acts = (float*)a.t[17];
  float* downp = (float*)a.t[18];
  bf* moeout = (bf*)a.t[19];
  float* qlat = (float*)a.t[20];
  float* scores = (float*)a.t[21];
  float* probs = (float*)a.t[22];
  float* latout = (float*)a.t[23];
  float* latpart = (float*)a.t[28];
  bf* qrot = (bf*)a.t[24];
  bf* krot = (bf*)a.t[25];
  const bf* cur = (const bf*)a.t[27];
  // t[27] aliases the original token and is replaced by the host before launch.

  for (int layer = 0; layer < 4; ++layer) {
    const int wb = layer < 3 ? layer * 38 : 114;
    const int is_kda = layer < 3;
    bf* next = (layer & 1) ? hB : hA;
    if (bid == 0) rms(cur, (const bf*)a.w[wb + 0], norm, red);
    grid.sync();

    if (is_kda) {
      constexpr int S2304 = 6;
      constexpr int S4096 = 8;
      constexpr int S1024 = 4;
      constexpr int QKV_BLOCKS = 4 * 16 * S2304;
      if (bid < QKV_BLOCKS) {
        const int p = bid / (16 * S2304);
        const int within = bid % (16 * S2304);
        const int tile = within / S2304;
        const int split = within % S2304;
        write_partial(a, wb + 2 + p * 3, norm, 2304, 4096, tile,
                      S2304, split, false, pp, p, 0, false);
      }
      if (bid >= QKV_BLOCKS && bid < QKV_BLOCKS + 32)
        beta_one(a, wb + 14, norm, beta, bid - QKV_BLOCKS, red);
      grid.sync();
      if (bid < 64) {
        const int p = bid / 16;
        const int tile = bid & 15;
        bf* out = p == 0 ? q : (p == 1 ? kbuf : (p == 2 ? vbuf : gbuf));
        reduce_gemv(pp + p * PST, 4096, tile, S2304, out);
      }
      grid.sync();
      if (bid < 48) {
        const int which = bid % 3;
        const int tile = bid / 3;
        const int layer_s = layer * 7;
        const bf* val = which == 0 ? q : (which == 1 ? kbuf : vbuf);
        bf* out = which == 0 ? q : (which == 1 ? kbuf : vbuf);
        const bf* old = (const bf*)a.s[layer_s + 1 + which];
        bf* win = (bf*)a.s[layer_s + 4 + which];
        conv_one(a, wb + 15, val, old, win, out, which, tile);
      }
      grid.sync();
      if (bid < 32)
        kda_state(a, layer * 7, q, kbuf, vbuf, gbuf, beta,
                  (float*)a.s[layer * 7], attin, bid);
      grid.sync();
      if (bid < 9 * S4096) {
        const int tile = bid / S4096, split = bid % S4096;
        write_partial(a, wb + 16, attin, 4096, 2304, tile,
                      S4096, split, false, pp2, 0, 0, false);
      }
      grid.sync();
      if (bid < 9) reduce_gemv(pp2, 2304, bid, S4096, attout);
      grid.sync();
    } else {
      constexpr int S2304 = 6;
      constexpr int S4096 = 8;
      constexpr int Q_BLOCKS = 24 * S2304;
      constexpr int KV_BLOCKS = 3 * S2304;
      if (bid < Q_BLOCKS + KV_BLOCKS) {
        int p, tile, split;
        if (bid < Q_BLOCKS) { p = 0; tile = bid / S2304; split = bid % S2304; }
        else { int z = bid - Q_BLOCKS; p = 1; tile = z / S2304; split = z % S2304; }
        const int N = p == 0 ? 6144 : 576;
        write_partial(a, wb + (p == 0 ? 2 : 5), norm, 2304, N,
                      tile, S2304, split, false, pp, p, 0, false);
      }
      grid.sync();
      if (bid < 24) reduce_gemv(pp, 6144, bid, S2304, q);
      if (bid >= 24 && bid < 27) reduce_gemv(pp + PST, 576, bid - 24, S2304, kbuf);
      grid.sync();
      if (bid == 0) rope_make(a, q, kbuf, qrot, krot);
      grid.sync();
      rope_cache(a, kbuf, krot);
      grid.sync();
      if (bid < 160) absorb_q(a, wb + 8, q, qlat, bid);
      grid.sync();
      if (bid < 32 * mla_parts) mla_score(a, qrot, (const bf*)a.s[24],
              (const bf*)a.s[23], qlat, scores, bid, mla_parts);
      grid.sync();
      if (bid < 32) mla_softmax(scores, probs, a.ctx + 1, bid, red);
      grid.sync();
      if (bid < 32 * mla_parts)
        latent_reduce_part((const bf*)a.s[23], probs, latpart, a.ctx + 1,
                           bid % 32, bid / 32, mla_parts);
      grid.sync();
      if (bid < 32) latent_merge(latpart, latout, bid, mla_parts);
      grid.sync();
      if (bid < 32) v_project(a, wb + 8, latout, attin, bid);
      grid.sync();
      if (bid < 9 * S4096) {
        const int tile = bid / S4096, split = bid % S4096;
        write_partial(a, wb + 11, attin, 4096, 2304, tile,
                      S4096, split, false, pp2, 0, 0, false);
      }
      grid.sync();
      if (bid < 9) reduce_gemv(pp2, 2304, bid, S4096, attout);
      grid.sync();
    }

    if (bid < 9) {
      for (int z = 0; z < 2; ++z) {
        const int i = bid * 256 + tid + z * 128;
        if (i < 2304) bfstore(hres, i, bfload(cur, i) + bfload(attout, i));
      }
    }
    grid.sync();
    if (bid == 0) rms(hres, (const bf*)a.w[wb + 1], norm, red);
    grid.sync();

    if (bid < 64) router_one(a, wb + (is_kda ? 19 : 14), norm, logits, bid, red);
    grid.sync();
    if (bid == 0 && tid == 0) choose_top8(logits, ids, routew);
    grid.sync();

    const int gatebase = wb + (is_kda ? 20 : 15);
    const int upbase = wb + (is_kda ? 23 : 18);
    const int downbase = wb + (is_kda ? 26 : 21);
    const int sgatebase = wb + (is_kda ? 29 : 24);
    const int supbase = wb + (is_kda ? 32 : 27);
    const int sdownbase = wb + (is_kda ? 35 : 30);
    constexpr int S2304 = 6;
    constexpr int S4096 = 4;
    constexpr int S1024 = 4;
    constexpr int GATE_BLOCKS = 9 * 4 * S2304;
    constexpr int DOWN_BLOCKS = 9 * 9 * S1024;
    if (bid < GATE_BLOCKS) {
      const int slot = bid / (4 * S2304);
      const int rem = bid % (4 * S2304);
      const int tile = rem / S2304;
      const int split = rem % S2304;
      expert_gate_up(a, slot == 8 ? sgatebase : gatebase,
                     slot == 8 ? supbase : upbase, norm, ids,
                     gatep, upp, slot, tile, S2304, split);
    }
    grid.sync();
    if (bid < 36) {
      const int slot = bid / 4;
      const int tile = bid & 3;
      for (int z = 0; z < 2; ++z) {
        const int n = tile * 256 + tid + z * 128;
        if (n >= 1024) continue;
        float gv = 0.0f, uv = 0.0f;
        for (int s = 0; s < S2304; ++s) {
          gv += gatep[(slot * S2304 + s) * 1024 + n];
          uv += upp[(slot * S2304 + s) * 1024 + n];
        }
        acts[slot * 1024 + n] = silu(gv) * uv;
      }
    }
    grid.sync();
    if (bid < DOWN_BLOCKS) {
      const int slot = bid / (9 * S1024);
      const int rem = bid % (9 * S1024);
      const int tile = rem / S1024;
      const int split = rem % S1024;
      const int db = slot == 8 ? sdownbase : downbase;
      expert_down(a, db, acts + slot * 1024, ids, downp, slot, tile, S1024, split);
    }
    grid.sync();
    if (bid < 9) {
      const int base = bid * 256 + tid;
      for (int z = 0; z < 2; ++z) {
        const int n = base + z * 128;
        if (n >= 2304) continue;
        float sum = 0.0f;
        for (int e = 0; e < 8; ++e) {
          float d = 0.0f;
          for (int s = 0; s < S1024; ++s) d += downp[(e * S1024 + s) * 2304 + n];
          sum = fmaf(routew[e], d, sum);
        }
        float shared = 0.0f;
        for (int s = 0; s < S1024; ++s) shared += downp[(8 * S1024 + s) * 2304 + n];
        bfstore(moeout, n, sum + shared);
      }
    }
    grid.sync();
    if (bid < 9) {
      for (int z = 0; z < 2; ++z) {
        const int i = bid * 256 + tid + z * 128;
        if (i < 2304) bfstore(next, i, bfload(hres, i) + bfload(moeout, i));
      }
    }
    grid.sync();
    cur = next;
  }
}

extern "C" cudaError_t launch_coop(const Args* a, cudaStream_t stream) {
  static int cached_grid = 0;
  if (cached_grid == 0) {
    int active = 0, sms = 0, dev = 0;
    cudaError_t e = cudaGetDevice(&dev);
    if (e != cudaSuccess) return e;
    e = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
    if (e != cudaSuccess) return e;
    e = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, decode_kernel, THREADS, 0);
    if (e != cudaSuccess) return e;
    cached_grid = active * sms;
    if (cached_grid > 448) cached_grid = 448;
    if (cached_grid < BLOCKS) cached_grid = BLOCKS;
  }
  void* args[] = {const_cast<Args*>(a)};
  return cudaLaunchCooperativeKernel((void*)decode_kernel,
      dim3(cached_grid), dim3(THREADS), args, 0, stream);
}
'''


class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = 128):
        super().__init__()
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16))


class QuantExperts(nn.Module):
    def __init__(self, n: int, in_f: int, out_f: int, group: int = 128):
        super().__init__()
        self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))


class KDA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, c = cfg.hidden, cfg.kda_heads * cfg.kda_head_dim
        self.q_proj = QuantLinear(d, c)
        self.k_proj = QuantLinear(d, c)
        self.v_proj = QuantLinear(d, c)
        self.g_proj = QuantLinear(d, c)
        self.beta_proj = nn.Linear(d, cfg.kda_heads, bias=False, dtype=cfg.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, c, cfg.short_conv, dtype=cfg.dtype))
        self.o_proj = QuantLinear(c, d)


class MLA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.q_proj = QuantLinear(cfg.hidden, cfg.mla_heads * (cfg.qk_nope + cfg.qk_rope))
        self.kv_a = QuantLinear(cfg.hidden, cfg.kv_lora + cfg.qk_rope)
        self.kv_b = QuantLinear(cfg.kv_lora, cfg.mla_heads * (cfg.qk_nope + cfg.v_head))
        self.o_proj = QuantLinear(cfg.mla_heads * cfg.v_head, cfg.hidden)


class MoE(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, m, e = cfg.hidden, cfg.moe_inter, cfg.n_experts
        self.router = nn.Linear(d, e, bias=False, dtype=cfg.dtype)
        self.gate = QuantExperts(e, d, m)
        self.up = QuantExperts(e, d, m)
        self.down = QuantExperts(e, m, d)
        self.s_gate = QuantExperts(cfg.n_shared, d, m)
        self.s_up = QuantExperts(cfg.n_shared, d, m)
        self.s_down = QuantExperts(cfg.n_shared, m, d)


class Block(nn.Module):
    def __init__(self, cfg, kind):
        super().__init__()
        self.kind = kind
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
        self.moe = MoE(cfg)


def _push_quant(dst, q):
    dst.extend((q.w_q, q.scales, q.zeros))


def _block_weights(block):
    out = [block.attn_norm, block.moe_norm]
    if block.kind == "K":
        a = block.attn
        for q in (a.q_proj, a.k_proj, a.v_proj, a.g_proj):
            _push_quant(out, q)
        out.extend((a.beta_proj.weight, a.conv_w))
        _push_quant(out, a.o_proj)
    else:
        a = block.attn
        for q in (a.q_proj, a.kv_a, a.kv_b, a.o_proj):
            _push_quant(out, q)
    m = block.moe
    out.append(m.router.weight)
    for q in (m.gate, m.up, m.down, m.s_gate, m.s_up, m.s_down):
        _push_quant(out, q)
    return out


def _state_buffers(state, device, cache_c, cache_r):
    ptrs, outputs = [], []
    for layer in range(4):
        st = state[layer]
        if layer < 3:
            ptrs.extend((st["S"], st["cq"], st["ck"], st["cv"]))
            news = (st["cq"], st["ck"], st["cv"])
            ptrs.extend(news)
            outputs.append(news)
        else:
            ptrs.extend((st["c_kv"], st["k_rope"]))
            ptrs.extend((cache_c, cache_r))
            length = st["c_kv"].shape[0] + 1
            outputs.append((cache_c.narrow(0, 0, length), cache_r.narrow(0, 0, length)))
    return ptrs, outputs


def _scratch(device, context, hidden):
    b = torch.bfloat16
    f = torch.float32
    return [
        torch.empty((2304,), device=device, dtype=b),
        torch.empty((2304,), device=device, dtype=b),
        torch.empty((2304,), device=device, dtype=b),
        torch.empty((6144,), device=device, dtype=b),
        torch.empty((6144,), device=device, dtype=b),
        torch.empty((4096,), device=device, dtype=b),
        torch.empty((4096,), device=device, dtype=b),
        torch.empty((32,), device=device, dtype=b),
        torch.empty((4096,), device=device, dtype=b),
        torch.empty((2304,), device=device, dtype=b),
        torch.empty((4 * 8 * 6144,), device=device, dtype=f),
        torch.empty((8 * 6144,), device=device, dtype=f),
        torch.empty((64,), device=device, dtype=b),
        torch.empty((8,), device=device, dtype=torch.int32),
        torch.empty((8,), device=device, dtype=f),
        torch.empty((9 * 6 * 1024,), device=device, dtype=f),
        torch.empty((9 * 6 * 1024,), device=device, dtype=f),
        torch.empty((9 * 1024,), device=device, dtype=f),
        torch.empty((9 * 4 * 2304,), device=device, dtype=f),
        torch.empty((2304,), device=device, dtype=b),
        torch.empty((32 * 512,), device=device, dtype=f),
        torch.empty((32 * (context + 1),), device=device, dtype=f),
        torch.empty((32 * (context + 1),), device=device, dtype=f),
        torch.empty((32 * 512,), device=device, dtype=f),
        torch.empty((32 * 64,), device=device, dtype=b),
        torch.empty((64,), device=device, dtype=b),
        torch.empty((2304,), device=device, dtype=b),
        hidden,
        torch.empty((32 * 10 * 512,), device=device, dtype=f),
    ]


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._kernel = None
        self._weight_refs = None
        self._weight_device = None
        self._cache_c = None
        self._cache_r = None

    def step(self, hidden, state):
        if self._kernel is None:
            self._kernel = load_inline(
                name="kimi_linear_w4a16_megakernel",
                cpp_sources=_CPP,
                cuda_sources=_CUDA,
                functions=None,
                extra_cflags=["-O3"],
                extra_cuda_cflags=["-O3", "-std=c++17", "-arch=sm_120"],
                with_cuda=True,
                no_implicit_headers=True,
                verbose=False,
            )
        if self._weight_refs is None or self._weight_device != hidden.device:
            weights = []
            for block in self.blocks:
                weights.extend(_block_weights(block))
            self._weight_refs = weights
            self._weight_device = hidden.device
        oldc = state[3]["c_kv"]
        oldr = state[3]["k_rope"]
        oldlen = oldc.shape[0]
        reusable = (
            self._cache_c is not None
            and self._cache_c.device == hidden.device
            and self._cache_c.data_ptr() == oldc.data_ptr()
            and self._cache_r.data_ptr() == oldr.data_ptr()
            and self._cache_c.shape[0] >= oldlen + 1
        )
        if not reusable:
            capacity = max(oldlen + 1, 2 * oldlen)
            self._cache_c = torch.empty((capacity, 512), device=hidden.device, dtype=oldc.dtype)
            self._cache_r = torch.empty((capacity, 64), device=hidden.device, dtype=oldr.dtype)
        state_ptrs, state_out = _state_buffers(state, hidden.device, self._cache_c, self._cache_r)
        temps = _scratch(hidden.device, oldlen, hidden)
        self._kernel.run(self._weight_refs, state_ptrs, temps, int(oldlen))
        for i in range(3):
            state[i]["cq"], state[i]["ck"], state[i]["cv"] = state_out[i]
        state[3]["c_kv"], state[3]["k_rope"] = state_out[3]
        return temps[26], state