Kimi-Linear Decode Claude Fable 5.1

"""Kimi-Linear W4A16 hybrid decode unit (batch 1) as ONE persistent dataflow megakernel.

Everything that happens per token -- RMSNorms, every int4 dequant-GEMV, the KDA short
conv + gated-delta recurrent state update, the MLA latent-cache attention (absorbed form),
the MoE router / top-8 / expert GEMVs, residual adds and the KV-cache append -- runs inside
a single `__global__` kernel that step() launches exactly once (cooperative launch, all
blocks co-resident).

Execution model inside the kernel
---------------------------------
* A global work queue (one atomic counter) hands out ~2.5k "items" in a fixed topological
  order (layer by layer, phase by phase).  Blocks are persistent: grab item, run it, repeat.
* Dependencies are per-tile arrival counters in global memory (release/acquire), NOT grid
  barriers: an item spins only on the specific counters it needs, so phases overlap.
* Split-K GEMV tiles write fp32 partials; the LAST arriving split (atomic ticket) runs the
  fused epilogue (residual add, SiLU*up, RoPE, cache append, router partials ...).
* The int4 weights are streamed exactly once, 4 bytes/lane (8 nibbles -> 8 FMAs) with the
  per-group asymmetric dequant folded into a per-16-row epilogue (no bf16 weight is ever
  materialised).
* The last block out resets the counters, so no memset/extra launch is needed per step.

No CUDA graphs, no torch.compile, no per-op kernel loops: the timed path is one launch.
"""
from __future__ import annotations

import math
import os
from dataclasses import dataclass, field

import torch
import torch.nn as nn

os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")

EPS = 1.0e-6
GROUP_SIZE = 128


@dataclass(frozen=True)
class Config:
    hidden: int = 2304
    kda_heads: int = 32
    kda_head_dim: int = 128
    short_conv: int = 4
    mla_heads: int = 32
    kv_lora: int = 512
    qk_nope: int = 128
    qk_rope: int = 64
    v_head: int = 128
    rope_theta: float = 10000.0
    n_experts: int = 64
    n_active: int = 8
    n_shared: int = 1
    moe_inter: int = 1024
    routed_scaling: float = 2.446
    group: int = 128
    pattern: tuple = ("K", "K", "K", "M")
    dtype: torch.dtype = field(default=torch.bfloat16)


def build_config(shape: dict) -> Config:
    return Config(n_experts=int(shape.get("n_experts", 64)))


# --------------------------------------------------------------------------- #
# CUDA source: the megakernel
# --------------------------------------------------------------------------- #
CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>
#include <cstring>
#include <cstdlib>
#include <vector>

namespace mk {

typedef __nv_bfloat16 bf16;

// ---- model constants (Kimi-Linear-48B-A3B motif) ----
constexpr int HID = 2304, NH = 32, DK = 128, CQ = 4096;
constexpr int KV_LORA = 512, QK_NOPE = 128, QK_ROPE = 64, V_HEAD = 128;
constexpr int Q_DIM = NH * (QK_NOPE + QK_ROPE);     // 6144
constexpr int KVA_DIM = KV_LORA + QK_ROPE;          // 576
constexpr int KVB_DIM = NH * (QK_NOPE + V_HEAD);    // 8192
constexpr int NE = 64, TOPK = 8, NSLOT = 9, INTER = 1024, GROUP = 128;
constexpr float ROUTED_SCALING = 2.446f, RMS_EPS = 1e-6f;
constexpr float KDA_SCALE = 0.08838834764831845f;   // 128^-0.5
constexpr float MLA_SCALE = 0.07216878364870323f;   // 192^-0.5

// ---- kernel geometry ----
constexpr int TILE = 128;        // GEMV tile = 128 output columns (4 per lane)
constexpr int NT = 256, NW = 8;  // threads / warps per block
constexpr int NTILE_HID = HID / TILE;          // 18
constexpr int NTILE_CQ = CQ / TILE;            // 32
constexpr int NTILE_Q = Q_DIM / TILE;          // 48
constexpr int NTILE_KVA = (KVA_DIM + TILE - 1) / TILE;  // 5 (last one half empty)
constexpr int NTILE_INTER = INTER / TILE;      // 8
constexpr int S_QKVG = 3, S_O = 4, S_GU = 3, S_Q = 3, S_KVA = 3;   // split-K factors
constexpr int G_ATTN = 4, HPG = NH / G_ATTN;   // attention head groups (8 heads = 8 warps per item)
constexpr int MAXCH = 130;                     // max attention chunks
constexpr int APITCH = 2 + KV_LORA;            // attention partial: m, l, o[512]
constexpr int SMEM_FLOATS = 4096 + NW * 128 + 1024;
constexpr int SMEM_BYTES = SMEM_FLOATS * 4;

// ---- counters ----
constexpr int CTL_NEXT = 0, CTL_EXIT = 1, CTL_PF = 8, CTL_BASE = 16, CTL_PER_LAYER = 512;
constexpr int C_HEAD = 0, C_O = 32, C_OTILE = 33, C_HRES = 51, C_GUT = 52, C_H1 = 124,
              C_DOWNT = 133, C_XNEXT = 151, C_QTILE = 152, C_QHEAD = 200, C_QREADY = 232,
              C_KVAT = 233, C_KVADONE = 238, C_ATTN = 239, C_OLAT = 243;
constexpr int N_CTL = CTL_BASE + 4 * CTL_PER_LAYER;

enum { IT_QKVG = 0, IT_STATE, IT_OPROJ, IT_GU, IT_DOWN, IT_MLAQ, IT_KVA, IT_ATTN, IT_COMB, IT_NONE };

struct QW { const uint8_t* w; const bf16* s; const bf16* z; int K; int N; };
struct KdaW { QW q, k, v, g, o; const bf16* beta_w; const bf16* conv_w; };
struct MlaW { QW q, kva, kvb, o; };
struct MoeW { const bf16* router; QW gate, up, down, s_gate, s_up, s_down; };
struct LayerW { int kind; const bf16* attn_norm; const bf16* moe_norm; KdaW kda; MlaW mla; MoeW moe; };

struct Params {
  LayerW L[4];
  // scratch (device, persistent)
  float* qkvg_part;   // [4][S_QKVG][4][CQ]
  float* q_part;      // [S_Q][Q_DIM]
  float* kva_part;    // [S_KVA][NTILE_KVA*TILE]
  float* q_abs;       // [NH][KV_LORA]
  float* q_rope;      // [NH][QK_ROPE]
  float* attn_part;   // [MAXCH][NH][APITCH]
  float* o_lat;       // [NH][KV_LORA]
  float* o_part;      // [4][S_O][HID]
  bf16*  attn_o;      // [4][CQ]
  bf16*  h_res;       // [4][HID]
  float* router_part; // [4][NTILE_HID][NE]
  float* gu_part;     // [4][NSLOT][2][S_GU][INTER]
  float* h1;          // [4][NSLOT][INTER]
  float* down_part;   // [4][NSLOT][HID]
  bf16*  xres;        // [3][HID]
  int*   ctl;
  unsigned long long* prof;   // optional per-item timeline (nullptr = off)
  // dynamic (per step)
  const bf16* hidden; bf16* out;
  float* S[3];
  const bf16* win_src[3][3];
  bf16* win_dst[3][3];
  const bf16* ckv_src; const bf16* kr_src; bf16* ckv_dst; bf16* kr_dst;
  int pos; int n_old; int n_chunks; int chunk; int n_items;
};

// ------------------------------------------------------------------ helpers
__device__ __forceinline__ float bf2f(bf16 v) { return __bfloat162float(v); }
__device__ __forceinline__ float bfr(float v) { return __bfloat162float(__float2bfloat16(v)); }
__device__ __forceinline__ float ldcg_bf(const bf16* p) {
  unsigned short u = __ldcg(reinterpret_cast<const unsigned short*>(p));
  return __uint_as_float(((uint32_t)u) << 16);
}
__device__ __forceinline__ float ld_bf(const bf16* p) {
  unsigned short u = *reinterpret_cast<const unsigned short*>(p);
  return __uint_as_float(((uint32_t)u) << 16);
}
__device__ __forceinline__ float lo_bf(uint32_t w) { return __uint_as_float(w << 16); }
__device__ __forceinline__ float hi_bf(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); }
__device__ __forceinline__ float nibf(uint32_t w, int sh) {   // 128 + nibble, exact fp32
  uint32_t b = (sh <= 16) ? ((w << (16 - sh)) & 0x000F0000u) : ((w >> (sh - 16)) & 0x000F0000u);
  return __uint_as_float(b | 0x43000000u);
}
__device__ __forceinline__ float silu(float x) { return x / (1.f + __expf(-x)); }
__device__ __forceinline__ float sigmoid(float x) { return 1.f / (1.f + __expf(-x)); }

__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;
}
// block-wide sum; red must hold >= NW floats; all threads get the total. Ends synced.
__device__ __forceinline__ float block_sum(float v, float* red) {
  v = warp_sum(v);
  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  __syncthreads();
  if (lane == 0) red[warp] = v;
  __syncthreads();
  float t = 0.f;
  #pragma unroll
  for (int i = 0; i < NW; i++) t += red[i];
  __syncthreads();
  return t;
}
__device__ __forceinline__ float block_max(float v, float* red) {
  v = warp_max(v);
  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  __syncthreads();
  if (lane == 0) red[warp] = v;
  __syncthreads();
  float t = red[0];
  #pragma unroll
  for (int i = 1; i < NW; i++) t = fmaxf(t, red[i]);
  __syncthreads();
  return t;
}

__device__ __forceinline__ int ld_acquire(const int* p) {
  int v;
  asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(v) : "l"(p) : "memory");
  return v;
}
__device__ __forceinline__ unsigned long long gtimer() {
  unsigned long long t;
  asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
  return t;
}
__shared__ unsigned long long s_wait_ns;   // per-item dependency-wait time (profiling)
// all threads call; returns when counter >= target
__device__ __forceinline__ void wait_ge(const int* p, int target) {
  if (threadIdx.x == 0) {
    if (ld_acquire(p) < target) {
      const unsigned long long t0 = gtimer();
      while (ld_acquire(p) < target) { __nanosleep(40); }
      s_wait_ns += gtimer() - t0;
    }
  }
  __syncthreads();
}
// all threads call after their global writes; bumps counter by 1
__device__ __forceinline__ void signal(int* p) {
  __threadfence();
  __syncthreads();
  if (threadIdx.x == 0) atomicAdd(p, 1);
}
// all threads call; returns the ticket (old value) to every thread. If the caller turns out to
// be the last arriver it may read the other arrivers' data (fence already done).
__device__ __forceinline__ int arrive(int* p, int* s_flag) {
  __threadfence();
  __syncthreads();
  if (threadIdx.x == 0) *s_flag = atomicAdd(p, 1);
  __syncthreads();
  int old = *s_flag;
  __threadfence();
  return old;
}

