"""Single-launch W4A16 Kimi-Linear decode megakernel for SM120.

The CUDA kernel is cooperative: all resident CTAs persist for the complete
four-block decode and use grid barriers between dependent stages.  Quantized
weights are unpacked and dequantized in registers as each GEMV consumes them.
The MLA path absorbs the key/value projection around latent attention instead
of constructing a context-by-head K/V tensor.
"""
from __future__ import annotations

import os

os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
os.environ.setdefault("MAX_JOBS", "8")

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

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <vector>

namespace cg = cooperative_groups;
using bf16 = __nv_bfloat16;

constexpr int D = 2304;
constexpr int H = 32;
constexpr int DK = 128;
constexpr int C = 4096;
constexpr int M = 1024;
constexpr int E = 64;
constexpr int TOP = 8;
constexpr int KL = 512;
constexpr int QR = 64;
constexpr int QN = 128;
constexpr int QD = 192;
constexpr int HDV = 128;
constexpr int THREADS = 256;

// Workspace regions.  SCORE has room for (16384 + benchmark decode steps)
// times 32 scores.  The Python module reserves 800000 float elements.
constexpr int XOFF = 0;
constexpr int NOFF = 2304;
constexpr int AOFF = 4608;
constexpr int BOFF = 20992;
constexpr int COFF = 37376;
constexpr int DOFF = 53760;
constexpr int SCOREOFF = 100000;
constexpr int ZOFF = 640000;

struct QW {
  const uint8_t* q;
  const bf16* s;
  const bf16* z;
};

struct KAttn {
  QW q, k, v, g, o;
  const bf16* beta;
  const bf16* conv;
};

struct MAttn {
  QW q, kva, kvb, o;
};

struct MoEP {
  const bf16* router;
  QW gate, up, down;
  QW sgate, sup, sdown;
};

struct BlockP {
  const bf16* anorm;
  const bf16* mnorm;
  int kind;
  KAttn ka;
  MAttn ma;
  MoEP moe;
};

struct Params { BlockP b[4]; };

__device__ __forceinline__ float fbf(float x) {
  return __bfloat162float(__float2bfloat16_rn(x));
}

// Reproduce the oracle's bf16 (uint4 - zero) * scale dequantization, including
// bf16 rounding after both elementary operations.
__device__ __forceinline__ float deq_at(const QW& w, int k, int n, int N) {
  uint8_t p = w.q[(k >> 1) * N + n];
  float qi = float((k & 1) ? (p >> 4) : (p & 15));
  int si = (k >> 7) * N + n;
  float d = fbf(qi - __bfloat162float(w.z[si]));
  return fbf(d * __bfloat162float(w.s[si]));
}

__device__ __forceinline__ float qgemv(const QW& w, const float* x,
                                        int K, int N, int n) {
  float acc = 0.0f;
  #pragma unroll 1
  for (int g = 0; g < K; g += 128) {
    int si = (g >> 7) * N + n;
    float sc = __bfloat162float(w.s[si]);
    float ze = __bfloat162float(w.z[si]);
    #pragma unroll 4
    for (int k = g; k < g + 128; k += 2) {
      uint8_t p = w.q[(k >> 1) * N + n];
      float w0 = fbf(fbf(float(p & 15) - ze) * sc);
      float w1 = fbf(fbf(float(p >> 4) - ze) * sc);
      acc = fmaf(x[k], w0, acc);
      acc = fmaf(x[k + 1], w1, acc);
    }
  }
  return acc;
}

__device__ __forceinline__ float qgemv_group(const QW& w, const float* x,
                                              int group, int N, int n) {
  int k0 = group * 128;
  int si = group * N + n;
  float sc = __bfloat162float(w.s[si]);
  float ze = __bfloat162float(w.z[si]);
  float acc = 0.0f;
  #pragma unroll 4
  for (int k = k0; k < k0 + 128; k += 2) {
    uint8_t p = w.q[(k >> 1) * N + n];
    float w0 = fbf(fbf(float(p & 15) - ze) * sc);
    float w1 = fbf(fbf(float(p >> 4) - ze) * sc);
    acc = fmaf(x[k], w0, acc);
    acc = fmaf(x[k + 1], w1, acc);
  }
  return acc;
}

