Kimi-Linear Decode Kimi K3 (256k)

"""Kimi-Linear W4A16 hybrid decode unit -- single-megakernel solution.

The entire per-token decode step (3x KDA + 1x MLA attention, 4x MoE FFN,
all int4 dequant GEMVs, convs, state updates, RMSNorms and residuals) runs
as ONE CUDA kernel launch inside step(). Grid-wide phase barriers are done
with a custom atomic barrier inside a cooperatively-launched persistent
kernel (188 blocks x 1024 threads on the RTX PRO 6000).

W4A16 dequant is fused into every GEMV: int4 nibbles are unpacked and
scaled per-128-group on the fly; no bf16 weight matrix is ever written out.

KIMI_EAGER=1 forces the slow-but-faithful eager path (debugging aid).
"""
from __future__ import annotations

import math
import os
from dataclasses import dataclass, field

import torch
import torch.nn as nn
import torch.nn.functional as F

OP_TYPE = "kimi_linear_w4a16_decode"
HARDWARE_REQUIRED = ["RTX_PRO_6000"]
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)))


# --------------------------------------------------------------------------- #
# W4A16 quantization (identical format to reference.py)
# --------------------------------------------------------------------------- #
def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    lo = w_q[0::2] & 0xF
    hi = w_q[1::2] & 0xF
    return (lo | (hi << 4)).contiguous()


def _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor:
    out = torch.empty((K, w_packed.shape[1]), dtype=torch.uint8, device=w_packed.device)
    out[0::2] = w_packed & 0xF
    out[1::2] = (w_packed >> 4) & 0xF
    return out


def dequant(w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor, K: int, group: int) -> torch.Tensor:
    wu = _unpack_int4(w_q, K).to(torch.bfloat16)
    s = scales.repeat_interleave(group, dim=0)
    z = zeros.repeat_interleave(group, dim=0)
    return (wu - z) * s


class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        assert in_f % group == 0 and in_f % 2 == 0
        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))

    def weight_bf(self) -> torch.Tensor:
        return dequant(self.w_q, self.scales, self.zeros, self.in_f, self.group)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return (x.float() @ self.weight_bf().float()).to(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))

    def weight_bf(self, e: int) -> torch.Tensor:
        return dequant(self.w_q[e], self.scales[e], self.zeros[e], self.in_f, self.group)


# --------------------------------------------------------------------------- #
# helpers (eager oracle path)
# --------------------------------------------------------------------------- #
def _rmsnorm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + EPS)
    return (xf * w.float()).to(x.dtype)


def _rope_cossin(pos: int, dim: int, theta: float, device):
    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    ang = pos * inv
    return torch.cos(ang), torch.sin(ang)


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    xf = x.float()
    even, odd = xf[..., 0::2], xf[..., 1::2]
    out = torch.empty_like(xf)
    out[..., 0::2] = even * cos - odd * sin
    out[..., 1::2] = odd * cos + even * sin
    return out.to(x.dtype)


# --------------------------------------------------------------------------- #
# layers -- eager oracle path (used for debugging / fallback)
# --------------------------------------------------------------------------- #
class KDA(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        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)
        self.scale = Dk ** -0.5

    def _short_conv(self, val, prev, idx):
        win = torch.cat([prev, val[None]], dim=0)
        w = self.conv_w[idx].float().transpose(0, 1)
        out = (win.float() * w).sum(0)
        return F.silu(out).to(val.dtype), win[1:]

    def step(self, x, st):
        H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
        q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        q, st["cq"] = self._short_conv(q, st["cq"], 0)
        k, st["ck"] = self._short_conv(k, st["ck"], 1)
        v, st["cv"] = self._short_conv(v, st["cv"], 2)
        q = q.view(H, Dk).float() * self.scale
        k = k.view(H, Dk).float()
        v = v.view(H, Dk).float()
        g = (-F.softplus(self.g_proj(x).float())).view(H, Dk)
        beta = torch.sigmoid(self.beta_proj(x).float())
        S = st["S"] * g.exp()[:, :, None]
        pred = (S * k[:, :, None]).sum(1)
        S = S + beta[:, None, None] * k[:, :, None] * (v - pred)[:, None, :]
        o = (S * q[:, :, None]).sum(1)
        st["S"] = S
        return self.o_proj(o.reshape(H * Dk).to(torch.bfloat16))


class MLA(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        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)
        self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5

    def step(self, x, st):
        cfg = self.cfg
        H = cfg.mla_heads
        pos = st["c_kv"].shape[0]
        q = self.q_proj(x).view(H, cfg.qk_nope + cfg.qk_rope)
        q_nope = q[:, : cfg.qk_nope].float()
        q_rope = q[:, cfg.qk_nope :]
        kv = self.kv_a(x)
        c_kv = kv[: cfg.kv_lora]
        k_rope = kv[cfg.kv_lora :]
        cos, sin = _rope_cossin(pos, cfg.qk_rope, cfg.rope_theta, x.device)
        q_rope = _apply_rope(q_rope, cos, sin).float()
        k_rope = _apply_rope(k_rope, cos, sin)
        st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0)
        st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)
        kvb = self.kv_b(st["c_kv"]).view(-1, H, cfg.qk_nope + cfg.v_head).float()
        k_nope = kvb[..., : cfg.qk_nope]
        v = kvb[..., cfg.qk_nope :]
        scores = (torch.einsum("hd,lhd->lh", q_nope, k_nope)
                  + torch.einsum("hd,ld->lh", q_rope, st["k_rope"].float())) * self.scale
        p = torch.softmax(scores, dim=0)
        o = torch.einsum("lh,lhd->hd", p, v)
        return self.o_proj(o.reshape(H * cfg.v_head).to(torch.bfloat16))


class MoE(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        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)

    def _ffn(self, x, experts_g, experts_u, experts_d, e):
        h = F.silu(x.float() @ experts_g.weight_bf(e).float()) * (x.float() @ experts_u.weight_bf(e).float())
        return h @ experts_d.weight_bf(e).float()

    def step(self, x):
        cfg = self.cfg
        probs = torch.softmax(self.router(x).float(), dim=-1)
        w, idx = torch.topk(probs, cfg.n_active)
        w = w / (w.sum() + 1e-9) * cfg.routed_scaling
        out = x.new_zeros(cfg.hidden, dtype=torch.float32)
        for j in range(cfg.n_active):
            out = out + w[j] * self._ffn(x, self.gate, self.up, self.down, int(idx[j]))
        for s in range(cfg.n_shared):
            out = out + self._ffn(x, self.s_gate, self.s_up, self.s_down, s)
        return out.to(torch.bfloat16)


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)

    def step(self, x, st):
        h = x + self.attn.step(_rmsnorm(x, self.attn_norm), st)
        return h + self.moe.step(_rmsnorm(h, self.moe_norm))


# --------------------------------------------------------------------------- #
# Megakernel integration
# --------------------------------------------------------------------------- #
from kimi_mega import MegaRunner  # local module in this directory