__device__ __forceinline__ QW expert_view(const QW base, int e) {
  QW v = base;
  v.w += (size_t)e * (size_t)(base.K / 2) * base.N;
  v.s += (size_t)e * (size_t)(base.K / GROUP) * base.N;
  v.z += (size_t)e * (size_t)(base.K / GROUP) * base.N;
  return v;
}

// ------------------------------------------------------------------ L2 prefetch
#ifdef MK_PF_HINT
#define MK_PF(ptr) asm volatile("prefetch.global.L2 [%0];" :: "l"(ptr))
#else
#define MK_PF(ptr) { unsigned _d; asm volatile("ld.global.cg.u32 %0, [%1];" : "=r"(_d) : "l"(ptr)); }
#endif
// pull the 128-column tile rows [r0, r1) of W (and their scales/zeros) into L2; all threads
__device__ __forceinline__ void prefetch_tile(const QW W, int n0, int r0, int r1) {
  const uint8_t* base = W.w + n0;
  const bool odd = (W.N & 127) != 0;
  for (int r = r0 + threadIdx.x; r < r1; r += NT) {
    const uint8_t* p = base + (size_t)r * W.N;
    MK_PF(p);
    if (odd) MK_PF(p + 127);
  }
  const int g0 = r0 >> 6, g1 = (r1 + 63) >> 6;
  for (int g = g0 + threadIdx.x; g < g1; g += NT) {
    const bf16* sp = W.s + (size_t)g * W.N + n0;
    const bf16* zp = W.z + (size_t)g * W.N + n0;
    MK_PF(sp); MK_PF(sp + 64); MK_PF(zp); MK_PF(zp + 64);
  }
}
// pull a [rows] x 128 B slice (row pitch in bytes) into L2
__device__ __forceinline__ void prefetch_rows(const void* base, int rows, size_t pitch_bytes) {
  for (int r = threadIdx.x; r < rows; r += NT) MK_PF(reinterpret_cast<const char*>(base) + (size_t)r * pitch_bytes);
}

// ------------------------------------------------------------------ vector loaders
// xs[k] = bf16 x[k] (fp32), K elements; all threads; synced at exit
__device__ __forceinline__ void load_vec(const bf16* __restrict__ x, int K, float* xs) {
  for (int i = threadIdx.x; i < K; i += NT) xs[i] = ldcg_bf(x + i);
  __syncthreads();
}
__device__ __forceinline__ void load_vec_f32(const float* __restrict__ x, int K, float* xs) {
  for (int i = threadIdx.x; i < K; i += NT) xs[i] = __ldcg(x + i);
  __syncthreads();
}
// xs = bf16( (x * rstd) * w ) as fp32   (reference _rmsnorm). all threads; synced.
__device__ __forceinline__ void load_norm(const bf16* __restrict__ x, const bf16* __restrict__ w, int K, float* xs, float* red) {
  float ss = 0.f;
  for (int i = threadIdx.x; i < K; i += NT) { float v = ldcg_bf(x + i); xs[i] = v; ss += v * v; }
  ss = block_sum(ss, red);
  const float rstd = rsqrtf(ss / (float)K + RMS_EPS);
  for (int i = threadIdx.x; i < K; i += NT) xs[i] = bfr((xs[i] * rstd) * ld_bf(w + i));
  __syncthreads();
}

// ------------------------------------------------------------------ fused int4 dequant GEMV tile
// Partial y for the 128 columns [n0, n0+128) of W over packed rows [r0, r1) (16-row units,
// round-robin over warps). x is fp32 in smem indexed by absolute k. Result in out[128] (smem).
__device__ __forceinline__ void gemv_tile(const QW W, int n0, int r0, int r1, const float* __restrict__ xs,
                          float* red, float* out) {
  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  const int N = W.N;
  const int n = n0 + lane * 4;
  const bool act = (n < N);
  const uint8_t* __restrict__ wbase = W.w + n;
  const int nunits = (r1 - r0) >> 4;
  float y0 = 0.f, y1 = 0.f, y2 = 0.f, y3 = 0.f;

  uint32_t w[16];
  uint2 sv = make_uint2(0u, 0u), zv = make_uint2(0u, 0u);
  int u = warp;
  if (u < nunits && act) {
    const int rb = r0 + (u << 4);
    #pragma unroll
    for (int i = 0; i < 16; i++) w[i] = __ldg(reinterpret_cast<const unsigned int*>(wbase + (size_t)(rb + i) * N));
    sv = __ldg(reinterpret_cast<const uint2*>(W.s + (size_t)(rb >> 6) * N + n));
    zv = __ldg(reinterpret_cast<const uint2*>(W.z + (size_t)(rb >> 6) * N + n));
  } else {
    #pragma unroll
    for (int i = 0; i < 16; i++) w[i] = 0u;
  }
  for (; u < nunits; u += NW) {
    const int rb = r0 + (u << 4);
    // prefetch next unit (weights + its group scales/zeros)
    uint32_t wn[16];
    uint2 svn = make_uint2(0u, 0u), zvn = make_uint2(0u, 0u);
    const int un = u + NW;
    if (un < nunits && act) {
      const int rbn = r0 + (un << 4);
      #pragma unroll
      for (int i = 0; i < 16; i++) wn[i] = __ldg(reinterpret_cast<const unsigned int*>(wbase + (size_t)(rbn + i) * N));
      svn = __ldg(reinterpret_cast<const uint2*>(W.s + (size_t)(rbn >> 6) * N + n));
      zvn = __ldg(reinterpret_cast<const uint2*>(W.z + (size_t)(rbn >> 6) * N + n));
    } else {
      #pragma unroll
      for (int i = 0; i < 16; i++) wn[i] = 0u;
    }
    const float s0 = lo_bf(sv.x), s1 = hi_bf(sv.x), s2 = lo_bf(sv.y), s3 = hi_bf(sv.y);
    const float z0 = lo_bf(zv.x), z1 = hi_bf(zv.x), z2 = lo_bf(zv.y), z3 = hi_bf(zv.y);
    float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f, ax = 0.f;
    #pragma unroll
    for (int i = 0; i < 16; i++) {
      const float2 xv = *reinterpret_cast<const float2*>(xs + 2 * (rb + i));
      ax += xv.x + xv.y;
      const uint32_t wi = w[i];
      a0 = fmaf(xv.x, nibf(wi, 0), a0);  a0 = fmaf(xv.y, nibf(wi, 4), a0);
      a1 = fmaf(xv.x, nibf(wi, 8), a1);  a1 = fmaf(xv.y, nibf(wi, 12), a1);
      a2 = fmaf(xv.x, nibf(wi, 16), a2); a2 = fmaf(xv.y, nibf(wi, 20), a2);
      a3 = fmaf(xv.x, nibf(wi, 24), a3); a3 = fmaf(xv.y, nibf(wi, 28), a3);
    }
    y0 = fmaf(s0, a0 - (128.f + z0) * ax, y0);
    y1 = fmaf(s1, a1 - (128.f + z1) * ax, y1);
    y2 = fmaf(s2, a2 - (128.f + z2) * ax, y2);
    y3 = fmaf(s3, a3 - (128.f + z3) * ax, y3);
    #pragma unroll
    for (int i = 0; i < 16; i++) w[i] = wn[i];
    sv = svn; zv = zvn;
  }
  __syncthreads();
  *reinterpret_cast<float4*>(red + warp * 128 + lane * 4) = make_float4(y0, y1, y2, y3);
  __syncthreads();
  if (threadIdx.x < 128) {
    float t = 0.f;
    #pragma unroll
    for (int k = 0; k < NW; k++) t += red[k * 128 + threadIdx.x];
    out[threadIdx.x] = t;
  }
  __syncthreads();
}

__device__ __forceinline__ void bf16x8_to_f32(const uint4& v, float* f) {
  f[0] = lo_bf(v.x); f[1] = hi_bf(v.x); f[2] = lo_bf(v.y); f[3] = hi_bf(v.y);
  f[4] = lo_bf(v.z); f[5] = hi_bf(v.z); f[6] = lo_bf(v.w); f[7] = hi_bf(v.w);
}

// transposed head GEMV for MLA absorb: out[c] = sum_d xq[d] * W[c][col0 + d], c in [0,512), d in [0,128).
// 16 B per lane (16 columns), 8 lanes per packed row, 4 rows per warp-load, 32 rows per warp in 2 batches.
__device__ __forceinline__ void gemv_t_head(const QW W, int col0, const float* __restrict__ xq, float* __restrict__ out) {
  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  const int N = W.N;
  const int sub = lane >> 3;
  const int c16 = (lane & 7) * 16;
  const int n = col0 + c16;
  const int g = (warp * 32) >> 6;    // one quant group per warp
  float xs_[16], cz = 0.f;
  {
    float sc[16], zr[16];
    const uint4* sp = reinterpret_cast<const uint4*>(W.s + (size_t)g * N + n);
    const uint4* zp = reinterpret_cast<const uint4*>(W.z + (size_t)g * N + n);
    const uint4 s0 = __ldg(sp), s1 = __ldg(sp + 1), z0 = __ldg(zp), z1 = __ldg(zp + 1);
    bf16x8_to_f32(s0, sc); bf16x8_to_f32(s1, sc + 8); bf16x8_to_f32(z0, zr); bf16x8_to_f32(z1, zr + 8);
    #pragma unroll
    for (int j = 0; j < 16; j++) { xs_[j] = xq[c16 + j] * sc[j]; cz = fmaf(xs_[j], 128.f + zr[j], cz); }
  }
  #pragma unroll
  for (int b = 0; b < 2; b++) {
    uint4 wv[4];
    #pragma unroll
    for (int i = 0; i < 4; i++) {
      const int r = warp * 32 + (b * 4 + i) * 4 + sub;
      wv[i] = __ldg(reinterpret_cast<const uint4*>(W.w + (size_t)r * N + n));
    }
    #pragma unroll
    for (int i = 0; i < 4; i++) {
      const int r = warp * 32 + (b * 4 + i) * 4 + sub;
      const uint32_t wd[4] = {wv[i].x, wv[i].y, wv[i].z, wv[i].w};
      float lo = 0.f, hi = 0.f;
      #pragma unroll
      for (int q = 0; q < 4; q++) {
        const uint32_t v = wd[q];
        lo = fmaf(xs_[4 * q + 0], nibf(v, 0), lo);  hi = fmaf(xs_[4 * q + 0], nibf(v, 4), hi);
        lo = fmaf(xs_[4 * q + 1], nibf(v, 8), lo);  hi = fmaf(xs_[4 * q + 1], nibf(v, 12), hi);
        lo = fmaf(xs_[4 * q + 2], nibf(v, 16), lo); hi = fmaf(xs_[4 * q + 2], nibf(v, 20), hi);
        lo = fmaf(xs_[4 * q + 3], nibf(v, 24), lo); hi = fmaf(xs_[4 * q + 3], nibf(v, 28), hi);
      }
      lo -= cz; hi -= cz;
      #pragma unroll
      for (int o = 1; o < 8; o <<= 1) {
        lo += __shfl_xor_sync(0xffffffffu, lo, o);
        hi += __shfl_xor_sync(0xffffffffu, hi, o);
      }
      if ((lane & 7) == 0) { out[2 * r] = lo; out[2 * r + 1] = hi; }
    }
  }
}