__device__ __forceinline__ QW expert(const QW& w, int e, int K, int N) {
  QW r;
  r.q = w.q + (size_t)e * (K / 2) * N;
  r.s = w.s + (size_t)e * (K / 128) * N;
  r.z = w.z + (size_t)e * (K / 128) * N;
  return r;
}

__device__ __forceinline__ int gtid() {
  return int(blockIdx.x) * int(blockDim.x) + int(threadIdx.x);
}
__device__ __forceinline__ int gstride() {
  return int(gridDim.x) * int(blockDim.x);
}

// One CTA owns one (256 output columns, one quantization group) partial.  This
// keeps a warp's packed-weight reads contiguous while exposing enough CTAs to
// occupy all 188 SMs, unlike a one-thread-per-output GEMV.
__device__ void parallel_qgemv(const QW& w, const float* x, int K, int N,
                               float* out, bool round_bf16,
                               cg::grid_group grid) {
  int t = gtid();
  for (int n = t; n < N; n += gstride()) out[n] = 0.0f;
  grid.sync();
  int tiles = (N + THREADS - 1) / THREADS;
  int groups = K / 128;
  int jobs = tiles * groups;
  for (int job = int(blockIdx.x); job < jobs; job += int(gridDim.x)) {
    int group = job % groups;
    int tile = job / groups;
    int n = tile * THREADS + int(threadIdx.x);
    if (n < N) atomicAdd(out + n, qgemv_group(w, x, group, N, n));
  }
  grid.sync();
  if (round_bf16) {
    for (int n = t; n < N; n += gstride()) out[n] = fbf(out[n]);
  }
  grid.sync();
}

__device__ void parallel_kda_inputs(const KAttn& p, const float* x,
                                    float* q, float* k, float* v, float* g,
                                    cg::grid_group grid) {
  int t = gtid();
  for (int n = t; n < C; n += gstride()) {
    q[n] = 0.0f; k[n] = 0.0f; v[n] = 0.0f; g[n] = 0.0f;
  }
  grid.sync();
  constexpr int GROUPS = D / 128;
  constexpr int TILES = C / THREADS;
  int jobs = 4 * TILES * GROUPS;
  for (int job = int(blockIdx.x); job < jobs; job += int(gridDim.x)) {
    int u = job;
    int group = u % GROUPS; u /= GROUPS;
    int tile = u % TILES; u /= TILES;
    int which = u;
    int n = tile * THREADS + int(threadIdx.x);
    const QW* w = which == 0 ? &p.q : (which == 1 ? &p.k : (which == 2 ? &p.v : &p.g));
    float* out = which == 0 ? q : (which == 1 ? k : (which == 2 ? v : g));
    atomicAdd(out + n, qgemv_group(*w, x, group, C, n));
  }
  grid.sync();
  for (int n = t; n < C; n += gstride()) {
    q[n] = fbf(q[n]); k[n] = fbf(k[n]); v[n] = fbf(v[n]); g[n] = fbf(g[n]);
  }
  grid.sync();
}

__device__ void parallel_mla_inputs(const MAttn& p, const float* x,
                                    float* q, float* kv,
                                    cg::grid_group grid) {
  int t = gtid();
  constexpr int QOUT = H * QD;
  constexpr int KOUT = KL + QR;
  for (int n = t; n < QOUT; n += gstride()) q[n] = 0.0f;
  for (int n = t; n < KOUT; n += gstride()) kv[n] = 0.0f;
  grid.sync();
  constexpr int GROUPS = D / 128;
  constexpr int QTILES = (QOUT + THREADS - 1) / THREADS;
  constexpr int KTILES = (KOUT + THREADS - 1) / THREADS;
  constexpr int QJOBS = QTILES * GROUPS;
  constexpr int JOBS = QJOBS + KTILES * GROUPS;
  for (int job = int(blockIdx.x); job < JOBS; job += int(gridDim.x)) {
    bool is_kv = job >= QJOBS;
    int u = is_kv ? job - QJOBS : job;
    int group = u % GROUPS;
    int tile = u / GROUPS;
    int n = tile * THREADS + int(threadIdx.x);
    int N = is_kv ? KOUT : QOUT;
    if (n < N) {
      const QW& w = is_kv ? p.kva : p.q;
      atomicAdd((is_kv ? kv : q) + n, qgemv_group(w, x, group, N, n));
    }
  }
  grid.sync();
  for (int n = t; n < QOUT; n += gstride()) q[n] = fbf(q[n]);
  for (int n = t; n < KOUT; n += gstride()) kv[n] = fbf(kv[n]);
  grid.sync();
}