class Model(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._runner = None
        self._eager_only = os.environ.get("KIMI_EAGER", "0") == "1"

    def step(self, hidden, state):
        if self._eager_only:
            return self._step_eager(hidden, state)
        if self._runner is None:
            dev = hidden.device
            self.to(dev)  # make sure all buffers are on the right device
            self._runner = MegaRunner(self, dev)
        return self._runner.step(hidden, state)

    def _step_eager(self, hidden, state):
        with torch.no_grad():
            for i, blk in enumerate(self.blocks):
                hidden = blk.step(hidden, state[i])
        return hidden, state


# --------------------------------------------------------------------------- #
# state / inputs (mirror reference helpers)
# --------------------------------------------------------------------------- #
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


# ==================================================================
# ===== sidecar: kimi_mega.py (49178 bytes, loaded by solution.py) =====
# ==================================================================

"""Megakernel runner for the Kimi-Linear W4A16 decode unit.

Builds one CUDA kernel (via torch.utils.cpp_extension.load_inline) that
executes the ENTIRE per-token decode step as a single cooperative launch:
3x [KDA attn + MoE] + 1x [MLA attn + MoE], all int4 dequant GEMVs fused.
Phase ordering inside the kernel is enforced by a custom atomic grid
barrier (the kernel is launched cooperatively, so all blocks are resident).
"""
from __future__ import annotations

import os

import torch

CPP_DECL = r"""
#include <vector>
#include <cstdint>
int64_t kimi_mega_setup(std::vector<int64_t> sp_host);
int64_t kimi_mega_step(int64_t sp, std::vector<int64_t> dyn, int64_t L,
                       int64_t gen_base, int64_t NCH, int64_t stop_after);
"""

CUDA_SRC = r"""
#include <cuda_bf16.h>
#include <cuda_pipeline_primitives.h>
#include <cstdint>
#include <cstdio>
#include <vector>

using bf16 = __nv_bfloat16;

// ------------------------------------------------------------------ consts
#define DD      2304
#define HC      4096
#define NHEADS  32
#define HDIM    128
#define GRP     128
#define MOE_M   1024
#define MLB_Q   6144
#define MLB_KVA 576
#define MLB_KVB 8192
#define MLB_QD  192
#define KVL     512
#define QK_NOPE 128
#define QK_ROPE 64
#define V_HEAD  128
#define KDA_SCALE 0.08838834764831845f
#define MLA_SCALE 0.07216878364870323f
#define RSCALE  2.446f
#define NEG_INF (-__int_as_float(0x7f800000))

// static param layout (must match python builder EXACTLY)
#define SPK(b)   ((b)*19)
#define K_NORMA 0
#define K_NORMM 1
#define K_QW 2
#define K_QS 3
#define K_QZ 4
#define K_KW 5
#define K_KS 6
#define K_KZ 7
#define K_VW 8
#define K_VS 9
#define K_VZ 10
#define K_GW 11
#define K_GS 12
#define K_GZ 13
#define K_OW 14
#define K_OS 15
#define K_OZ 16
#define K_CONV 17
#define K_BETA 18

#define SPE(b)   (57 + (b)*19)
#define E_ROUT 0
#define E_GW 1
#define E_GS 2
#define E_GZ 3
#define E_UW 4
#define E_US 5
#define E_UZ 6
#define E_DW 7
#define E_DS 8
#define E_DZ 9
#define E_SGW 10
#define E_SGS 11
#define E_SGZ 12
#define E_SUW 13
#define E_SUS 14
#define E_SUZ 15
#define E_SDW 16
#define E_SDS 17
#define E_SDZ 18

#define SPM      133
#define M_NORMA 0
#define M_NORMM 1
#define M_QW 2
#define M_QS 3
#define M_QZ 4
#define M_KVAW 5
#define M_KVAS 6
#define M_KVAZ 7
#define M_KVBW 8
#define M_KVBS 9
#define M_KVBZ 10
#define M_OW 11
#define M_OS 12
#define M_OZ 13

#define SPW      147
#define W_Q 0
#define W_K 1
#define W_V 2
#define W_G 3
#define W_O 4
#define W_XN 5
#define W_BX 6
#define W_A 7
#define W_D 8
#define W_HH 9
#define W_TKI 10
#define W_TKW 11
#define W_QM 12
#define W_QF 13
#define W_KRR 14
#define W_AP 15
#define W_CK 16
#define W_KR 17
#define W_CTR 18
#define W_INV 19
#define SP_TOTAL 167

struct DynArgs {
  long long x_in;
  long long h_out;
  long long S[3];
  long long cq[3];
  long long ck[3];
  long long cv[3];
  int L;
  int NCH;
  int stop_after;
  unsigned gen_base;
};

// ------------------------------------------------------------------ helpers
__device__ __forceinline__ float bf2f(bf16 v) { return __bfloat162float(v); }
__device__ __forceinline__ bf16 f2bf(float v) { return __float2bfloat16_rn(v); }

union B8 { uint4 v; bf16 f[8]; };
union B4 { uint2 v; bf16 f[4]; };

__device__ __forceinline__ float siluf(float x) { return x / (1.f + expf(-x)); }
__device__ __forceinline__ float softplusf(float x) { return x > 20.f ? x : log1pf(expf(x)); }

// tensor-core mma for the MLA attention phase (bf16, fp32 accumulate)
__device__ __forceinline__ void mma_bf16(float* c, const unsigned* a, const unsigned* b) {
  asm volatile(
    "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
    "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
    : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
    : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
}
#define M3_RT 48
#define M3_STP 584
#define M3_QP 584

__device__ __forceinline__ float block_sum(float v, float* scratch) {
  int w = threadIdx.x >> 5, l = threadIdx.x & 31;
  #pragma unroll
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(~0u, v, o);
  if (l == 0) scratch[w] = v;
  __syncthreads();
  if (w == 0) {
    float a = (l < (int)(blockDim.x >> 5)) ? scratch[l] : 0.f;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) a += __shfl_xor_sync(~0u, a, o);
    if (l == 0) scratch[0] = a;
  }
  __syncthreads();
  float r = scratch[0];
  __syncthreads();
  return r;
}

__device__ __forceinline__ void compute_smx(const float* sx, float* smx, int ng) {
  int w = threadIdx.x >> 5, l = threadIdx.x & 31;
  __syncthreads();
  if (w < ng) {
    const float* g = sx + w * 128;
    float a = g[l] + g[l + 32] + g[l + 64] + g[l + 96];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) a += __shfl_xor_sync(~0u, a, o);
    if (l == 0) smx[w] = a;
  }
  __syncthreads();
}

__device__ __forceinline__ void load_sx_bf16(const bf16* __restrict__ x, float* sx, int K) {
  int nv = K >> 3;
  for (int i = threadIdx.x; i < nv; i += blockDim.x) {
    B8 u; u.v = __ldcg((const uint4*)(x + 8 * i));
    #pragma unroll
    for (int j = 0; j < 8; ++j) sx[8 * i + j] = bf2f(u.f[j]);
  }
}

__device__ __forceinline__ void load_sx_f32(const float* __restrict__ x, float* sx, int K) {
  int nv = K >> 2;
  float4* sx4 = (float4*)sx;
  for (int i = threadIdx.x; i < nv; i += blockDim.x) {
    sx4[i] = __ldcg((const float4*)(x + 4 * i));
  }
}

__device__ __forceinline__ void gbar(unsigned* cnt, unsigned* gen, unsigned target) {
  __syncthreads();
  if (threadIdx.x == 0) {
    __threadfence();
    unsigned a = atomicAdd(cnt, 1u);
    if (a == gridDim.x - 1) {
      atomicExch(cnt, 0u);
      __threadfence();
      atomicAdd(gen, 1u);
    }
    while (atomicAdd(gen, 0u) < target) __nanosleep(32);
  }
  __syncthreads();
}

__device__ __forceinline__ int grab(unsigned* ctr_slot, int ntasks, int* ismall) {
  __syncthreads();
  if (threadIdx.x == 0) ismall[0] = (int)atomicAdd(ctr_slot, 1u);
  __syncthreads();
  return ismall[0];
}

// ------------------------------------------------------------------ gemv64
__device__ __forceinline__ void gemv64(
    const unsigned char* __restrict__ wq, const bf16* __restrict__ sc, const bf16* __restrict__ zr,
    int N, int g0, int ng, int col0,
    const float* __restrict__ sx, const float* __restrict__ smx,
    float* team_part, float* out64) {
  const int team = threadIdx.x >> 7;
  const int lane = threadIdx.x & 31;
  const int wrp = (threadIdx.x >> 5) & 3;
  const int base = ng >> 3, rem = ng & 7;
  const int gstart = team * base + min(team, rem);
  const int gcount = base + (team < rem ? 1 : 0);
  const int cb = (lane & 7) * 8;
  const int lsub = lane >> 3;
  const int ldb = N;

  float accf[8];
  #pragma unroll
  for (int j = 0; j < 8; ++j) accf[j] = 0.f;

  const unsigned char* wbase = wq + (size_t)(g0 + gstart) * 64 * (size_t)ldb + col0 + cb;
  for (int gl = 0; gl < gcount; ++gl) {
    float dotq[8];
    #pragma unroll
    for (int j = 0; j < 8; ++j) dotq[j] = 0.f;
    const unsigned char* wr = wbase + (size_t)(gl * 64 + wrp * 4) * (size_t)ldb;
    uint2 wv[4];
    #pragma unroll
    for (int it = 0; it < 4; ++it) wv[it] = __ldg((const uint2*)(wr + (size_t)(it * 16 + lsub) * ldb));
    #pragma unroll
    for (int it = 0; it < 4; ++it) {
      int rloc = (gstart + gl) * 64 + wrp * 4 + it * 16 + lsub;
      float x0 = sx[2 * rloc], x1 = sx[2 * rloc + 1];
      const unsigned char* pb = (const unsigned char*)&wv[it];
      #pragma unroll
      for (int j = 0; j < 8; ++j) {
        int qv = pb[j];
        dotq[j] += x0 * (float)(qv & 15) + x1 * (float)(qv >> 4);
      }
    }
    int ga = g0 + gstart + gl;
    uint4 sv = __ldg((const uint4*)(sc + (size_t)ga * N + col0 + cb));
    uint4 zv = __ldg((const uint4*)(zr + (size_t)ga * N + col0 + cb));
    const bf16* sp = (const bf16*)&sv;
    const bf16* zp = (const bf16*)&zv;
    float sxg = smx[gstart + gl];
    // zero-point correction must be applied exactly once per group/column:
    // only one lane of the team's k-split applies it.
    bool zleader = (wrp == 0 && lsub == 0);
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
      float sf = bf2f(sp[j]);
      accf[j] += sf * dotq[j] - (zleader ? sf * bf2f(zp[j]) * sxg : 0.f);
    }
  }
  // reduce the 4 lsub lanes' k-partials (lanes l, l+8, l+16, l+24 share cols)
  #pragma unroll
  for (int j = 0; j < 8; ++j) {
    accf[j] += __shfl_xor_sync(~0u, accf[j], 8);
    accf[j] += __shfl_xor_sync(~0u, accf[j], 16);
  }
  if (lsub == 0) {
    float* warp_part = team_part + (threadIdx.x >> 5) * 64 + cb;
    #pragma unroll
    for (int j = 0; j < 8; ++j) warp_part[j] = accf[j];
  }
  __syncthreads();
  if (threadIdx.x < 64) {
    float a = 0.f;
    #pragma unroll 8
    for (int t = 0; t < 32; ++t) a += team_part[t * 64 + threadIdx.x];
    out64[threadIdx.x] = a;
  }
  __syncthreads();
}

__device__ __forceinline__ void rmsn_apply(float* xf, const bf16* __restrict__ w,
                                           float* sx, float* scratch, bf16* ws_xn) {
  float part = 0.f;
  for (int i = threadIdx.x; i < DD; i += blockDim.x) part += xf[i] * xf[i];
  float tot = block_sum(part, scratch);
  float rs = 1.f / sqrtf(tot * (1.f / (float)DD) + 1e-6f);
  for (int i = threadIdx.x; i < DD; i += blockDim.x) {
    bf16 nb = f2bf(xf[i] * rs * bf2f(__ldg(w + i)));
    sx[i] = bf2f(nb);
    if (ws_xn) ws_xn[i] = nb;
  }
  __syncthreads();
}

// recombine x_b (block b>0 input) from previous-block partials, into h1x
__device__ __forceinline__ void recomb_xb(const bf16* xp, const float* A0, const float* Dm,
                                          float* h1x, int tid, int bdim) {
  for (int i = tid; i < DD; i += bdim) {
    float xb_ = bf2f(__ldcg(xp + i));
    float att = bf2f(f2bf(__ldcg(A0 + i) + __ldcg(A0 + DD + i)));
    float h1 = bf2f(f2bf(xb_ + att));
    float macc = __ldcg(Dm + 8 * DD + i);
    #pragma unroll
    for (int s = 0; s < 8; ++s) macc += __ldcg(Dm + s * DD + i);
    h1x[i] = bf2f(f2bf(h1 + bf2f(f2bf(macc))));
  }
  __syncthreads();
}

// router + top8 on smem xn vector sx; results in ismall[80..88), w in small[96..104)
__device__ __forceinline__ void router_topk(const bf16* rw, const float* sx, float* small, int* ismall, int tid) {
  int w = tid >> 5, l = tid & 31;
  if (w < 32) {
    float a0 = 0.f, a1 = 0.f;
    const bf16* r0 = rw + w * DD;
    const bf16* r1 = rw + (w + 32) * DD;
    for (int i = 0; i < DD / 32; i += 4) {
      B4 u0, u1;
      u0.v = __ldg((const uint2*)(r0 + l * 4 + i * 32));
      u1.v = __ldg((const uint2*)(r1 + l * 4 + i * 32));
      #pragma unroll
      for (int j = 0; j < 4; ++j) {
        a0 += bf2f(u0.f[j]) * sx[l * 4 + i * 32 + j];
        a1 += bf2f(u1.f[j]) * sx[l * 4 + i * 32 + j];
      }
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) { a0 += __shfl_xor_sync(~0u, a0, o); a1 += __shfl_xor_sync(~0u, a1, o); }
    if (l == 0) { small[16 + w] = bf2f(f2bf(a0)); small[16 + w + 32] = bf2f(f2bf(a1)); }
  }
  __syncthreads();
  if (w == 0) {
    float v0 = small[16 + l], v1 = small[16 + l + 32];
    int m0 = l, m1 = l + 32;
    float gmx = fmaxf(v0, v1);
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) gmx = fmaxf(gmx, __shfl_xor_sync(~0u, gmx, o));
    float den = expf(v0 - gmx) + expf(v1 - gmx);
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) den += __shfl_xor_sync(~0u, den, o);
    for (int i = 0; i < 8; ++i) {
      float bv = fmaxf(v0, v1);
      int bi = (v0 >= v1) ? m0 : m1;
      #pragma unroll
      for (int o = 16; o > 0; o >>= 1) {
        float ov = __shfl_xor_sync(~0u, bv, o);
        int oi = __shfl_xor_sync(~0u, bi, o);
        if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; }
      }
      float p = expf(small[16 + bi] - gmx) / den;
      if (l == 0) { ismall[80 + i] = bi; small[96 + i] = p; }
      if (m0 == bi) { v0 = NEG_INF; }
      if (m1 == bi) { v1 = NEG_INF; }
    }
    __syncwarp();
    float wsum = small[96 + (l & 7)];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) wsum += __shfl_xor_sync(~0u, wsum, o);
    wsum = wsum * 0.25f;
    if (l < 8) small[96 + l] = small[96 + l] / (wsum + 1e-9f) * RSCALE;
  }
  __syncthreads();
}

// gate+up task loop body (shared by K4/M6)
__device__ __forceinline__ void moe_gate_up(const long long* SE, unsigned* ctrs, int* ismall,
                                            const int* topk8, const float* sx, const float* smx,
                                            float* tp0, float* tp1, float* out64, float* small,
                                            float* hh) {
  for (;;) {
    int t = grab(ctrs, 144, ismall);
    if (t >= 144) break;
    int slot_e = t >> 4, tile = t & 15;
    const unsigned char* wqg; const bf16* scg; const bf16* zrg;
    const unsigned char* wqu; const bf16* scu; const bf16* zru;
    if (slot_e < 8) {
      long long e = topk8[slot_e];
      wqg = (const unsigned char*)(const void*)SE[E_GW] + e * (1152 * 1024);
      scg = (const bf16*)(const void*)SE[E_GS] + e * (18 * 1024);
      zrg = (const bf16*)(const void*)SE[E_GZ] + e * (18 * 1024);
      wqu = (const unsigned char*)(const void*)SE[E_UW] + e * (1152 * 1024);
      scu = (const bf16*)(const void*)SE[E_US] + e * (18 * 1024);
      zru = (const bf16*)(const void*)SE[E_UZ] + e * (18 * 1024);
    } else {
      wqg = (const unsigned char*)(const void*)SE[E_SGW];
      scg = (const bf16*)(const void*)SE[E_SGS];
      zrg = (const bf16*)(const void*)SE[E_SGZ];
      wqu = (const unsigned char*)(const void*)SE[E_SUW];
      scu = (const bf16*)(const void*)SE[E_SUS];
      zru = (const bf16*)(const void*)SE[E_SUZ];
    }
    gemv64(wqg, scg, zrg, 1024, 0, 18, tile * 64, sx, smx, tp0, out64);
    if (threadIdx.x < 64) tp1[threadIdx.x] = out64[threadIdx.x];
    __syncthreads();
    gemv64(wqu, scu, zru, 1024, 0, 18, tile * 64, sx, smx, tp0, out64);
    if (threadIdx.x < 64) hh[slot_e * MOE_M + tile * 64 + threadIdx.x] = siluf(tp1[threadIdx.x]) * out64[threadIdx.x];
  }
}

// down task loop body (shared by K5/M7)
__device__ __forceinline__ void moe_down(const long long* SE, unsigned* ctrs, int* ismall,
                                         const int* tki, const float* tkw, const float* hh,
                                         float* Dm, float* sx, float* smx, float* tp0,
                                         float* out64, float* red) {
  for (;;) {
    int t = grab(ctrs, 324, ismall);
    if (t >= 324) break;
    int slot_e = t / 36, tile = t % 36;
    float wgt = (slot_e < 8) ? __ldcg(tkw + slot_e) : 1.f;
    const unsigned char* wqd; const bf16* scd; const bf16* zrd;
    if (slot_e < 8) {
      long long e = __ldcg(tki + slot_e);
      wqd = (const unsigned char*)(const void*)SE[E_DW] + e * (512 * 2304);
      scd = (const bf16*)(const void*)SE[E_DS] + e * (8 * 2304);
      zrd = (const bf16*)(const void*)SE[E_DZ] + e * (8 * 2304);
    } else {
      wqd = (const unsigned char*)(const void*)SE[E_SDW];
      scd = (const bf16*)(const void*)SE[E_SDS];
      zrd = (const bf16*)(const void*)SE[E_SDZ];
    }
    load_sx_f32(hh + slot_e * MOE_M, sx, MOE_M);
    compute_smx(sx, smx, 8);
    gemv64(wqd, scd, zrd, 2304, 0, 8, tile * 64, sx, smx, tp0, out64);
    if (threadIdx.x < 64) Dm[slot_e * DD + tile * 64 + threadIdx.x] = wgt * out64[threadIdx.x];
  }
}

// ------------------------------------------------------------------ kernel
extern "C" __global__ void __launch_bounds__(1024, 1)
kimi_mega(const long long* __restrict__ SP, const DynArgs dyn) {
  extern __shared__ char smem_raw[];
  float* h1x   = (float*)smem_raw;          // 2304
  float* sx    = h1x + 2304;                // 2304
  float* smx   = sx + 2304;                 // 32
  float* tp0   = smx + 32;                  // 2048 (32 warps x 64 partials)
  float* tp1   = tp0 + 2048;                // 2048
  float* out64 = tp1 + 2048;                // 128
  float* red   = out64 + 128;               // 64
  float* small = red + 64;                  // 256
  int*   ismall = (int*)small;

  unsigned* ctr = (unsigned*)(const void*)SP[SPW + W_CTR];
  unsigned* barc = ctr + 30;
  unsigned* gen = ctr + 31;

  const int POS = dyn.L;
  const int Lp1 = dyn.L + 1;
  const int NCH = dyn.NCH;
  const int stopa = dyn.stop_after;

  unsigned target = dyn.gen_base;
  int slot = 0;
  const int tid = threadIdx.x;

  // ===================================================================
  for (int b = 0; b < 3; ++b) {
    const long long* SK = SP + SPK(b);
    const long long* SE = SP + SPE(b);

    // ---------------- K1 --------------------------------------------
    if (slot <= stopa) {
      const bf16* xin = (const bf16*)(const void*)dyn.x_in;
      bf16* xbuf = (bf16*)(const void*)SP[SPW + W_BX];
      if (b == 0) {
        for (int i = tid; i < DD; i += blockDim.x) h1x[i] = bf2f(__ldcg(xin + i));
        __syncthreads();
      } else {
        const bf16* xp = (b == 1) ? xin : (xbuf + ((b - 1) & 1) * DD);
        recomb_xb(xp, (const float*)(const void*)SP[SPW + W_A], (const float*)(const void*)SP[SPW + W_D], h1x, tid, blockDim.x);
        if (blockIdx.x == 0) {
          bf16* ob = xbuf + (b & 1) * DD;
          for (int i = tid; i < DD; i += blockDim.x) ob[i] = f2bf(h1x[i]);
        }
      }
      bf16* ws_xn = (blockIdx.x == 0) ? (bf16*)(const void*)SP[SPW + W_XN] : nullptr;
      rmsn_apply(h1x, (const bf16*)(const void*)SK[K_NORMA], sx, red, ws_xn);
      compute_smx(sx, smx, 18);
      unsigned* ctrs = ctr + slot;
      for (;;) {
        int t = grab(ctrs, 256, ismall);
        if (t >= 256) break;
        int c0 = t * 64;
        int m = c0 >> 12, n0 = c0 & 4095;
        const unsigned char* wq; const bf16* sc; const bf16* zr;
        if (m == 0)      { wq = (const unsigned char*)(const void*)SK[K_QW]; sc = (const bf16*)(const void*)SK[K_QS]; zr = (const bf16*)(const void*)SK[K_QZ]; }
        else if (m == 1) { wq = (const unsigned char*)(const void*)SK[K_KW]; sc = (const bf16*)(const void*)SK[K_KS]; zr = (const bf16*)(const void*)SK[K_KZ]; }
        else if (m == 2) { wq = (const unsigned char*)(const void*)SK[K_VW]; sc = (const bf16*)(const void*)SK[K_VS]; zr = (const bf16*)(const void*)SK[K_VZ]; }
        else             { wq = (const unsigned char*)(const void*)SK[K_GW]; sc = (const bf16*)(const void*)SK[K_GS]; zr = (const bf16*)(const void*)SK[K_GZ]; }
        gemv64(wq, sc, zr, 4096, 0, 18, n0, sx, smx, tp0, out64);
        if (tid < 64) {
          int cc = n0 + tid;
          float raw = out64[tid];
          if (m < 3) {
            long long stptr = (m == 0) ? dyn.cq[b] : (m == 1) ? dyn.ck[b] : dyn.cv[b];
            bf16* win = (bf16*)(const void*)stptr;
            bf16 p0 = __ldcg(win + cc);
            bf16 p1 = __ldcg(win + HC + cc);
            bf16 p2 = __ldcg(win + 2 * HC + cc);
            bf16 rawb = f2bf(raw);
            const bf16* cw = (const bf16*)(const void*)SK[K_CONV] + m * (HC * 4) + cc * 4;
            float acc = bf2f(p0) * bf2f(__ldg(cw + 0)) + bf2f(p1) * bf2f(__ldg(cw + 1))
                      + bf2f(p2) * bf2f(__ldg(cw + 2)) + bf2f(rawb) * bf2f(__ldg(cw + 3));
            long long wsptr = (m == 0) ? SP[SPW + W_Q] : (m == 1) ? SP[SPW + W_K] : SP[SPW + W_V];
            ((bf16*)(const void*)wsptr)[cc] = f2bf(siluf(acc));
            win[cc] = p1;
            win[HC + cc] = p2;
            win[2 * HC + cc] = rawb;
          } else {
            ((float*)(const void*)SP[SPW + W_G])[cc] = -softplusf(bf2f(f2bf(raw)));
          }
        }
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- K2 --------------------------------------------
    if (slot <= stopa) {
      const bf16* ws_xn = (const bf16*)(const void*)SP[SPW + W_XN];
      float* fk = h1x; float* fq = h1x + 128; float* fg = h1x + 256;
      float* fv = sx; float* prd = sx + 64; float* oscr = tp0;   // 1024 floats: tp0+tp1
      unsigned* ctrs = ctr + slot;
      for (;;) {
        int t = grab(ctrs, 64, ismall);
        if (t >= 64) break;
        int h = t >> 1, dv0 = (t & 1) * 64;
        float part = 0.f;
        const bf16* bw = (const bf16*)(const void*)SK[K_BETA] + h * DD;
        for (int i = tid; i < DD; i += blockDim.x) part += bf2f(__ldcg(ws_xn + i)) * bf2f(__ldg(bw + i));
        float tot = block_sum(part, red);
        if (tid == 0) small[104] = 1.f / (1.f + expf(-bf2f(f2bf(tot))));
        if (tid < 128) {
          fk[tid] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_K] + h * HDIM + tid));
          fg[tid] = __ldcg((const float*)(const void*)SP[SPW + W_G] + h * HDIM + tid);
          fq[tid] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_Q] + h * HDIM + tid)) * KDA_SCALE;
        } else if (tid < 192) {
          fv[tid - 128] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_V] + h * HDIM + dv0 + tid - 128));
        }
        __syncthreads();
        int dv = tid & 63, dkg = tid >> 6;
        float* Sg = (float*)(const void*)dyn.S[b];
        size_t soff = ((size_t)h * HDIM + dkg * 8) * HDIM + dv0 + dv;
        float sreg[8];
        float pp = 0.f;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
          sreg[j] = __ldcg(Sg + soff + (size_t)j * HDIM);
          pp += sreg[j] * fk[dkg * 8 + j];
        }
        oscr[dkg * 64 + dv] = pp;
        __syncthreads();
        if (tid < 64) {
          float a = 0.f;
          #pragma unroll
          for (int tt = 0; tt < 16; ++tt) a += oscr[tt * 64 + tid];
          prd[tid] = a;
        }
        __syncthreads();
        float bb = small[104];
        float fvd = fv[dv], pd = prd[dv];
        float op = 0.f;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
          float sn = sreg[j] * expf(fg[dkg * 8 + j]) + bb * fk[dkg * 8 + j] * (fvd - pd);
          sreg[j] = sn;
          op += sn * fq[dkg * 8 + j];
        }
        #pragma unroll
        for (int j = 0; j < 8; ++j) *(Sg + soff + (size_t)j * HDIM) = sreg[j];
        __syncthreads();
        oscr[dkg * 64 + dv] = op;
        __syncthreads();
        if (tid < 64) {
          float a = 0.f;
          #pragma unroll
          for (int tt = 0; tt < 16; ++tt) a += oscr[tt * 64 + tid];
          ((bf16*)(const void*)SP[SPW + W_O])[h * HDIM + dv0 + tid] = f2bf(a);
        }
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- K3 --------------------------------------------
    if (slot <= stopa) {
      unsigned* ctrs = ctr + slot;
      float* A = (float*)(const void*)SP[SPW + W_A];
      load_sx_bf16((const bf16*)(const void*)SP[SPW + W_O], sx, 4096);
      compute_smx(sx, smx, 32);
      for (;;) {
        int t = grab(ctrs, 72, ismall);
        if (t >= 72) break;
        int kh = t / 36, tile = t % 36;
        gemv64((const unsigned char*)(const void*)SK[K_OW], (const bf16*)(const void*)SK[K_OS], (const bf16*)(const void*)SK[K_OZ],
               2304, kh * 16, 16, tile * 64, sx + kh * 2048, smx + kh * 16, tp0, out64);
        if (tid < 64) A[kh * DD + tile * 64 + tid] = out64[tid];
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- K4 --------------------------------------------
    int topk_i_sh[8];
    if (slot <= stopa) {
      const bf16* xin = (const bf16*)(const void*)dyn.x_in;
      const bf16* xbuf = (const bf16*)(const void*)SP[SPW + W_BX];
      const bf16* xb = (b == 0) ? xin : (xbuf + (b & 1) * DD);
      const float* A = (const float*)(const void*)SP[SPW + W_A];
      for (int i = tid; i < DD; i += blockDim.x) {
        float att = bf2f(f2bf(__ldcg(A + i) + __ldcg(A + DD + i)));
        h1x[i] = bf2f(f2bf(bf2f(__ldcg(xb + i)) + att));
      }
      __syncthreads();
      rmsn_apply(h1x, (const bf16*)(const void*)SK[K_NORMM], sx, red, nullptr);
      compute_smx(sx, smx, 18);
      router_topk((const bf16*)(const void*)SE[E_ROUT], sx, small, ismall, tid);
      if (blockIdx.x == 0 && tid < 8) {
        ((int*)(const void*)SP[SPW + W_TKI])[tid] = ismall[80 + tid];
        ((float*)(const void*)SP[SPW + W_TKW])[tid] = small[96 + tid];
      }
      #pragma unroll
      for (int i = 0; i < 8; ++i) topk_i_sh[i] = ismall[80 + i];
      moe_gate_up(SE, ctr + slot, ismall, topk_i_sh, sx, smx, tp0, tp1, out64, small,
                  (float*)(const void*)SP[SPW + W_HH]);
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- K5 --------------------------------------------
    if (slot <= stopa) {
      moe_down(SE, ctr + slot, ismall,
               (const int*)(const void*)SP[SPW + W_TKI], (const float*)(const void*)SP[SPW + W_TKW],
               (const float*)(const void*)SP[SPW + W_HH], (float*)(const void*)SP[SPW + W_D],
               sx, smx, tp0, out64, red);
    }
    gbar(barc, gen, ++target); slot++;
  }

  // ===================================================================
  {
    const int b = 3;
    const long long* SM = SP + SPM;
    const long long* SE = SP + SPE(3);

    // ---------------- M1 --------------------------------------------
    if (slot <= stopa) {
      bf16* xbuf = (bf16*)(const void*)SP[SPW + W_BX];
      recomb_xb(xbuf + ((b - 1) & 1) * DD, (const float*)(const void*)SP[SPW + W_A],
                (const float*)(const void*)SP[SPW + W_D], h1x, tid, blockDim.x);
      if (blockIdx.x == 0) {
        bf16* ob = xbuf + (b & 1) * DD;
        for (int i = tid; i < DD; i += blockDim.x) ob[i] = f2bf(h1x[i]);
      }
      __syncthreads();
      rmsn_apply(h1x, (const bf16*)(const void*)SM[M_NORMA], sx, red, nullptr);
      compute_smx(sx, smx, 18);
      unsigned* ctrs = ctr + slot;
      for (;;) {
        int t = grab(ctrs, 105, ismall);
        if (t >= 105) break;
        if (t < 96) {
          gemv64((const unsigned char*)(const void*)SM[M_QW], (const bf16*)(const void*)SM[M_QS], (const bf16*)(const void*)SM[M_QZ],
                 MLB_Q, 0, 18, t * 64, sx, smx, tp0, out64);
          if (tid < 64) ((bf16*)(const void*)SP[SPW + W_QM])[t * 64 + tid] = f2bf(out64[tid]);
        } else {
          int c0 = (t - 96) * 64;
          gemv64((const unsigned char*)(const void*)SM[M_KVAW], (const bf16*)(const void*)SM[M_KVAS], (const bf16*)(const void*)SM[M_KVAZ],
                 MLB_KVA, 0, 18, c0, sx, smx, tp0, out64);
          if (tid < 64) {
            int c = c0 + tid;
            if (c < KVL) ((bf16*)(const void*)SP[SPW + W_CK])[(size_t)POS * KVL + c] = f2bf(out64[tid]);
            else ((bf16*)(const void*)SP[SPW + W_KRR])[c - KVL] = f2bf(out64[tid]);
          }
        }
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M2 --------------------------------------------
    if (slot <= stopa) {
      const bf16* ws_qm = (const bf16*)(const void*)SP[SPW + W_QM];
      bf16* ws_qf = (bf16*)(const void*)SP[SPW + W_QF];
      const float* invf = (const float*)(const void*)SP[SPW + W_INV];
      const unsigned char* wq = (const unsigned char*)(const void*)SM[M_KVBW];
      const bf16* sc = (const bf16*)(const void*)SM[M_KVBS];
      const bf16* zr = (const bf16*)(const void*)SM[M_KVBZ];
      float* qn_eff = h1x;      // 128
      float* zred = tp1;        // 128
      unsigned* ctrs = ctr + slot;
      int w = tid >> 5, l = tid & 31;
      for (;;) {
        int t = grab(ctrs, 128, ismall);
        if (t >= 128) break;
        int h = t >> 2, st = t & 3;
        if (tid < 128) {
          float qn = bf2f(__ldcg(ws_qm + h * MLB_QD + tid));
          float sf = bf2f(__ldg(sc + (size_t)st * MLB_KVB + h * (QK_NOPE + V_HEAD) + tid));
          float zf = bf2f(__ldg(zr + (size_t)st * MLB_KVB + h * (QK_NOPE + V_HEAD) + tid));
          qn_eff[tid] = qn * sf;
          zred[tid] = qn * sf * zf;
        }
        __syncthreads();
        if (tid == 0) {
          float a = 0.f;
          for (int i = 0; i < 128; ++i) a += zred[i];
          small[108] = a;
        }
        __syncthreads();
        float zterm = small[108];
        for (int it = 0; it < 2; ++it) {
          int r = st * 64 + w * 2 + it;
          unsigned u = __ldg((const unsigned*)(wq + (size_t)r * MLB_KVB + h * (QK_NOPE + V_HEAD) + l * 4));
          float a_lo = 0.f, a_hi = 0.f;
          const unsigned char* pb = (const unsigned char*)&u;
          #pragma unroll
          for (int j = 0; j < 4; ++j) {
            a_lo += qn_eff[l * 4 + j] * (float)(pb[j] & 15);
            a_hi += qn_eff[l * 4 + j] * (float)(pb[j] >> 4);
          }
          #pragma unroll
          for (int o = 16; o > 0; o >>= 1) { a_lo += __shfl_xor_sync(~0u, a_lo, o); a_hi += __shfl_xor_sync(~0u, a_hi, o); }
          if (l == 0) {
            ws_qf[h * M3_QP + 2 * r] = f2bf(a_lo - zterm);
            ws_qf[h * M3_QP + 2 * r + 1] = f2bf(a_hi - zterm);
          }
        }
        __syncthreads();
        if (st == 0 && tid < 64) {
          int i = tid & 31;
          float ang = (float)POS * __ldcg(invf + i);
          float cs = cosf(ang), sn = sinf(ang);
          if (tid < 32) {
            float x0 = bf2f(__ldcg(ws_qm + h * MLB_QD + QK_NOPE + 2 * i));
            float x1 = bf2f(__ldcg(ws_qm + h * MLB_QD + QK_NOPE + 2 * i + 1));
            ws_qf[h * M3_QP + 512 + 2 * i] = f2bf(x0 * cs - x1 * sn);
            ws_qf[h * M3_QP + 512 + 2 * i + 1] = f2bf(x1 * cs + x0 * sn);
          } else if (h == 0) {
            const bf16* kr = (const bf16*)(const void*)SP[SPW + W_KRR];
            float x0 = bf2f(__ldcg(kr + 2 * i));
            float x1 = bf2f(__ldcg(kr + 2 * i + 1));
            bf16* kbf = (bf16*)(const void*)SP[SPW + W_KR] + (size_t)POS * QK_ROPE;
            kbf[2 * i] = f2bf(x0 * cs - x1 * sn);
            kbf[2 * i + 1] = f2bf(x1 * cs + x0 * sn);
          }
        }
        __syncthreads();
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M3 (tensor-core flash-decode over latent cache) ----
    if (slot <= stopa) {
      const bf16* ckv = (const bf16*)(const void*)SP[SPW + W_CK];
      const bf16* krb = (const bf16*)(const void*)SP[SPW + W_KR];
      float* ap = (float*)(const void*)SP[SPW + W_AP];
      const bf16* ws_qf = (const bf16*)(const void*)SP[SPW + W_QF];
      unsigned* ctrs = ctr + slot;
      const int rch = (Lp1 + NCH - 1) / NCH;
      const int w = tid >> 5, l = tid & 31;
      bf16* stage_s = (bf16*)smem_raw;                       // [M3_RT][M3_STP]
      float* sbuf_s = (float*)(smem_raw + M3_RT * M3_STP * 2); // 2 x [32][M3_RT] fp32
      __shared__ int stask;
      __shared__ float sh_rs[32];
      for (;;) {
        __syncthreads();
        if (tid == 0) stask = (int)atomicAdd(ctrs, 1u);
        __syncthreads();
        int t = (int)stask;
        if (t >= NCH) break;
        int r0 = t * rch, r1 = min(Lp1, r0 + rch);
        int nrows = r1 - r0;
        float m = NEG_INF, lsum = 0.f;
        int mm = w >> 4, ng = w & 15;
        float accf[4][4];
        #pragma unroll
        for (int i = 0; i < 4; ++i) for (int j = 0; j < 4; ++j) accf[i][j] = 0.f;
        __syncthreads();
        int nb = (nrows + M3_RT - 1) / M3_RT;
        for (int bb = 0; bb < nb; ++bb) {
          int rows_b = min(M3_RT, nrows - bb * M3_RT);
          int rbase = r0 + bb * M3_RT;
          for (int rr = w; rr < M3_RT; rr += 32) {
            if (rr < rows_b) {
              const bf16* crow = ckv + (size_t)(rbase + rr) * 512;
              const bf16* krow = krb + (size_t)(rbase + rr) * 64;
              bf16* dst = stage_s + rr * M3_STP;
              *(uint4*)(dst + l * 16)      = __ldcg((const uint4*)crow + l * 2);
              *(uint4*)(dst + l * 16 + 8) = __ldcg((const uint4*)crow + l * 2 + 1);
              *(unsigned short*)(dst + 512 + 2 * l) = __ldcg((const unsigned short*)(krow + 2 * l));
              *(unsigned short*)(dst + 512 + 2 * l + 1) = __ldcg((const unsigned short*)(krow + 2 * l + 1));
            } else {
              bf16* dst = stage_s + rr * M3_STP;
              for (int i = l * 8; i < 576; i += 32 * 8) *(uint4*)(dst + i) = make_uint4(0, 0, 0, 0);
            }
          }
          __syncthreads();
          // mma1: S[32,RT] = q' @ K^T (24 warps: 2m x 6n x 2 khalf)
          if (w < 24) {
            int kg = w / 12, m1 = (w % 12) & 1, nt = (w % 12) >> 1;
            float c[4] = {0, 0, 0, 0};
            const bf16* qa = ws_qf + (m1 * 16) * M3_QP;
            int rr2 = l >> 2, cc = (l & 3) * 2;
            int kend = kg ? 36 : 18;
            for (int ks = kg * 18; ks < kend; ++ks) {
              unsigned a[4], b[2];
              const bf16* ap_ = qa + ks * 16;
              a[0] = __ldg((const unsigned*)&ap_[rr2 * M3_QP + cc]);
              a[1] = __ldg((const unsigned*)&ap_[(rr2 + 8) * M3_QP + cc]);
              a[2] = __ldg((const unsigned*)&ap_[rr2 * M3_QP + cc + 8]);
              a[3] = __ldg((const unsigned*)&ap_[(rr2 + 8) * M3_QP + cc + 8]);
              const bf16* bk = stage_s + (nt * 8) * M3_STP + ks * 16;
              b[0] = *(unsigned*)&bk[rr2 * M3_STP + cc];
              b[1] = *(unsigned*)&bk[rr2 * M3_STP + cc + 8];
              mma_bf16(c, a, b);
            }
            float* sb = sbuf_s + kg * (32 * M3_RT);
            sb[(m1 * 16 + rr2) * M3_RT + nt * 8 + cc] = c[0];
            sb[(m1 * 16 + rr2) * M3_RT + nt * 8 + cc + 1] = c[1];
            sb[(m1 * 16 + rr2 + 8) * M3_RT + nt * 8 + cc] = c[2];
            sb[(m1 * 16 + rr2 + 8) * M3_RT + nt * 8 + cc + 1] = c[3];
          }
          __syncthreads();
          // softmax (warp = head), reads both sbuf layers
          {
            int h = w;
            float vv[2];
            vv[0] = sbuf_s[h * M3_RT + l] + sbuf_s[32 * M3_RT + h * M3_RT + l];
            vv[1] = (l + 32 < M3_RT) ? (sbuf_s[h * M3_RT + l + 32] + sbuf_s[32 * M3_RT + h * M3_RT + l + 32]) : NEG_INF;
            vv[0] *= MLA_SCALE; vv[1] *= MLA_SCALE;
            if (l >= rows_b) vv[0] = NEG_INF;
            if (l + 32 >= rows_b) vv[1] = NEG_INF;
            float tmax = fmaxf(vv[0], vv[1]);
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) tmax = fmaxf(tmax, __shfl_xor_sync(~0u, tmax, o));
            float mnew = fmaxf(m, tmax);
            float rs = (m == NEG_INF) ? 0.f : __expf(m - mnew);
            float p0 = __expf(vv[0] - mnew);
            float p1 = __expf(vv[1] - mnew);
            float wsum = p0 + p1;
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) wsum += __shfl_xor_sync(~0u, wsum, o);
            lsum = lsum * rs + wsum;
            m = mnew;
            __syncwarp();
            bf16* prow = (bf16*)(sbuf_s + 32 * M3_RT);
            prow[h * 96 + l] = __float2bfloat16_rn(p0);
            if (l + 32 < M3_RT) prow[h * 96 + l + 32] = __float2bfloat16_rn(p1);
            if (l == 0) sh_rs[h] = rs;
          }
          __syncthreads();
          // acc rescale + mma2: acc[32,512] += P @ V
          {
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
              #pragma unroll
              for (int j = 0; j < 4; ++j) accf[i][j] *= sh_rs[mm * 16 + ((j >= 2) ? ((l >> 2) + 8) : (l >> 2))];
            }
            const bf16* pa16 = (const bf16*)(sbuf_s + 32 * M3_RT);
            int rr2 = l >> 2, cc = (l & 3) * 2;
            for (int nt2 = 0; nt2 < 4; ++nt2) {
              for (int ks = 0; ks < 3; ++ks) {
                unsigned a[4], b[2];
                const bf16* ap_ = pa16 + (mm * 16) * 96 + ks * 16;
                a[0] = *(unsigned*)&ap_[rr2 * 96 + cc];
                a[1] = *(unsigned*)&ap_[(rr2 + 8) * 96 + cc];
                a[2] = *(unsigned*)&ap_[rr2 * 96 + cc + 8];
                a[3] = *(unsigned*)&ap_[(rr2 + 8) * 96 + cc + 8];
                unsigned lo = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc) * M3_STP + ng * 32 + nt2 * 8 + rr2];
                unsigned hi = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 1) * M3_STP + ng * 32 + nt2 * 8 + rr2];
                b[0] = lo | (hi << 16);
                lo = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 8) * M3_STP + ng * 32 + nt2 * 8 + rr2];
                hi = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 9) * M3_STP + ng * 32 + nt2 * 8 + rr2];
                b[1] = lo | (hi << 16);
                mma_bf16(accf[nt2], a, b);
              }
            }
          }
          __syncthreads();
        }
        // partial write
        if (w < 32 && l == 0) {
          float* op = ap + ((size_t)t * 32 + w) * 514;
          op[512] = m;
          op[513] = lsum;
        }
        {
          int rr2 = l >> 2, cc2 = (l & 3) * 2;
          float* base = ap + ((size_t)t * 32) * 514;
          #pragma unroll
          for (int nt2 = 0; nt2 < 4; ++nt2) {
            int hrow = mm * 16 + rr2;
            int col = ng * 32 + nt2 * 8 + cc2;
            base[hrow * 514 + col] = accf[nt2][0];
            base[hrow * 514 + col + 1] = accf[nt2][1];
            base[(hrow + 8) * 514 + col] = accf[nt2][2];
            base[(hrow + 8) * 514 + col + 1] = accf[nt2][3];
          }
        }
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M4 --------------------------------------------
    if (slot <= stopa) {
      const float* ap = (const float*)(const void*)SP[SPW + W_AP];
      unsigned* ctrs = ctr + slot;
      float* mc = red;          // 48
      float* wc = tp1;          // 48
      float* lc = tp1 + 48;     // 48
      for (;;) {
        int t = grab(ctrs, 64, ismall);
        if (t >= 64) break;
        int h = t >> 1, dv0 = (t & 1) * 64;
        size_t pbase = (size_t)h * 514;
        size_t chstride = (size_t)32 * 514;
        if (tid < 32) {
          for (int c = tid; c < NCH; c += 32) {
            mc[c] = __ldcg(ap + (size_t)c * chstride + pbase + 512);
            lc[c] = __ldcg(ap + (size_t)c * chstride + pbase + 513);
          }
        }
        __syncthreads();
        if (tid == 0) {
          float M = NEG_INF;
          for (int c = 0; c < NCH; ++c) M = fmaxf(M, mc[c]);
          float lacc = 0.f;
          for (int c = 0; c < NCH; ++c) {
            float wgt = (M == NEG_INF) ? 0.f : expf(mc[c] - M);
            wc[c] = wgt;
            lacc += wgt * lc[c];
          }
          small[108] = lacc;
        }
        __syncthreads();
        float ltot = small[108];
        if (tid < 512) {
          float a = 0.f;
          for (int c = 0; c < NCH; ++c)
            a += wc[c] * __ldcg(ap + (size_t)c * chstride + pbase + tid);
          sx[tid] = a / ltot;
        }
        __syncthreads();
        compute_smx(sx, smx, 4);
        gemv64((const unsigned char*)(const void*)SM[M_KVBW], (const bf16*)(const void*)SM[M_KVBS], (const bf16*)(const void*)SM[M_KVBZ],
               MLB_KVB, 0, 4, h * (QK_NOPE + V_HEAD) + QK_NOPE + dv0, sx, smx, tp0, out64);
        if (tid < 64) ((bf16*)(const void*)SP[SPW + W_O])[h * V_HEAD + dv0 + tid] = f2bf(out64[tid]);
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M5 --------------------------------------------
    if (slot <= stopa) {
      unsigned* ctrs = ctr + slot;
      float* A = (float*)(const void*)SP[SPW + W_A];
      load_sx_bf16((const bf16*)(const void*)SP[SPW + W_O], sx, 4096);
      compute_smx(sx, smx, 32);
      for (;;) {
        int t = grab(ctrs, 72, ismall);
        if (t >= 72) break;
        int kh = t / 36, tile = t % 36;
        gemv64((const unsigned char*)(const void*)SM[M_OW], (const bf16*)(const void*)SM[M_OS], (const bf16*)(const void*)SM[M_OZ],
               2304, kh * 16, 16, tile * 64, sx + kh * 2048, smx + kh * 16, tp0, out64);
        if (tid < 64) A[kh * DD + tile * 64 + tid] = out64[tid];
      }
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M6 --------------------------------------------
    int topk_i_sh[8];
    if (slot <= stopa) {
      const bf16* xbuf = (const bf16*)(const void*)SP[SPW + W_BX];
      const bf16* xb = xbuf + (b & 1) * DD;
      const float* A = (const float*)(const void*)SP[SPW + W_A];
      for (int i = tid; i < DD; i += blockDim.x) {
        float att = bf2f(f2bf(__ldcg(A + i) + __ldcg(A + DD + i)));
        h1x[i] = bf2f(f2bf(bf2f(__ldcg(xb + i)) + att));
      }
      __syncthreads();
      rmsn_apply(h1x, (const bf16*)(const void*)SM[M_NORMM], sx, red, nullptr);
      compute_smx(sx, smx, 18);
      router_topk((const bf16*)(const void*)SE[E_ROUT], sx, small, ismall, tid);
      if (blockIdx.x == 0 && tid < 8) {
        ((int*)(const void*)SP[SPW + W_TKI])[tid] = ismall[80 + tid];
        ((float*)(const void*)SP[SPW + W_TKW])[tid] = small[96 + tid];
      }
      #pragma unroll
      for (int i = 0; i < 8; ++i) topk_i_sh[i] = ismall[80 + i];
      moe_gate_up(SE, ctr + slot, ismall, topk_i_sh, sx, smx, tp0, tp1, out64, small,
                  (float*)(const void*)SP[SPW + W_HH]);
    }
    gbar(barc, gen, ++target); slot++;

    // ---------------- M7 --------------------------------------------
    if (slot <= stopa) {
      moe_down(SE, ctr + slot, ismall,
               (const int*)(const void*)SP[SPW + W_TKI], (const float*)(const void*)SP[SPW + W_TKW],
               (const float*)(const void*)SP[SPW + W_HH], (float*)(const void*)SP[SPW + W_D],
               sx, smx, tp0, out64, red);
    }
    gbar(barc, gen, ++target); slot++;
  }

  // ===================================================================
  // epilogue (no barrier after; resets task counters for the next step)
  // ===================================================================
  if (slot <= stopa) {
    const bf16* xp = (const bf16*)(const void*)SP[SPW + W_BX] + DD;   // x_3
    const float* A0 = (const float*)(const void*)SP[SPW + W_A];
    const float* Dm = (const float*)(const void*)SP[SPW + W_D];
    bf16* out = (bf16*)(const void*)dyn.h_out;
    int cpb = (DD + gridDim.x - 1) / gridDim.x;
    int c0 = blockIdx.x * cpb, c1 = min(DD, c0 + cpb);
    for (int i = c0 + tid; i < c1; i += blockDim.x) {
      float xb_ = bf2f(__ldcg(xp + i));
      float att = bf2f(f2bf(__ldcg(A0 + i) + __ldcg(A0 + DD + i)));
      float h1 = bf2f(f2bf(xb_ + att));
      float macc = __ldcg(Dm + 8 * DD + i);
      #pragma unroll
      for (int s = 0; s < 8; ++s) macc += __ldcg(Dm + s * DD + i);
      out[i] = f2bf(h1 + bf2f(f2bf(macc)));
    }
    if (blockIdx.x == 0 && tid < 30) ctr[tid] = 0;
  } else if (blockIdx.x == 0 && tid < 30) {
    ctr[tid] = 0;
  }
}

// ------------------------------------------------------------------ host
static bool g_smem_set = false;

int64_t kimi_mega_setup(std::vector<int64_t> sp_host) {
  if (sp_host.size() != SP_TOTAL) {
    printf("kimi_mega_setup: bad SP size %zu want %d\n", sp_host.size(), SP_TOTAL);
    return 0;
  }
  void* dev = nullptr;
  cudaError_t e = cudaMalloc(&dev, SP_TOTAL * sizeof(int64_t));
  if (e != cudaSuccess) { printf("cudaMalloc: %s\n", cudaGetErrorString(e)); return 0; }
  e = cudaMemcpy(dev, sp_host.data(), SP_TOTAL * sizeof(int64_t), cudaMemcpyHostToDevice);
  if (e != cudaSuccess) { printf("cudaMemcpy: %s\n", cudaGetErrorString(e)); return 0; }
  if (!g_smem_set) {
    e = cudaFuncSetAttribute((const void*)kimi_mega, cudaFuncAttributeMaxDynamicSharedMemorySize, 99328);
    if (e != cudaSuccess) { printf("cudaFuncSetAttribute: %s\n", cudaGetErrorString(e)); return 0; }
    g_smem_set = true;
  }
  int coop = 0;
  cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, 0);
  if (!coop) { printf("cooperative launch unsupported!\n"); return 0; }
  int occ = 0;
  e = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, (const void*)kimi_mega, 1024, 99328);
  if (e != cudaSuccess || occ < 1) { printf("occupancy: %d %s\n", occ, cudaGetErrorString(e)); return 0; }
  return (int64_t)dev;
}

int64_t kimi_mega_step(int64_t sp, std::vector<int64_t> dyn, int64_t L,
                       int64_t gen_base, int64_t NCH, int64_t stop_after) {
  if (dyn.size() != 14) { printf("bad dyn size\n"); return 1; }
  DynArgs a;
  a.x_in = dyn[0];
  a.h_out = dyn[1];
  for (int i = 0; i < 3; ++i) { a.S[i] = dyn[2 + i]; a.cq[i] = dyn[5 + i]; a.ck[i] = dyn[8 + i]; a.cv[i] = dyn[11 + i]; }
  a.L = (int)L;
  a.NCH = (int)NCH;
  a.stop_after = (int)stop_after;
  a.gen_base = (unsigned)gen_base;
  int nsm = 188;
  cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, 0);
  const long long* SPp = (const long long*)sp;
  void* params[2] = { (void*)&SPp, (void*)&a };
  cudaError_t e = cudaLaunchCooperativeKernel((const void*)kimi_mega, dim3(nsm), dim3(1024), params, 99328, 0);
  if (e != cudaSuccess) { printf("launch: %s\n", cudaGetErrorString(e)); return 2; }
  return 0;
}
"""

_EXT = None
SP_TOTAL = 167


def _get_ext():
    global _EXT
    if _EXT is None:
        from torch.utils.cpp_extension import load_inline
        import shutil

        here = os.path.dirname(os.path.abspath(__file__))
        os.environ.setdefault("TORCH_EXTENSIONS_DIR", os.path.join(here, ".torch_extensions"))
        flags = [
            "-O3",
            "-std=c++17",
            "--generate-code=arch=compute_120,code=sm_120",
        ]
        gpp = shutil.which("g++")
        if gpp:
            flags.append(f"--compiler-bindir={gpp}")
        _EXT = load_inline(
            name="kimi_mega_v1",
            cpp_sources=[CPP_DECL],
            cuda_sources=[CUDA_SRC],
            functions=["kimi_mega_setup", "kimi_mega_step"],
            extra_cuda_cflags=flags,
            verbose=False,
        )
    return _EXT


class MegaRunner:
    RINGL = 32768

    def __init__(self, model: nn.Module, dev: torch.device):
        cfg = model.cfg
        assert cfg.pattern[3] == "M" and all(k == "K" for k in cfg.pattern[:3]), \
            "megakernel fast path assumes pattern K,K,K,M"
        self.model = model
        self.dev = dev
        d = cfg.hidden
        bf = torch.bfloat16
        t = lambda *shape, dtype: torch.zeros(*shape, dtype=dtype, device=dev)
        C = cfg.kda_heads * cfg.kda_head_dim
        ws = {}
        ws["ws_q"] = t(C, dtype=bf)
        ws["ws_k"] = t(C, dtype=bf)
        ws["ws_v"] = t(C, dtype=bf)
        ws["ws_g"] = t(C, dtype=torch.float32)
        ws["ws_o"] = t(C, dtype=bf)
        ws["ws_xn"] = t(d, dtype=bf)
        ws["buf_x"] = t(2, d, dtype=bf)
        ws["A"] = t(2, d, dtype=torch.float32)
        ws["D"] = t(9, d, dtype=torch.float32)
        ws["hh"] = t(9, cfg.moe_inter, dtype=torch.float32)
        ws["topk_i"] = t(8, dtype=torch.int32)
        ws["topk_w"] = t(8, dtype=torch.float32)
        ws["ws_qm"] = t(cfg.mla_heads * (cfg.qk_nope + cfg.qk_rope), dtype=bf)
        ws["ws_qf"] = t(cfg.mla_heads, 584, dtype=bf)
        ws["kr_raw"] = t(cfg.qk_rope, dtype=bf)
        ws["attn_part"] = t(128, 32, 514, dtype=torch.float32)
        ws["c_buf"] = t(self.RINGL, cfg.kv_lora, dtype=bf)
        ws["k_buf"] = t(self.RINGL, cfg.qk_rope, dtype=bf)
        ws["ctr"] = t(32, dtype=torch.int32)
        invf = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.qk_rope, 2, dtype=torch.float32) / cfg.qk_rope))
        ws["invf"] = invf.to(dev)
        ws["h_out"] = t(d, dtype=bf)
        self.ws = ws
        self.h_out = ws["h_out"]
        self.c_buf = ws["c_buf"]
        self.k_buf = ws["k_buf"]

        sp = []
        def add(x): sp.append(x.data_ptr())
        for b in range(3):
            blk = model.blocks[b]; at = blk.attn
            add(blk.attn_norm); add(blk.moe_norm)
            for pr in (at.q_proj, at.k_proj, at.v_proj, at.g_proj, at.o_proj):
                add(pr.w_q); add(pr.scales); add(pr.zeros)
            add(at.conv_w); add(at.beta_proj.weight)
        for b in range(4):
            moe = model.blocks[b].moe
            add(moe.router.weight)
            for qe in (moe.gate, moe.up, moe.down, moe.s_gate, moe.s_up, moe.s_down):
                add(qe.w_q); add(qe.scales); add(qe.zeros)
        mla = model.blocks[3]; at = mla.attn
        add(mla.attn_norm); add(mla.moe_norm)
        for pr in (at.q_proj, at.kv_a, at.kv_b, at.o_proj):
            add(pr.w_q); add(pr.scales); add(pr.zeros)
        add(ws["ws_q"]); add(ws["ws_k"]); add(ws["ws_v"]); add(ws["ws_g"])
        add(ws["ws_o"]); add(ws["ws_xn"]); add(ws["buf_x"]); add(ws["A"])
        add(ws["D"]); add(ws["hh"]); add(ws["topk_i"]); add(ws["topk_w"])
        add(ws["ws_qm"]); add(ws["ws_qf"]); add(ws["kr_raw"])
        add(ws["attn_part"]); add(ws["c_buf"]); add(ws["k_buf"])
        add(ws["ctr"]); add(ws["invf"])
        assert len(sp) == SP_TOTAL, len(sp)

        self.ext = _get_ext()
        self.sp = self.ext.kimi_mega_setup(sp)
        assert self.sp != 0, "megakernel setup failed"
        self.gen_base = 0
        self.mla_idx = 3
        self.stop_after = int(os.environ.get("KIMI_STOP_AFTER", "999"))

    def step(self, hidden, state):
        st = state[self.mla_idx]
        c_kv = st["c_kv"]
        L = c_kv.shape[0]
        l1 = L + 1
        if c_kv.data_ptr() != self.c_buf.data_ptr():
            assert l1 <= self.RINGL
            self.c_buf[:L].copy_(c_kv, non_blocking=True)
            self.k_buf[:L].copy_(st["k_rope"], non_blocking=True)
        st["c_kv"] = self.c_buf[:l1]
        st["k_rope"] = self.k_buf[:l1]
        dyn = [hidden.data_ptr(), self.h_out.data_ptr()]
        for kd in ("S", "cq", "ck", "cv"):
            dyn += [state[b][kd].data_ptr() for b in range(3)]
        NCH = min(128, max(1, (l1 + 111) // 112))
        rc = self.ext.kimi_mega_step(self.sp, dyn, L, self.gen_base, NCH, self.stop_after)
        self.gen_base += 22
        assert rc == 0, f"megakernel step failed rc={rc}"
        return self.h_out, state