// ------------------------------------------------------------------ MoE routing (recomputed per item)
// Needs h_res[l] (bf16) and router partials. Fills s_idx[8], s_wgt[8]; if xs != nullptr also
// writes hn = rmsnorm(h_res, moe_norm) (bf16-rounded, fp32) into xs. Ends synced.
__device__ __forceinline__ void moe_route(const Params& P, int l, float* xs, float* red, float* tmp, int* s_idx, float* s_wgt) {
  const bf16* h = P.h_res + (size_t)l * HID;
  float ss = 0.f;
  for (int i = threadIdx.x; i < HID; i += NT) { float v = ldcg_bf(h + i); if (xs) xs[i] = v; ss += v * v; }
  ss = block_sum(ss, red);
  const float rstd = rsqrtf(ss / (float)HID + RMS_EPS);
  if (xs) {
    const bf16* w = P.L[l].moe_norm;
    for (int i = threadIdx.x; i < HID; i += NT) xs[i] = bfr((xs[i] * rstd) * ld_bf(w + i));
  }
  if (threadIdx.x < NE) {
    const float* rp = P.router_part + (size_t)l * NTILE_HID * NE + threadIdx.x;
    float acc = 0.f;
    #pragma unroll
    for (int t = 0; t < NTILE_HID; t++) acc += __ldcg(rp + t * NE);
    tmp[threadIdx.x] = bfr(acc * rstd);
  }
  __syncthreads();
  if (threadIdx.x < 32) {
    const int lane = threadIdx.x;
    float l0 = tmp[lane], l1 = tmp[lane + 32];
    float m = warp_max(fmaxf(l0, l1));
    float p0 = __expf(l0 - m), p1 = __expf(l1 - m);
    float den = warp_sum(p0 + p1);
    p0 /= den; p1 /= den;
    float wsum = 0.f;
    #pragma unroll
    for (int j = 0; j < TOPK; j++) {
      float best = fmaxf(p0, p1);
      float bm = warp_max(best);
      int cand = (p0 == bm) ? lane : ((p1 == bm) ? lane + 32 : 1000);
      int idx = cand;
      #pragma unroll
      for (int o = 16; o > 0; o >>= 1) idx = min(idx, __shfl_xor_sync(0xffffffffu, idx, o));
      if (lane == 0) { s_idx[j] = idx; s_wgt[j] = bm; }
      wsum += bm;
      if (idx == lane) p0 = -1.f;
      if (idx == lane + 32) p1 = -1.f;
    }
    if (lane == 0) {
      const float inv = ROUTED_SCALING / (wsum + 1e-9f);
      #pragma unroll
      for (int j = 0; j < TOPK; j++) s_wgt[j] = s_wgt[j] * inv;
    }
  }
  __syncthreads();
}

// ------------------------------------------------------------------ item decode
struct Item { int type, layer, i; };
__device__ __forceinline__ int layer_items(int kind, int n_chunks, int phase) {
  if (kind == 0) {  // KDA
    switch (phase) {
      case 0: return NH * 4 * S_QKVG;          // IT_QKVG
      case 1: return NH * 4;                   // IT_STATE (head, slice)
      case 2: return NTILE_HID * S_O;          // IT_OPROJ
      case 3: return NSLOT * NTILE_INTER * 2 * S_GU;   // IT_GU
      case 4: return NSLOT * NTILE_HID;        // IT_DOWN
      default: return 0;
    }
  } else {
    switch (phase) {
      case 0: return NTILE_Q * S_Q;            // IT_MLAQ
      case 1: return NTILE_KVA * S_KVA;        // IT_KVA
      case 2: return n_chunks * G_ATTN;        // IT_ATTN
      case 3: return NH * 4;                   // IT_COMB (head, quarter)
      case 4: return NTILE_HID * S_O;          // IT_OPROJ
      case 5: return NSLOT * NTILE_INTER * 2 * S_GU;
      case 6: return NSLOT * NTILE_HID;
      default: return 0;
    }
  }
}
__device__ __forceinline__ int phase_type(int kind, int phase) {
  if (kind == 0) { const int t[5] = {IT_QKVG, IT_STATE, IT_OPROJ, IT_GU, IT_DOWN}; return t[phase]; }
  const int t[7] = {IT_MLAQ, IT_KVA, IT_ATTN, IT_COMB, IT_OPROJ, IT_GU, IT_DOWN};
  return t[phase];
}
__device__ __forceinline__ Item decode_item(const Params& P, int idx) {
  Item it; it.type = IT_NONE; it.layer = 0; it.i = 0;
  for (int l = 0; l < 4; l++) {
    const int kind = P.L[l].kind;
    const int np = (kind == 0) ? 5 : 7;
    for (int ph = 0; ph < np; ph++) {
      const int c = layer_items(kind, P.n_chunks, ph);
      if (idx < c) { it.type = phase_type(kind, ph); it.layer = l; it.i = idx; return it; }
      idx -= c;
    }
  }
  return it;
}
__host__ __device__ inline int total_items(const int* kinds, int n_chunks) {
  int t = 0;
  for (int l = 0; l < 4; l++) {
    const int kind = kinds[l];
    if (kind == 0) t += NH * 4 * S_QKVG + NH * 4 + NTILE_HID * S_O + NSLOT * NTILE_INTER * 2 * S_GU + NSLOT * NTILE_HID;
    else t += NTILE_Q * S_Q + NTILE_KVA * S_KVA + n_chunks * G_ATTN + NH * 4 + NTILE_HID * S_O + NSLOT * NTILE_INTER * 2 * S_GU + NSLOT * NTILE_HID;
  }
  return t;
}

__device__ __forceinline__ const bf16* layer_in(const Params& P, int l) {
  return (l == 0) ? P.hidden : (P.xres + (size_t)(l - 1) * HID);
}
__device__ __forceinline__ bf16* layer_out(const Params& P, int l) {
  return (l == 3) ? P.out : (P.xres + (size_t)l * HID);
}
__device__ __forceinline__ int* lctl(const Params& P, int l) { return P.ctl + CTL_BASE + l * CTL_PER_LAYER; }

// ------------------------------------------------------------------ static-data prefetch queues
// Queue q holds the static (input-independent) data of layer (q+1)&3: weight tiles, KDA state
// slices, shared expert, router, and (for the MLA layer) the latent cache. Blocks that wait on a
// dependency drain the queue into L2 so DRAM keeps streaming while the critical path runs.
constexpr int PF_QKVG = NH * 4 * S_QKVG, PF_S = NH * 4, PF_O = NTILE_HID * S_O, PF_SGU = NTILE_INTER * 2 * S_GU,
              PF_SD = NTILE_HID, PF_RT = (NE * HID * 2) / (384 * 128);   // 6 router pieces of 48 KB