__device__ void norm_stage(const float* x, const bf16* w, float* out,
                           float* scratch, cg::grid_group grid) {
  int t = gtid();
  if (t == 0) {
    float ss = 0.0f;
    for (int i = 0; i < D; ++i) ss = fmaf(x[i], x[i], ss);
    scratch[0] = rsqrtf(ss * (1.0f / float(D)) + 1.0e-6f);
  }
  grid.sync();
  float r = scratch[0];
  for (int i = t; i < D; i += gstride())
    out[i] = fbf(x[i] * r * __bfloat162float(w[i]));
  grid.sync();
}

__device__ void moe_stage(const MoEP& p, float* ws, cg::grid_group grid) {
  int t = gtid(), stride = gstride();
  float* x = ws + XOFF;
  float* xn = ws + NOFF;
  float* a = ws + AOFF;
  float* b = ws + BOFF;
  float* route = ws + DOFF;

  for (int e = t; e < E; e += stride) {
    float s = 0.0f;
    for (int k = 0; k < D; ++k)
      s = fmaf(xn[k], __bfloat162float(p.router[(size_t)e * D + k]), s);
    a[e] = fbf(s);
  }
  grid.sync();

  // Only 64 logits are involved.  A scalar selection avoids another shared
  // reduction structure and is lost in the quantized weight stream latency.
  if (t == 0) {
    float vmax[TOP];
    int vidx[TOP];
    for (int j = 0; j < TOP; ++j) {
      float best = -1.0e30f;
      int bi = -1;
      for (int e = 0; e < E; ++e) {
        bool used = false;
        for (int u = 0; u < j; ++u) used |= (vidx[u] == e);
        if (!used && a[e] > best) { best = a[e]; bi = e; }
      }
      vmax[j] = best;
      vidx[j] = bi;
    }
    float mx = vmax[0], den = 0.0f;
    for (int j = 0; j < TOP; ++j) den += __expf(vmax[j] - mx);
    for (int j = 0; j < TOP; ++j) {
      route[j] = float(vidx[j]);
      route[TOP + j] = __expf(vmax[j] - mx) / den * 2.446f;
    }
  }
  grid.sync();

  for (int z = t; z < (TOP + 1) * M; z += stride) { a[z] = 0.0f; b[z] = 0.0f; }
  grid.sync();
  int etiles = M / THREADS;
  int egroups = D / 128;
  int ejobs = 2 * (TOP + 1) * etiles * egroups;
  for (int job = int(blockIdx.x); job < ejobs; job += int(gridDim.x)) {
    int u = job;
    int group = u % egroups; u /= egroups;
    int tile = u % etiles; u /= etiles;
    int j = u % (TOP + 1); u /= (TOP + 1);
    int which = u;
    int n = tile * THREADS + int(threadIdx.x);
    QW w;
    if (j < TOP) {
      int ei = int(route[j]);
      w = expert(which ? p.up : p.gate, ei, D, M);
    } else {
      w = which ? p.sup : p.sgate;
    }
    atomicAdd((which ? b : a) + j * M + n,
              qgemv_group(w, xn, group, M, n));
  }
  grid.sync();

  for (int z = t; z < (TOP + 1) * M; z += stride) {
    float gv = a[z];
    a[z] = (gv / (1.0f + __expf(-gv))) * b[z];
    if (z < D) b[z] = 0.0f;
  }
  grid.sync();

  int dtiles = (D + THREADS - 1) / THREADS;
  int dgroups = M / 128;
  int djobs = (TOP + 1) * dtiles * dgroups;
  for (int job = int(blockIdx.x); job < djobs; job += int(gridDim.x)) {
    int u = job;
    int group = u % dgroups; u /= dgroups;
    int tile = u % dtiles; u /= dtiles;
    int j = u;
    int n = tile * THREADS + int(threadIdx.x);
    if (n < D) {
      QW wd;
      float rw;
      if (j < TOP) {
        wd = expert(p.down, int(route[j]), M, D);
        rw = route[TOP + j];
      } else {
        wd = p.sdown;
        rw = 1.0f;
      }
      atomicAdd(b + n, rw * qgemv_group(wd, a + j * M, group, D, n));
    }
  }
  grid.sync();
  for (int n = t; n < D; n += stride) x[n] = fbf(x[n] + fbf(b[n]));
  grid.sync();
}

