"""Single-kernel megakernel decode for the Kimi-Linear W4A16 hybrid unit.

The whole per-token forward -- 3x(KDA + MoE) + (MLA + MoE), every int4
dequant-GEMV, short conv, KDA state update, MLA latent attention, MoE router +
expert FFNs, all RMSNorms and residuals -- runs as ONE cooperative CUDA kernel
launch per step(). Phases inside the kernel are separated by grid.sync()
(cooperative groups), so sequential dependencies are honored without extra
launches.

Design notes:
  * int4 GEMV: weights stay packed in global memory. Each (128-col tile,
    k-chunk) unit is one CTA pass: 16 warps each take 4 consecutive packed rows
    per 128-row quant group, lanes take 4 consecutive columns (uint32 loads,
    fully coalesced). Nibbles are expanded 2-at-a-time with the fp16 magic
    trick ((w>>s & 0x000F000F) | 0x64006400 -> half2(1024+n), then -1024) and
    accumulated in half2 against a /8-prescaled activation vector staged in
    shared memory; per-group fold applies scale, and warp 0 alone applies the
    -scale*zero*groupsum(x) correction. Cross-warp reduce through smem, then
    store or fp32 atomicAdd (split-K / expert accumulation directly into the
    residual-preloaded destination -- no zeroing or extra reduce phases).
  * MLA uses the absorbed-latent form: q_nope is folded through kv_b once
    (q_abs = q . W_knope per latent channel), attention runs directly on the
    bf16 c_kv cache, and the context vector is folded through the W_v half of
    kv_b afterwards -- the cache is never expanded to per-head K/V.
  * The MLA cache lives in a capacity buffer owned by the model; step() returns
    views [:L+1], so cat() disappears. The first step of a fresh state copies
    the fed cache into the buffer inside the same kernel (phase Z).
  * KDA conv windows double-buffer (ping-pong) between two internal banks so
    the shifted window can be written while other CTAs still read the old one.
  * MoE: the router runs as two phases -- 16 CTAs compute the 64 expert
    logits in parallel (4 dot products each), then one warp does softmax +
    top-8; gate/up/down expert GEMVs read the selected ids from scratch.
  * Idle CTAs in compute-light phases issue prefetch.global.L2::evict_last
    on weights used a phase or two later (shared experts, next attention
    layer, the growing c_kv cache), hiding DRAM latency behind grid syncs.
    Warming only pays in bandwidth-idle windows; in streaming phases it
    competes with demand reads and was measured to hurt.

A slow eager path is NOT kept around: this file's step() always uses the fused
kernel. Debugging happens via the `debug_max_phase` attribute (early-exit after
N grid barriers) used by scratch tests during development.
"""
from __future__ import annotations

import os

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

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

EPS = 1.0e-6