constexpr int PF_KDA = PF_QKVG + PF_S + PF_O + PF_SGU + PF_SD + PF_RT;
constexpr int PF_Q = NTILE_Q * S_Q, PF_KVA = NTILE_KVA * S_KVA, PF_KVB = NH * 2;
constexpr int PF_MLA = PF_Q + PF_KVA + PF_KVB + PF_O + PF_SGU + PF_SD + PF_RT;
__device__ __forceinline__ int n_pieces(const Params& P, int q) {
  const int L = (q + 1) & 3;
  if (P.L[L].kind == 0) return PF_KDA;
  return PF_MLA + (P.pos + 47) / 48 + (P.pos + 383) / 384;
}
__device__ __forceinline__ void issue_piece(const Params& P, int q, int piece) {
  const int L = (q + 1) & 3;
  const LayerW& W = P.L[L];
  if (W.kind == 0) {
    constexpr int R = (HID / 2) / S_QKVG;
    if (piece < PF_QKVG) {
      const int m = piece / (NH * S_QKVG), rem = piece - m * (NH * S_QKVG), h = rem / S_QKVG, sp = rem - h * S_QKVG;
      const QW M = (m == 0) ? W.kda.q : (m == 1) ? W.kda.k : (m == 2) ? W.kda.v : W.kda.g;
      prefetch_tile(M, h * TILE, sp * R, (sp + 1) * R); return;
    }
    piece -= PF_QKVG;
    if (piece < PF_S) { prefetch_rows(P.S[L] + (size_t)(piece >> 2) * DK * DK + (piece & 3) * 32, DK, DK * sizeof(float)); return; }
    piece -= PF_S;
    if (piece < PF_O) { constexpr int RO = (CQ / 2) / S_O; prefetch_tile(W.kda.o, (piece / S_O) * TILE, (piece % S_O) * RO, (piece % S_O + 1) * RO); return; }
    piece -= PF_O;
    if (piece < PF_SGU) {
      constexpr int RG = (HID / 2) / S_GU;
      const int which = piece / (NTILE_INTER * S_GU), rem = piece - which * (NTILE_INTER * S_GU), tile = rem / S_GU, sp = rem - tile * S_GU;
      prefetch_tile(which ? W.moe.s_up : W.moe.s_gate, tile * TILE, sp * RG, (sp + 1) * RG); return;
    }
    piece -= PF_SGU;
    if (piece < PF_SD) { prefetch_tile(W.moe.s_down, piece * TILE, 0, INTER / 2); return; }
    piece -= PF_SD;
    prefetch_rows(reinterpret_cast<const char*>(W.moe.router) + (size_t)piece * 384 * 128, 384, 128);
  } else {
    constexpr int R = (HID / 2) / S_Q, RK = (HID / 2) / S_KVA, RO = (CQ / 2) / S_O, RG = (HID / 2) / S_GU;
    if (piece < PF_Q) { prefetch_tile(W.mla.q, (piece / S_Q) * TILE, (piece % S_Q) * R, (piece % S_Q + 1) * R); return; }
    piece -= PF_Q;
    if (piece < PF_KVA) { prefetch_tile(W.mla.kva, (piece / S_KVA) * TILE, (piece % S_KVA) * RK, (piece % S_KVA + 1) * RK); return; }
    piece -= PF_KVA;
    if (piece < PF_KVB) { prefetch_tile(W.mla.kvb, (piece >> 1) * (QK_NOPE + V_HEAD) + (piece & 1) * TILE, 0, KV_LORA / 2); return; }
    piece -= PF_KVB;
    if (piece < PF_O) { prefetch_tile(W.mla.o, (piece / S_O) * TILE, (piece % S_O) * RO, (piece % S_O + 1) * RO); return; }
    piece -= PF_O;
    if (piece < PF_SGU) {
      const int which = piece / (NTILE_INTER * S_GU), rem = piece - which * (NTILE_INTER * S_GU), tile = rem / S_GU, sp = rem - tile * S_GU;
      prefetch_tile(which ? W.moe.s_up : W.moe.s_gate, tile * TILE, sp * RG, (sp + 1) * RG); return;
    }
    piece -= PF_SGU;
    if (piece < PF_SD) { prefetch_tile(W.moe.s_down, piece * TILE, 0, INTER / 2); return; }
    piece -= PF_SD;
    if (piece < PF_RT) { prefetch_rows(reinterpret_cast<const char*>(W.moe.router) + (size_t)piece * 384 * 128, 384, 128); return; }
    piece -= PF_RT;
    const int nck = (P.pos + 47) / 48;
    if (piece < nck) {
      const int r0 = piece * 48, r1 = min(r0 + 48, P.pos);
      const bf16* src = (P.n_old > 0) ? P.ckv_src : P.ckv_dst;
      prefetch_rows(src + (size_t)r0 * KV_LORA, (r1 - r0) * (KV_LORA * 2 / 128), 128); return;
    }
    piece -= nck;
    const int r0 = piece * 384, r1 = min(r0 + 384, P.pos);
    const bf16* src = (P.n_old > 0) ? P.kr_src : P.kr_dst;
    if (r1 > r0) prefetch_rows(src + (size_t)r0 * QK_ROPE, r1 - r0, 128);
  }
}
__shared__ int s_pf;
// all threads call; returns when *p >= target. While waiting, the block drains prefetch queue q.
__device__ __forceinline__ void wait_pf(const Params& P, int q, const int* p, int target) {
  if (threadIdx.x == 0) s_pf = (ld_acquire(p) < target) ? 1 : 0;
  __syncthreads();
  if (s_pf == 0) return;
  unsigned long long t0 = 0;
  if (threadIdx.x == 0) t0 = gtimer();
  const int np = n_pieces(P, q);
  int* pfc = P.ctl + CTL_PF + q;
  while (true) {
    __syncthreads();
    if (threadIdx.x == 0) {
      int piece = -1;
      if (ld_acquire(p) < target) { piece = atomicAdd(pfc, 1); if (piece >= np) piece = -2; }
      s_pf = piece;
    }
    __syncthreads();
    const int piece = s_pf;
    if (piece == -1) break;
    if (piece >= 0) { issue_piece(P, q, piece); continue; }
    if (threadIdx.x == 0) { while (ld_acquire(p) < target) { __nanosleep(60); } }
    __syncthreads();
    break;
  }
  if (threadIdx.x == 0) s_wait_ns += gtimer() - t0;
}

// ------------------------------------------------------------------ items
__device__ __forceinline__ void item_qkvg(const Params& P, int l, int i, float* xs, float* red, float* out, int* s_flag) {
  const int h = i / (4 * S_QKVG);
  const int rem = i - h * 4 * S_QKVG;
  const int m = rem / S_QKVG;
  const int sp = rem - m * S_QKVG;
  const LayerW& L = P.L[l];
  const QW W = (m == 0) ? L.kda.q : (m == 1) ? L.kda.k : (m == 2) ? L.kda.v : L.kda.g;
  const int rows = (HID / 2) / S_QKVG;
  prefetch_tile(W, h * TILE, sp * rows, (sp + 1) * rows);
  if (l > 0) wait_pf(P, l - 1, lctl(P, l - 1) + C_XNEXT, NTILE_HID);
  load_norm(layer_in(P, l), L.attn_norm, HID, xs, red);
  gemv_tile(W, h * TILE, sp * rows, (sp + 1) * rows, xs, red, out);
  float* dst = P.qkvg_part + (((size_t)l * S_QKVG + sp) * 4 + m) * CQ + h * TILE;
  if (threadIdx.x < TILE) dst[threadIdx.x] = out[threadIdx.x];
  signal(lctl(P, l) + C_HEAD + h);
}

__device__ __forceinline__ void item_state(const Params& P, int l, int i, float* xs, float* red, float* aux, int* s_flag) {
  const int h = i >> 2, js = i & 3, j0 = js * 32;
  const LayerW& L = P.L[l];
  prefetch_rows(P.S[l] + (size_t)h * DK * DK + j0, DK, DK * sizeof(float));
  wait_pf(P, l, lctl(P, l) + C_HEAD + h, 4 * S_QKVG);
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
  // beta_h = sigmoid(bf16(xn . beta_w[h]))
  load_norm(layer_in(P, l), L.attn_norm, HID, xs, red);
  float bacc = 0.f;
  {
    const bf16* bw = L.kda.beta_w + (size_t)h * HID;
    for (int k = tid; k < HID; k += NT) bacc += xs[k] * ld_bf(bw + k);
  }
  const float beta = sigmoid(bfr(block_sum(bacc, red)));
  float* qc = aux;            // [128]  q after conv/silu/scale
  float* kc = aux + 128;      // [128]
  float* ge = aux + 256;      // [128]  exp(g)
  float* vc = aux + 384;      // [32]
  float* qraw = aux + 416;    // [128] bf16-rounded raw q (for window write)
  float* kraw = aux + 544;    // [128]
  float* vraw = aux + 672;    // [32]
  // raw projections (sum split partials, bf16 round)
  const float* base = P.qkvg_part + (size_t)l * S_QKVG * 4 * CQ + h * DK;
  if (tid < 128) {
    float aq = 0.f, ak = 0.f, ag = 0.f;
    #pragma unroll
    for (int sp = 0; sp < S_QKVG; sp++) {
      const float* b = base + (size_t)sp * 4 * CQ;
      aq += __ldcg(b + tid); ak += __ldcg(b + CQ + tid); ag += __ldcg(b + 3 * CQ + tid);
    }
    qraw[tid] = bfr(aq); kraw[tid] = bfr(ak);
    const float graw = bfr(ag);
    // exp(-softplus(g)) ; softplus threshold 20 like torch
    const float spl = (graw > 20.f) ? graw : log1pf(__expf(graw));
    ge[tid] = __expf(-spl);
  } else if (tid < 160) {
    const int j = tid - 128;
    float av = 0.f;
    #pragma unroll
    for (int sp = 0; sp < S_QKVG; sp++) av += __ldcg(base + (size_t)sp * 4 * CQ + 2 * CQ + j0 + j);
    vraw[j] = bfr(av);
  }
  __syncthreads();
  // short conv (window rows: 0 oldest .. 2 newest, then current) + silu
  {
    const bf16* cw = L.kda.conv_w;
    if (tid < 128) {          // q channel h*128+tid
      const int c = h * DK + tid;
      const bf16* ws = P.win_src[l][0];
      const uint2 wv = *reinterpret_cast<const uint2*>(cw + (size_t)c * 4);
      float acc = ldcg_bf(ws + c) * lo_bf(wv.x) + ldcg_bf(ws + CQ + c) * hi_bf(wv.x)
                + ldcg_bf(ws + 2 * CQ + c) * lo_bf(wv.y) + qraw[tid] * hi_bf(wv.y);
      qc[tid] = bfr(silu(acc)) * KDA_SCALE;
    } else {                  // k channel
      const int t = tid - 128;
      const int c = h * DK + t;
      const bf16* ws = P.win_src[l][1];
      const uint2 wv = *reinterpret_cast<const uint2*>(cw + ((size_t)CQ + c) * 4);
      float acc = ldcg_bf(ws + c) * lo_bf(wv.x) + ldcg_bf(ws + CQ + c) * hi_bf(wv.x)
                + ldcg_bf(ws + 2 * CQ + c) * lo_bf(wv.y) + kraw[t] * hi_bf(wv.y);
      kc[t] = bfr(silu(acc));
    }
    if (tid < 32) {           // v channel of the slice
      const int c = h * DK + j0 + tid;
      const bf16* ws = P.win_src[l][2];
      const uint2 wv = *reinterpret_cast<const uint2*>(cw + ((size_t)2 * CQ + c) * 4);
      float acc = ldcg_bf(ws + c) * lo_bf(wv.x) + ldcg_bf(ws + CQ + c) * hi_bf(wv.x)
                + ldcg_bf(ws + 2 * CQ + c) * lo_bf(wv.y) + vraw[tid] * hi_bf(wv.y);
      vc[tid] = bfr(silu(acc));
    }
    // window update for the 32 owned channels of q,k,v : new = [old1, old2, cur]
    if (tid < 96) {
      const int m = tid >> 5, t = tid & 31;
      const int c = h * DK + j0 + t;
      const bf16* ws = P.win_src[l][m];
      bf16* wd = P.win_dst[l][m];
      wd[c] = ws[CQ + c];
      wd[CQ + c] = ws[2 * CQ + c];
      const float cur = (m == 0) ? qraw[j0 + t] : (m == 1) ? kraw[j0 + t] : vraw[t];
      wd[2 * CQ + c] = __float2bfloat16(cur);
    }
  }
  __syncthreads();
  // recurrent state update for columns [j0, j0+32) of head h
  {
    float* Sp = P.S[l] + ((size_t)h * DK + warp * 16) * DK + j0 + lane;
    float sv[16];
    float pp = 0.f;
    #pragma unroll
    for (int r = 0; r < 16; r++) {
      const int ii = warp * 16 + r;
      sv[r] = __ldcg(Sp + r * DK) * ge[ii];
      pp = fmaf(sv[r], kc[ii], pp);
    }
    red[warp * 32 + lane] = pp;
    __syncthreads();
    float pred = 0.f;
    #pragma unroll
    for (int w = 0; w < NW; w++) pred += red[w * 32 + lane];
    const float delta = beta * (vc[lane] - pred);
    float oo = 0.f;
    #pragma unroll
    for (int r = 0; r < 16; r++) {
      const int ii = warp * 16 + r;
      sv[r] = fmaf(kc[ii], delta, sv[r]);
      oo = fmaf(sv[r], qc[ii], oo);
      Sp[r * DK] = sv[r];
    }
    __syncthreads();
    red[warp * 32 + lane] = oo;
    __syncthreads();
    if (tid < 32) {
      float o = 0.f;
      #pragma unroll
      for (int w = 0; w < NW; w++) o += red[w * 32 + tid];
      P.attn_o[(size_t)l * CQ + h * DK + j0 + tid] = __float2bfloat16(o);
    }
  }
  signal(lctl(P, l) + C_O);
}