__device__ void kda_stage(const BlockP& p, float* ws, float* S,
                          bf16* cq, bf16* ck, bf16* cv,
                          cg::grid_group grid) {
  int t = gtid(), stride = gstride();
  float* x = ws + XOFF;
  float* xn = ws + NOFF;
  float* q = ws + AOFF;
  float* k = ws + BOFF;
  float* v = ws + COFF;
  float* g = ws + DOFF;
  float* pred = g + C;

  norm_stage(x, p.anorm, xn, q + C * 2, grid);

  parallel_kda_inputs(p.ka, xn, q, k, v, g, grid);

  for (int n = t; n < C; n += stride) {
    float vals[3] = {q[n], k[n], v[n]};
    bf16* hist[3] = {cq, ck, cv};
    for (int j = 0; j < 3; ++j) {
      bf16* hp = hist[j];
      float z0 = __bfloat162float(hp[n]);
      float z1 = __bfloat162float(hp[C + n]);
      float z2 = __bfloat162float(hp[2 * C + n]);
      const bf16* cw = p.ka.conv + (size_t)j * C * 4 + n * 4;
      float y = z0 * __bfloat162float(cw[0])
              + z1 * __bfloat162float(cw[1])
              + z2 * __bfloat162float(cw[2])
              + vals[j] * __bfloat162float(cw[3]);
      y = y / (1.0f + __expf(-y));
      vals[j] = fbf(y);
      hp[n] = __float2bfloat16_rn(z1);
      hp[C + n] = __float2bfloat16_rn(z2);
      hp[2 * C + n] = __float2bfloat16_rn((j == 0 ? q[n] : (j == 1 ? k[n] : v[n])));
    }
    q[n] = vals[0] * 0.08838834764831845f;
    k[n] = vals[1];
    v[n] = vals[2];
    g[n] = 1.0f / (1.0f + __expf(g[n]));
  }
  grid.sync();

  if (t < H) {
    float z = 0.0f;
    for (int i = 0; i < D; ++i)
      z = fmaf(xn[i], __bfloat162float(p.ka.beta[(size_t)t * D + i]), z);
    pred[C + t] = 1.0f / (1.0f + __expf(-fbf(z)));
  }
  for (int z = t; z < H * DK * DK; z += stride) {
    int hi = z / DK;
    int i = hi % DK;
    S[z] *= g[hi];
  }
  grid.sync();

  for (int z = t; z < H * DK; z += stride) {
    int h = z / DK, j = z % DK;
    float s = 0.0f;
    const float* sp = S + (size_t)h * DK * DK + j;
    for (int i = 0; i < DK; ++i) s = fmaf(sp[(size_t)i * DK], k[h * DK + i], s);
    pred[z] = s;
  }
  grid.sync();

  for (int z = t; z < H * DK * DK; z += stride) {
    int h = z / (DK * DK);
    int rem = z - h * DK * DK;
    int i = rem / DK, j = rem % DK;
    S[z] += pred[C + h] * k[h * DK + i] * (v[h * DK + j] - pred[h * DK + j]);
  }
  grid.sync();

  for (int z = t; z < H * DK; z += stride) {
    int h = z / DK, j = z % DK;
    float s = 0.0f;
    const float* sp = S + (size_t)h * DK * DK + j;
    for (int i = 0; i < DK; ++i) s = fmaf(sp[(size_t)i * DK], q[h * DK + i], s);
    v[z] = fbf(s);
  }
  grid.sync();

  parallel_qgemv(p.ka.o, v, C, D, k, true, grid);
  for (int n = t; n < D; n += stride) x[n] = fbf(x[n] + k[n]);
  grid.sync();

  norm_stage(x, p.mnorm, xn, q + C * 2, grid);
  moe_stage(p.moe, ws, grid);
}