# ---------------------------------------------------------------------------
# CUDA source: one cooperative kernel, 33 grid.sync-separated phases.
# ---------------------------------------------------------------------------
_CUDA_SRC = r"""
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#include <vector>

namespace cg = cooperative_groups;

#define D_HID 2304
#define C_KDA 4096
#define NH 32
#define DK 128
#define M_INTER 1024
#define HQ_MLA 6144
#define KVL 512
#define QR 64
#define KVA 576
#define NKVB 8192
#define NBLK 512
#define KDA_SCALE 0.08838834764831845f
#define MLA_SCALE 0.07216878364870323f
#define SMEM_BYTES 81920

typedef __nv_bfloat16 bf16;
typedef __nv_bfloat162 bf162;

struct KWeights {
    const bf16* attn_norm[4];
    const bf16* moe_norm[4];
    const uint8_t* kda_wq[3][5];       // q,k,v,g,o
    const bf16* kda_sc[3][5];
    const bf16* kda_zr[3][5];
    const bf16* beta_w[3];             // [32,2304]
    const bf16* conv_w[3];             // [3,4096,4]
    const uint8_t* mla_wq[4];          // q, kv_a, kv_b, o
    const bf16* mla_sc[4];
    const bf16* mla_zr[4];
    const bf16* router[4];             // [64,2304]
    const uint8_t* moe_wq[4][6];       // gate,up,down,s_gate,s_up,s_down
    const bf16* moe_sc[4][6];
    const bf16* moe_zr[4][6];
};

struct KScratch {
    float* x;        // [5][D]
    float* hattn;    // [4][D]
    float* qkvg;     // [4*C]
    float* betaout;  // [32]
    float* obuf;     // [4096]
    __half* xn_moe;  // [D]
    int* ids;        // [8]
    float* wts;      // [8]
    float* logits;   // [64]
    float* hg;       // [9*1024]
    float* hu;       // [9*1024]
    float* q6144;    // [6144]
    float* kva;      // [576]
    __half* qabs;    // [512*32]  (c-major: qabs[c*32+h])
    __half* qr;      // [64*32]
    float* scores;   // [LCAP*32]
    float* mx;       // [NCHMAX*32]
    float* sm;       // [NCHMAX*32]
    float* ctx;      // [32*512]
    bf16* outb;      // [D]
    bf16* ckv;       // [LCAP*512]
    bf16* krope;     // [LCAP*64]
};

struct KStep {
    const bf16* hidden;
    float* S[3];
    const bf16* winr[9];   // [layer*3+kind] -> [3][4096]
    bf16* winw[9];
    const bf16* ckv_src;
    const bf16* kr_src;
    int copy_len;
    int L;
    int max_phase;
};

struct KArg {
    KWeights w;
    KScratch s;
    KStep a;
    float routed;
    float theta;
};

// ---------------------------------------------------------------- int4 GEMV
__device__ __forceinline__ __half2 deq2(uint32_t w, int sh) {
    uint32_t v = ((w >> sh) & 0x000F000Fu) | 0x64006400u;
    __half2 h = *reinterpret_cast<__half2*>(&v);
    return __hsub2(h, __float2half2_rn(1024.f));
}

// y[col0+t] (for t < ncols) = 8 * sum_k xs[k] * deq(w[k, col0+t]) over groups
// [g0, g0+ng). xs is the FULL-K activation staged in smem, prescaled by 1/8.
// gsum[g] = sum of xs over group g. smem_red needs 16*128 floats.
__device__ void gemv_i4(const uint8_t* __restrict__ wq,
                        const bf16* __restrict__ sc,
                        const bf16* __restrict__ zr,
                        int N, int g0, int ng, int col0, int ncols,
                        const __half* __restrict__ xs,
                        const float* __restrict__ gsum,
                        float* __restrict__ out, float wmul, bool atomic,
                        float* smem_red) {
    const int lane = threadIdx.x & 31;
    const int w = threadIdx.x >> 5;
    const int c4 = lane * 4;
    const bool active = (c4 < ncols);
    float y0 = 0.f, y1 = 0.f, y2 = 0.f, y3 = 0.f;
    if (active) {
        #pragma unroll 2
        for (int g = g0; g < g0 + ng; ++g) {
            const uint8_t* rp = wq + (size_t)(g * 64 + w * 4) * N + col0 + c4;
            const __half2* xp = reinterpret_cast<const __half2*>(xs + (g * 64 + w * 4) * 2);
            __half2 a02 = __float2half2_rn(0.f);
            __half2 a13 = __float2half2_rn(0.f);
            #pragma unroll
            for (int r = 0; r < 4; ++r) {
                uint32_t wrd = *reinterpret_cast<const uint32_t*>(rp);
                rp += N;
                __half2 xpair = xp[r];
                __half2 xlo = __half2half2(__low2half(xpair));
                __half2 xhi = __half2half2(__high2half(xpair));
                a02 = __hfma2(deq2(wrd, 0), xlo, a02);
                a13 = __hfma2(deq2(wrd, 8), xlo, a13);
                a02 = __hfma2(deq2(wrd, 4), xhi, a02);
                a13 = __hfma2(deq2(wrd, 12), xhi, a13);
            }
            float2 f02 = __half22float2(a02);
            float2 f13 = __half22float2(a13);
            const bf16* sp = sc + (size_t)g * N + col0 + c4;
            const bf16* zp = zr + (size_t)g * N + col0 + c4;
            bf162 sA = *reinterpret_cast<const bf162*>(sp);
            bf162 sB = *reinterpret_cast<const bf162*>(sp + 2);
            bf162 zA = *reinterpret_cast<const bf162*>(zp);
            bf162 zB = *reinterpret_cast<const bf162*>(zp + 2);
            float xg = (w == 0) ? gsum[g] : 0.f;
            y0 += __bfloat162float(sA.x) * (f02.x - __bfloat162float(zA.x) * xg);
            y1 += __bfloat162float(sA.y) * (f13.x - __bfloat162float(zA.y) * xg);
            y2 += __bfloat162float(sB.x) * (f02.y - __bfloat162float(zB.x) * xg);
            y3 += __bfloat162float(sB.y) * (f13.y - __bfloat162float(zB.y) * xg);
        }
    }
    __syncthreads();
    if (active) {
        smem_red[w * 128 + c4 + 0] = y0;
        smem_red[w * 128 + c4 + 1] = y1;
        smem_red[w * 128 + c4 + 2] = y2;
        smem_red[w * 128 + c4 + 3] = y3;
    }
    __syncthreads();
    const int t = threadIdx.x;
    if (t < ncols) {
        float acc = 0.f;
        #pragma unroll
        for (int i = 0; i < 16; ++i) acc += smem_red[i * 128 + t];
        acc = acc * 8.f * wmul;
        if (atomic) atomicAdd(out + t, acc);
        else out[t] = acc;
    }
}

// ------------------------------------------------------------- x preparation
__device__ float block_sum(float v, float* red) {
    for (int o = 16; o; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
    if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = v;
    __syncthreads();
    if (threadIdx.x == 0) {
        float tot = 0.f;
        for (int i = 0; i < 16; ++i) tot += red[i];
        red[16] = tot;
    }
    __syncthreads();
    return red[16];
}

__device__ void fill_gsum(const __half* xs, int n, float* gsum) {
    const int w = threadIdx.x >> 5, lane = threadIdx.x & 31;
    for (int g = w; g < n / 128; g += 16) {
        float s = 0.f;
        for (int j = lane; j < 128; j += 32) s += __half2float(xs[g * 128 + j]);
        for (int o = 16; o; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
        if (lane == 0) gsum[g] = s;
    }
    __syncthreads();
}

// rmsnorm(src)*normw -> xs (half, /8), gsum
__device__ void prep_norm(const float* src, const bf16* normw, int n,
                          __half* xs, float* gsum, float* red) {
    float ss = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float v = src[i];
        ss += v * v;
    }
    float tot = block_sum(ss, red);
    float scale = rsqrtf(tot / (float)n + 1.0e-6f);
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float v = src[i] * scale * __bfloat162float(normw[i]);
        v = __bfloat162float(__float2bfloat16(v));   // reference emits bf16
        xs[i] = __float2half(v * 0.125f);
    }
    __syncthreads();
    fill_gsum(xs, n, gsum);
}

// bf16-round(src [+ src2]) -> xs (half, /8), gsum
__device__ void prep_bf16(const float* src, const float* src2, int n,
                          __half* xs, float* gsum) {
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        float v = src[i] + (src2 ? src2[i] : 0.f);
        v = __bfloat162float(__float2bfloat16(v));
        xs[i] = __float2half(v * 0.125f);
    }
    __syncthreads();
    fill_gsum(xs, n, gsum);
}

// -------------------------------------------------------------------- phases
__device__ void phase_Z(const KArg& A) {
    const int T = blockIdx.x * blockDim.x + threadIdx.x;
    const int NT = gridDim.x * blockDim.x;
    for (int i = T; i < D_HID; i += NT) A.s.x[i] = __bfloat162float(A.a.hidden[i]);
    for (int i = T; i < NH * KVL; i += NT) A.s.ctx[i] = 0.f;
    if (A.a.copy_len > 0) {
        const uint4* s1 = reinterpret_cast<const uint4*>(A.a.ckv_src);
        uint4* d1 = reinterpret_cast<uint4*>(A.s.ckv);
        int n1 = A.a.copy_len * 64;   // 512 bf16 = 64 uint4 per row
        for (int i = T; i < n1; i += NT) d1[i] = s1[i];
        const uint4* s2 = reinterpret_cast<const uint4*>(A.a.kr_src);
        uint4* d2 = reinterpret_cast<uint4*>(A.s.krope);
        int n2 = A.a.copy_len * 8;
        for (int i = T; i < n2; i += NT) d2[i] = s2[i];
    }
}

__device__ void l2_warm(const uint8_t* const* bases, const size_t* sizes,
                        int nreg, int pu, int np) {
    size_t total = 0;
    for (int i = 0; i < nreg; ++i) total += sizes[i];
    size_t b0 = total * pu / np, b1 = total * (pu + 1) / np;
    for (size_t off = b0 + (size_t)threadIdx.x * 128; off < b1;
         off += (size_t)blockDim.x * 128) {
        size_t o = off;
        for (int i = 0; i < nreg; ++i) {
            if (o < sizes[i]) {
                asm volatile("prefetch.global.L2::evict_last [%0];" :: "l"(bases[i] + o));
                break;
            }
            o -= sizes[i];
        }
    }
}

__device__ void kda_phaseA(const KArg& A, int layer, char* SM) {
    __half* xs = (__half*)SM;                       // 2304 half
    float* gsum = (float*)(SM + 8192);              // 18
    float* red = (float*)(SM + 8320);               // 17
    float* ysred = (float*)(SM + 8448);             // 16*128
    prep_norm(A.s.x + layer * D_HID, A.w.attn_norm[layer], D_HID, xs, gsum, red);
    for (int u = blockIdx.x; u < 133; u += gridDim.x) {
        if (u < 128) {
            int mat = u >> 5, tile = u & 31;
            gemv_i4(A.w.kda_wq[layer][mat], A.w.kda_sc[layer][mat], A.w.kda_zr[layer][mat],
                    C_KDA, 0, 18, tile * 128, 128, xs, gsum,
                    A.s.qkvg + mat * C_KDA + tile * 128, 1.f, false, ysred);
        } else if (u == 128) {
            int hh = threadIdx.x >> 4;
            int ss = threadIdx.x & 15;
            float acc = 0.f;
            const bf16* bw = A.w.beta_w[layer] + (size_t)hh * D_HID;
            for (int d = ss; d < D_HID; d += 16)
                acc += __half2float(xs[d]) * 8.f * __bfloat162float(bw[d]);
            for (int o = 8; o; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o, 16);
            if (ss == 0) {
                acc = __bfloat162float(__float2bfloat16(acc));
                A.s.betaout[hh] = 1.f / (1.f + expf(-acc));
            }
        } else {
            int j0 = (u - 129) * 576;
            for (int i = threadIdx.x; i < 576; i += blockDim.x)
                A.s.hattn[layer * D_HID + j0 + i] = A.s.x[layer * D_HID + j0 + i];
        }
    }
}

__device__ void kda_phaseB(const KArg& A, int layer, char* SM) {
    float* qc = (float*)SM;             // 128
    float* kc = qc + 128;
    float* vc = kc + 128;
    float* dec = vc + 128;
    float* red = dec + 128;             // 16*32
    float* predv = red + 512;           // 32
    for (int u = blockIdx.x; u < 128; u += gridDim.x) {
        __syncthreads();
        const int h = u >> 2, j0 = (u & 3) * 32;
        const int t = threadIdx.x;
        if (t < 384) {
            const int kind = t >> 7, c = t & 127;
            const int ch = h * 128 + c;
            const bf16* cw = A.w.conv_w[layer] + (size_t)(kind * C_KDA + ch) * 4;
            const bf16* wr = A.a.winr[layer * 3 + kind];
            bf16 p1b = wr[1 * C_KDA + ch];
            bf16 p2b = wr[2 * C_KDA + ch];
            bf16 valb = __float2bfloat16(A.s.qkvg[kind * C_KDA + ch]);
            float s = __bfloat162float(wr[ch]) * __bfloat162float(cw[0])
                    + __bfloat162float(p1b) * __bfloat162float(cw[1])
                    + __bfloat162float(p2b) * __bfloat162float(cw[2])
                    + __bfloat162float(valb) * __bfloat162float(cw[3]);
            s = s / (1.f + expf(-s));
            s = __bfloat162float(__float2bfloat16(s));
            if (kind == 0) qc[c] = s * KDA_SCALE;
            else if (kind == 1) kc[c] = s;
            else vc[c] = s;
            if (j0 == 0) {
                bf16* ww = A.a.winw[layer * 3 + kind];
                ww[ch] = p1b;
                ww[C_KDA + ch] = p2b;
                ww[2 * C_KDA + ch] = valb;
            }
        } else {
            const int c = t - 384;
            float graw = __bfloat162float(__float2bfloat16(A.s.qkvg[3 * C_KDA + h * 128 + c]));
            dec[c] = 1.f / (1.f + expf(graw));   // exp(-softplus(g)) == sigmoid(-g)
        }
        __syncthreads();
        const int j = t & 31, ic = t >> 5;
        const int jc = j0 + j;
        const int i0 = ic * 8;
        float* Sg = A.a.S[layer] + (size_t)h * DK * DK;
        float sreg[8];
        float pp = 0.f;
        #pragma unroll
        for (int r = 0; r < 8; ++r) {
            float sv = Sg[(i0 + r) * DK + jc] * dec[i0 + r];
            sreg[r] = sv;
            pp += sv * kc[i0 + r];
        }
        red[ic * 32 + j] = pp;
        __syncthreads();
        if (t < 32) {
            float s = 0.f;
            for (int i = 0; i < 16; ++i) s += red[i * 32 + t];
            predv[t] = s;
        }
        __syncthreads();
        const float dv = A.s.betaout[h] * (vc[jc] - predv[j]);
        float oo = 0.f;
        #pragma unroll
        for (int r = 0; r < 8; ++r) {
            float sv = sreg[r] + kc[i0 + r] * dv;
            Sg[(i0 + r) * DK + jc] = sv;
            oo += sv * qc[i0 + r];
        }
        __syncthreads();
        red[ic * 32 + j] = oo;
        __syncthreads();
        if (t < 32) {
            float s = 0.f;
            for (int i = 0; i < 16; ++i) s += red[i * 32 + t];
            A.s.obuf[h * 128 + j0 + t] = s;
        }
    }
}

__device__ void oproj_phase(const KArg& A, const uint8_t* wq, const bf16* sc,
                            const bf16* zr, float* dst, const float* src2,
                            char* SM) {
    __half* xs = (__half*)SM;                       // 4096 half
    float* gsum = (float*)(SM + 8192);              // 32
    float* ysred = (float*)(SM + 8448);
    prep_bf16(A.s.obuf, src2, C_KDA, xs, gsum);
    for (int u = blockIdx.x; u < 144; u += gridDim.x) {
        int tile = u % 18, kchunk = u / 18;
        gemv_i4(wq, sc, zr, D_HID, kchunk * 4, 4, tile * 128, 128,
                xs, gsum, dst + tile * 128, 1.f, true, ysred);
    }
}

// Warm L2 with [b0,b1) of the concatenated regions (fire-and-forget).

// Router phase A: 16 CTAs compute 4 expert logits each (bf16-rounded, into
// global scratch) while four more preload x[l+1] with the attention residual.
__device__ void moe_logits(const KArg& A, int layer, char* SM) {
    const int u = blockIdx.x;
    if (u < 16) {
        float* xh = (float*)SM;             // 2304
        float* red = xh + 2304;             // 17
        float* er = red + 32;               // 16 partials (4 experts x 4 warps)
        const float* hsrc = A.s.hattn + layer * D_HID;
        float ss = 0.f;
        for (int i = threadIdx.x; i < D_HID; i += blockDim.x) {
            float v = __bfloat162float(__float2bfloat16(hsrc[i]));
            xh[i] = v;
            ss += v * v;
        }
        float tot = block_sum(ss, red);
        float scale = rsqrtf(tot / (float)D_HID + 1.0e-6f);
        for (int i = threadIdx.x; i < D_HID; i += blockDim.x) {
            float v = xh[i] * scale * __bfloat162float(A.w.moe_norm[layer][i]);
            v = __bfloat162float(__float2bfloat16(v));
            xh[i] = v;
            if (u == 0) A.s.xn_moe[i] = __float2half(v);
        }
        __syncthreads();
        const int el = threadIdx.x >> 7, t7 = threadIdx.x & 127;
        const bf16* rw = A.w.router[layer] + (size_t)(u * 4 + el) * D_HID;
        float acc = 0.f;
        for (int d = t7; d < D_HID; d += 128) acc += xh[d] * __bfloat162float(rw[d]);
        for (int o = 16; o; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o);
        if ((t7 & 31) == 0) er[el * 4 + (t7 >> 5)] = acc;
        __syncthreads();
        if (t7 == 0) {
            float sacc = er[el * 4] + er[el * 4 + 1] + er[el * 4 + 2] + er[el * 4 + 3];
            A.s.logits[u * 4 + el] = __bfloat162float(__float2bfloat16(sacc));
        }
    } else if (u < 20) {
        int j0 = (u - 16) * 576;
        for (int i = threadIdx.x; i < 576; i += blockDim.x)
            A.s.x[(layer + 1) * D_HID + j0 + i] = A.s.hattn[layer * D_HID + j0 + i];
    } else {
        // bandwidth-idle: warm L2 with the shared-expert down-proj plus a
        // slice of the next attention layer's weight stream.
        const uint8_t* bases[3];
        size_t sizes[3];
        int nreg = 0;
        bases[nreg] = A.w.moe_wq[layer][5];
        sizes[nreg++] = (size_t)512 * 2304;
        const size_t SZ = (size_t)1152 * C_KDA;
        if (layer < 2) {
            bases[nreg] = A.w.kda_wq[layer + 1][0];
            sizes[nreg++] = SZ;
            bases[nreg] = A.w.kda_wq[layer + 1][1];
            sizes[nreg++] = SZ;
        } else if (layer == 2) {
            bases[nreg] = A.w.mla_wq[0];
            sizes[nreg++] = (size_t)1152 * HQ_MLA;
        } else {
            bases[nreg] = A.w.kda_wq[0][0];
            sizes[nreg++] = SZ;
            bases[nreg] = A.w.kda_wq[0][1];
            sizes[nreg++] = SZ;
        }
        l2_warm(bases, sizes, nreg, u - 20, gridDim.x - 20);
    }
}

// Router phase B: warp 0 of CTA 0 does softmax + top-8 (min-index tie-break,
// matching torch.topk) over the 64 logits into global ids/wts; every other
// CTA warms L2 with the shared-expert gate/up weights the next phase reads.
__device__ void moe_pick(const KArg& A, int layer) {
    if (blockIdx.x != 0) {
        const size_t BLK = (size_t)1152 * 1024;
        const size_t SZ = (size_t)1152 * C_KDA;
        const uint8_t* bases[4] = { A.w.moe_wq[layer][3], A.w.moe_wq[layer][4],
                                    nullptr, nullptr };
        size_t sizes[4] = { BLK, BLK, 0, 0 };
        int nreg = 2;
        if (layer == 2) {
            bases[nreg] = A.w.mla_wq[1];
            sizes[nreg++] = (size_t)1152 * KVA;
        } else {
            int nl = (layer + 1) & 3;   // layer 3 -> next step's layer 0
            bases[nreg] = A.w.kda_wq[nl][2];
            sizes[nreg++] = SZ;
            bases[nreg] = A.w.kda_wq[nl][3];
            sizes[nreg++] = SZ;
        }
        l2_warm(bases, sizes, nreg, blockIdx.x - 1, (int)gridDim.x - 1);
        return;
    }
    if (threadIdx.x >= 32) return;
    const int lane = threadIdx.x;
    const float* lg = A.s.logits;
    float v0 = lg[lane], v1 = lg[lane + 32];
    float mx = fmaxf(v0, v1);
    for (int o = 16; o; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
    float sum = expf(v0 - mx) + expf(v1 - mx);
    for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
    float inv = 1.f / sum;
    float wsum = 0.f;
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        float best = fmaxf(v0, v1);
        for (int o = 16; o; o >>= 1) best = fmaxf(best, __shfl_xor_sync(0xffffffffu, best, o));
        int win = (v0 == best) ? lane : ((v1 == best) ? lane + 32 : 64);
        for (int o = 16; o; o >>= 1) win = min(win, __shfl_xor_sync(0xffffffffu, win, o));
        float p = expf(best - mx) * inv;
        wsum += p;
        if (lane == 0) { A.s.ids[j] = win; A.s.wts[j] = p; }
        if (win == lane) v0 = -1e30f;
        else if (win == lane + 32) v1 = -1e30f;
    }
    __syncwarp();
    float sc2 = A.routed / (wsum + 1e-9f);
    if (lane < 8) A.s.wts[lane] *= sc2;
}

__device__ void moe_gateup(const KArg& A, int layer, char* SM) {
    __half* xs = (__half*)SM;                       // 2304 half
    float* gsum = (float*)(SM + 8192);              // 18
    float* ysred = (float*)(SM + 8448);
    for (int i = threadIdx.x; i < D_HID; i += blockDim.x)
        xs[i] = __hmul(A.s.xn_moe[i], __float2half(0.125f));
    __syncthreads();
    fill_gsum(xs, D_HID, gsum);
    for (int u = blockIdx.x; u < 144; u += gridDim.x) {
        int slot = u >> 4, gu = (u >> 3) & 1, tile = u & 7;
        const uint8_t* wq;
        const bf16 *sc, *zr;
        if (slot < 8) {
            int e = A.s.ids[slot];
            wq = A.w.moe_wq[layer][gu] + (size_t)e * (1152 * 1024);
            sc = A.w.moe_sc[layer][gu] + (size_t)e * (18 * 1024);
            zr = A.w.moe_zr[layer][gu] + (size_t)e * (18 * 1024);
        } else {
            wq = A.w.moe_wq[layer][3 + gu];
            sc = A.w.moe_sc[layer][3 + gu];
            zr = A.w.moe_zr[layer][3 + gu];
        }
        gemv_i4(wq, sc, zr, M_INTER, 0, 18, tile * 128, 128, xs, gsum,
                (gu ? A.s.hu : A.s.hg) + slot * M_INTER + tile * 128, 1.f, false, ysred);
    }
}

__device__ void moe_down(const KArg& A, int layer, char* SM) {
    __half* xs = (__half*)SM;                       // 1024 half
    float* gsum = (float*)(SM + 8192);              // 8
    float* ysred = (float*)(SM + 8448);
    for (int u = blockIdx.x; u < 162; u += gridDim.x) {
        __syncthreads();
        int slot = u / 18, tile = u % 18;
        for (int i = threadIdx.x; i < M_INTER; i += blockDim.x) {
            float g = A.s.hg[slot * M_INTER + i];
            float up = A.s.hu[slot * M_INTER + i];
            float act = g / (1.f + expf(-g)) * up;
            xs[i] = __float2half(act * 0.125f);
        }
        __syncthreads();
        fill_gsum(xs, M_INTER, gsum);
        const uint8_t* wq;
        const bf16 *sc, *zr;
        float wmul;
        if (slot < 8) {
            int e = A.s.ids[slot];
            wq = A.w.moe_wq[layer][2] + (size_t)e * (512 * 2304);
            sc = A.w.moe_sc[layer][2] + (size_t)e * (8 * 2304);
            zr = A.w.moe_zr[layer][2] + (size_t)e * (8 * 2304);
            wmul = A.s.wts[slot];
        } else {
            wq = A.w.moe_wq[layer][5];
            sc = A.w.moe_sc[layer][5];
            zr = A.w.moe_zr[layer][5];
            wmul = 1.f;
        }
        gemv_i4(wq, sc, zr, D_HID, 0, 8, tile * 128, 128, xs, gsum,
                A.s.x + (layer + 1) * D_HID + tile * 128, wmul, true, ysred);
    }
}

__device__ void mla_phaseA(const KArg& A, char* SM) {
    __half* xs = (__half*)SM;
    float* gsum = (float*)(SM + 8192);
    float* red = (float*)(SM + 8320);
    float* ysred = (float*)(SM + 8448);
    prep_norm(A.s.x + 3 * D_HID, A.w.attn_norm[3], D_HID, xs, gsum, red);
    for (int u = blockIdx.x; u < 57; u += gridDim.x) {
        if (u < 48) {
            gemv_i4(A.w.mla_wq[0], A.w.mla_sc[0], A.w.mla_zr[0], HQ_MLA,
                    0, 18, u * 128, 128, xs, gsum, A.s.q6144 + u * 128,
                    1.f, false, ysred);
        } else if (u < 53) {
            int tile = u - 48;
            int nc = min(128, KVA - tile * 128);
            gemv_i4(A.w.mla_wq[1], A.w.mla_sc[1], A.w.mla_zr[1], KVA,
                    0, 18, tile * 128, nc, xs, gsum, A.s.kva + tile * 128,
                    1.f, false, ysred);
        } else if (u < 57) {
            int j0 = (u - 53) * 576;
            for (int i = threadIdx.x; i < 576; i += blockDim.x)
                A.s.hattn[3 * D_HID + j0 + i] = A.s.x[3 * D_HID + j0 + i];
        } else {
            // idle slots: warm L2 with the absorbed kv_b matrix (next phase)
            // and the ckv/krope cache rows the scores phase streams.
            const int L1 = A.a.L + 1;
            const uint8_t* bases[3];
            size_t sizes[3];
            bases[0] = A.w.mla_wq[2];
            sizes[0] = (size_t)512 * NKVB;
            bases[1] = (const uint8_t*)A.s.ckv;
            sizes[1] = (size_t)L1 * KVL * 2;
            bases[2] = (const uint8_t*)A.s.krope;
            sizes[2] = (size_t)L1 * QR * 2;
            l2_warm(bases, sizes, 3, u - 57, gridDim.x - 57);
        }
    }
}

__device__ void mla_phaseB(const KArg& A, char* SM) {
    float* qn = (float*)SM;   // [32][128] scaled q_nope
    for (int idx = threadIdx.x; idx < NH * 128; idx += blockDim.x) {
        int hh = idx >> 7, dd = idx & 127;
        float v = __bfloat162float(__float2bfloat16(
            A.s.q6144[hh * 192 + dd]));
        qn[idx] = v * MLA_SCALE;
    }
    __syncthreads();
    const int L = A.a.L;
    for (int u = blockIdx.x; u < 257; u += gridDim.x) {
        if (u < 256) {
            const int c2 = u, gc = c2 >> 6;
            const int w = threadIdx.x >> 5, lane = threadIdx.x & 31;
            const uint8_t* wr = A.w.mla_wq[2] + (size_t)c2 * NKVB;
            const bf16* scb = A.w.mla_sc[2] + (size_t)gc * NKVB;
            const bf16* zrb = A.w.mla_zr[2] + (size_t)gc * NKVB;
            for (int hh = w * 2; hh < w * 2 + 2; ++hh) {
                const int d0 = lane * 4;
                uint32_t wrd = *reinterpret_cast<const uint32_t*>(wr + hh * 256 + d0);
                float alo = 0.f, ahi = 0.f;
                #pragma unroll
                for (int jj = 0; jj < 4; ++jj) {
                    uint32_t b = (wrd >> (8 * jj)) & 0xFFu;
                    int n = hh * 256 + d0 + jj;
                    float s = __bfloat162float(scb[n]);
                    float z = __bfloat162float(zrb[n]);
                    float qv = qn[hh * 128 + d0 + jj];
                    alo += qv * ((float)(b & 15u) - z) * s;
                    ahi += qv * ((float)(b >> 4) - z) * s;
                }
                for (int o = 16; o; o >>= 1) {
                    alo += __shfl_down_sync(0xffffffffu, alo, o);
                    ahi += __shfl_down_sync(0xffffffffu, ahi, o);
                }
                if (lane == 0) {
                    A.s.qabs[(2 * c2) * NH + hh] = __float2half(alo);
                    A.s.qabs[(2 * c2 + 1) * NH + hh] = __float2half(ahi);
                }
            }
        } else {
            for (int i = threadIdx.x; i < KVL; i += blockDim.x)
                A.s.ckv[(size_t)L * KVL + i] = __float2bfloat16(A.s.kva[i]);
            if (threadIdx.x < 32) {
                const int r = threadIdx.x;
                float inv = powf(A.theta, -((float)(2 * r)) / 64.f);
                float ang = (float)L * inv;
                float cc = cosf(ang), sn = sinf(ang);
                float ke = __bfloat162float(__float2bfloat16(
                    A.s.kva[KVL + 2 * r]));
                float ko = __bfloat162float(__float2bfloat16(
                    A.s.kva[KVL + 2 * r + 1]));
                A.s.krope[(size_t)L * QR + 2 * r] = __float2bfloat16(ke * cc - ko * sn);
                A.s.krope[(size_t)L * QR + 2 * r + 1] = __float2bfloat16(ko * cc + ke * sn);
                for (int hh = 0; hh < NH; ++hh) {
                    int b0 = hh * 192 + 128 + 2 * r;
                    float qe = __bfloat162float(__float2bfloat16(A.s.q6144[b0]));
                    float qo = __bfloat162float(__float2bfloat16(A.s.q6144[b0 + 1]));
                    A.s.qr[(2 * r) * NH + hh] = __float2half((qe * cc - qo * sn) * MLA_SCALE);
                    A.s.qr[(2 * r + 1) * NH + hh] = __float2half((qo * cc + qe * sn) * MLA_SCALE);
                }
            }
        }
    }
}

__device__ void mla_scores(const KArg& A, char* SM) {
    const int L1 = A.a.L + 1;
    // adaptive chunk rows: smallest of 16/32/64/... so one round covers L1
    int crows = 16;
    while ((L1 + crows - 1) / crows > (int)gridDim.x) crows <<= 1;
    const int NCH = (L1 + crows - 1) / crows;
    if (blockIdx.x >= NCH && (int)gridDim.x > NCH) {
        // idle chunks: warm L2 with the o-proj weights used two phases later
        const uint8_t* bases[1] = { A.w.mla_wq[3] };
        size_t sizes[1] = { (size_t)2048 * 2304 };
        l2_warm(bases, sizes, 1, blockIdx.x - NCH, (int)gridDim.x - NCH);
    }
    __half* qa_s = (__half*)SM;                 // 512*32
    __half* qr_s = (__half*)(SM + 32768);       // 64*32
    __half* ckv_s = (__half*)(SM + 36864);      // 32*512
    float* red = (float*)(SM + 69632);          // 32*32 per-batch dots
    __half* kr_s = (__half*)(SM + 77824);       // 32*64
    if (blockIdx.x >= NCH) return;              // idle CTA: skip straight to
                                                // the grid barrier in the caller
    {
        const uint32_t* qsrc = reinterpret_cast<const uint32_t*>(A.s.qabs);
        uint32_t* qdst = reinterpret_cast<uint32_t*>(qa_s);
        for (int i = threadIdx.x; i < KVL * 16; i += blockDim.x) qdst[i] = qsrc[i];
        const uint32_t* rsrc = reinterpret_cast<const uint32_t*>(A.s.qr);
        uint32_t* rdst = reinterpret_cast<uint32_t*>(qr_s);
        for (int i = threadIdx.x; i < QR * 16; i += blockDim.x) rdst[i] = rsrc[i];
    }
    __syncthreads();
    for (int u = blockIdx.x; u < NCH; u += gridDim.x) {
        __syncthreads();
        const int l0 = u * crows;
        const int cnt = min(crows, L1 - l0);
        float m_run = -1e30f, s_run = 0.f;   // online softmax stats (threads<32)
        for (int b = 0; b * 32 < cnt; ++b) {
            const int bc = min(32, cnt - b * 32);
            {
                const int lb = l0 + b * 32;
                for (int idx = threadIdx.x; idx < 32 * 256; idx += blockDim.x) {
                    int row = idx >> 8, cp = (idx & 255) * 2;
                    if (row < bc) {
                        bf162 v = *reinterpret_cast<const bf162*>(A.s.ckv + (size_t)(lb + row) * KVL + cp);
                        *reinterpret_cast<__half2*>(ckv_s + row * KVL + cp) =
                            __floats2half2_rn(__bfloat162float(v.x), __bfloat162float(v.y));
                    }
                }
                for (int idx = threadIdx.x; idx < 32 * 32; idx += blockDim.x) {
                    int row = idx >> 5, cp = (idx & 31) * 2;
                    if (row < bc) {
                        bf162 v = *reinterpret_cast<const bf162*>(A.s.krope + (size_t)(lb + row) * QR + cp);
                        *reinterpret_cast<__half2*>(kr_s + row * QR + cp) =
                            __floats2half2_rn(__bfloat162float(v.x), __bfloat162float(v.y));
                    }
                }
            }
            __syncthreads();
            {
                const int lloc = threadIdx.x >> 4;
                const int h2 = threadIdx.x & 15;
                if (lloc < bc) {
                    float sx = 0.f, sy = 0.f;
                    const __half* crow = ckv_s + lloc * KVL;
                    const __half2* qa2 = reinterpret_cast<const __half2*>(qa_s) + h2;
                    #pragma unroll 1
                    for (int seg = 0; seg < 8; ++seg) {
                        __half2 a0 = __float2half2_rn(0.f);
                        __half2 a1 = __float2half2_rn(0.f);
                        #pragma unroll 8
                        for (int c = seg * 64; c < seg * 64 + 64; c += 2) {
                            a0 = __hfma2(__half2half2(crow[c]), qa2[c * 16], a0);
                            a1 = __hfma2(__half2half2(crow[c + 1]), qa2[(c + 1) * 16], a1);
                        }
                        float2 f0 = __half22float2(a0);
                        float2 f1 = __half22float2(a1);
                        sx += f0.x + f1.x;
                        sy += f0.y + f1.y;
                    }
                    const __half* krow = kr_s + lloc * QR;
                    const __half2* qr2 = reinterpret_cast<const __half2*>(qr_s) + h2;
                    __half2 a0 = __float2half2_rn(0.f);
                    #pragma unroll 8
                    for (int r = 0; r < QR; ++r)
                        a0 = __hfma2(__half2half2(krow[r]), qr2[r * 16], a0);
                    float2 f0 = __half22float2(a0);
                    sx += f0.x;
                    sy += f0.y;
                    red[lloc * 32 + h2 * 2] = sx;
                    red[lloc * 32 + h2 * 2 + 1] = sy;
                    float2 sv = make_float2(sx, sy);
                    *reinterpret_cast<float2*>(A.s.scores + (size_t)(l0 + b * 32 + lloc) * NH + h2 * 2) = sv;
                }
            }
            __syncthreads();
            if (threadIdx.x < 32) {
                const int hh = threadIdx.x;
                float bm = -1e30f;
                for (int l = 0; l < bc; ++l) bm = fmaxf(bm, red[l * 32 + hh]);
                float bs = 0.f;
                for (int l = 0; l < bc; ++l) bs += expf(red[l * 32 + hh] - bm);
                float nm = fmaxf(m_run, bm);
                s_run = s_run * expf(m_run - nm) + bs * expf(bm - nm);
                m_run = nm;
            }
        }
        if (threadIdx.x < 32) {
            A.s.mx[u * 32 + threadIdx.x] = m_run;
            A.s.sm[u * 32 + threadIdx.x] = s_run;
        }
    }
}

__device__ void mla_ctx(const KArg& A, char* SM) {
    const int L1 = A.a.L + 1;
    int crows = 16;                     // same adaptive chunking as scores
    while ((L1 + crows - 1) / crows > (int)gridDim.x) crows <<= 1;
    const int NCH = (L1 + crows - 1) / crows;
    const int NCHS = NCH;
    float* red = (float*)SM;            // 16*32
    float* Msm = red + 512;             // 32
    float* Sinv = Msm + 32;             // 32
    __half* p_s = (__half*)(SM + 4096); // crows*32 (<=128*32)
    if (blockIdx.x >= NCH) return;
    const int part = threadIdx.x >> 5, hh = threadIdx.x & 31;
    if (part < 16) {
        float m = -1e30f;
        for (int ch = part; ch < NCHS; ch += 16) m = fmaxf(m, A.s.mx[ch * 32 + hh]);
        red[part * 32 + hh] = m;
    }
    __syncthreads();
    if (threadIdx.x < 32) {
        float m = -1e30f;
        for (int p = 0; p < 16; ++p) m = fmaxf(m, red[p * 32 + threadIdx.x]);
        Msm[threadIdx.x] = m;
    }
    __syncthreads();
    if (part < 16) {
        float s = 0.f;
        for (int ch = part; ch < NCHS; ch += 16)
            s += A.s.sm[ch * 32 + hh] * expf(A.s.mx[ch * 32 + hh] - Msm[hh]);
        red[part * 32 + hh] = s;
    }
    __syncthreads();
    if (threadIdx.x < 32) {
        float s = 0.f;
        for (int p = 0; p < 16; ++p) s += red[p * 32 + threadIdx.x];
        Sinv[threadIdx.x] = 1.f / s;
    }
    __syncthreads();
    for (int u = blockIdx.x; u < NCH; u += gridDim.x) {
        __syncthreads();
        const int l0 = u * crows;
        const int cnt = min(crows, L1 - l0);
        for (int idx = threadIdx.x; idx < crows * 32; idx += blockDim.x) {
            int ll = idx >> 5, h5 = idx & 31;
            float p = 0.f;
            if (ll < cnt)
                p = expf(A.s.scores[(size_t)(l0 + ll) * NH + h5] - Msm[h5]) * Sinv[h5];
            p_s[ll * 32 + h5] = __float2half(p);
        }
        __syncthreads();
        const int c = threadIdx.x;
        __half2 acc[16];
        #pragma unroll
        for (int i = 0; i < 16; ++i) acc[i] = __float2half2_rn(0.f);
        for (int ll = 0; ll < cnt; ++ll) {
            __half2 cv = __half2half2(__float2half(
                __bfloat162float(A.s.ckv[(size_t)(l0 + ll) * KVL + c])));
            const __half2* pr = reinterpret_cast<const __half2*>(p_s + ll * 32);
            #pragma unroll
            for (int i = 0; i < 16; ++i) acc[i] = __hfma2(pr[i], cv, acc[i]);
        }
        #pragma unroll
        for (int i = 0; i < 16; ++i) {
            float2 f = __half22float2(acc[i]);
            atomicAdd(A.s.ctx + (size_t)(2 * i) * KVL + c, f.x);
            atomicAdd(A.s.ctx + (size_t)(2 * i + 1) * KVL + c, f.y);
        }
    }
}

__device__ void mla_wv(const KArg& A, char* SM) {
    __half* xs = (__half*)SM;                       // 512
    float* gsum = (float*)(SM + 8192);              // 4
    float* ysred = (float*)(SM + 8448);
    for (int u = blockIdx.x; u < NH; u += gridDim.x) {
        __syncthreads();
        for (int i = threadIdx.x; i < KVL; i += blockDim.x)
            xs[i] = __float2half(A.s.ctx[u * KVL + i] * 0.125f);
        __syncthreads();
        fill_gsum(xs, KVL, gsum);
        gemv_i4(A.w.mla_wq[2], A.w.mla_sc[2], A.w.mla_zr[2], NKVB, 0, 4,
                u * 256 + 128, 128, xs, gsum, A.s.obuf + u * 128, 1.f, false,
                ysred);
    }
}

__device__ void tail_phase(const KArg& A) {
    const int T = blockIdx.x * blockDim.x + threadIdx.x;
    const int NT = gridDim.x * blockDim.x;
    for (int i = T; i < D_HID; i += NT)
        A.s.outb[i] = __float2bfloat16(A.s.x[4 * D_HID + i]);
}

// ------------------------------------------------------------------- kernel
extern "C" __global__ void __launch_bounds__(NBLK, 1) mega_kernel(KArg A) {
    cg::grid_group grid = cg::this_grid();
    extern __shared__ char SM[];
    int ph = 0;
    #define PHASE_END() do { grid.sync(); if (++ph >= A.a.max_phase) return; } while (0)
    phase_Z(A);
    PHASE_END();
    for (int layer = 0; layer < 3; ++layer) {
        kda_phaseA(A, layer, SM);
        PHASE_END();
        kda_phaseB(A, layer, SM);
        PHASE_END();
        oproj_phase(A, A.w.kda_wq[layer][4], A.w.kda_sc[layer][4],
                    A.w.kda_zr[layer][4], A.s.hattn + layer * D_HID, nullptr, SM);
        PHASE_END();
        moe_logits(A, layer, SM);
        PHASE_END();
        moe_pick(A, layer);
        PHASE_END();
        moe_gateup(A, layer, SM);
        PHASE_END();
        moe_down(A, layer, SM);
        PHASE_END();
    }
    mla_phaseA(A, SM);
    PHASE_END();
    mla_phaseB(A, SM);
    PHASE_END();
    mla_scores(A, SM);
    PHASE_END();
    mla_ctx(A, SM);
    PHASE_END();
    mla_wv(A, SM);
    PHASE_END();
    oproj_phase(A, A.w.mla_wq[3], A.w.mla_sc[3], A.w.mla_zr[3],
                A.s.hattn + 3 * D_HID, nullptr, SM);
    PHASE_END();
    moe_logits(A, 3, SM);
    PHASE_END();
    moe_pick(A, 3);
    PHASE_END();
    moe_gateup(A, 3, SM);
    PHASE_END();
    moe_down(A, 3, SM);
    PHASE_END();
    tail_phase(A);
    #undef PHASE_END
}

// --------------------------------------------------------------------- host
static KArg g_arg;
static bf16* g_winbank = nullptr;   // [2][3][3][3][C]
static int g_grid = 0;
static bool g_attr_set = false;

static void launch(cudaStream_t stream) {
    void* p = (void*)&g_arg;
    void* args[] = {p};
    cudaError_t err = cudaLaunchCooperativeKernel(
        (void*)mega_kernel, dim3(g_grid), dim3(NBLK), args, SMEM_BYTES, stream);
    TORCH_CHECK(err == cudaSuccess, "coop launch failed: ", cudaGetErrorString(err));
}

void begin(std::vector<torch::Tensor> ws, std::vector<torch::Tensor> scr,
           torch::Tensor winbank,
           torch::Tensor S0, torch::Tensor S1, torch::Tensor S2,
           std::vector<torch::Tensor> winr,
           torch::Tensor ckv_src, torch::Tensor kr_src,
           torch::Tensor hidden, int64_t L, double routed, double theta,
           int64_t max_phase) {
    if (!g_attr_set) {
        cudaFuncSetAttribute((void*)mega_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_BYTES);
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        int per_sm = 0;
        cudaOccupancyMaxActiveBlocksPerMultiprocessor(
            &per_sm, (void*)mega_kernel, NBLK, SMEM_BYTES);
        TORCH_CHECK(per_sm >= 1, "megakernel: zero occupancy");
        g_grid = prop.multiProcessorCount * per_sm;
        g_attr_set = true;
    }
    // ---- weights, fixed order built in Python (see _weight_list) ----
    size_t i = 0;
    for (int l = 0; l < 3; ++l) {
        g_arg.w.attn_norm[l] = (const bf16*)ws[i++].data_ptr();
        g_arg.w.moe_norm[l] = (const bf16*)ws[i++].data_ptr();
        for (int m = 0; m < 5; ++m) {
            g_arg.w.kda_wq[l][m] = (const uint8_t*)ws[i++].data_ptr();
            g_arg.w.kda_sc[l][m] = (const bf16*)ws[i++].data_ptr();
            g_arg.w.kda_zr[l][m] = (const bf16*)ws[i++].data_ptr();
        }
        g_arg.w.beta_w[l] = (const bf16*)ws[i++].data_ptr();
        g_arg.w.conv_w[l] = (const bf16*)ws[i++].data_ptr();
        g_arg.w.router[l] = (const bf16*)ws[i++].data_ptr();
        for (int m = 0; m < 6; ++m) {
            g_arg.w.moe_wq[l][m] = (const uint8_t*)ws[i++].data_ptr();
            g_arg.w.moe_sc[l][m] = (const bf16*)ws[i++].data_ptr();
            g_arg.w.moe_zr[l][m] = (const bf16*)ws[i++].data_ptr();
        }
    }
    g_arg.w.attn_norm[3] = (const bf16*)ws[i++].data_ptr();
    g_arg.w.moe_norm[3] = (const bf16*)ws[i++].data_ptr();
    for (int m = 0; m < 4; ++m) {
        g_arg.w.mla_wq[m] = (const uint8_t*)ws[i++].data_ptr();
        g_arg.w.mla_sc[m] = (const bf16*)ws[i++].data_ptr();
        g_arg.w.mla_zr[m] = (const bf16*)ws[i++].data_ptr();
    }
    g_arg.w.router[3] = (const bf16*)ws[i++].data_ptr();
    for (int m = 0; m < 6; ++m) {
        g_arg.w.moe_wq[3][m] = (const uint8_t*)ws[i++].data_ptr();
        g_arg.w.moe_sc[3][m] = (const bf16*)ws[i++].data_ptr();
        g_arg.w.moe_zr[3][m] = (const bf16*)ws[i++].data_ptr();
    }
    TORCH_CHECK(i == ws.size(), "weight list size mismatch");
    // ---- scratch, fixed order (see _scratch_list) ----
    size_t j = 0;
    g_arg.s.x = (float*)scr[j++].data_ptr();
    g_arg.s.hattn = (float*)scr[j++].data_ptr();
    g_arg.s.qkvg = (float*)scr[j++].data_ptr();
    g_arg.s.betaout = (float*)scr[j++].data_ptr();
    g_arg.s.obuf = (float*)scr[j++].data_ptr();
    g_arg.s.xn_moe = (__half*)scr[j++].data_ptr();
    g_arg.s.ids = (int*)scr[j++].data_ptr();
    g_arg.s.wts = (float*)scr[j++].data_ptr();
    g_arg.s.logits = (float*)scr[j++].data_ptr();
    g_arg.s.hg = (float*)scr[j++].data_ptr();
    g_arg.s.hu = (float*)scr[j++].data_ptr();
    g_arg.s.q6144 = (float*)scr[j++].data_ptr();
    g_arg.s.kva = (float*)scr[j++].data_ptr();
    g_arg.s.qabs = (__half*)scr[j++].data_ptr();
    g_arg.s.qr = (__half*)scr[j++].data_ptr();
    g_arg.s.scores = (float*)scr[j++].data_ptr();
    g_arg.s.mx = (float*)scr[j++].data_ptr();
    g_arg.s.sm = (float*)scr[j++].data_ptr();
    g_arg.s.ctx = (float*)scr[j++].data_ptr();
    g_arg.s.outb = (bf16*)scr[j++].data_ptr();
    g_arg.s.ckv = (bf16*)scr[j++].data_ptr();
    g_arg.s.krope = (bf16*)scr[j++].data_ptr();
    TORCH_CHECK(j == scr.size(), "scratch list size mismatch");

    g_winbank = (bf16*)winbank.data_ptr();
    g_arg.a.S[0] = (float*)S0.data_ptr();
    g_arg.a.S[1] = (float*)S1.data_ptr();
    g_arg.a.S[2] = (float*)S2.data_ptr();
    for (int k = 0; k < 9; ++k) {
        g_arg.a.winr[k] = (const bf16*)winr[k].data_ptr();
        g_arg.a.winw[k] = g_winbank + (size_t)k * 3 * C_KDA;   // bank 0
    }
    g_arg.a.ckv_src = (const bf16*)ckv_src.data_ptr();
    g_arg.a.kr_src = (const bf16*)kr_src.data_ptr();
    g_arg.a.hidden = (const bf16*)hidden.data_ptr();
    g_arg.a.copy_len = (int)L;
    g_arg.a.L = (int)L;
    g_arg.a.max_phase = (int)max_phase;
    g_arg.routed = (float)routed;
    g_arg.theta = (float)theta;
    launch(at::cuda::getCurrentCUDAStream());
}

void nextstep(torch::Tensor hidden, int64_t L, int64_t widx_w, int64_t max_phase) {
    const size_t bank = (size_t)9 * 3 * C_KDA;
    for (int k = 0; k < 9; ++k) {
        g_arg.a.winr[k] = g_winbank + (size_t)(1 - widx_w) * bank + (size_t)k * 3 * C_KDA;
        g_arg.a.winw[k] = g_winbank + (size_t)widx_w * bank + (size_t)k * 3 * C_KDA;
    }
    g_arg.a.hidden = (const bf16*)hidden.data_ptr();
    g_arg.a.copy_len = 0;
    g_arg.a.L = (int)L;
    g_arg.a.max_phase = (int)max_phase;
    launch(at::cuda::getCurrentCUDAStream());
}
"""