__device__ __forceinline__ void item_oproj(const Params& P, int l, int i, float* xs, float* red, float* out, int* s_flag) {
  const int tile = i / S_O, sp = i - tile * S_O;
  const LayerW& L = P.L[l];
  const QW W = (L.kind == 0) ? L.kda.o : L.mla.o;
  const int rows = (CQ / 2) / S_O;
  const int n0 = tile * TILE;
  prefetch_tile(W, n0, sp * rows, (sp + 1) * rows);
  wait_pf(P, l, lctl(P, l) + C_O, (L.kind == 0) ? NH * 4 : NH);
  load_vec(P.attn_o + (size_t)l * CQ, CQ, xs);
  gemv_tile(W, n0, sp * rows, (sp + 1) * rows, xs, red, out);
  float* part = P.o_part + ((size_t)l * S_O + sp) * HID + n0;
  if (threadIdx.x < TILE) part[threadIdx.x] = out[threadIdx.x];
  const int old = arrive(lctl(P, l) + C_OTILE + tile, s_flag);
  if (old != S_O - 1) return;
  // last split of this tile: residual + h_res + router partials
  float* hs = out;   // reuse [128]
  const bf16* xin = layer_in(P, l);
  if (threadIdx.x < TILE) {
    const int n = n0 + threadIdx.x;
    float acc = 0.f;
    #pragma unroll
    for (int s = 0; s < S_O; s++) acc += __ldcg(P.o_part + ((size_t)l * S_O + s) * HID + n);
    const float hv = bfr(ldcg_bf(xin + n) + bfr(acc));
    hs[threadIdx.x] = hv;
    P.h_res[(size_t)l * HID + n] = __float2bfloat16(hv);
  }
  __syncthreads();
  {
    const int e = threadIdx.x >> 2, q = threadIdx.x & 3;
    const bf16* wr = L.moe.router + (size_t)e * HID + n0 + q * 32;
    const bf16* mn = L.moe_norm + n0 + q * 32;
    float acc = 0.f;
    #pragma unroll 8
    for (int t = 0; t < 32; t++) acc += hs[q * 32 + t] * ld_bf(mn + t) * ld_bf(wr + t);
    acc += __shfl_xor_sync(0xffffffffu, acc, 1);
    acc += __shfl_xor_sync(0xffffffffu, acc, 2);
    if (q == 0) P.router_part[((size_t)l * NTILE_HID + tile) * NE + e] = acc;
  }
  signal(lctl(P, l) + C_HRES);
}

__device__ __forceinline__ void item_gu(const Params& P, int l, int i, float* xs, float* red, float* out, float* tmp,
                        int* s_idx, float* s_wgt, int* s_flag) {
  const int per_slot = NTILE_INTER * 2 * S_GU;
  const int slot = i / per_slot;
  int rem = i - slot * per_slot;
  const int tile = rem / (2 * S_GU); rem -= tile * 2 * S_GU;
  const int which = rem / S_GU;             // 0 gate, 1 up
  const int sp = rem - which * S_GU;
  const LayerW& L = P.L[l];
  const int rows = (HID / 2) / S_GU;
  const int n0 = tile * TILE;
  if (slot == TOPK) prefetch_tile(which ? L.moe.s_up : L.moe.s_gate, n0, sp * rows, (sp + 1) * rows);
  wait_pf(P, l, lctl(P, l) + C_HRES, NTILE_HID);
  moe_route(P, l, xs, red, tmp, s_idx, s_wgt);
  QW W;
  if (slot < TOPK) { W = expert_view(which ? L.moe.up : L.moe.gate, s_idx[slot]); prefetch_tile(W, n0, sp * rows, (sp + 1) * rows); }
  else W = which ? L.moe.s_up : L.moe.s_gate;
  gemv_tile(W, n0, sp * rows, (sp + 1) * rows, xs, red, out);
  float* part = P.gu_part + ((((size_t)l * NSLOT + slot) * 2 + which) * S_GU + sp) * INTER + n0;
  if (threadIdx.x < TILE) part[threadIdx.x] = out[threadIdx.x];
  const int old = arrive(lctl(P, l) + C_GUT + slot * NTILE_INTER + tile, s_flag);
  if (old != 2 * S_GU - 1) return;
  if (threadIdx.x < TILE) {
    const int n = n0 + threadIdx.x;
    const float* gb = P.gu_part + (((size_t)l * NSLOT + slot) * 2 + 0) * S_GU * INTER + n;
    const float* ub = P.gu_part + (((size_t)l * NSLOT + slot) * 2 + 1) * S_GU * INTER + n;
    float g = 0.f, u = 0.f;
    #pragma unroll
    for (int s = 0; s < S_GU; s++) { g += __ldcg(gb + (size_t)s * INTER); u += __ldcg(ub + (size_t)s * INTER); }
    P.h1[((size_t)l * NSLOT + slot) * INTER + n] = silu(g) * u;
  }
  signal(lctl(P, l) + C_H1 + slot);
}

__device__ __forceinline__ void item_down(const Params& P, int l, int i, float* xs, float* red, float* out, float* tmp,
                          int* s_idx, float* s_wgt, int* s_flag) {
  const int slot = i / NTILE_HID, tile = i - slot * NTILE_HID;
  const LayerW& L = P.L[l];
  const int n0 = tile * TILE;
  if (slot == TOPK) prefetch_tile(L.moe.s_down, n0, 0, INTER / 2);
  wait_pf(P, l, lctl(P, l) + C_H1 + slot, NTILE_INTER);
  moe_route(P, l, nullptr, red, tmp, s_idx, s_wgt);
  QW W;
  if (slot < TOPK) { W = expert_view(L.moe.down, s_idx[slot]); prefetch_tile(W, n0, 0, INTER / 2); }
  else W = L.moe.s_down;
  load_vec_f32(P.h1 + ((size_t)l * NSLOT + slot) * INTER, INTER, xs);
  gemv_tile(W, n0, 0, INTER / 2, xs, red, out);
  float* part = P.down_part + ((size_t)l * NSLOT + slot) * HID + n0;
  if (threadIdx.x < TILE) part[threadIdx.x] = out[threadIdx.x];
  const int old = arrive(lctl(P, l) + C_DOWNT + tile, s_flag);
  if (old != NSLOT - 1) return;
  if (threadIdx.x < TILE) {
    const int n = n0 + threadIdx.x;
    const float* dp = P.down_part + (size_t)l * NSLOT * HID + n;
    float acc = 0.f;
    #pragma unroll
    for (int j = 0; j < TOPK; j++) acc += s_wgt[j] * __ldcg(dp + (size_t)j * HID);
    acc += __ldcg(dp + (size_t)TOPK * HID);
    const float hv = ldcg_bf(P.h_res + (size_t)l * HID + n);
    layer_out(P, l)[n] = __float2bfloat16(hv + bfr(acc));
  }
  signal(lctl(P, l) + C_XNEXT);
}

// ---- MLA ----
__device__ __forceinline__ void mla_head_epilogue(const Params& P, int l, int hh, float* xs, float* red, float* aux) {
  // xs free to use (size >= 512+...)
  float* qn = aux;          // [128]
  float* qrr = aux + 128;   // [64]
  const int tid = threadIdx.x;
  if (tid < QK_NOPE + QK_ROPE) {
    float acc = 0.f;
    #pragma unroll
    for (int s = 0; s < S_Q; s++) acc += __ldcg(P.q_part + (size_t)s * Q_DIM + hh * (QK_NOPE + QK_ROPE) + tid);
    acc = bfr(acc);
    if (tid < QK_NOPE) qn[tid] = acc; else qrr[tid - QK_NOPE] = acc;
  }
  __syncthreads();
  if (tid < QK_ROPE / 2) {
    const float e = qrr[2 * tid], o = qrr[2 * tid + 1];
    const float inv = 1.0f / powf(10000.0f, (float)(2 * tid) / (float)QK_ROPE);
    const float ang = (float)P.pos * inv;
    float sn, cs;
    sincosf(ang, &sn, &cs);
    P.q_rope[(size_t)hh * QK_ROPE + 2 * tid] = bfr(e * cs - o * sn);
    P.q_rope[(size_t)hh * QK_ROPE + 2 * tid + 1] = bfr(o * cs + e * sn);
  }
  gemv_t_head(P.L[l].mla.kvb, hh * (QK_NOPE + V_HEAD), qn, P.q_abs + (size_t)hh * KV_LORA);
  signal(lctl(P, l) + C_QREADY);
}