__device__ void mla_stage(const BlockP& p, float* ws,
                          const bf16* cin, const bf16* rin,
                          bf16* cbuf, bf16* rbuf, int L, bool copy_cache,
                          cg::grid_group grid) {
  int t = gtid(), stride = gstride();
  float* x = ws + XOFF;
  float* xn = ws + NOFF;
  float* q = ws + AOFF;
  float* kv = ws + BOFF;
  float* ab = ws + COFF;
  float* aux = ws + DOFF;
  float* scores = ws + SCOREOFF;
  float* zbuf = ws + ZOFF;

  norm_stage(x, p.anorm, xn, aux + 1024, grid);

  parallel_mla_inputs(p.ma, xn, q, kv, grid);

  if (copy_cache) {
    for (int i = t; i < L * KL; i += stride) cbuf[i] = cin[i];
    for (int i = t; i < L * QR; i += stride) rbuf[i] = rin[i];
  }
  for (int i = t; i < KL; i += stride) cbuf[(size_t)L * KL + i] = __float2bfloat16_rn(kv[i]);

  // Decoupled RoPE.  One thread handles each pair so reads and writes cannot
  // race.  Position is the pre-append cache length.
  for (int z = t; z < H * (QR / 2); z += stride) {
    int h = z / (QR / 2), j = z % (QR / 2);
    float inv = __expf((-logf(10000.0f) / float(QR)) * float(2 * j));
    float co = cosf(float(L) * inv), si = sinf(float(L) * inv);
    int i0 = h * QD + QN + 2 * j;
    float a0 = q[i0], a1 = q[i0 + 1];
    q[i0] = fbf(a0 * co - a1 * si);
    q[i0 + 1] = fbf(a1 * co + a0 * si);
  }
  for (int j = t; j < QR / 2; j += stride) {
    float inv = __expf((-logf(10000.0f) / float(QR)) * float(2 * j));
    float co = cosf(float(L) * inv), si = sinf(float(L) * inv);
    float a0 = kv[KL + 2 * j], a1 = kv[KL + 2 * j + 1];
    rbuf[(size_t)L * QR + 2 * j] = __float2bfloat16_rn(a0 * co - a1 * si);
    rbuf[(size_t)L * QR + 2 * j + 1] = __float2bfloat16_rn(a1 * co + a0 * si);
  }
  grid.sync();
  int LL = L + 1;

  // q_nope @ W_k^T, stored k-major so a warp evaluating one token reads the
  // 32 heads contiguously while the latent value is broadcast.
  for (int zz = t; zz < KL * H; zz += stride) {
    int k = zz / H, h = zz % H;
    float s = 0.0f;
    for (int d = 0; d < QN; ++d)
      s = fmaf(q[h * QD + d], deq_at(p.ma.kvb, k, h * (QN + HDV) + d, H * (QN + HDV)), s);
    ab[k * H + h] = s;
  }
  grid.sync();

  // Each warp owns a (token, all-heads) score row.  Latent/cache values are
  // warp broadcasts and the absorbed query coefficients are coalesced.
  int lane = threadIdx.x & 31;
  int warp = threadIdx.x >> 5;
  int warps_per_grid = int(gridDim.x) * (THREADS / 32);
  int first_warp = int(blockIdx.x) * (THREADS / 32) + warp;
  for (int l = first_warp; l < LL; l += warps_per_grid) {
    float s = 0.0f;
    for (int k = 0; k < KL; ++k)
      s = fmaf(__bfloat162float(cbuf[(size_t)l * KL + k]), ab[k * H + lane], s);
    for (int d = 0; d < QR; ++d)
      s = fmaf(q[lane * QD + QN + d], __bfloat162float(rbuf[(size_t)l * QR + d]), s);
    scores[(size_t)l * H + lane] = s * 0.07216878364870322f;
  }
  grid.sync();

  __shared__ float softmax_red[THREADS];
  if (int(blockIdx.x) < H) {
    int h = int(blockIdx.x), tx = int(threadIdx.x);
    float mx = -1.0e30f;
    for (int l = tx; l < LL; l += THREADS)
      mx = fmaxf(mx, scores[(size_t)l * H + h]);
    softmax_red[tx] = mx;
    __syncthreads();
    for (int d = THREADS / 2; d; d >>= 1) {
      if (tx < d) softmax_red[tx] = fmaxf(softmax_red[tx], softmax_red[tx + d]);
      __syncthreads();
    }
    mx = softmax_red[0];
    float den = 0.0f;
    for (int l = tx; l < LL; l += THREADS)
      den += __expf(scores[(size_t)l * H + h] - mx);
    softmax_red[tx] = den;
    __syncthreads();
    for (int d = THREADS / 2; d; d >>= 1) {
      if (tx < d) softmax_red[tx] += softmax_red[tx + d];
      __syncthreads();
    }
    if (tx == 0) {
      aux[h] = mx;
      aux[H + h] = 1.0f / softmax_red[0];
    }
  }
  grid.sync();

  // Materialize normalized probabilities once.  Recomputing exp(score) for
  // every one of the 512 latent channels would issue H*KL*L exponentials.
  for (int zz = t; zz < LL * H; zz += stride) {
    int h = zz % H;
    scores[zz] = __expf(scores[zz] - aux[h]) * aux[H + h];
  }
  grid.sync();

  for (int zz = t; zz < H * KL; zz += stride) zbuf[zz] = 0.0f;
  grid.sync();
  constexpr int LPARTS = 8;
  int ztiles = (H * KL) / THREADS;
  int zjobs = ztiles * LPARTS;
  int lchunk = (LL + LPARTS - 1) / LPARTS;
  for (int job = int(blockIdx.x); job < zjobs; job += int(gridDim.x)) {
    int part = job % LPARTS;
    int tile = job / LPARTS;
    int zz = tile * THREADS + int(threadIdx.x);
    int h = zz / KL, k = zz % KL;
    int l0 = part * lchunk, l1 = min(LL, l0 + lchunk);
    float s = 0.0f;
    for (int l = l0; l < l1; ++l) {
      s = fmaf(scores[(size_t)l * H + h],
               __bfloat162float(cbuf[(size_t)l * KL + k]), s);
    }
    atomicAdd(zbuf + zz, s);
  }
  grid.sync();

  for (int zz = t; zz < H * HDV; zz += stride) {
    int h = zz / HDV, d = zz % HDV;
    float s = 0.0f;
    int n = h * (QN + HDV) + QN + d;
    for (int k = 0; k < KL; ++k)
      s = fmaf(zbuf[h * KL + k], deq_at(p.ma.kvb, k, n, H * (QN + HDV)), s);
    q[zz] = fbf(s);
  }
  grid.sync();

  parallel_qgemv(p.ma.o, q, H * HDV, D, kv, true, grid);
  for (int n = t; n < D; n += stride) x[n] = fbf(x[n] + kv[n]);
  grid.sync();

  norm_stage(x, p.mnorm, xn, aux + 1024, grid);
  moe_stage(p.moe, ws, grid);
}