_CPP_SRC = r"""
#include <torch/extension.h>
#include <vector>
void begin(std::vector<torch::Tensor> ws, std::vector<torch::Tensor> scr,
           torch::Tensor winbank,
           torch::Tensor S0, torch::Tensor S1, torch::Tensor S2,
           std::vector<torch::Tensor> winr,
           torch::Tensor ckv_src, torch::Tensor kr_src,
           torch::Tensor hidden, int64_t L, double routed, double theta,
           int64_t max_phase);
void nextstep(torch::Tensor hidden, int64_t L, int64_t widx_w, int64_t max_phase);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("begin", &begin);
    m.def("next", &nextstep);
}
"""

_ext = None


def _get_ext():
    global _ext
    if _ext is None:
        _ext = load_inline(
            name="kimi_mega_v1",
            cpp_sources=[_CPP_SRC],
            cuda_sources=[_CUDA_SRC],
            extra_cuda_cflags=["-O3", "-gencode=arch=compute_120,code=sm_120"],
            verbose=False,
        )
    return _ext


# ---------------------------------------------------------------------------
# Module tree mirroring reference.py's state_dict layout exactly.
# ---------------------------------------------------------------------------
class _QL(nn.Module):
    def __init__(self, in_f, out_f, group=128):
        super().__init__()
        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 _QE(nn.Module):
    def __init__(self, n, in_f, out_f, group=128):
        super().__init__()
        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):
        super().__init__()
        H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
        self.q_proj = _QL(d, H * Dk)
        self.k_proj = _QL(d, H * Dk)
        self.v_proj = _QL(d, H * Dk)
        self.g_proj = _QL(d, H * Dk)
        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 = _QL(H * Dk, d)


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


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