__device__ __forceinline__ void item_mlaq(const Params& P, int l, int i, float* xs, float* red, float* out, float* aux, int* s_flag) {
  const int tile = i / S_Q, sp = i - tile * S_Q;
  const LayerW& L = P.L[l];
  const int rows = (HID / 2) / S_Q;
  const int n0 = tile * TILE;
  prefetch_tile(L.mla.q, n0, sp * rows, (sp + 1) * rows);
  {
    const int ph0 = n0 / (QK_NOPE + QK_ROPE), ph1 = (n0 + TILE - 1) / (QK_NOPE + QK_ROPE);
    prefetch_tile(L.mla.kvb, ph0 * (QK_NOPE + V_HEAD), 0, KV_LORA / 2);
    if (ph1 != ph0) prefetch_tile(L.mla.kvb, ph1 * (QK_NOPE + V_HEAD), 0, KV_LORA / 2);
  }
  if (l > 0) wait_pf(P, l - 1, lctl(P, l - 1) + C_XNEXT, NTILE_HID);
  load_norm(layer_in(P, l), L.attn_norm, HID, xs, red);
  gemv_tile(L.mla.q, n0, sp * rows, (sp + 1) * rows, xs, red, out);
  float* part = P.q_part + (size_t)sp * Q_DIM + n0;
  if (threadIdx.x < TILE) part[threadIdx.x] = out[threadIdx.x];
  int old = arrive(lctl(P, l) + C_QTILE + tile, s_flag);
  if (old != S_Q - 1) return;
  const int h0 = (tile * TILE) / (QK_NOPE + QK_ROPE);
  const int h1 = (tile * TILE + TILE - 1) / (QK_NOPE + QK_ROPE);
  old = arrive(lctl(P, l) + C_QHEAD + h0, s_flag);
  if (old == 1) mla_head_epilogue(P, l, h0, xs, red, aux);
  if (h1 != h0) {
    old = arrive(lctl(P, l) + C_QHEAD + h1, s_flag);
    if (old == 1) mla_head_epilogue(P, l, h1, xs, red, aux);
  }
}

__device__ __forceinline__ void item_kva(const Params& P, int l, int i, float* xs, float* red, float* out, float* aux, int* s_flag) {
  const int tile = i / S_KVA, sp = i - tile * S_KVA;
  const LayerW& L = P.L[l];
  const int rows = (HID / 2) / S_KVA;
  const int n0 = tile * TILE;
  prefetch_tile(L.mla.kva, n0, sp * rows, (sp + 1) * rows);
  if (l > 0) wait_pf(P, l - 1, lctl(P, l - 1) + C_XNEXT, NTILE_HID);
  load_norm(layer_in(P, l), L.attn_norm, HID, xs, red);
  gemv_tile(L.mla.kva, n0, sp * rows, (sp + 1) * rows, xs, red, out);
  float* part = P.kva_part + (size_t)sp * (NTILE_KVA * TILE) + n0;
  if (threadIdx.x < TILE) part[threadIdx.x] = out[threadIdx.x];
  const int old = arrive(lctl(P, l) + C_KVAT + tile, s_flag);
  if (old != S_KVA - 1) return;
  float* krr = aux;  // [64]
  const int tid = threadIdx.x;
  if (tid < TILE) {
    const int n = n0 + tid;
    if (n < KVA_DIM) {
      float acc = 0.f;
      #pragma unroll
      for (int s = 0; s < S_KVA; s++) acc += __ldcg(P.kva_part + (size_t)s * (NTILE_KVA * TILE) + n);
      if (n < KV_LORA) P.ckv_dst[(size_t)P.pos * KV_LORA + n] = __float2bfloat16(acc);
      else krr[n - KV_LORA] = bfr(acc);
    }
  }
  __syncthreads();
  if (n0 + TILE > KV_LORA && tid < QK_ROPE / 2) {
    const float e = krr[2 * tid], o = krr[2 * tid + 1];
    const float inv = 1.0f / powf(10000.0f, (float)(2 * tid) / (float)QK_ROPE);
    const float ang = (float)P.pos * inv;
    float sn, cs;
    sincosf(ang, &sn, &cs);
    P.kr_dst[(size_t)P.pos * QK_ROPE + 2 * tid] = __float2bfloat16(e * cs - o * sn);
    P.kr_dst[(size_t)P.pos * QK_ROPE + 2 * tid + 1] = __float2bfloat16(o * cs + e * sn);
  }
  signal(lctl(P, l) + C_KVADONE);
}

// attention over one chunk of positions for 8 heads (one head per warp), flash-decoding partials
__device__ __forceinline__ void item_attn(const Params& P, int l, int i, float* xs, float* red) {
  const int c = i / G_ATTN, g = i - c * G_ATTN;
  wait_pf(P, l, lctl(P, l) + C_QREADY, NH);
  wait_pf(P, l, lctl(P, l) + C_KVADONE, NTILE_KVA);
  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  const int hh = g * HPG + warp;
  const int p_begin = c * P.chunk;
  const int p_end = min(p_begin + P.chunk, P.pos + 1);
  const bool do_copy = (P.n_old > 0) && (g == 0) && (warp == 0);
  float qa[16];
  {
    const float4* qp = reinterpret_cast<const float4*>(P.q_abs + (size_t)hh * KV_LORA + lane * 16);
    #pragma unroll
    for (int t = 0; t < 4; t++) { float4 v = __ldcg(qp + t); qa[4*t] = v.x; qa[4*t+1] = v.y; qa[4*t+2] = v.z; qa[4*t+3] = v.w; }
  }
  const float2 qr = __ldcg(reinterpret_cast<const float2*>(P.q_rope + (size_t)hh * QK_ROPE + lane * 2));
  float m = __int_as_float(0xff800000), lsum = 0.f;
  float o[16];
  #pragma unroll
  for (int t = 0; t < 16; t++) o[t] = 0.f;

  for (int p0 = p_begin; p0 < p_end; p0 += 4) {
    uint4 r0[4], r1[4]; uint32_t kr[4];
    #pragma unroll
    for (int u = 0; u < 4; u++) {
      const int p = p0 + u;
      if (p < p_end) {
        const bool old = (p < P.n_old);
        const bf16* rowp = (old ? P.ckv_src : P.ckv_dst) + (size_t)p * KV_LORA + lane * 16;
        const bf16* krp = (old ? P.kr_src : P.kr_dst) + (size_t)p * QK_ROPE + lane * 2;
        if (p == P.pos) {
          r0[u] = __ldcg(reinterpret_cast<const uint4*>(rowp));
          r1[u] = __ldcg(reinterpret_cast<const uint4*>(rowp) + 1);
          kr[u] = __ldcg(reinterpret_cast<const unsigned int*>(krp));
        } else {
          r0[u] = *reinterpret_cast<const uint4*>(rowp);
          r1[u] = *(reinterpret_cast<const uint4*>(rowp) + 1);
          kr[u] = *reinterpret_cast<const unsigned int*>(krp);
        }
      } else {
        r0[u] = make_uint4(0, 0, 0, 0); r1[u] = r0[u]; kr[u] = 0u;
      }
    }
    #pragma unroll
    for (int u = 0; u < 4; u++) {
      const int p = p0 + u;
      if (p < p_end) {
        float f[16];
        bf16x8_to_f32(r0[u], f); bf16x8_to_f32(r1[u], f + 8);
        if (do_copy && p < P.n_old) {
          bf16* d = P.ckv_dst + (size_t)p * KV_LORA + lane * 16;
          *reinterpret_cast<uint4*>(d) = r0[u];
          *(reinterpret_cast<uint4*>(d) + 1) = r1[u];
          *reinterpret_cast<unsigned int*>(P.kr_dst + (size_t)p * QK_ROPE + lane * 2) = kr[u];
        }
        float s = qr.x * lo_bf(kr[u]) + qr.y * hi_bf(kr[u]);
        #pragma unroll
        for (int t = 0; t < 16; t++) s = fmaf(qa[t], f[t], s);
        s = warp_sum(s) * MLA_SCALE;
        if (s > m) {
          const float alpha = __expf(m - s);
          lsum *= alpha;
          #pragma unroll
          for (int t = 0; t < 16; t++) o[t] *= alpha;
          m = s;
        }
        const float pr = __expf(s - m);
        lsum += pr;
        #pragma unroll
        for (int t = 0; t < 16; t++) o[t] = fmaf(pr, f[t], o[t]);
      }
    }
  }
  float* dst = P.attn_part + ((size_t)c * NH + hh) * APITCH;
  if (lane == 0) { dst[0] = m; dst[1] = lsum; }
  #pragma unroll
  for (int t = 0; t < 16; t++) dst[2 + lane * 16 + t] = o[t];
  signal(lctl(P, l) + C_ATTN + g);
}

// combine partials for (head, quarter of the latent dim); last quarter does the kv_b (v part) GEMV
__device__ __forceinline__ void item_comb(const Params& P, int l, int i, float* xs, float* red, float* out, int* s_flag) {
  const int hh = i >> 2, q = i & 3;
  prefetch_tile(P.L[l].mla.kvb, hh * (QK_NOPE + V_HEAD) + QK_NOPE, 0, KV_LORA / 2);
  wait_pf(P, l, lctl(P, l) + C_ATTN + (hh / HPG), P.n_chunks);
  const int tid = threadIdx.x;
  const int nc = P.n_chunks;
  const float* base = P.attn_part + (size_t)hh * APITCH;
  float mloc = __int_as_float(0xff800000);
  for (int c = tid; c < nc; c += NT) mloc = fmaxf(mloc, __ldcg(base + (size_t)c * NH * APITCH));
  const float M = block_max(mloc, red);
  const int cc = tid & 127, half = tid >> 7;
  const int c_lo = half ? (nc >> 1) : 0, c_hi = half ? nc : (nc >> 1);
  float acc = 0.f, lacc = 0.f;
  for (int c = c_lo; c < c_hi; c++) {
    const float* bp = base + (size_t)c * NH * APITCH;
    const float sc = __expf(__ldcg(bp) - M);
    lacc = fmaf(__ldcg(bp + 1), sc, lacc);
    acc = fmaf(__ldcg(bp + 2 + q * 128 + cc), sc, acc);
  }
  __syncthreads();
  red[tid] = acc; red[NT + tid] = lacc;
  __syncthreads();
  if (tid < 128) {
    const float L = red[NT] + red[NT + 128];
    P.o_lat[(size_t)hh * KV_LORA + q * 128 + tid] = (red[tid] + red[128 + tid]) / L;
  }
  const int old = arrive(lctl(P, l) + C_OLAT + hh, s_flag);
  if (old != 3) return;
  load_vec_f32(P.o_lat + (size_t)hh * KV_LORA, KV_LORA, xs);
  gemv_tile(P.L[l].mla.kvb, hh * (QK_NOPE + V_HEAD) + QK_NOPE, 0, KV_LORA / 2, xs, red, out);
  if (tid < V_HEAD) P.attn_o[(size_t)l * CQ + hh * V_HEAD + tid] = __float2bfloat16(out[tid]);
  signal(lctl(P, l) + C_O);
}