__global__ void kimi_mega(Params p, const bf16* hidden, bf16* out, float* ws,
                          float* S0, bf16* cq0, bf16* ck0, bf16* cv0,
                          float* S1, bf16* cq1, bf16* ck1, bf16* cv1,
                          float* S2, bf16* cq2, bf16* ck2, bf16* cv2,
                          const bf16* cin, const bf16* rin,
                          bf16* cbuf, bf16* rbuf, int L, bool copy_cache) {
  cg::grid_group grid = cg::this_grid();
  int t = gtid();
  for (int i = t; i < D; i += gstride()) ws[XOFF + i] = __bfloat162float(hidden[i]);
  grid.sync();

  kda_stage(p.b[0], ws, S0, cq0, ck0, cv0, grid);
  kda_stage(p.b[1], ws, S1, cq1, ck1, cv1, grid);
  kda_stage(p.b[2], ws, S2, cq2, ck2, cv2, grid);
  mla_stage(p.b[3], ws, cin, rin, cbuf, rbuf, L, copy_cache, grid);

  for (int i = t; i < D; i += gstride()) out[i] = __float2bfloat16_rn(ws[XOFF + i]);
}

static inline const bf16* bp(const torch::Tensor& t) {
  return reinterpret_cast<const bf16*>(t.data_ptr<at::BFloat16>());
}
static inline bf16* bpm(torch::Tensor& t) {
  return reinterpret_cast<bf16*>(t.data_ptr<at::BFloat16>());
}