class _Block(nn.Module):
    def __init__(self, cfg, kind):
        super().__init__()
        self.kind = kind
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = _KDA(cfg) if kind == "K" else _MLA(cfg)
        self.moe = _MoE(cfg)


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        assert tuple(cfg.pattern) == ("K", "K", "K", "M"), "kernel is specialized"
        assert cfg.n_experts == 64 and cfg.moe_inter == 1024 and cfg.hidden == 2304
        self.blocks = nn.ModuleList(_Block(cfg, k) for k in cfg.pattern)
        self._tr = None          # (expected_len, widx, S0_ptr) for continuation detect
        self._cap = 0
        self._dev = None
        self.debug_max_phase = 1 << 30

    # ---- internal buffers (plain attributes: must NOT enter state_dict) ----
    def _alloc(self, dev, L):
        cap = L + 4096
        if self._dev == dev and self._cap >= L + 64:
            return
        C = 4096
        nch = cap // 32 + 2
        f32 = dict(device=dev, dtype=torch.float32)
        f16 = dict(device=dev, dtype=torch.float16)
        bf = dict(device=dev, dtype=torch.bfloat16)
        self._ckv = torch.zeros(cap, 512, **bf)
        self._kr = torch.zeros(cap, 64, **bf)
        self._winbank = torch.zeros(2, 9, 3, C, **bf)
        self._scr = [
            torch.zeros(5 * 2304, **f32),        # x
            torch.zeros(4 * 2304, **f32),        # hattn
            torch.zeros(4 * C, **f32),           # qkvg
            torch.zeros(32, **f32),              # betaout
            torch.zeros(C, **f32),               # obuf
            torch.zeros(2304, **f16),            # xn_moe
            torch.zeros(8, dtype=torch.int32, device=dev),  # ids
            torch.zeros(8, **f32),               # wts
            torch.zeros(64, **f32),              # logits
            torch.zeros(9 * 1024, **f32),        # hg
            torch.zeros(9 * 1024, **f32),        # hu
            torch.zeros(6144, **f32),            # q6144
            torch.zeros(576, **f32),             # kva
            torch.zeros(512 * 32, **f16),        # qabs
            torch.zeros(64 * 32, **f16),         # qr
            torch.zeros(cap * 32, **f32),        # scores
            torch.zeros(nch * 32, **f32),        # mx
            torch.zeros(nch * 32, **f32),        # sm
            torch.zeros(32 * 512, **f32),        # ctx
            torch.zeros(2304, **bf),             # outb
            self._ckv.view(-1),
            self._kr.view(-1),
        ]
        self._out = self._scr[19].view(2304)
        self._cap = cap
        self._dev = dev

    def _weight_list(self):
        ws = []
        for i in range(3):
            b = self.blocks[i]
            ws += [b.attn_norm, b.moe_norm]
            for p in (b.attn.q_proj, b.attn.k_proj, b.attn.v_proj, b.attn.g_proj, b.attn.o_proj):
                ws += [p.w_q, p.scales, p.zeros]
            ws += [b.attn.beta_proj.weight, b.attn.conv_w, b.moe.router.weight]
            for p in (b.moe.gate, b.moe.up, b.moe.down, b.moe.s_gate, b.moe.s_up, b.moe.s_down):
                ws += [p.w_q, p.scales, p.zeros]
        b = self.blocks[3]
        ws += [b.attn_norm, b.moe_norm]
        for p in (b.attn.q_proj, b.attn.kv_a, b.attn.kv_b, b.attn.o_proj):
            ws += [p.w_q, p.scales, p.zeros]
        ws += [b.moe.router.weight]
        for p in (b.moe.gate, b.moe.up, b.moe.down, b.moe.s_gate, b.moe.s_up, b.moe.s_down):
            ws += [p.w_q, p.scales, p.zeros]
        return [w.contiguous() for w in ws]

    @torch.no_grad()
    def step(self, hidden, state):
        ext = _get_ext()
        L = state[3]["c_kv"].shape[0]
        cont = (
            self._tr is not None
            and state[3]["c_kv"].data_ptr() == self._ckv.data_ptr()
            and L == self._tr[0]
            and state[0]["S"].data_ptr() == self._tr[2]
        )
        if not cont:
            dev = hidden.device
            self._alloc(dev, L)
            winr = []
            for i in range(3):
                for k in ("cq", "ck", "cv"):
                    t = state[i][k]
                    assert t.is_contiguous() and t.dtype == torch.bfloat16
                    winr.append(t)
                assert state[i]["S"].is_contiguous() and state[i]["S"].dtype == torch.float32
            ckv_src = state[3]["c_kv"].contiguous()
            kr_src = state[3]["k_rope"].contiguous()
            assert L + 1 <= self._cap
            ext.begin(
                self._weight_list(), self._scr, self._winbank.view(-1),
                state[0]["S"], state[1]["S"], state[2]["S"], winr,
                ckv_src, kr_src, hidden.contiguous(), L,
                float(self.cfg.routed_scaling), float(self.cfg.rope_theta),
                self.debug_max_phase,
            )
            widx = 0
        else:
            widx = 1 - self._tr[1]
            ext.next(hidden, L, widx, self.debug_max_phase)
        # hand back views into our capacity buffers / window bank
        state[3]["c_kv"] = self._ckv[: L + 1]
        state[3]["k_rope"] = self._kr[: L + 1]
        for i in range(3):
            state[i]["cq"] = self._winbank[widx, i * 3 + 0]
            state[i]["ck"] = self._winbank[widx, i * 3 + 1]
            state[i]["cv"] = self._winbank[widx, i * 3 + 2]
        self._tr = (L + 1, widx, state[0]["S"].data_ptr())
        return self._out, state