// ------------------------------------------------------------------ the megakernel
__global__ void __launch_bounds__(NT, 2) megakernel(const Params P) {
  extern __shared__ __align__(16) float smem[];
  __shared__ int s_item;
  __shared__ int s_flag;
  __shared__ int s_idx[TOPK];
  __shared__ float s_wgt[TOPK];
  float* xs = smem;                    // [4096]
  float* red = smem + 4096;            // [NW*128]
  float* aux = red + NW * 128;         // [1024]
  float* out = aux + 768;              // [128] gemv tile result
  float* tmp = aux + 896;              // [128]

  int* ctl = P.ctl;
  unsigned long long t_entry = 0, t0 = 0;
  if (P.prof != nullptr && threadIdx.x == 0) t_entry = gtimer();
  if (threadIdx.x == 0) s_item = atomicAdd(&ctl[CTL_NEXT], 1);
  __syncthreads();
  int item = s_item;
  while (item < P.n_items) {
    __syncthreads();
    if (threadIdx.x == 0) { s_item = atomicAdd(&ctl[CTL_NEXT], 1); s_wait_ns = 0; }   // prefetch next ticket
    if (P.prof != nullptr && threadIdx.x == 0) t0 = gtimer();
    const Item it = decode_item(P, item);
    switch (it.type) {
      case IT_QKVG:  item_qkvg(P, it.layer, it.i, xs, red, out, &s_flag); break;
      case IT_STATE: item_state(P, it.layer, it.i, xs, red, aux, &s_flag); break;
      case IT_OPROJ: item_oproj(P, it.layer, it.i, xs, red, out, &s_flag); break;
      case IT_GU:    item_gu(P, it.layer, it.i, xs, red, out, tmp, s_idx, s_wgt, &s_flag); break;
      case IT_DOWN:  item_down(P, it.layer, it.i, xs, red, out, tmp, s_idx, s_wgt, &s_flag); break;
      case IT_MLAQ:  item_mlaq(P, it.layer, it.i, xs, red, out, aux, &s_flag); break;
      case IT_KVA:   item_kva(P, it.layer, it.i, xs, red, out, aux, &s_flag); break;
      case IT_ATTN:  item_attn(P, it.layer, it.i, xs, red); break;
      case IT_COMB:  item_comb(P, it.layer, it.i, xs, red, out, &s_flag); break;
      default: break;
    }
    __syncthreads();
    if (P.prof != nullptr && threadIdx.x == 0) {
      unsigned long long* r = P.prof + (size_t)item * 4;
      unsigned smid;
      asm("mov.u32 %0, %%smid;" : "=r"(smid));
      r[0] = t0; r[1] = gtimer(); r[2] = s_wait_ns; r[3] = ((unsigned long long)smid << 32) | blockIdx.x;
    }
    item = s_item;
  }
  // last block out resets all counters for the next launch
  __syncthreads();
  if (threadIdx.x == 0) {
    if (P.prof != nullptr) {
      P.prof[(size_t)P.n_items * 4 + (size_t)blockIdx.x * 2] = t_entry;
      P.prof[(size_t)P.n_items * 4 + (size_t)blockIdx.x * 2 + 1] = gtimer();
    }
    __threadfence();
    s_flag = atomicAdd(&ctl[CTL_EXIT], 1);
  }
  __syncthreads();
  if (s_flag == (int)gridDim.x - 1) {
    __threadfence();
    for (int i = threadIdx.x; i < N_CTL; i += NT) ctl[i] = 0;
    __threadfence();
  }
}

// ------------------------------------------------------------------ host side
static std::vector<Params> g_models;
static int g_grid = 0;

static QW take_qw(const int64_t* v, size_t& p, int K, int N) {
  QW q;
  q.w = reinterpret_cast<const uint8_t*>(v[p++]);
  q.s = reinterpret_cast<const bf16*>(v[p++]);
  q.z = reinterpret_cast<const bf16*>(v[p++]);
  q.K = K; q.N = N;
  return q;
}

int64_t mk_setup(torch::Tensor kinds_t, torch::Tensor ptrs_t) {
  TORCH_CHECK(kinds_t.dtype() == torch::kInt64 && ptrs_t.dtype() == torch::kInt64 && !ptrs_t.is_cuda());
  const int64_t* kinds = kinds_t.data_ptr<int64_t>();
  const int64_t* ptrs = ptrs_t.data_ptr<int64_t>();
  const size_t nptrs = (size_t)ptrs_t.numel();
  Params P;
  memset(&P, 0, sizeof(P));
  size_t p = 0;
  for (int l = 0; l < 4; l++) {
    LayerW& L = P.L[l];
    L.kind = (int)kinds[l];
    L.attn_norm = reinterpret_cast<const bf16*>(ptrs[p++]);
    L.moe_norm = reinterpret_cast<const bf16*>(ptrs[p++]);
    if (L.kind == 0) {
      L.kda.q = take_qw(ptrs, p, HID, CQ);
      L.kda.k = take_qw(ptrs, p, HID, CQ);
      L.kda.v = take_qw(ptrs, p, HID, CQ);
      L.kda.g = take_qw(ptrs, p, HID, CQ);
      L.kda.o = take_qw(ptrs, p, CQ, HID);
      L.kda.beta_w = reinterpret_cast<const bf16*>(ptrs[p++]);
      L.kda.conv_w = reinterpret_cast<const bf16*>(ptrs[p++]);
    } else {
      L.mla.q = take_qw(ptrs, p, HID, Q_DIM);
      L.mla.kva = take_qw(ptrs, p, HID, KVA_DIM);
      L.mla.kvb = take_qw(ptrs, p, KV_LORA, KVB_DIM);
      L.mla.o = take_qw(ptrs, p, CQ, HID);
    }
    L.moe.router = reinterpret_cast<const bf16*>(ptrs[p++]);
    L.moe.gate = take_qw(ptrs, p, HID, INTER);
    L.moe.up = take_qw(ptrs, p, HID, INTER);
    L.moe.down = take_qw(ptrs, p, INTER, HID);
    L.moe.s_gate = take_qw(ptrs, p, HID, INTER);
    L.moe.s_up = take_qw(ptrs, p, HID, INTER);
    L.moe.s_down = take_qw(ptrs, p, INTER, HID);
  }
  P.qkvg_part = reinterpret_cast<float*>(ptrs[p++]);
  P.q_part = reinterpret_cast<float*>(ptrs[p++]);
  P.kva_part = reinterpret_cast<float*>(ptrs[p++]);
  P.q_abs = reinterpret_cast<float*>(ptrs[p++]);
  P.q_rope = reinterpret_cast<float*>(ptrs[p++]);
  P.attn_part = reinterpret_cast<float*>(ptrs[p++]);
  P.o_lat = reinterpret_cast<float*>(ptrs[p++]);
  P.o_part = reinterpret_cast<float*>(ptrs[p++]);
  P.attn_o = reinterpret_cast<bf16*>(ptrs[p++]);
  P.h_res = reinterpret_cast<bf16*>(ptrs[p++]);
  P.router_part = reinterpret_cast<float*>(ptrs[p++]);
  P.gu_part = reinterpret_cast<float*>(ptrs[p++]);
  P.h1 = reinterpret_cast<float*>(ptrs[p++]);
  P.down_part = reinterpret_cast<float*>(ptrs[p++]);
  P.xres = reinterpret_cast<bf16*>(ptrs[p++]);
  P.ctl = reinterpret_cast<int*>(ptrs[p++]);
  TORCH_CHECK(p == nptrs, "pointer table size mismatch: ", p, " vs ", nptrs);
  if (g_grid == 0) {
    int dev = 0;
    cudaGetDevice(&dev);
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, dev);
    int nb = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&nb, megakernel, NT, SMEM_BYTES);
    TORCH_CHECK(nb >= 1, "megakernel cannot be resident");
    if (nb > 2) nb = 2;
    g_grid = nb * prop.multiProcessorCount;
    if (const char* e = getenv("MK_GRID")) { const int g = atoi(e); if (g > 0 && g < g_grid) g_grid = g; }
  }
  g_models.push_back(P);
  return (int64_t)g_models.size() - 1;
}

void mk_step(int64_t id, torch::Tensor dyn_t, int64_t pos, int64_t n_old, int64_t n_chunks, int64_t chunk) {
  Params P = g_models.at((size_t)id);
  const int64_t* dyn = dyn_t.data_ptr<int64_t>();
  size_t p = 0;
  P.hidden = reinterpret_cast<const bf16*>(dyn[p++]);
  P.out = reinterpret_cast<bf16*>(dyn[p++]);
  for (int l = 0; l < 3; l++) P.S[l] = reinterpret_cast<float*>(dyn[p++]);
  for (int l = 0; l < 3; l++) for (int m = 0; m < 3; m++) P.win_src[l][m] = reinterpret_cast<const bf16*>(dyn[p++]);
  for (int l = 0; l < 3; l++) for (int m = 0; m < 3; m++) P.win_dst[l][m] = reinterpret_cast<bf16*>(dyn[p++]);
  P.ckv_src = reinterpret_cast<const bf16*>(dyn[p++]);
  P.kr_src = reinterpret_cast<const bf16*>(dyn[p++]);
  P.ckv_dst = reinterpret_cast<bf16*>(dyn[p++]);
  P.kr_dst = reinterpret_cast<bf16*>(dyn[p++]);
  P.prof = reinterpret_cast<unsigned long long*>(dyn[p++]);
  P.pos = (int)pos; P.n_old = (int)n_old; P.n_chunks = (int)n_chunks; P.chunk = (int)chunk;
  int kinds[4];
  for (int l = 0; l < 4; l++) kinds[l] = P.L[l].kind;
  P.n_items = total_items(kinds, P.n_chunks);
  void* args[] = { (void*)&P };
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  cudaError_t err = cudaLaunchCooperativeKernel((const void*)megakernel, dim3(g_grid), dim3(NT), args, SMEM_BYTES, stream);
  TORCH_CHECK(err == cudaSuccess, "megakernel launch failed: ", cudaGetErrorString(err));
}

int64_t mk_grid() { return g_grid; }
int64_t mk_n_ctl() { return N_CTL; }
int64_t mk_maxch() { return MAXCH; }

}  // namespace mk