torch::Tensor mega_step(torch::Tensor hidden, torch::Tensor out,
                        torch::Tensor workspace,
                        std::vector<torch::Tensor> states,
                        torch::Tensor cin, torch::Tensor rin,
                        torch::Tensor cbuf, torch::Tensor rbuf,
                        int64_t L64, bool copy_cache,
                        std::vector<torch::Tensor> weights) {
  TORCH_CHECK(hidden.is_cuda() && hidden.scalar_type() == at::kBFloat16,
              "hidden must be CUDA bf16");
  TORCH_CHECK(states.size() == 12, "expected 12 KDA state tensors");
  TORCH_CHECK(weights.size() == 147, "expected 147 model tensors, got ", weights.size());
  TORCH_CHECK(L64 + 1 <= 16800, "context exceeds megakernel score workspace");

  Params p{};
  size_t i = 0;
  auto qw = [&]() -> QW {
    QW r{weights[i].data_ptr<uint8_t>(), bp(weights[i + 1]), bp(weights[i + 2])};
    i += 3;
    return r;
  };
  for (int b = 0; b < 4; ++b) {
    p.b[b].anorm = bp(weights[i++]);
    p.b[b].mnorm = bp(weights[i++]);
    p.b[b].kind = (b == 3);
    if (b < 3) {
      p.b[b].ka.q = qw(); p.b[b].ka.k = qw(); p.b[b].ka.v = qw();
      p.b[b].ka.g = qw();
      p.b[b].ka.beta = bp(weights[i++]);
      p.b[b].ka.conv = bp(weights[i++]);
      p.b[b].ka.o = qw();
    } else {
      p.b[b].ma.q = qw(); p.b[b].ma.kva = qw();
      p.b[b].ma.kvb = qw(); p.b[b].ma.o = qw();
    }
    p.b[b].moe.router = bp(weights[i++]);
    p.b[b].moe.gate = qw(); p.b[b].moe.up = qw(); p.b[b].moe.down = qw();
    p.b[b].moe.sgate = qw(); p.b[b].moe.sup = qw(); p.b[b].moe.sdown = qw();
  }
  TORCH_CHECK(i == weights.size(), "internal model tensor layout mismatch");

  static int blocks = 0;
  if (blocks == 0) {
    int blocks_per_sm = 0, sms = 0;
    cudaError_t query_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &blocks_per_sm, kimi_mega, THREADS, 0);
    TORCH_CHECK(query_err == cudaSuccess, "occupancy query failed: ", cudaGetErrorString(query_err));
    query_err = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, hidden.get_device());
    TORCH_CHECK(query_err == cudaSuccess, "SM query failed: ", cudaGetErrorString(query_err));
    blocks = sms * blocks_per_sm;
  }
  int L = int(L64);

  Params* pp = &p;
  const bf16* hp = bp(hidden);
  bf16* op = bpm(out);
  float* wsp = workspace.data_ptr<float>();
  float* S0 = states[0].data_ptr<float>(); bf16* cq0 = bpm(states[1]); bf16* ck0 = bpm(states[2]); bf16* cv0 = bpm(states[3]);
  float* S1 = states[4].data_ptr<float>(); bf16* cq1 = bpm(states[5]); bf16* ck1 = bpm(states[6]); bf16* cv1 = bpm(states[7]);
  float* S2 = states[8].data_ptr<float>(); bf16* cq2 = bpm(states[9]); bf16* ck2 = bpm(states[10]); bf16* cv2 = bpm(states[11]);
  const bf16* cip = bp(cin); const bf16* rip = bp(rin);
  bf16* cbp = bpm(cbuf); bf16* rbp = bpm(rbuf);

  void* args[] = {pp, &hp, &op, &wsp,
                  &S0, &cq0, &ck0, &cv0, &S1, &cq1, &ck1, &cv1,
                  &S2, &cq2, &ck2, &cv2, &cip, &rip, &cbp, &rbp,
                  &L, &copy_cache};
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  cudaError_t err = cudaLaunchCooperativeKernel((void*)kimi_mega, dim3(blocks), dim3(THREADS),
                                                args, 0, stream);
  TORCH_CHECK(err == cudaSuccess, "cooperative launch failed: ", cudaGetErrorString(err));
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("mega_step", &mega_step, "Kimi-Linear single cooperative megakernel");
}
"""


_ext = load_inline(
    name="kimi_linear_sm120_mega_v8",
    cpp_sources="",
    cuda_sources=CUDA_SRC,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
    with_cuda=True,
    verbose=False,
)


class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = 128):
        super().__init__()
        self.in_f, self.out_f, self.group = in_f, out_f, group
        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.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group
        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__()
        h, dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
        self.q_proj = QuantLinear(d, h * dk, cfg.group)
        self.k_proj = QuantLinear(d, h * dk, cfg.group)
        self.v_proj = QuantLinear(d, h * dk, cfg.group)
        self.g_proj = QuantLinear(d, h * dk, cfg.group)
        self.beta_proj = nn.Linear(d, h, bias=False, dtype=cfg.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, h * dk, cfg.short_conv, dtype=cfg.dtype))
        self.o_proj = QuantLinear(h * dk, d, cfg.group)


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


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, cfg.group)
        self.up = QuantExperts(e, d, m, cfg.group)
        self.down = QuantExperts(e, m, d, cfg.group)
        self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group)


class Block(nn.Module):
    def __init__(self, cfg, kind: str):
        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)


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, kind) for kind in cfg.pattern)
        self.register_buffer("_workspace", torch.empty(800000, dtype=torch.float32), persistent=False)
        self.register_buffer("_output", torch.empty(cfg.hidden, dtype=torch.bfloat16), persistent=False)
        self._weight_tensors = None

    @staticmethod
    def _q(dst, q):
        dst.extend((q.w_q, q.scales, q.zeros))

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

    def step(self, hidden, state):
        if self._weight_tensors is None or self._weight_tensors[0].device != hidden.device:
            self._weight_tensors = self._collect_weights()

        mla = state[3]
        c_in = mla["c_kv"]
        r_in = mla["k_rope"]
        length = c_in.shape[0]
        cbuf = mla.get("_c_kv_buffer")
        rbuf = mla.get("_k_rope_buffer")
        copy_cache = cbuf is None or cbuf.shape[0] <= length
        if copy_cache:
            capacity = max(length + 256, 16640)
            cbuf = torch.empty((capacity, 512), device=hidden.device, dtype=torch.bfloat16)
            rbuf = torch.empty((capacity, 64), device=hidden.device, dtype=torch.bfloat16)
            mla["_c_kv_buffer"] = cbuf
            mla["_k_rope_buffer"] = rbuf

        states = []
        for i in range(3):
            st = state[i]
            states.extend((st["S"], st["cq"], st["ck"], st["cv"]))

        out = _ext.mega_step(
            hidden,
            self._output,
            self._workspace,
            states,
            c_in,
            r_in,
            cbuf,
            rbuf,
            length,
            copy_cache,
            self._weight_tensors,
        )
        mla["c_kv"] = cbuf[: length + 1]
        mla["k_rope"] = rbuf[: length + 1]
        return out, state