// Python bindings live in this (nvcc-compiled) translation unit.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("mk_setup", &mk::mk_setup, "register weights/scratch pointer table; returns model id");
  m.def("mk_step", &mk::mk_step, "launch the decode megakernel once");
  m.def("mk_grid", &mk::mk_grid);
  m.def("mk_n_ctl", &mk::mk_n_ctl);
  m.def("mk_maxch", &mk::mk_maxch);
}
"""

# The host stub is empty: GCC 15 rejects torch's python headers, so the pybind module is
# defined in the nvcc-compiled unit above and the stub is built without TORCH_API_INCLUDE_EXTENSION_H.
CPP_SRC = "// bindings are defined in the CUDA translation unit\n"

_EXT = None


def _ext():
    global _EXT
    if _EXT is None:
        from torch.utils.cpp_extension import load_inline
        _EXT = load_inline(
            name="kimi_linear_megakernel_v1",
            cpp_sources=CPP_SRC,
            cuda_sources=CUDA_SRC,
            extra_cflags=["-UTORCH_API_INCLUDE_EXTENSION_H"],
            extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "-std=c++20", "-Xptxas", "-v"]
            + (["-DMK_PF_HINT"] if os.environ.get("MK_PF_HINT") else []),
            verbose=bool(os.environ.get("MK_VERBOSE")),
        )
    return _EXT


# --------------------------------------------------------------------------- #
# Modules (same parameter / buffer names as the reference so its state_dict loads)
# --------------------------------------------------------------------------- #
class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        self.in_f, self.out_f, self.group = in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16))


class QuantExperts(nn.Module):
    def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group
        ng = in_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, ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16))


class KDA(nn.Module):
    def __init__(self, cfg: Config):
        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: Config):
        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: Config):
        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: Config, 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):
    """Same interface as reference.Model; step() = one cooperative megakernel launch."""

    CACHE_SLACK = 256  # extra cache rows allocated beyond the ingested context

    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        assert tuple(cfg.pattern) == ("K", "K", "K", "M"), "kernel is specialised for [K,K,K,M]"
        assert cfg.hidden == 2304 and cfg.n_experts == 64 and cfg.n_active == 8 and cfg.n_shared == 1
        assert cfg.moe_inter == 1024 and cfg.kda_heads == 32 and cfg.kda_head_dim == 128
        assert cfg.mla_heads == 32 and cfg.kv_lora == 512 and cfg.qk_nope == 128 and cfg.qk_rope == 64
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._id = None
        self._dev = None
        self._parity = 0
        self._ckv_buf = None
        self._kr_buf = None
        self._cap = 0

    # ---- one-time device setup (after weights are loaded) ----
    def _qw(self, q):
        return [q.w_q.data_ptr(), q.scales.data_ptr(), q.zeros.data_ptr()]

    def _setup(self, dev: torch.device):
        ext = _ext()
        S_QKVG, S_O, S_GU, S_Q, S_KVA = 3, 4, 3, 3, 3
        maxch = int(ext.mk_maxch())
        f32 = dict(dtype=torch.float32, device=dev)
        bf = dict(dtype=torch.bfloat16, device=dev)
        self._scr = dict(
            qkvg_part=torch.zeros(4, S_QKVG, 4, 4096, **f32),
            q_part=torch.zeros(S_Q, 6144, **f32),
            kva_part=torch.zeros(S_KVA, 5 * 128, **f32),
            q_abs=torch.zeros(32, 512, **f32),
            q_rope=torch.zeros(32, 64, **f32),
            attn_part=torch.zeros(maxch, 32, 514, **f32),
            o_lat=torch.zeros(32, 512, **f32),
            o_part=torch.zeros(4, S_O, 2304, **f32),
            attn_o=torch.zeros(4, 4096, **bf),
            h_res=torch.zeros(4, 2304, **bf),
            router_part=torch.zeros(4, 18, 64, **f32),
            gu_part=torch.zeros(4, 9, 2, S_GU, 1024, **f32),
            h1=torch.zeros(4, 9, 1024, **f32),
            down_part=torch.zeros(4, 9, 2304, **f32),
            xres=torch.zeros(3, 2304, **bf),
            ctl=torch.zeros(int(ext.mk_n_ctl()), dtype=torch.int32, device=dev),
        )
        kinds, ptrs = [], []
        for blk in self.blocks:
            kinds.append(0 if blk.kind == "K" else 1)
            ptrs += [blk.attn_norm.data_ptr(), blk.moe_norm.data_ptr()]
            a = blk.attn
            if blk.kind == "K":
                for q in (a.q_proj, a.k_proj, a.v_proj, a.g_proj, a.o_proj):
                    ptrs += self._qw(q)
                ptrs += [a.beta_proj.weight.data_ptr(), a.conv_w.data_ptr()]
            else:
                for q in (a.q_proj, a.kv_a, a.kv_b, a.o_proj):
                    ptrs += self._qw(q)
            m = blk.moe
            ptrs.append(m.router.weight.data_ptr())
            for q in (m.gate, m.up, m.down, m.s_gate, m.s_up, m.s_down):
                ptrs += self._qw(q)
        for k in ("qkvg_part", "q_part", "kva_part", "q_abs", "q_rope", "attn_part", "o_lat", "o_part",
                  "attn_o", "h_res", "router_part", "gu_part", "h1", "down_part", "xres", "ctl"):
            ptrs.append(self._scr[k].data_ptr())
        # sanity: every weight tensor must be contiguous on this device
        for t in list(self.parameters()) + list(self.buffers()):
            assert t.is_contiguous() and t.device == dev, "weights must be contiguous on the kernel device"
        self._id = int(ext.mk_setup(torch.tensor(kinds, dtype=torch.int64), torch.tensor(ptrs, dtype=torch.int64)))
        self._dev = dev
        # optional per-item timeline buffer (profiling aid; off unless MK_PROF is set)
        self._prof = torch.zeros(1 << 16, dtype=torch.int64, device=dev) if os.environ.get("MK_PROF") else None
        self._out = [torch.zeros(2304, **bf), torch.zeros(2304, **bf)]
        # ping-pong short-conv windows for the three KDA layers: [layer][q/k/v][parity]
        self._win = [[[torch.zeros(3, 4096, **bf) for _ in range(2)] for _ in range(3)] for _ in range(3)]

    def _ensure_cache(self, n_rows: int, dev: torch.device):
        if self._ckv_buf is None or self._cap < n_rows:
            cap = n_rows + self.CACHE_SLACK
            self._ckv_buf = torch.zeros(cap, 512, dtype=torch.bfloat16, device=dev)
            self._kr_buf = torch.zeros(cap, 64, dtype=torch.bfloat16, device=dev)
            self._cap = cap

    @torch.no_grad()
    def step(self, hidden: torch.Tensor, state: list):
        dev = hidden.device
        if self._id is None:
            self._setup(dev)
        ext = _EXT
        if hidden.dtype != torch.bfloat16 or not hidden.is_contiguous():
            hidden = hidden.contiguous().to(torch.bfloat16)
        mla = state[3]
        ckv, kr = mla["c_kv"], mla["k_rope"]
        pos = ckv.shape[0]
        # KV cache: keep a capacity buffer; ingest a foreign cache inside the kernel (no copy launch)
        n_old = 0
        src_ckv, src_kr = ckv, kr
        if self._ckv_buf is None or ckv.data_ptr() != self._ckv_buf.data_ptr() or self._cap < pos + 1:
            old_ckv, old_kr = self._ckv_buf, self._kr_buf
            if self._ckv_buf is not None and ckv.data_ptr() == self._ckv_buf.data_ptr():
                src_ckv, src_kr = old_ckv, old_kr   # growing our own buffer: re-ingest from it
                self._ckv_buf = None
            self._ensure_cache(pos + 1, dev)
            n_old = pos
            if not src_ckv.is_contiguous():
                src_ckv = src_ckv.contiguous()
            if not src_kr.is_contiguous():
                src_kr = src_kr.contiguous()
        self._parity ^= 1
        par = self._parity
        out = self._out[par]
        dyn = [hidden.data_ptr(), out.data_ptr()]
        win_dst = []
        for l in range(3):
            st = state[l]
            dyn.append(st["S"].data_ptr())
        for l in range(3):
            st = state[l]
            wd = []
            for m, key in enumerate(("cq", "ck", "cv")):
                src = st[key]
                dyn.append(src.data_ptr())
                d = self._win[l][m][par]
                if d.data_ptr() == src.data_ptr():
                    d = self._win[l][m][par ^ 1]
                wd.append(d)
            win_dst.append(wd)
        for l in range(3):
            for m in range(3):
                dyn.append(win_dst[l][m].data_ptr())
        dyn += [src_ckv.data_ptr(), src_kr.data_ptr(), self._ckv_buf.data_ptr(), self._kr_buf.data_ptr(),
                self._prof.data_ptr() if self._prof is not None else 0]
        n_tok = pos + 1
        chunk = 32
        while (n_tok + chunk - 1) // chunk > 130:
            chunk *= 2
        n_chunks = (n_tok + chunk - 1) // chunk
        ext.mk_step(self._id, torch.tensor(dyn, dtype=torch.int64), pos, n_old, n_chunks, chunk)
        # state bookkeeping (views only; no kernels)
        for l in range(3):
            st = state[l]
            st["cq"], st["ck"], st["cv"] = win_dst[l][0], win_dst[l][1], win_dst[l][2]
        mla["c_kv"] = self._ckv_buf[: pos + 1]
        mla["k_rope"] = self._kr_buf[: pos + 1]
        return out, state


# --------------------------------------------------------------------------- #
# state / inputs (same as reference)
# --------------------------------------------------------------------------- #
def init_state(cfg: Config, context_len: int, seed: int) -> list:
    dev = torch.device("cuda:0")
    g = torch.Generator(device=dev).manual_seed(seed)
    H, Dk = cfg.kda_heads, cfg.kda_head_dim
    C = H * Dk
    state = []
    for kind in cfg.pattern:
        if kind == "K":
            state.append({
                "S": torch.randn(H, Dk, Dk, device=dev, generator=g) * 0.05,
                "cq": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "ck": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "cv": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
            })
        else:
            state.append({
                "c_kv": torch.randn(context_len, cfg.kv_lora, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "k_rope": torch.randn(context_len, cfg.qk_rope, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
            })
    return state


def init_token(cfg: Config, seed: int) -> torch.Tensor:
    dev = torch.device("cuda:0")
    g = torch.Generator(device=dev).manual_seed(seed + 1)
    return torch.randn(cfg.hidden, device=dev, generator=g, dtype=cfg.dtype) * 0.25