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

The whole per-token forward of one [KDA, KDA, KDA, MLA] motif (each layer followed by a 64-expert top-8 + shared MoE FFN,
RMSNorms and residual adds) runs inside a single `__global__` launch (`mega_kernel`, built with `load_inline`): 188 persistent
CTAs (one per SM, 512 threads) walk the 26 dependent phases of the token and synchronize with a grid barrier between them.

* int4 weights are repacked once, when the state dict is loaded, into 16-column x 128-k "units" (1088 B: nibbles laid out in
  mma A-fragment order + per-column scale/zero). A unit is streamed once from DRAM into registers/shared memory, dequantized
  in registers exactly like the reference ((q - z) * s in bf16, via packed bf16x2 ops), and multiplied on the tensor cores
  (mma.m16n8k16, fp32 accumulate). No bf16 weight is ever materialized.
* The first units of the next phase are staged into shared memory with bulk async copies (one mbarrier per warp), issued right
  after this CTA has arrived at the phase barrier (arrive / prefetch / wait), so DRAM streams through barrier and prologue
  gaps without delaying the arrival itself; state write-back that nothing downstream needs is issued in the same window.
* KDA: q/k/v/g projections -> short causal conv + gates + recurrent-state update (one CTA per head x 32 value columns)
  -> o_proj. MLA: absorbed formulation (kv_b folded into the query and the output) so attention runs over the 512+64-d latent
  cache with tensor cores (split over cache tokens, online softmax, partial combine); the new latent row is appended in place
  (the cache is handed over into a capacity buffer on the first step, copied while it is read).
* MoE: router + top-8 (ties -> lower index, like torch.topk), gate/up and down projections for the 8 routed + 1 shared
  expert, down-proj input passed as a bf16 hi/lo pair so it behaves like the reference's fp32.
* Split-K partial sums use fp32 atomics into accumulators that are cleaned (zeroed) by the kernel itself once dead.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field

import torch
import torch.nn as nn

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

_CUDA_SRC = r'''// ==== mk_common.cuh ====
// Common device helpers for the Kimi-Linear decode megakernel.
#include <cstdint>
#include <cuda_runtime.h>
#include <cuda_bf16.h>

#define DEVI __device__ __forceinline__

// ---------------------------------------------------------------- bf16 helpers
DEVI float bf2f(uint32_t bits16) { return __uint_as_float(bits16 << 16); }
DEVI float bf2f_lo(uint32_t w) { return __uint_as_float(w << 16); }
DEVI float bf2f_hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); }
DEVI uint32_t f2bf_bits(float x) {           // round-to-nearest-even, returns 16-bit pattern
    __nv_bfloat16 h = __float2bfloat16_rn(x);
    return (uint32_t)__bfloat16_as_ushort(h);
}
DEVI float round_bf16(float x) { return __uint_as_float(f2bf_bits(x) << 16); }
DEVI uint32_t pack_bf16x2(float lo, float hi) {
    __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi);
    return *reinterpret_cast<uint32_t*>(&v);
}

// ---------------------------------------------------------------- memory helpers
DEVI uint4 ldg_stream16(const void* p) {
    uint4 v;
    asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];"
                 : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p));
    return v;
}
DEVI uint2 ldg_stream8(const void* p) {
    uint2 v;
    asm volatile("ld.global.nc.L1::no_allocate.v2.u32 {%0,%1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(p));
    return v;
}
// loads of data produced by other CTAs inside the same kernel: bypass L1 (.cg)
DEVI float ldcg_f32(const float* p) { float v; asm volatile("ld.global.cg.f32 %0, [%1];" : "=f"(v) : "l"(p)); return v; }
DEVI uint32_t ldcg_u32(const void* p) { uint32_t v; asm volatile("ld.global.cg.u32 %0, [%1];" : "=r"(v) : "l"(p)); return v; }
DEVI void stcg_f32(float* p, float v) { asm volatile("st.global.cg.f32 [%0], %1;" :: "l"(p), "f"(v) : "memory"); }
DEVI void red_add_f32(float* p, float v) { asm volatile("red.global.add.f32 [%0], %1;" :: "l"(p), "f"(v) : "memory"); }

// ---------------------------------------------------------------- MMA
DEVI void mma_bf16_16816(float (&c)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) {
    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};"
                 : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
                 : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
DEVI uint32_t bf16x2_sub(uint32_t a, uint32_t b) {
    uint32_t d;
    asm("sub.rn.bf16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b));
    return d;
}
DEVI uint32_t bf16x2_mul(uint32_t a, uint32_t b) {
    uint32_t d;
    asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b));
    return d;
}

// ---------------------------------------------------------------- warp / block reductions
DEVI float warp_sum(float v) {
#pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
    return v;
}
DEVI float warp_max(float v) {
#pragma unroll
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o));
    return v;
}

// ---------------------------------------------------------------- grid barrier
// Monotonic-counter barrier. `epoch` is per-thread-0 state; each CTA adds exactly 1 per barrier.
DEVI void grid_barrier(unsigned int* ctr, unsigned int nblocks, unsigned int& target_base) {
    __syncthreads();
    if (threadIdx.x == 0) {
        target_base += nblocks;
        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" :: "l"(ctr), "r"(1) : "memory");
        unsigned int v;
        unsigned int spins = 0;
        do {
            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(ctr) : "memory");
            if (++spins > (1u << 26)) __trap();          // a lost CTA would otherwise hang the GPU forever
        } while ((int)(v - target_base) < 0);
    }
    __syncthreads();
}

DEVI void barrier_arrive(unsigned int* ctr, unsigned int nblocks, unsigned int& target_base) {
    __syncthreads();
    if (threadIdx.x == 0) {
        target_base += nblocks;
        asm volatile("red.release.gpu.global.add.u32 [%0], %1;" :: "l"(ctr), "r"(1) : "memory");
    }
}
DEVI void barrier_wait(unsigned int* ctr, unsigned int target_base) {
    if (threadIdx.x == 0) {
        unsigned int v;
        unsigned int spins = 0;
        do {
            asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(ctr) : "memory");
            if (++spins > (1u << 26)) __trap();
        } while ((int)(v - target_base) < 0);
    }
    __syncthreads();
}

// ---------------------------------------------------------------- packed weight format
// Unit = 16 output columns x 128 k (one quant group) = 1088 bytes:
//   [0,512)   chunk0: lane l -> 16 B = 4 words (mma k-steps 0..3)
//   [512,1024) chunk1: k-steps 4..7
//   [1024,1088) sz block: for gq in 0..7 : 4 x bf16 = {scale(gq), scale(gq+8), 128+zero(gq), 128+zero(gq+8)}
// word (lane=(gq,t), kstep s): nibble for (a-reg r in 0..3, half h in 0..1) at bit 4r+16h holds
//   q[k = 16 s + 2t + 8(r>>1) + h][n = 16 m + gq + 8(r&1)]
constexpr int UNIT_BYTES = 1088;

// ---------------------------------------------------------------- cp.async / ldmatrix
DEVI uint32_t smem_u32(const void* p) { return (uint32_t)__cvta_generic_to_shared(p); }
DEVI void cp_async16(uint32_t dst_smem, const void* src) {
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(dst_smem), "l"(src) : "memory");
}
DEVI void cp_async_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); }
template <int N> DEVI void cp_async_wait() { asm volatile("cp.async.wait_group %0;" :: "n"(N) : "memory"); }
DEVI void ldmatrix_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
}
DEVI void ldmatrix_x4_trans(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
}

// ---------------------------------------------------------------- mbarrier + bulk async copy (global -> shared)
DEVI void mbar_init(uint32_t addr, uint32_t count) { asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(addr), "r"(count) : "memory"); }
DEVI void fence_proxy_async() { asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); }
DEVI void mbar_expect_tx(uint32_t addr, uint32_t bytes) {
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" :: "r"(addr), "r"(bytes) : "memory");
}
DEVI void bulk_g2s(uint32_t dst, const void* src, uint32_t bytes, uint32_t mbar) {
    asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];"
                 :: "r"(dst), "l"(src), "r"(bytes), "r"(mbar) : "memory");
}
DEVI uint32_t mbar_test_wait(uint32_t addr, uint32_t parity) {
    uint32_t ok;
    asm volatile("{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n\tselp.u32 %0, 1, 0, p;\n}"
                 : "=r"(ok) : "r"(addr), "r"(parity) : "memory");
    return ok;
}
DEVI uint4 lds128(const void* p) {
    uint4 v; uint32_t a = (uint32_t)__cvta_generic_to_shared(p);
    asm volatile("ld.shared.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "r"(a) : "memory");
    return v;
}
DEVI uint2 lds64(const void* p) {
    uint2 v; uint32_t a = (uint32_t)__cvta_generic_to_shared(p);
    asm volatile("ld.shared.v2.u32 {%0,%1}, [%2];" : "=r"(v.x), "=r"(v.y) : "r"(a) : "memory");
    return v;
}

// ==== repack.cuh ====

// Repack normal-orientation W4 weights ( (K/2, Nsrc) uint8, scales/zeros (K/128, Nsrc) bf16 ) into unit stream.
// Matrix used = columns [col_off, col_off+N) of the source. Batched via blockIdx.y (batch b):
//   src bytes advance by b*src_batch_stride (w_q), src_sz_batch_stride (scales/zeros), col_off_batch (columns),
//   dst advances by b*dst_batch_stride bytes.
__global__ void repack_normal_kernel(const uint8_t* __restrict__ w_q, const uint16_t* __restrict__ sc, const uint16_t* __restrict__ zr,
                                     uint8_t* __restrict__ dst, int K, int N, int Nsrc, int col_off,
                                     long long src_batch_stride, long long src_sz_batch_stride, int col_off_batch, long long dst_batch_stride) {
    int b = blockIdx.y;
    w_q += (long long)b * src_batch_stride;
    sc += (long long)b * src_sz_batch_stride;
    zr += (long long)b * src_sz_batch_stride;
    dst += (long long)b * dst_batch_stride;
    int col0 = col_off + b * col_off_batch;
    int KG = K / 128;
    long long units = (long long)(N / 16) * KG;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= units * 256) return;
    int s = idx & 7;
    int lane = (idx >> 3) & 31;
    long long u = idx >> 8;
    int m = (int)(u / KG), g = (int)(u % KG);
    int gq = lane >> 2, t = lane & 3;
    uint32_t word = 0;
#pragma unroll
    for (int r = 0; r < 4; r++) {
#pragma unroll
        for (int h = 0; h < 2; h++) {
            int row = gq + 8 * (r & 1);
            int n = col0 + m * 16 + row;
            int k = g * 128 + 16 * s + 2 * t + 8 * (r >> 1) + h;
            uint32_t byte = w_q[(long long)(k >> 1) * Nsrc + n];
            uint32_t q = (k & 1) ? (byte >> 4) : (byte & 0xF);
            word |= q << (4 * r + 16 * h);
        }
    }
    uint8_t* up = dst + u * UNIT_BYTES;
    *reinterpret_cast<uint32_t*>(up + (s >> 2) * 512 + lane * 16 + (s & 3) * 4) = word;
    if (s == 0 && lane < 8) {
        int n0 = col0 + m * 16 + lane, n1 = n0 + 8;
        uint32_t s0 = sc[(long long)g * Nsrc + n0], s1 = sc[(long long)g * Nsrc + n1];
        float z0 = bf2f(zr[(long long)g * Nsrc + n0]), z1 = bf2f(zr[(long long)g * Nsrc + n1]);
        uint32_t zz0 = f2bf_bits(128.f + z0), zz1 = f2bf_bits(128.f + z1);      // exact in bf16; used as (128+q) - (128+z)
        uint32_t* sp = reinterpret_cast<uint32_t*>(up + 1024 + lane * 8);
        sp[0] = s0 | (s1 << 16);
        sp[1] = zz0 | (zz1 << 16);
    }
}

// kv_b K-part, transposed use: y[c] = sum_d W[c, h*256 + d] * q[d].  Unit (h, mc) = 16 c-rows x 128 d, nibbles only (1024 B).
// word (lane=(gq,t), kstep s): nibble (r,hh) = q[c = 16 mc + gq + 8(r&1)][d = 16 s + 2t + 8(r>>1) + hh]
// Also emits szKb[h][g][d] = {scale bf16 | (8-zero) bf16 << 16}, g = c/128.
__global__ void repack_kbt_kernel(const uint8_t* __restrict__ w_q, const uint16_t* __restrict__ sc, const uint16_t* __restrict__ zr,
                                  uint8_t* __restrict__ dst, uint32_t* __restrict__ szdst, int Nsrc) {
    int h = blockIdx.y;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    // words: 32 mc * 32 lanes * 8 s = 8192 per head
    if (idx < 32 * 256) {
        int s = idx & 7;
        int lane = (idx >> 3) & 31;
        int mc = (int)(idx >> 8);
        int gq = lane >> 2, t = lane & 3;
        uint32_t word = 0;
#pragma unroll
        for (int r = 0; r < 4; r++) {
#pragma unroll
            for (int hh = 0; hh < 2; hh++) {
                int c = mc * 16 + gq + 8 * (r & 1);
                int d = 16 * s + 2 * t + 8 * (r >> 1) + hh;
                uint32_t byte = w_q[(long long)(c >> 1) * Nsrc + h * 256 + d];
                uint32_t q = (c & 1) ? (byte >> 4) : (byte & 0xF);
                word |= q << (4 * r + 16 * hh);
            }
        }
        uint8_t* up = dst + ((long long)h * 32 + mc) * 1024;
        *reinterpret_cast<uint32_t*>(up + (s >> 2) * 512 + lane * 16 + (s & 3) * 4) = word;
    }
    // sz table: 4 groups * 128 d per head
    if (idx < 4 * 128) {
        int g = (int)(idx >> 7), d = (int)(idx & 127);
        int n = h * 256 + d;
        uint32_t s0 = sc[(long long)g * Nsrc + n];
        float z0 = bf2f(zr[(long long)g * Nsrc + n]);
        szdst[h * 512 + idx] = s0 | (f2bf_bits(8.f - z0) << 16);
    }
}

// ==== gemv.cuh ====

struct WUnit { uint4 a, b; uint2 sz; };

DEVI WUnit load_unit(const uint8_t* p, int lane) {
    WUnit u;
    u.a = ldg_stream16(p + lane * 16);
    u.b = ldg_stream16(p + 512 + lane * 16);
    u.sz = ldg_stream8(p + 1024 + (lane >> 2) * 8);
    return u;
}

// Dequantize one unit exactly like the reference (bf16 (q - z) * s, one rounding) and accumulate x . w into ca/cb (fp32,
// two interleaved chains). xf_t: smem uint2 array for THIS lane's t and this unit's first k-step: xf_t[s*4] = {b0,b1}.
// LO = true also multiplies a second x fragment set (bf16 low part of an fp32 activation) located `lo_stride` uint2's later.
template <bool LO = false>
DEVI void unit_mma(const WUnit& u, const uint2* xf_t, float (&ca)[4], float (&cb)[4], int lo_stride = 0) {
    const uint32_t S0 = (u.sz.x & 0xFFFFu) * 0x10001u, S1 = (u.sz.x >> 16) * 0x10001u;      // scales of rows g, g+8 (bf16x2 broadcast)
    const uint32_t Z0 = (u.sz.y & 0xFFFFu) * 0x10001u, Z1 = (u.sz.y >> 16) * 0x10001u;      // 128 + zero point
    const uint32_t words[8] = {u.a.x, u.a.y, u.a.z, u.a.w, u.b.x, u.b.y, u.b.z, u.b.w};
#pragma unroll
    for (int s = 0; s < 8; s++) {
        const uint32_t w = words[s];
        const uint32_t a0 = bf16x2_mul(bf16x2_sub((w & 0x000F000Fu) | 0x43004300u, Z0), S0);
        const uint32_t a1 = bf16x2_mul(bf16x2_sub(((w >> 4) & 0x000F000Fu) | 0x43004300u, Z1), S1);
        const uint32_t a2 = bf16x2_mul(bf16x2_sub(((w >> 8) & 0x000F000Fu) | 0x43004300u, Z0), S0);
        const uint32_t a3 = bf16x2_mul(bf16x2_sub(((w >> 12) & 0x000F000Fu) | 0x43004300u, Z1), S1);
        const uint2 b = xf_t[s * 4];
        if (s & 1) mma_bf16_16816(cb, a0, a1, a2, a3, b.x, b.y);
        else mma_bf16_16816(ca, a0, a1, a2, a3, b.x, b.y);
        if (LO) {
            const uint2 bl = xf_t[s * 4 + lo_stride];
            if (s & 1) mma_bf16_16816(cb, a0, a1, a2, a3, bl.x, bl.y);
            else mma_bf16_16816(ca, a0, a1, a2, a3, bl.x, bl.y);
        }
    }
}

struct UnitDesc { const uint8_t* ptr; int key; int xoff; };

// Process units [wb, we) (warp-level). Dec(u) -> UnitDesc; Flush(key, row_g_sum, row_g8_sum) is called when the M-tile key
// changes and at the end. Units of one M-tile accumulate in the tensor-core accumulators without intermediate epilogues.
template <int UR, class Dec, class Flush, bool LO = false>
DEVI void gemv_units(int wb, int we, const uint2* xf_t, Dec dec, Flush flush, int lo_stride = 0) {
    const int lane = threadIdx.x & 31;
    float ca[4] = {0.f, 0.f, 0.f, 0.f}, cb[4] = {0.f, 0.f, 0.f, 0.f};
    int cur_key = -1;
    auto step = [&](const WUnit& w, const UnitDesc& d) {
        if (d.key != cur_key) {
            if (cur_key >= 0) flush(cur_key, ca[0] + cb[0], ca[2] + cb[2]);
#pragma unroll
            for (int i = 0; i < 4; i++) { ca[i] = 0.f; cb[i] = 0.f; }
            cur_key = d.key;
        }
        unit_mma<LO>(w, xf_t + d.xoff * 4, ca, cb, lo_stride);
    };
    int u = wb;
    for (; u + UR <= we; u += UR) {
        WUnit wu[UR];
        UnitDesc d[UR];
#pragma unroll
        for (int i = 0; i < UR; i++) {
            d[i] = dec(u + i);
            wu[i] = load_unit(d[i].ptr, lane);
        }
#pragma unroll
        for (int i = 0; i < UR; i++) step(wu[i], d[i]);
    }
    for (; u < we; u++) {
        UnitDesc d = dec(u);
        WUnit w = load_unit(d.ptr, lane);
        step(w, d);
    }
    if (cur_key >= 0) flush(cur_key, ca[0] + cb[0], ca[2] + cb[2]);
}

// ---- per-warp shared-memory staging of a phase's first units (bulk async copies, one mbarrier per warp)
struct WarpStage {
    const uint8_t* ptr;   // generic pointer to this warp's slot (NSU units)
    uint32_t base;        // smem address of the slot
    uint32_t mbar;        // smem address of the warp's mbarrier
    uint32_t parity;      // completed uses so far
    int n;                // units currently staged (0 = nothing in flight)
};
DEVI WUnit load_unit_smem(const uint8_t* p, int lane) {
    WUnit u;
    u.a = lds128(p + lane * 16);
    u.b = lds128(p + 512 + lane * 16);
    u.sz = lds64(p + 1024 + (lane >> 2) * 8);
    return u;
}
template <int NS, class PtrOf>
DEVI void stage_issue(WarpStage& st, int wb, int we, PtrOf ptr_of) {
    const int n = we - wb < NS ? (we - wb > 0 ? we - wb : 0) : NS;
    st.n = n;
    if (n > 0 && (threadIdx.x & 31) == 0) {
        fence_proxy_async();
        mbar_expect_tx(st.mbar, (uint32_t)n * UNIT_BYTES);
#pragma unroll
        for (int i = 0; i < NS; i++)
            if (i < n) bulk_g2s(st.base + (uint32_t)i * UNIT_BYTES, ptr_of(wb + i), UNIT_BYTES, st.mbar);
    }
    __syncwarp();
}
DEVI void stage_wait(WarpStage& st) {
    if (st.n > 0) {
        unsigned int spins = 0;
        while (!mbar_test_wait(st.mbar, st.parity & 1u)) { if (++spins > (1u << 26)) __trap(); }
        st.parity++;
    }
}

// Like gemv_units, but the first st.n units of [wb,we) are read from the warp's staged slot (already waited on).
template <int UR, class Dec, class Flush, bool LO = false>
DEVI void gemv_units_stage(int wb, int we, const uint2* xf_t, Dec dec, Flush flush, const WarpStage& st, int lo_stride = 0) {
    const int lane = threadIdx.x & 31;
    float ca[4] = {0.f, 0.f, 0.f, 0.f}, cb[4] = {0.f, 0.f, 0.f, 0.f};
    int cur_key = -1;
    auto step = [&](const WUnit& w, const UnitDesc& d) {
        if (d.key != cur_key) {
            if (cur_key >= 0) flush(cur_key, ca[0] + cb[0], ca[2] + cb[2]);
#pragma unroll
            for (int i = 0; i < 4; i++) { ca[i] = 0.f; cb[i] = 0.f; }
            cur_key = d.key;
        }
        unit_mma<LO>(w, xf_t + d.xoff * 4, ca, cb, lo_stride);
    };
    const int n0 = st.n;
    int u = wb + n0;
    // issue the first global batch (clamped indices; only used if a full batch exists)
    const bool have = (u + UR <= we);
    WUnit wu[UR];
    UnitDesc d[UR];
    {
        const int last = we > wb ? we - 1 : wb;
#pragma unroll
        for (int i = 0; i < UR; i++) {
            d[i] = dec(min(u + i, last));
            wu[i] = load_unit(d[i].ptr, lane);
        }
    }
    for (int i = 0; i < n0; i++) {
        const UnitDesc ds = dec(wb + i);
        const WUnit w = load_unit_smem(st.ptr + i * UNIT_BYTES, lane);
        step(w, ds);
    }
    if (have) {
#pragma unroll
        for (int i = 0; i < UR; i++) step(wu[i], d[i]);
        u += UR;
    }
    for (; u + UR <= we; u += UR) {
        WUnit wv[UR];
        UnitDesc dd[UR];
#pragma unroll
        for (int i = 0; i < UR; i++) {
            dd[i] = dec(u + i);
            wv[i] = load_unit(dd[i].ptr, lane);
        }
#pragma unroll
        for (int i = 0; i < UR; i++) step(wv[i], dd[i]);
    }
    for (; u < we; u++) {
        const UnitDesc dt = dec(u);
        const WUnit w = load_unit(dt.ptr, lane);
        step(w, dt);
    }
    if (cur_key >= 0) flush(cur_key, ca[0] + cb[0], ca[2] + cb[2]);
}

// Build x fragments from a bf16 smem array xs[K] (raw 16-bit). xf[(K/16)*4] uint2. Block-wide (all threads), ends with __syncthreads.
DEVI void build_xfrag(const uint16_t* xs, int K, uint2* xf) {
    const int tid = threadIdx.x;
    const int nent = (K / 16) * 4;
    for (int e = tid; e < nent; e += blockDim.x) {
        int ks = e >> 2, t = e & 3;
        int base = ks * 16 + 2 * t;
        uint32_t b0 = *reinterpret_cast<const uint32_t*>(xs + base);
        uint32_t b1 = *reinterpret_cast<const uint32_t*>(xs + base + 8);
        xf[e] = make_uint2(b0, b1);
    }
    __syncthreads();
}


// ==== attn.cuh ====
// MLA absorbed attention (tensor cores). One CTA = one chunk of cache tokens, all 32 heads.
//   scores S^T[token][head] = sum_c K[token][c] * Q[head][c]   (M=tokens, N=heads, K=576)  -> mma with A=K rows (ldmatrix), B=Q (global bf16)
//   O[head][c] += P[head][token] * V[token][c]                 (M=heads,  N=c,     K=tokens) -> A=P (smem), B=V^T (ldmatrix.trans)
// Rounds of 32 tokens, double-buffered with cp.async, online softmax across rounds. Writes an unnormalized partial (m, l, O).

namespace attn {
constexpr int RT = 32;                    // tokens per round
constexpr int ROWB = 1168;                // smem row pitch in bytes (1152 + 16 pad => conflict-free ldmatrix)
constexpr int SS_STRIDE = 33;             // floats
constexpr int P_STRIDE = 20;              // 32-bit words per head row (40 bf16)
constexpr size_t BUF_BYTES = (size_t)RT * ROWB;                           // 37376
constexpr size_t OFF_BUF0 = 0;
constexpr size_t OFF_BUF1 = BUF_BYTES;
constexpr size_t OFF_SS = 2 * BUF_BYTES;                                  // float [32][33]
constexpr size_t OFF_P = OFF_SS + 32 * SS_STRIDE * 4;                     // u32 [32][20]
constexpr size_t OFF_ML = OFF_P + 32 * P_STRIDE * 4;                      // float m[32], l[32], alpha[32]
constexpr size_t SMEM_BYTES = OFF_ML + 3 * 32 * 4;

struct Args {
    const uint16_t* qb;        // bf16 [32][576] : q_abs (512) | q_rope (64), already rotated
    const uint16_t* ckv_src; const uint16_t* krope_src;    // old rows
    uint16_t* ckv_dst; uint16_t* krope_dst;                // capacity buffers (row L holds the new token)
    int L; int migrate;
    float scale;
};

// Load round r (tokens [tb, tb+nt)) into smem buffer via cp.async; zero-fill unused rows.
DEVI void load_round(const Args& A, uint8_t* buf, int tb, int nt) {
    const uint32_t bsm = smem_u32(buf);
    for (int e = threadIdx.x; e < RT * 72; e += blockDim.x) {
        const int j = e / 72, q = e - j * 72;
        const uint32_t d = bsm + j * ROWB + q * 16;
        if (j < nt) {
            const int l = tb + j;
            const bool old = l < A.L;
            const void* src;
            if (q < 64) src = (old ? A.ckv_src : A.ckv_dst) + (size_t)l * 512 + q * 8;
            else src = (old ? A.krope_src : A.krope_dst) + (size_t)l * 64 + (q - 64) * 8;
            cp_async16(d, src);
        } else {
            *reinterpret_cast<uint4*>(buf + j * ROWB + q * 16) = make_uint4(0, 0, 0, 0);
        }
    }
}

// Write old rows to their new home (first step after the state was handed over).
DEVI void migrate_round(const Args& A, const uint8_t* buf, int tb, int nt) {
    for (int e = threadIdx.x; e < RT * 72; e += blockDim.x) {
        const int j = e / 72, q = e - j * 72;
        const int l = tb + j;
        if (j < nt && l < A.L) {
            uint4 v = *reinterpret_cast<const uint4*>(buf + j * ROWB + q * 16);
            if (q < 64) *reinterpret_cast<uint4*>(A.ckv_dst + (size_t)l * 512 + q * 8) = v;
            else *reinterpret_cast<uint4*>(A.krope_dst + (size_t)l * 64 + (q - 64) * 8) = v;
        }
    }
}

// part: [32][512] O (fp32) + m[32] at +16384, l[32] at +16416
DEVI void run(const Args& A, float* part, int t0, int t1, uint8_t* sm) {
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    const int g = lane >> 2, t = lane & 3;
    uint8_t* buf[2] = {sm + OFF_BUF0, sm + OFF_BUF1};
    float* Ss = reinterpret_cast<float*>(sm + OFF_SS);
    uint32_t* Ps = reinterpret_cast<uint32_t*>(sm + OFF_P);
    float* m_run = reinterpret_cast<float*>(sm + OFF_ML);
    float* l_run = m_run + 32;
    float* alpha_s = l_run + 32;
    const int T = t1 - t0;
    const int R = (T + RT - 1) / RT;
    if (tid < 32) { m_run[tid] = -INFINITY; l_run[tid] = 0.f; }
    float oacc[2][4][4];
#pragma unroll
    for (int a = 0; a < 2; a++)
#pragma unroll
        for (int b = 0; b < 4; b++)
#pragma unroll
            for (int c = 0; c < 4; c++) oacc[a][b][c] = 0.f;

    if (R > 0) { load_round(A, buf[0], t0, min(RT, T)); }
    cp_async_commit();
    for (int r = 0; r < R; r++) {
        const int tb = t0 + r * RT;
        const int nt = min(RT, t1 - tb);
        if (r + 1 < R) load_round(A, buf[(r + 1) & 1], tb + RT, min(RT, t1 - (tb + RT)));
        cp_async_commit();
        cp_async_wait<1>();
        __syncthreads();
        const uint8_t* kb = buf[r & 1];
        if (A.migrate) migrate_round(A, kb, tb, nt);
        // ---- scores: warps 0..3 each own 8 heads and both 16-token tiles
        if (warp < 4) {
            float ca[4] = {0.f, 0.f, 0.f, 0.f}, cb[4] = {0.f, 0.f, 0.f, 0.f};
            const uint32_t* qrow = reinterpret_cast<const uint32_t*>(A.qb + (size_t)(warp * 8 + g) * 576);
            const uint32_t base = smem_u32(kb) + ((lane & 7) + ((lane >> 3) & 1) * 8) * ROWB + (lane >> 4) * 16;
#pragma unroll
            for (int blk = 0; blk < 3; blk++) {
                uint32_t qw[24];
#pragma unroll
                for (int i = 0; i < 12; i++) {
                    qw[2 * i] = __ldg(qrow + (blk * 12 + i) * 8 + t);
                    qw[2 * i + 1] = __ldg(qrow + (blk * 12 + i) * 8 + t + 4);
                }
#pragma unroll
                for (int i = 0; i < 12; i++) {
                    const int ks = blk * 12 + i;
                    uint32_t a0, a1, a2, a3;
                    ldmatrix_x4(a0, a1, a2, a3, base + ks * 32);
                    mma_bf16_16816(ca, a0, a1, a2, a3, qw[2 * i], qw[2 * i + 1]);
                    ldmatrix_x4(a0, a1, a2, a3, base + 16 * ROWB + ks * 32);
                    mma_bf16_16816(cb, a0, a1, a2, a3, qw[2 * i], qw[2 * i + 1]);
                }
            }
            const int h0 = warp * 8 + 2 * t;
            Ss[(h0) * SS_STRIDE + g] = ca[0] * A.scale;
            Ss[(h0 + 1) * SS_STRIDE + g] = ca[1] * A.scale;
            Ss[(h0) * SS_STRIDE + g + 8] = ca[2] * A.scale;
            Ss[(h0 + 1) * SS_STRIDE + g + 8] = ca[3] * A.scale;
            Ss[(h0) * SS_STRIDE + 16 + g] = cb[0] * A.scale;
            Ss[(h0 + 1) * SS_STRIDE + 16 + g] = cb[1] * A.scale;
            Ss[(h0) * SS_STRIDE + 16 + g + 8] = cb[2] * A.scale;
            Ss[(h0 + 1) * SS_STRIDE + 16 + g + 8] = cb[3] * A.scale;
        }
        __syncthreads();
        // ---- online softmax: warp w handles heads 2w, 2w+1; lane = token
#pragma unroll
        for (int hh = 0; hh < 2; hh++) {
            const int h = warp * 2 + hh;
            const float s = (lane < nt) ? Ss[h * SS_STRIDE + lane] : -INFINITY;
            const float mr = warp_max(s);
            const float mo = m_run[h];
            const float mn = fmaxf(mo, mr);
            const float al = __expf(mo - mn);
            const float p = (lane < nt) ? __expf(s - mn) : 0.f;
            const uint32_t pb = f2bf_bits(p);
            const float pr = bf2f(pb);
            const float ps = warp_sum(pr);
            uint16_t* prow = reinterpret_cast<uint16_t*>(Ps + h * P_STRIDE);
            prow[lane] = (uint16_t)pb;
            if (lane == 0) { m_run[h] = mn; l_run[h] = l_run[h] * al + ps; alpha_s[h] = al; }
        }
        __syncthreads();
        // ---- PV: warp w owns c in [32w, 32w+32), both head m-tiles
#pragma unroll
        for (int mt = 0; mt < 2; mt++) {
            const float al0 = alpha_s[mt * 16 + g], al1 = alpha_s[mt * 16 + g + 8];
#pragma unroll
            for (int nt_ = 0; nt_ < 4; nt_++) {
                oacc[mt][nt_][0] *= al0; oacc[mt][nt_][1] *= al0;
                oacc[mt][nt_][2] *= al1; oacc[mt][nt_][3] *= al1;
            }
        }
        {
            const uint32_t vbase = smem_u32(kb) + ((lane & 7) + ((lane >> 3) & 1) * 8) * ROWB + (warp * 32 + (lane >> 4) * 8) * 2;
#pragma unroll
            for (int kt = 0; kt < 2; kt++) {
                uint32_t bv[4][2];
                ldmatrix_x4_trans(bv[0][0], bv[0][1], bv[1][0], bv[1][1], vbase + kt * 16 * ROWB);
                ldmatrix_x4_trans(bv[2][0], bv[2][1], bv[3][0], bv[3][1], vbase + kt * 16 * ROWB + 32);
#pragma unroll
                for (int mt = 0; mt < 2; mt++) {
                    const uint32_t a0 = Ps[(mt * 16 + g) * P_STRIDE + kt * 8 + t];
                    const uint32_t a1 = Ps[(mt * 16 + g + 8) * P_STRIDE + kt * 8 + t];
                    const uint32_t a2 = Ps[(mt * 16 + g) * P_STRIDE + kt * 8 + t + 4];
                    const uint32_t a3 = Ps[(mt * 16 + g + 8) * P_STRIDE + kt * 8 + t + 4];
#pragma unroll
                    for (int nt_ = 0; nt_ < 4; nt_++) mma_bf16_16816(oacc[mt][nt_], a0, a1, a2, a3, bv[nt_][0], bv[nt_][1]);
                }
            }
        }
        __syncthreads();
    }
    // ---- write partial
#pragma unroll
    for (int mt = 0; mt < 2; mt++) {
#pragma unroll
        for (int nt_ = 0; nt_ < 4; nt_++) {
            const int c = warp * 32 + nt_ * 8 + 2 * t;
            const int h0 = mt * 16 + g;
            *reinterpret_cast<float2*>(part + (size_t)h0 * 512 + c) = make_float2(oacc[mt][nt_][0], oacc[mt][nt_][1]);
            *reinterpret_cast<float2*>(part + (size_t)(h0 + 8) * 512 + c) = make_float2(oacc[mt][nt_][2], oacc[mt][nt_][3]);
        }
    }
    if (tid < 32) { part[16384 + tid] = m_run[tid]; part[16416 + tid] = l_run[tid]; }
}
}  // namespace attn

// ==== mk.cuh ====
// Kimi-Linear decode megakernel: device code.

namespace mk {

constexpr int GEMV_UR = 1;       // units in flight per warp in the weight-streaming loop

constexpr int D = 2304, KGD = 18, NTD = 144;
constexpr int NH = 32, DK = 128, CH = 4096;
constexpr int MI = 1024;
constexpr int TOPK = 8, NSLOT = 9;
constexpr int NTHREADS = 512, NWARPS = 16;
constexpr int KVL = 512, ROPE = 64, ATT = 576;
constexpr float EPS = 1e-6f;
constexpr float ROUTED_SCALE = 2.446f;

// ------------------------------------------------------------------ parameters
struct LayerW {
    const uint8_t* wA;          // KDA: q|k|v|g packs.  MLA: q_proj|kv_a packs
    const uint8_t* wO;          // o_proj pack
    const uint16_t* beta_w;     // KDA (32, 2304) bf16
    const uint16_t* conv_w;     // KDA (3, 4096, 4) bf16
    const uint8_t* wKb;         // MLA kv_b K-part (transposed units)
    const uint32_t* szKb;       // MLA kv_b K-part sz table
    const uint8_t* wVb;         // MLA kv_b V-part packs
    const uint16_t* attn_norm;
    const uint16_t* moe_norm;
    const uint16_t* router_w;   // (64, 2304) bf16
    const uint8_t* wGU;         // [65][2][1152 units]
    const uint8_t* wD;          // [65][1152 units]
};
struct LayerS {
    float* S; uint16_t* cq; uint16_t* ck; uint16_t* cv;                  // KDA state
    const uint16_t* ckv_src; const uint16_t* krope_src;                   // MLA cache (old rows)
    uint16_t* ckv_dst; uint16_t* krope_dst;                               // MLA cache (capacity buffers)
};
struct Params {
    LayerW w[4];
    LayerS s[4];
    const uint16_t* hidden_in;
    uint16_t* hidden_out;
    float* scratch;
    unsigned* bar;
    unsigned bar_base;
    int L;          // cached tokens before this step
    int migrate;    // 1: copy old cache rows src->dst while attending
};

// ------------------------------------------------------------------ scratch layout (floats)
constexpr int ACC_A = 0;          // KDA: q|k|v|g raw sums [16384]; MLA: q [6144] | kv [576]
constexpr int ACC_O = 16384;      // [2304]
constexpr int ACC_GU = 18688;     // [9][2][1024]
constexpr int ACC_DN = 37120;     // [2304]
constexpr int ACC_OV = 39424;     // [4096] (MLA absorbed V output)
constexpr int ACC_L = 43520;
constexpr int ACC_TOTAL = 4 * ACC_L;

constexpr int NZ_LOGITS = 0;      // [64]
constexpr int NZ_BETA = 64;       // [32]
constexpr int NZ_OBUF = 128;      // bf16 [4096]  (2048 floats)
constexpr int NZ_QABS = 2176;     // [32][512]
constexpr int NZ_L = 20608;
constexpr int NZ_BASE = ACC_TOTAL;
constexpr int PART_BASE = NZ_BASE + 4 * NZ_L;
constexpr int PART_STRIDE = 16384 + 64;
constexpr int NAMAX = 188;
constexpr int SCRATCH_FLOATS = PART_BASE + NAMAX * PART_STRIDE;

// ------------------------------------------------------------------ shared memory layout (bytes)
constexpr int NSU = 3;                                   // shared-memory stage slot capacity (units per warp)
constexpr int NS_K1 = 3;                                 // staged units: next layer's K1 / M1 weights
constexpr int NS_O = 1;                                  // staged units: o_proj weights (the rest streams)
constexpr int NS_E2 = 2;                                 // staged units: MoE down-proj weights
constexpr int SM_XIN = 0;                                // bf16 [2304]
constexpr int SM_HS = 4608;                              // bf16 [2304]
constexpr int SM_MBAR = 9216;                            // 16 x 8 B mbarriers (one per warp)
constexpr int SM_XS = 9472;                              // bf16 [4096] (8192)
constexpr int SM_XF = SM_XS + 8192;                      // uint2 [256*4] (8192); also phase scratch for K2 / M2 / M5
constexpr int SM_RED = SM_XF + 8192;                     // float [64]
constexpr int SM_TK = SM_RED + 256;                      // int idx[16], float w[16]
constexpr int SM_STAGE = SM_TK + 256;                    // per-warp staged units: 16 x NSU x 1088 B
constexpr int SM_RROW = SM_STAGE + NWARPS * NSU * UNIT_BYTES;   // router row (4608 B), cp.async target in E0
constexpr int SMEM_BYTES = 101376;
static_assert(SM_RROW + 4608 <= SMEM_BYTES, "smem overflow");
static_assert(SM_STAGE % 16 == 0 && SM_RROW % 16 == 0, "alignment");

DEVI int min_i(int a, int b) { return a < b ? a : b; }

// ------------------------------------------------------------------ small block-level helpers
DEVI void cta_range(int U, int& ub, int& ue) {
    ub = (int)((long long)U * blockIdx.x / gridDim.x);
    ue = (int)((long long)U * (blockIdx.x + 1) / gridDim.x);
}
DEVI void warp_range(int ub, int ue, int& wb, int& we) {
    int n = ue - ub, w = threadIdx.x >> 5;
    wb = ub + (int)((long long)n * w / NWARPS);
    we = ub + (int)((long long)n * (w + 1) / NWARPS);
}

DEVI void rmsnorm_block(const uint16_t* src, const uint16_t* w, uint16_t* dst, float* red) {
    const int tid = threadIdx.x;
    uint32_t wv[5];                                   // norm weights: issue the (cold) loads before the reduction
#pragma unroll
    for (int i = 0; i < 5; i++) { const int idx = tid + i * NTHREADS; wv[i] = idx < D ? (uint32_t)__ldg(w + idx) : 0u; }
    float ss = 0.f;
    for (int i = tid; i < D; i += NTHREADS) { float v = bf2f(src[i]); ss += v * v; }
    ss = warp_sum(ss);
    if ((tid & 31) == 0) red[tid >> 5] = ss;
    __syncthreads();
    float tot = 0.f;
#pragma unroll
    for (int i = 0; i < NWARPS; i++) tot += red[i];
    float inv = rsqrtf(tot * (1.0f / D) + EPS);
#pragma unroll
    for (int i = 0; i < 5; i++) {
        const int idx = tid + i * NTHREADS;
        if (idx < D) dst[idx] = (uint16_t)f2bf_bits((bf2f(src[idx]) * inv) * bf2f(wv[i]));
    }
    __syncthreads();
}

// ------------------------------------------------------------------ KDA
// Phase K1: rmsnorm + q/k/v/g projections (atomics into acc_a) + beta.
DEVI void kda_in_proj(const Params& P, int layer, const LayerW& W, float* acc_a, float* beta_out, uint8_t* sm, WarpStage& st) {
    uint16_t* xin = reinterpret_cast<uint16_t*>(sm + SM_XIN);
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
    float* red = reinterpret_cast<float*>(sm + SM_RED);
    const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    rmsnorm_block(xin, W.attn_norm, xs, red);
    build_xfrag(xs, D, xf);
    constexpr int U = 4 * 256 * KGD;
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    warp_range(ub, ue, wb, we);
    const uint2* xf_t = xf + (lane & 3);
    auto dec = [&](int u) -> UnitDesc {
        UnitDesc d;
        int mg = u / KGD, g = u - mg * KGD;
        d.ptr = W.wA + (size_t)u * UNIT_BYTES; d.key = mg; d.xoff = g * 8;
        return d;
    };
    auto flush = [&](int key, float a0, float a1) {
        int t = lane & 3, g = lane >> 2;
        if (t == 0) red_add_f32(acc_a + key * 16 + g, a0);
        else if (t == 1) red_add_f32(acc_a + key * 16 + g + 8, a1);
    };
    stage_wait(st);
    gemv_units_stage<GEMV_UR>(wb, we, xf_t, dec, flush, st);
    __syncthreads();
    // beta = bf16(dot(xn, beta_w[h]))
    if (blockIdx.x < NH && warp == 0) {
        const uint16_t* row = W.beta_w + (size_t)blockIdx.x * D;
        float acc = 0.f;
#pragma unroll
        for (int i = 0; i < 9; i++) {
            int base = i * 256 + lane * 8;
            uint4 wv = ldg_stream16(row + base);
            uint4 xv = *reinterpret_cast<const uint4*>(xs + base);
            acc += bf2f_lo(wv.x) * bf2f_lo(xv.x) + bf2f_hi(wv.x) * bf2f_hi(xv.x)
                 + bf2f_lo(wv.y) * bf2f_lo(xv.y) + bf2f_hi(wv.y) * bf2f_hi(xv.y)
                 + bf2f_lo(wv.z) * bf2f_lo(xv.z) + bf2f_hi(wv.z) * bf2f_hi(xv.z)
                 + bf2f_lo(wv.w) * bf2f_lo(xv.w) + bf2f_hi(wv.w) * bf2f_hi(xv.w);
        }
        acc = warp_sum(acc);
        if (lane == 0) stcg_f32(beta_out + blockIdx.x, round_bf16(acc));
    }
}

DEVI float silu_f(float x) { return x / (1.0f + expf(-x)); }
DEVI float softplus_f(float x) { return x > 20.f ? x : log1pf(expf(x)); }

// Phase K2: conv + gates + S update. CTAs 0..127: head h=c/4, dv slice j=c%4.
DEVI void kda_state_preload(const LayerS& S, float (&sv)[8]) {
    if (blockIdx.x >= NH * 4) return;
    const int h = blockIdx.x >> 2, j = blockIdx.x & 3, lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    const float* Sp = S.S + (size_t)h * DK * DK + j * 32 + lane;
#pragma unroll
    for (int i = 0; i < 8; i++) sv[i] = Sp[(size_t)(warp * 8 + i) * DK];
}

DEVI void kda_state_update(const LayerW& W, const LayerS& S, const float* acc_a, const float* beta_raw, uint16_t* obuf, uint8_t* sm, float (&spre)[8]) {
    if (blockIdx.x >= NH * 4) return;
    const int h = blockIdx.x >> 2, j = blockIdx.x & 3;
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    float* qs = reinterpret_cast<float*>(sm + SM_XF);          // [128] scaled q
    float* ks = qs + 128;                                       // [128]
    float* gs = ks + 128;                                       // [128] decay
    float* vs = gs + 128;                                       // [32]
    float* predp = vs + 32;                                     // [16][32]
    float* op = predp + NWARPS * 32;                            // [16][32]
    float* bt = op + NWARPS * 32;                               // [1]
    const float scale = 0.08838834764831845f;
    if (tid < 128) {
        int ch = h * DK + tid;
        // q (conv idx 0), k (conv idx 1)
#pragma unroll
        for (int which = 0; which < 2; which++) {
            uint16_t* st = which == 0 ? S.cq : S.ck;
            float raw = round_bf16(ldcg_f32(acc_a + which * CH + ch));
            uint32_t c0 = st[0 * CH + ch], c1 = st[1 * CH + ch], c2 = st[2 * CH + ch];
            uint2 wv = *reinterpret_cast<const uint2*>(W.conv_w + ((size_t)which * CH + ch) * 4);
            float o = bf2f(c0) * bf2f_lo(wv.x) + bf2f(c1) * bf2f_hi(wv.x) + bf2f(c2) * bf2f_lo(wv.y) + raw * bf2f_hi(wv.y);
            float y = round_bf16(silu_f(o));
            if (which == 0) qs[tid] = y * scale; else ks[tid] = y;
        }
        float graw = round_bf16(ldcg_f32(acc_a + 3 * CH + ch));
        gs[tid] = expf(-softplus_f(graw));
    } else if (tid < 160) {
        int dv = tid - 128;
        int ch = h * DK + j * 32 + dv;
        float raw = round_bf16(ldcg_f32(acc_a + 2 * CH + ch));
        uint32_t c0 = S.cv[0 * CH + ch], c1 = S.cv[1 * CH + ch], c2 = S.cv[2 * CH + ch];
        uint2 wv = *reinterpret_cast<const uint2*>(W.conv_w + ((size_t)2 * CH + ch) * 4);
        float o = bf2f(c0) * bf2f_lo(wv.x) + bf2f(c1) * bf2f_hi(wv.x) + bf2f(c2) * bf2f_lo(wv.y) + raw * bf2f_hi(wv.y);
        vs[dv] = round_bf16(silu_f(o));
    } else if (tid == 160) {
        float braw = ldcg_f32(beta_raw + h);
        *bt = 1.0f / (1.0f + expf(-braw));
    }
    __syncthreads();
    float sv[8];
    float pp = 0.f;
#pragma unroll
    for (int i = 0; i < 8; i++) {
        int dk = warp * 8 + i;
        float v = spre[i] * gs[dk];
        sv[i] = v;
        pp += v * ks[dk];
    }
    predp[warp * 32 + lane] = pp;
    __syncthreads();
    float pred = 0.f;
#pragma unroll
    for (int w = 0; w < NWARPS; w++) pred += predp[w * 32 + lane];
    const float delta = (*bt) * (vs[lane] - pred);
    float o = 0.f;
#pragma unroll
    for (int i = 0; i < 8; i++) {
        int dk = warp * 8 + i;
        float v = sv[i] + ks[dk] * delta;
        spre[i] = v;                                  // written back after this CTA has arrived at the barrier
        o += v * qs[dk];
    }
    op[warp * 32 + lane] = o;
    __syncthreads();
    if (warp == 0) {
        float t = 0.f;
#pragma unroll
        for (int w = 0; w < NWARPS; w++) t += op[w * 32 + lane];
        obuf[h * DK + j * 32 + lane] = (uint16_t)f2bf_bits(t);
    }
}

// Deferred S write-back (issued after the K2 barrier arrival, so the arrival's release fence does not wait for these stores).
DEVI void kda_state_store(const LayerS& S, const float (&sv)[8]) {
    if (blockIdx.x >= NH * 4) return;
    const int h = blockIdx.x >> 2, j = blockIdx.x & 3, lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    float* Sp = S.S + (size_t)h * DK * DK + j * 32 + lane;
#pragma unroll
    for (int i = 0; i < 8; i++) Sp[(size_t)(warp * 8 + i) * DK] = sv[i];
}

// Conv window shift for q/k/v (done after the K2 barrier so no CTA can still be reading the old window).
DEVI void kda_conv_state_update(const LayerS& S, const float* acc_a) {
    for (int idx = blockIdx.x * NTHREADS + threadIdx.x; idx < 3 * CH; idx += gridDim.x * NTHREADS) {
        int which = idx / CH, ch = idx - which * CH;
        uint16_t* st = which == 0 ? S.cq : (which == 1 ? S.ck : S.cv);
        float raw = round_bf16(ldcg_f32(acc_a + which * CH + ch));
        uint16_t c1 = st[1 * CH + ch], c2 = st[2 * CH + ch];
        st[0 * CH + ch] = c1; st[1 * CH + ch] = c2; st[2 * CH + ch] = (uint16_t)f2bf_bits(raw);
    }
}

// Generic o_proj: K=4096, N=2304. Input bf16 vector in smem XS already. Accumulates into acc_o.
DEVI void oproj_gemv(const uint8_t* wO, float* acc_o, uint8_t* sm, WarpStage& st) {
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
    const int lane = threadIdx.x & 31;
    build_xfrag(xs, 4096, xf);
    constexpr int KG = 32;
    constexpr int U = NTD * KG;
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    warp_range(ub, ue, wb, we);
    const uint2* xf_t = xf + (lane & 3);
    auto dec = [&](int u) -> UnitDesc {
        UnitDesc d;
        int mg = u / KG, g = u - mg * KG;
        d.ptr = wO + (size_t)u * UNIT_BYTES; d.key = mg; d.xoff = g * 8;
        return d;
    };
    auto flush = [&](int key, float a0, float a1) {
        int t = lane & 3, g = lane >> 2;
        if (t == 0) red_add_f32(acc_o + key * 16 + g, a0);
        else if (t == 1) red_add_f32(acc_o + key * 16 + g + 8, a1);
    };
    stage_wait(st);
    gemv_units_stage<GEMV_UR>(wb, we, xf_t, dec, flush, st);
}

// ------------------------------------------------------------------ MoE
// E0: h = x_in + bf16(attn_out); moe rmsnorm; router logits (CTAs 0..63).
DEVI void moe_pre(const LayerW& W, const float* acc_o, float* logits, uint8_t* sm) {
    uint16_t* xin = reinterpret_cast<uint16_t*>(sm + SM_XIN);
    uint16_t* hs = reinterpret_cast<uint16_t*>(sm + SM_HS);
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
    float* red = reinterpret_cast<float*>(sm + SM_RED);
    uint16_t* rrow = reinterpret_cast<uint16_t*>(sm + SM_RROW);        // router row staged by cp.async (4608 B)
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    if (blockIdx.x < 64 && warp == 0) {
        const uint16_t* row = W.router_w + (size_t)blockIdx.x * D;
#pragma unroll
        for (int i = 0; i < 9; i++) cp_async16(smem_u32(rrow + i * 256 + lane * 8), row + i * 256 + lane * 8);
        cp_async_commit();
    }
    for (int i = tid; i < D; i += NTHREADS) {
        float a = round_bf16(ldcg_f32(acc_o + i));
        hs[i] = (uint16_t)f2bf_bits(bf2f(xin[i]) + a);
    }
    __syncthreads();
    rmsnorm_block(hs, W.moe_norm, xs, red);
    build_xfrag(xs, D, xf);
    if (blockIdx.x < 64 && warp == 0) {
        cp_async_wait<0>();
        __syncwarp();
        float acc = 0.f;
#pragma unroll
        for (int i = 0; i < 9; i++) {
            int base = i * 256 + lane * 8;
            uint4 wv = *reinterpret_cast<const uint4*>(rrow + base);
            uint4 xv = *reinterpret_cast<const uint4*>(xs + base);
            acc += bf2f_lo(wv.x) * bf2f_lo(xv.x) + bf2f_hi(wv.x) * bf2f_hi(xv.x)
                 + bf2f_lo(wv.y) * bf2f_lo(xv.y) + bf2f_hi(wv.y) * bf2f_hi(xv.y)
                 + bf2f_lo(wv.z) * bf2f_lo(xv.z) + bf2f_hi(wv.z) * bf2f_hi(xv.z)
                 + bf2f_lo(wv.w) * bf2f_lo(xv.w) + bf2f_hi(wv.w) * bf2f_hi(xv.w);
        }
        acc = warp_sum(acc);
        if (lane == 0) stcg_f32(logits + blockIdx.x, round_bf16(acc));
    }
}

// Softmax over the 64 router logits + top-8 by rank (value desc, index asc on ties). One warp: logits go to a small smem
// array and each lane ranks its two experts against all 64 (broadcast reads), instead of 8 dependent argmax rounds.
DEVI void router_topk(const float* logits, int* idx_s, float* w_s, float* sv) {
    const int lane = threadIdx.x & 31;
    const float x0 = ldcg_f32(logits + lane), x1 = ldcg_f32(logits + lane + 32);
    sv[lane] = x0; sv[lane + 32] = x1;
    __syncwarp();
    const float m = warp_max(fmaxf(x0, x1));
    const float e0 = expf(x0 - m), e1 = expf(x1 - m);
    const float sum = warp_sum(e0 + e1);
    int r0 = 0, r1 = 0;
#pragma unroll 16
    for (int j = 0; j < 64; j++) {
        const float vj = sv[j];
        r0 += (vj > x0) || (vj == x0 && j < lane);
        r1 += (vj > x1) || (vj == x1 && j < lane + 32);
    }
    const float p0 = e0 / sum, p1 = e1 / sum;
    float sel = 0.f;
    if (r0 < TOPK) sel += p0;
    if (r1 < TOPK) sel += p1;
    sel = warp_sum(sel);
    const float denom = sel + 1e-9f;
    if (r0 < TOPK) { idx_s[r0] = lane; w_s[r0] = p0 / denom * ROUTED_SCALE; }
    if (r1 < TOPK) { idx_s[r1] = lane + 32; w_s[r1] = p1 / denom * ROUTED_SCALE; }
    if (lane == 0) { idx_s[TOPK] = 64; w_s[TOPK] = 1.0f; }
}

// E1: gate/up projections for the 9 experts (atomics into acc_gu).
DEVI void moe_gate_up(const LayerW& W, float* acc_gu, uint8_t* sm) {
    int* tk_idx = reinterpret_cast<int*>(sm + SM_TK);
    const uint2* xf = reinterpret_cast<const uint2*>(sm + SM_XF);
    const int lane = threadIdx.x & 31;
    constexpr int UPS = 2 * 64 * KGD;                // units per slot = 2304
    constexpr int U = NSLOT * UPS;                   // 20736
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    warp_range(ub, ue, wb, we);
    const uint2* xf_t = xf + (lane & 3);
    auto dec = [&](int u) -> UnitDesc {
        UnitDesc d;
        int slot = u / UPS, r = u - slot * UPS;          // r in [0, 2304): [mat][m][g]
        int mat = r / (64 * KGD), r2 = r - mat * (64 * KGD);
        int m = r2 / KGD, g = r2 - m * KGD;
        int e = tk_idx[slot];
        size_t unit = ((size_t)(e * 2 + mat) * 1152) + m * KGD + g;
        d.ptr = W.wGU + unit * UNIT_BYTES;
        d.key = (slot * 2 + mat) * 64 + m; d.xoff = g * 8;
        return d;
    };
    auto flush = [&](int key, float a0, float a1) {
        int t = lane & 3, g = lane >> 2;
        if (t == 0) red_add_f32(acc_gu + key * 16 + g, a0);
        else if (t == 1) red_add_f32(acc_gu + key * 16 + g + 8, a1);
    };
    gemv_units<GEMV_UR>(wb, we, xf_t, dec, flush);
}

// E2: down projections. Prologue builds h_e = silu(gate)*up for the slots this CTA touches.
DEVI void moe_down(const LayerW& W, const float* acc_gu, float* acc_dn, uint8_t* sm, WarpStage& st) {
    int* tk_idx = reinterpret_cast<int*>(sm + SM_TK);
    float* tk_w = reinterpret_cast<float*>(sm + SM_TK + 64);
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
    const int tid = threadIdx.x, lane = tid & 31;
    constexpr int UPS = NTD * 8;                     // 1152
    constexpr int U = NSLOT * UPS;                   // 10368
    constexpr int LO_STRIDE = (MI / 16) * 4;         // uint2's between the hi and lo fragment sets of one slot (256)
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    const int s_lo = ub / UPS, s_hi = (ue - 1) / UPS;   // s_hi - s_lo <= 1
    // h_e = silu(gate) * up in fp32, handed to the tensor core as a bf16 hi/lo pair (~fp32 fidelity, like the reference)
    for (int s = s_lo; s <= s_hi; s++) {
        uint16_t* xh = xs + (size_t)(s - s_lo) * 2 * MI;
        uint16_t* xl = xh + MI;
        for (int i = tid; i < MI; i += NTHREADS) {
            float g = ldcg_f32(acc_gu + (size_t)(s * 2 + 0) * MI + i);
            float u_ = ldcg_f32(acc_gu + (size_t)(s * 2 + 1) * MI + i);
            float hv = silu_f(g) * u_;
            uint32_t hb = f2bf_bits(hv);
            xh[i] = (uint16_t)hb;
            xl[i] = (uint16_t)f2bf_bits(hv - bf2f(hb));
        }
    }
    __syncthreads();
    for (int s = s_lo; s <= s_hi; s++) {
        build_xfrag(xs + (size_t)(s - s_lo) * 2 * MI, MI, xf + (size_t)(s - s_lo) * 2 * LO_STRIDE);
        build_xfrag(xs + (size_t)(s - s_lo) * 2 * MI + MI, MI, xf + (size_t)(s - s_lo) * 2 * LO_STRIDE + LO_STRIDE);
    }
    warp_range(ub, ue, wb, we);
    const uint2* xf_t = xf + (lane & 3);
    auto dec = [&](int u) -> UnitDesc {
        UnitDesc d;
        int slot = u / UPS, r = u - slot * UPS;
        int m = r >> 3, g = r & 7;
        int e = tk_idx[slot];
        size_t unit = (size_t)e * UPS + r;
        d.ptr = W.wD + unit * UNIT_BYTES;
        d.key = slot * NTD + m;
        d.xoff = (slot - s_lo) * 2 * (MI / 16) + g * 8;
        return d;
    };
    auto flush = [&](int key, float a0, float a1) {
        int slot = key / NTD, m = key - slot * NTD;
        float wgt = tk_w[slot];
        int t = lane & 3, g = lane >> 2;
        if (t == 0) red_add_f32(acc_dn + m * 16 + g, a0 * wgt);
        else if (t == 1) red_add_f32(acc_dn + m * 16 + g + 8, a1 * wgt);
    };
    stage_wait(st);
    gemv_units_stage<GEMV_UR, decltype(dec), decltype(flush), true>(wb, we, xf_t, dec, flush, st, LO_STRIDE);
}

// ------------------------------------------------------------------ MLA
// M2: unpack q, rope, cache append (CTA 0), q_abs via transposed kv_b K-part.
DEVI void mla_qabs(const Params& P, const LayerW& W, const LayerS& S, const float* acc_a, uint16_t* qb, uint8_t* sm) {
    constexpr int U = 32 * 32;
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    int ub, ue;
    cta_range(U, ub, ue);
    const int h_lo = ub >> 5, h_hi = (ue - 1) >> 5;      // h_hi - h_lo <= 1
    float* qn = reinterpret_cast<float*>(sm + SM_XF);    // [2][128]
    float* cs = qn + 256;                                 // [32] cos, [32] sin
    const int L = P.L;
    // rope tables
    if (tid < 32) {
        float e = (float)(2 * tid) / 64.f;
        float inv = 1.0f / powf(10000.0f, e);
        float ang = (float)L * inv;
        float sn, cn;
        sincosf(ang, &sn, &cn);
        cs[tid] = cn; cs[32 + tid] = sn;
    }
    __syncthreads();
    // q_nope for needed heads (bf16-rounded q)
    for (int hh = h_lo; hh <= h_hi; hh++) {
        for (int i = tid; i < 128; i += NTHREADS) qn[(hh - h_lo) * 128 + i] = round_bf16(ldcg_f32(acc_a + hh * 192 + i));
    }
    // q_rope for heads whose first unit is in this CTA
    for (int hh = h_lo; hh <= h_hi; hh++) {
        if (hh * 32 >= ub && hh * 32 < ue) {
            if (tid < 32) {
                float e_ = round_bf16(ldcg_f32(acc_a + hh * 192 + 128 + 2 * tid));
                float o_ = round_bf16(ldcg_f32(acc_a + hh * 192 + 128 + 2 * tid + 1));
                float c = cs[tid], s = cs[32 + tid];
                float re = round_bf16(e_ * c - o_ * s);
                float ro = round_bf16(o_ * c + e_ * s);
                qb[hh * 576 + 512 + 2 * tid] = (uint16_t)f2bf_bits(re); qb[hh * 576 + 512 + 2 * tid + 1] = (uint16_t)f2bf_bits(ro);
            }
        }
    }
    // cache append: c_kv_new = bf16(acc_kv[0:512]); k_rope_new = rope(bf16(acc_kv[512:576]))
    if (blockIdx.x == 0) {
        const float* kv = acc_a + 6144;
        for (int i = tid; i < 512; i += NTHREADS) S.ckv_dst[(size_t)L * 512 + i] = (uint16_t)f2bf_bits(ldcg_f32(kv + i));
        if (tid < 32) {
            float e_ = round_bf16(ldcg_f32(kv + 512 + 2 * tid));
            float o_ = round_bf16(ldcg_f32(kv + 512 + 2 * tid + 1));
            float c = cs[tid], s = cs[32 + tid];
            S.krope_dst[(size_t)L * 64 + 2 * tid] = (uint16_t)f2bf_bits(e_ * c - o_ * s);
            S.krope_dst[(size_t)L * 64 + 2 * tid + 1] = (uint16_t)f2bf_bits(o_ * c + e_ * s);
        }
    }
    __syncthreads();
    // absorbed q: one unit per warp
    int u = ub + warp;
    if (u < ue) {
        const int h = u >> 5, mc = u & 31, g = mc >> 3;
        const int t = lane & 3, gq = lane >> 2;
        const uint8_t* up = W.wKb + (size_t)u * 1024;
        uint4 w0 = ldg_stream16(up + lane * 16), w1 = ldg_stream16(up + 512 + lane * 16);
        const float* qh = qn + (h - h_lo) * 128;
        const uint32_t* sz = W.szKb + (size_t)h * 512 + g * 128;
        uint32_t b0[8], b1[8];
        float Cc = 0.f;
#pragma unroll
        for (int s = 0; s < 8; s++) {
            int d = 16 * s + 2 * t;
            uint2 z01 = *reinterpret_cast<const uint2*>(sz + d);
            uint2 z89 = *reinterpret_cast<const uint2*>(sz + d + 8);
            float sc0 = bf2f_lo(z01.x), zz0 = bf2f_hi(z01.x), sc1 = bf2f_lo(z01.y), zz1 = bf2f_hi(z01.y);
            float sc2 = bf2f_lo(z89.x), zz2 = bf2f_hi(z89.x), sc3 = bf2f_lo(z89.y), zz3 = bf2f_hi(z89.y);
            uint32_t p01 = pack_bf16x2(sc0 * qh[d], sc1 * qh[d + 1]);
            uint32_t p89 = pack_bf16x2(sc2 * qh[d + 8], sc3 * qh[d + 9]);
            b0[s] = p01; b1[s] = p89;
            Cc += zz0 * bf2f_lo(p01) + zz1 * bf2f_hi(p01) + zz2 * bf2f_lo(p89) + zz3 * bf2f_hi(p89);
        }
        Cc += __shfl_xor_sync(0xffffffffu, Cc, 1);
        Cc += __shfl_xor_sync(0xffffffffu, Cc, 2);
        float c[4] = {0.f, 0.f, 0.f, 0.f};
        const uint32_t words[8] = {w0.x, w0.y, w0.z, w0.w, w1.x, w1.y, w1.z, w1.w};
#pragma unroll
        for (int s = 0; s < 8; s++) {
            uint32_t w = words[s];
            uint32_t a0 = bf16x2_sub((w & 0x000F000Fu) | 0x43004300u, 0x43084308u);
            uint32_t a1 = bf16x2_sub(((w >> 4) & 0x000F000Fu) | 0x43004300u, 0x43084308u);
            uint32_t a2 = bf16x2_sub(((w >> 8) & 0x000F000Fu) | 0x43004300u, 0x43084308u);
            uint32_t a3 = bf16x2_sub(((w >> 12) & 0x000F000Fu) | 0x43004300u, 0x43084308u);
            mma_bf16_16816(c, a0, a1, a2, a3, b0[s], b1[s]);
        }
        if (t == 0) {
            qb[h * 576 + mc * 16 + gq] = (uint16_t)f2bf_bits(c[0] + Cc);
            qb[h * 576 + mc * 16 + gq + 8] = (uint16_t)f2bf_bits(c[2] + Cc);
        }
    }
}

// M5: combine partials + Wv absorb. CTAs 0..127: head h=c/4, c-group g=c%4.
DEVI void mla_combine_wv(const LayerW& W, const float* part, int na, float* acc_ov, uint8_t* sm) {
    if (blockIdx.x >= NH * 4) return;
    const int h = blockIdx.x >> 2, g = blockIdx.x & 3;
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    float* wi = reinterpret_cast<float*>(sm + SM_XF + 512);    // [188]
    float* rd = wi + 192;                                       // [4][128] partial sums
    float* misc = rd + 512;                                     // [2]
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
    float* red = reinterpret_cast<float*>(sm + SM_RED);
    // global max
    float mloc = -INFINITY;
    for (int i = tid; i < na; i += NTHREADS) mloc = fmaxf(mloc, ldcg_f32(part + (size_t)i * PART_STRIDE + 16384 + h));
    mloc = warp_max(mloc);
    if (lane == 0) red[warp] = mloc;
    __syncthreads();
    float M = -INFINITY;
    for (int i = 0; i < NWARPS; i++) M = fmaxf(M, red[i]);
    __syncthreads();
    float zloc = 0.f;
    for (int i = tid; i < na; i += NTHREADS) {
        float mi = ldcg_f32(part + (size_t)i * PART_STRIDE + 16384 + h);
        float li = ldcg_f32(part + (size_t)i * PART_STRIDE + 16416 + h);
        float w = expf(mi - M);
        wi[i] = w;
        zloc += w * li;
    }
    zloc = warp_sum(zloc);
    if (lane == 0) red[warp] = zloc;
    __syncthreads();
    float Z = 0.f;
    for (int i = 0; i < NWARPS; i++) Z += red[i];
    // o_lat slice
    {
        const int cl = tid & 127, pt = tid >> 7;
        float a = 0.f;
        for (int i = pt; i < na; i += 4) a += wi[i] * ldcg_f32(part + (size_t)i * PART_STRIDE + h * 512 + g * 128 + cl);
        rd[pt * 128 + cl] = a;
    }
    __syncthreads();
    if (tid < 128) {
        float v = (rd[tid] + rd[128 + tid] + rd[256 + tid] + rd[384 + tid]) / Z;
        xs[tid] = (uint16_t)f2bf_bits(v);
    }
    __syncthreads();
    build_xfrag(xs, 128, xf);
    if (warp < 8) {
        const int m = warp;
        const uint8_t* up = W.wVb + (size_t)((h * 8 + m) * 4 + g) * UNIT_BYTES;
        WUnit wu = load_unit(up, lane);
        float ca[4] = {0.f, 0.f, 0.f, 0.f}, cb[4] = {0.f, 0.f, 0.f, 0.f};
        unit_mma(wu, xf + (lane & 3), ca, cb);
        const float v0 = ca[0] + cb[0], v1 = ca[2] + cb[2];
        int t = lane & 3, gq = lane >> 2;
        if (t == 0) red_add_f32(acc_ov + h * 128 + m * 16 + gq, v0);
        else if (t == 1) red_add_f32(acc_ov + h * 128 + m * 16 + gq + 8, v1);
    }
}

// Accumulator regions are self-cleaning: a layer's ACC block is zeroed (by all CTAs) once nothing can still read it,
// so the next launch starts with all-zero accumulators and no initial zeroing barrier is needed.
DEVI void zero_acc_block(float* scr, int layer, int skip_dn) {
    float4* a4 = reinterpret_cast<float4*>(scr + layer * ACC_L);
    constexpr int N4 = ACC_L / 4;
    constexpr int DN0 = ACC_DN / 4, DN1 = (ACC_DN + D) / 4;
    for (int i = blockIdx.x * NTHREADS + threadIdx.x; i < N4; i += gridDim.x * NTHREADS) {
        if (skip_dn && i >= DN0 && i < DN1) continue;
        a4[i] = make_float4(0.f, 0.f, 0.f, 0.f);
    }
}

// ------------------------------------------------------------------ staging targets
template <int NS>
DEVI void stage_contig(WarpStage& st, int U, const uint8_t* base) {
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    warp_range(ub, ue, wb, we);
    stage_issue<NS>(st, wb, we, [&](int u) { return base + (size_t)u * UNIT_BYTES; });
}
DEVI void stage_e2(WarpStage& st, const LayerW& W, const int* tk_idx) {
    constexpr int UPS = NTD * 8, U = NSLOT * UPS;
    int ub, ue, wb, we;
    cta_range(U, ub, ue);
    warp_range(ub, ue, wb, we);
    stage_issue<NS_E2>(st, wb, we, [&](int u) {
        const int slot = u / UPS, r = u - slot * UPS;
        return W.wD + ((size_t)tk_idx[slot] * UPS + r) * UNIT_BYTES;
    });
}

// ------------------------------------------------------------------ the kernel
__global__ void __launch_bounds__(NTHREADS, 1) mega_kernel(const __grid_constant__ Params P) {
    extern __shared__ __align__(16) uint8_t sm[];
    unsigned bar_t = P.bar_base;
    const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5;
    float* scr = P.scratch;
    // per-warp staging slot + mbarrier
    WarpStage st;
    st.ptr = sm + SM_STAGE + warp * (NSU * UNIT_BYTES);
    st.base = smem_u32(st.ptr);
    st.mbar = smem_u32(sm + SM_MBAR + warp * 8);
    st.parity = 0;
    st.n = 0;
    if (lane == 0) mbar_init(st.mbar, 1);
    fence_proxy_async();
    __syncwarp();
    stage_contig<NS_K1>(st, 4 * 256 * KGD, P.w[0].wA);           // layer-0 K1 weights start streaming immediately
    uint16_t* xin = reinterpret_cast<uint16_t*>(sm + SM_XIN);
    uint16_t* hs = reinterpret_cast<uint16_t*>(sm + SM_HS);
    uint16_t* xs = reinterpret_cast<uint16_t*>(sm + SM_XS);
    int* tk_idx = reinterpret_cast<int*>(sm + SM_TK);
    float* tk_w = reinterpret_cast<float*>(sm + SM_TK + 64);
#define PHASE_END() grid_barrier(P.bar, gridDim.x, bar_t)
#define PHASE_END_S(STMT) do { barrier_arrive(P.bar, gridDim.x, bar_t); STMT; barrier_wait(P.bar, bar_t); } while (0)
    const int na = min_i(NAMAX, min_i((int)gridDim.x, (P.L + 1 + attn::RT - 1) / attn::RT));
    for (int layer = 0; layer < 4; layer++) {
        const LayerW& W = P.w[layer];
        const LayerS& S = P.s[layer];
        float* acc = scr + layer * ACC_L;
        float* nz = scr + NZ_BASE + layer * NZ_L;
        // ---- layer input
        if (layer == 0) {
            for (int i = tid; i < D; i += NTHREADS) xin[i] = __ldg(P.hidden_in + i);
        } else {
            const float* prev_dn = scr + (layer - 1) * ACC_L + ACC_DN;
            for (int i = tid; i < D; i += NTHREADS) xin[i] = (uint16_t)f2bf_bits(bf2f(hs[i]) + round_bf16(ldcg_f32(prev_dn + i)));
        }
        __syncthreads();
        if (layer == 2) zero_acc_block(scr, 0, 0);
        if (layer == 3) zero_acc_block(scr, 1, 0);
        if (layer < 3) {
            // ---- KDA
            kda_in_proj(P, layer, W, acc + ACC_A, nz + NZ_BETA, sm, st);
            float s_pre[8];
            kda_state_preload(S, s_pre);                             // state rows load while the barrier resolves
            PHASE_END();
            kda_state_update(W, S, acc + ACC_A, nz + NZ_BETA, reinterpret_cast<uint16_t*>(nz + NZ_OBUF), sm, s_pre);
            PHASE_END_S(stage_contig<NS_O>(st, NTD * 32, W.wO); kda_state_store(S, s_pre));   // o_proj weights + S write-back while the barrier resolves
            {
                const uint16_t* ob = reinterpret_cast<const uint16_t*>(nz + NZ_OBUF);
                for (int i = tid; i < CH / 2; i += NTHREADS) reinterpret_cast<uint32_t*>(xs)[i] = ldcg_u32(ob + 2 * i);
                __syncthreads();
                oproj_gemv(W.wO, acc + ACC_O, sm, st);
                kda_conv_state_update(S, acc + ACC_A);
            }
            PHASE_END();
        } else {
            // ---- MLA
            {
                float* red = reinterpret_cast<float*>(sm + SM_RED);
                uint2* xf = reinterpret_cast<uint2*>(sm + SM_XF);
                            rmsnorm_block(xin, W.attn_norm, xs, red);
                build_xfrag(xs, D, xf);
                constexpr int U = (384 + 36) * KGD;
                int ub, ue, wb, we;
                cta_range(U, ub, ue);
                warp_range(ub, ue, wb, we);
                const uint2* xf_t = xf + (lane & 3);
                float* acc_a = acc + ACC_A;
                auto dec = [&](int u) -> UnitDesc {
                    UnitDesc d;
                    int mg = u / KGD, g = u - mg * KGD;
                    d.ptr = W.wA + (size_t)u * UNIT_BYTES; d.key = mg; d.xoff = g * 8;
                    return d;
                };
                auto flush = [&](int key, float a0, float a1) {
                    int t = lane & 3, g = lane >> 2;
                    if (t == 0) red_add_f32(acc_a + key * 16 + g, a0);
                    else if (t == 1) red_add_f32(acc_a + key * 16 + g + 8, a1);
                };
                stage_wait(st);
                gemv_units_stage<GEMV_UR>(wb, we, xf_t, dec, flush, st);
            }
            PHASE_END();
            zero_acc_block(scr, 2, 0);                                   // layer-2 residual was consumed before the M1 barrier
            mla_qabs(P, W, S, acc + ACC_A, reinterpret_cast<uint16_t*>(nz + NZ_QABS), sm);
            PHASE_END();
            if (blockIdx.x < na) {
                attn::Args AA;
                AA.qb = reinterpret_cast<const uint16_t*>(nz + NZ_QABS);
                AA.ckv_src = S.ckv_src; AA.krope_src = S.krope_src; AA.ckv_dst = S.ckv_dst; AA.krope_dst = S.krope_dst;
                AA.L = P.L; AA.migrate = P.migrate; AA.scale = 0.07216878364870322f;
                const int ntok = P.L + 1;
                const int t0 = (int)((long long)ntok * blockIdx.x / na), t1 = (int)((long long)ntok * (blockIdx.x + 1) / na);
                attn::run(AA, scr + PART_BASE + (size_t)blockIdx.x * PART_STRIDE, t0, t1, sm + SM_XS);
            }
            __syncthreads();                                         // attention smem is dead from here on
            stage_contig<NS_O>(st, NTD * 32, W.wO);                        // (stage region overlaps attention smem, so issue after it)
            PHASE_END();
            mla_combine_wv(W, scr + PART_BASE, na, acc + ACC_OV, sm);
            PHASE_END();
            {
                for (int i = tid; i < CH; i += NTHREADS) xs[i] = (uint16_t)f2bf_bits(ldcg_f32(acc + ACC_OV + i));
                __syncthreads();
                oproj_gemv(W.wO, acc + ACC_O, sm, st);
            }
            PHASE_END();
        }
        // ---- MoE
        moe_pre(W, acc + ACC_O, nz + NZ_LOGITS, sm);
        PHASE_END();
        if (tid < 32) router_topk(nz + NZ_LOGITS, tk_idx, tk_w, reinterpret_cast<float*>(sm + SM_RED));
        __syncthreads();
        moe_gate_up(W, acc + ACC_GU, sm);
        PHASE_END_S(stage_e2(st, W, tk_idx));                        // down-proj weights stream while the barrier resolves
        moe_down(W, acc + ACC_GU, acc + ACC_DN, sm, st);
        PHASE_END_S(if (layer < 2) stage_contig<NS_K1>(st, 4 * 256 * KGD, P.w[layer + 1].wA);       // next layer's K1 weights
                    else if (layer == 2) stage_contig<NS_K1>(st, (384 + 36) * KGD, P.w[3].wA));       // MLA q/kv_a weights
    }
    // ---- final output: out = bf16(h + bf16(acc_dn)); each element is read (and re-zeroed) by exactly one CTA thread
    {
        float* dn = scr + 3 * ACC_L + ACC_DN;
        zero_acc_block(scr, 3, 1);
        for (int i = blockIdx.x * NTHREADS + tid; i < D; i += gridDim.x * NTHREADS) {
            P.hidden_out[i] = (uint16_t)f2bf_bits(bf2f(hs[i]) + round_bf16(ldcg_f32(dn + i)));
            stcg_f32(dn + i, 0.f);
        }
    }
#undef PHASE_END
}

}  // namespace mk

// ==== mk_host.cu ====
// Host-side wrappers for the megakernel (built via load_inline / torch cpp_extension).
#include <vector>
#include <cstdio>
#include <cstring>
#include <c10/cuda/CUDAStream.h>

struct ModelState { mk::Params p; unsigned epoch; };
static std::vector<ModelState> g_models;      // one entry per Model instance (own weights, scratch, barrier counter)
static int g_nsm = 0;
constexpr int NBARRIERS = 26;      // grid barriers per launch (must match mk.cuh)
int64_t scratch_floats() { return mk::SCRATCH_FLOATS; }

void repack_normal(int64_t w_q, int64_t sc, int64_t zr, int64_t dst, int64_t K, int64_t N, int64_t Nsrc, int64_t col_off,
                   int64_t batch, int64_t src_bs, int64_t src_sz_bs, int64_t col_off_batch, int64_t dst_bs) {
    long long units = (N / 16) * (K / 128);
    long long thr = units * 256;
    repack_normal_kernel<<<dim3((unsigned)((thr + 255) / 256), (unsigned)batch), 256, 0, c10::cuda::getCurrentCUDAStream()>>>(
        (const uint8_t*)w_q, (const uint16_t*)sc, (const uint16_t*)zr, (uint8_t*)dst, (int)K, (int)N, (int)Nsrc, (int)col_off,
        src_bs, src_sz_bs, (int)col_off_batch, dst_bs);
}

void repack_kbt(int64_t w_q, int64_t sc, int64_t zr, int64_t dst, int64_t szdst, int64_t Nsrc) {
    repack_kbt_kernel<<<dim3(32, 32), 256, 0, c10::cuda::getCurrentCUDAStream()>>>(
        (const uint8_t*)w_q, (const uint16_t*)sc, (const uint16_t*)zr, (uint8_t*)dst, (uint32_t*)szdst, (int)Nsrc);
}

// v: per layer 12 pointers (wA, wO, beta_w, conv_w, wKb, szKb, wVb, attn_norm, moe_norm, router_w, wGU, wD) x4, then scratch, bar
// Returns a handle used by run().
int64_t make_static(std::vector<int64_t> v) {
    ModelState ms;
    memset(&ms, 0, sizeof(ms));
    size_t i = 0;
    for (int l = 0; l < 4; l++) {
        mk::LayerW& w = ms.p.w[l];
        w.wA = (const uint8_t*)v[i++]; w.wO = (const uint8_t*)v[i++];
        w.beta_w = (const uint16_t*)v[i++]; w.conv_w = (const uint16_t*)v[i++];
        w.wKb = (const uint8_t*)v[i++]; w.szKb = (const uint32_t*)v[i++]; w.wVb = (const uint8_t*)v[i++];
        w.attn_norm = (const uint16_t*)v[i++]; w.moe_norm = (const uint16_t*)v[i++]; w.router_w = (const uint16_t*)v[i++];
        w.wGU = (const uint8_t*)v[i++]; w.wD = (const uint8_t*)v[i++];
    }
    ms.p.scratch = (float*)v[i++];
    ms.p.bar = (unsigned*)v[i++];
    ms.epoch = 0;
    if (g_nsm == 0) {
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, 0);
        g_nsm = prop.multiProcessorCount;
        cudaFuncSetAttribute(mk::mega_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, mk::SMEM_BYTES);
    }
    g_models.push_back(ms);
    return (int64_t)g_models.size() - 1;
}

// v: per layer 8 pointers (S, cq, ck, cv, ckv_src, krope_src, ckv_dst, krope_dst) x4, hidden_in, hidden_out, L, migrate
void run(int64_t handle, std::vector<int64_t> v) {
    ModelState& ms = g_models[handle];
    mk::Params P = ms.p;
    size_t i = 0;
    for (int l = 0; l < 4; l++) {
        mk::LayerS& s = P.s[l];
        s.S = (float*)v[i++]; s.cq = (uint16_t*)v[i++]; s.ck = (uint16_t*)v[i++]; s.cv = (uint16_t*)v[i++];
        s.ckv_src = (const uint16_t*)v[i++]; s.krope_src = (const uint16_t*)v[i++];
        s.ckv_dst = (uint16_t*)v[i++]; s.krope_dst = (uint16_t*)v[i++];
    }
    P.hidden_in = (const uint16_t*)v[i++];
    P.hidden_out = (uint16_t*)v[i++];
    P.L = (int)v[i++];
    P.migrate = (int)v[i++];
    P.bar_base = ms.epoch;
    ms.epoch += (unsigned)NBARRIERS * (unsigned)g_nsm;
    mk::mega_kernel<<<g_nsm, mk::NTHREADS, mk::SMEM_BYTES, c10::cuda::getCurrentCUDAStream()>>>(P);
}

'''


_EXT = None


def _get_ext():
    global _EXT
    if _EXT is None:
        from torch.utils.cpp_extension import load_inline
        cpp = """
#include <vector>
#include <cstdint>
int64_t scratch_floats();
void repack_normal(int64_t w_q, int64_t sc, int64_t zr, int64_t dst, int64_t K, int64_t N, int64_t Nsrc, int64_t col_off,
                   int64_t batch, int64_t src_bs, int64_t src_sz_bs, int64_t col_off_batch, int64_t dst_bs);
void repack_kbt(int64_t w_q, int64_t sc, int64_t zr, int64_t dst, int64_t szdst, int64_t Nsrc);
int64_t make_static(std::vector<int64_t> v);
void run(int64_t handle, std::vector<int64_t> v);
"""
        _EXT = load_inline(
            name="kimi_linear_megakernel", cpp_sources=cpp, cuda_sources=_CUDA_SRC,
            functions=["scratch_floats", "repack_normal", "repack_kbt", "make_static", "run"],
            extra_cuda_cflags=["-O3", "-std=c++20"], verbose=False)
    return _EXT


GROUP_SIZE = 128
EPS = 1.0e-6

@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)))


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


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


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


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


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


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


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


UNIT = 1088   # bytes per packed weight unit (see mk_common.cuh)


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        if (cfg.hidden, cfg.kda_heads, cfg.kda_head_dim, cfg.mla_heads, cfg.kv_lora, cfg.qk_nope, cfg.qk_rope, cfg.v_head,
                cfg.n_experts, cfg.n_active, cfg.n_shared, cfg.moe_inter, cfg.group, tuple(cfg.pattern)) != \
                (2304, 32, 128, 32, 512, 128, 64, 128, 64, 8, 1, 1024, 128, ("K", "K", "K", "M")):
            raise NotImplementedError("the megakernel is specialized to the Kimi-Linear-48B-A3B motif shapes of this benchmark")
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._packed = None
        self.register_load_state_dict_post_hook(self._after_load)

    @staticmethod
    def _after_load(module, incompatible_keys):
        # weights changed: drop the packed copy and rebuild it right away (not inside step) when the model is on the GPU
        module._packed = None
        if module.blocks[0].attn_norm.is_cuda:
            module._pack()

    # ---------------------------------------------------------------- weight repack (one time per weight load)
    def _pack(self):
        ext = _get_ext()
        dev = self.blocks[0].attn_norm.device
        keep = []
        static = []

        def pk(mod, dst, dst_off, K, N, Nsrc=None, col_off=0):
            ext.repack_normal(mod.w_q.data_ptr(), mod.scales.data_ptr(), mod.zeros.data_ptr(), dst.data_ptr() + dst_off,
                              K, N, Nsrc or N, col_off, 1, 0, 0, 0, 0)

        for l, blk in enumerate(self.blocks):
            a, m = blk.attn, blk.moe
            if blk.kind == "K":
                wA = torch.empty(4 * 256 * 18 * UNIT, dtype=torch.uint8, device=dev)
                for i, mod in enumerate([a.q_proj, a.k_proj, a.v_proj, a.g_proj]):
                    pk(mod, wA, i * 256 * 18 * UNIT, 2304, 4096)
                wO = torch.empty(144 * 32 * UNIT, dtype=torch.uint8, device=dev)
                pk(a.o_proj, wO, 0, 4096, 2304)
                keep += [wA, wO]
                fields = [wA.data_ptr(), wO.data_ptr(), a.beta_proj.weight.data_ptr(), a.conv_w.data_ptr(), 0, 0, 0]
            else:
                wA = torch.empty((384 + 36) * 18 * UNIT, dtype=torch.uint8, device=dev)
                pk(a.q_proj, wA, 0, 2304, 6144)
                pk(a.kv_a, wA, 384 * 18 * UNIT, 2304, 576)
                wO = torch.empty(144 * 32 * UNIT, dtype=torch.uint8, device=dev)
                pk(a.o_proj, wO, 0, 4096, 2304)
                wKb = torch.empty(32 * 32 * 1024, dtype=torch.uint8, device=dev)
                szKb = torch.empty(32 * 512, dtype=torch.int32, device=dev)
                ext.repack_kbt(a.kv_b.w_q.data_ptr(), a.kv_b.scales.data_ptr(), a.kv_b.zeros.data_ptr(), wKb.data_ptr(), szKb.data_ptr(), 8192)
                wVb = torch.empty(32 * 8 * 4 * UNIT, dtype=torch.uint8, device=dev)
                ext.repack_normal(a.kv_b.w_q.data_ptr(), a.kv_b.scales.data_ptr(), a.kv_b.zeros.data_ptr(), wVb.data_ptr(),
                                  512, 128, 8192, 128, 32, 0, 0, 256, 8 * 4 * UNIT)
                keep += [wA, wO, wKb, szKb, wVb]
                fields = [wA.data_ptr(), wO.data_ptr(), 0, 0, wKb.data_ptr(), szKb.data_ptr(), wVb.data_ptr()]
            # MoE experts: 65 slots (64 routed + shared)
            per_mat = 1152 * UNIT
            wGU = torch.empty(65 * 2 * per_mat, dtype=torch.uint8, device=dev)
            wD = torch.empty(65 * per_mat, dtype=torch.uint8, device=dev)
            for mat, (ex, exs) in enumerate([(m.gate, m.s_gate), (m.up, m.s_up)]):
                ext.repack_normal(ex.w_q.data_ptr(), ex.scales.data_ptr(), ex.zeros.data_ptr(), wGU.data_ptr() + mat * per_mat,
                                  2304, 1024, 1024, 0, 64, ex.w_q[0].numel(), ex.scales[0].numel(), 0, 2 * per_mat)
                ext.repack_normal(exs.w_q.data_ptr(), exs.scales.data_ptr(), exs.zeros.data_ptr(), wGU.data_ptr() + 64 * 2 * per_mat + mat * per_mat,
                                  2304, 1024, 1024, 0, 1, 0, 0, 0, 0)
            ext.repack_normal(m.down.w_q.data_ptr(), m.down.scales.data_ptr(), m.down.zeros.data_ptr(), wD.data_ptr(),
                              1024, 2304, 2304, 0, 64, m.down.w_q[0].numel(), m.down.scales[0].numel(), 0, per_mat)
            ext.repack_normal(m.s_down.w_q.data_ptr(), m.s_down.scales.data_ptr(), m.s_down.zeros.data_ptr(), wD.data_ptr() + 64 * per_mat,
                              1024, 2304, 2304, 0, 1, 0, 0, 0, 0)
            keep += [wGU, wD]
            fields += [blk.attn_norm.data_ptr(), blk.moe_norm.data_ptr(), m.router.weight.data_ptr(), wGU.data_ptr(), wD.data_ptr()]
            static += fields
        scratch = torch.zeros(ext.scratch_floats(), dtype=torch.float32, device=dev)
        bar = torch.zeros(4, dtype=torch.int32, device=dev)
        static += [scratch.data_ptr(), bar.data_ptr()]
        handle = ext.make_static(static)
        torch.cuda.synchronize()
        self._packed = dict(keep=keep, scratch=scratch, bar=bar, handle=handle)

    # ---------------------------------------------------------------- one decode step (a single kernel launch)
    def step(self, hidden, state):
        if self._packed is None:
            self._pack()
        ext = _get_ext()
        dev = hidden.device
        if hidden.dtype != torch.bfloat16 or not hidden.is_contiguous():
            hidden = hidden.to(torch.bfloat16).contiguous()
        v = []
        for l, kind in enumerate(self.cfg.pattern):
            st = state[l]
            if kind == "K":
                S, cq, ck, cv = st["S"], st["cq"], st["ck"], st["cv"]
                if S.dtype != torch.float32 or not S.is_contiguous():
                    S = st["S"] = S.float().contiguous()
                for key, cv_ in (("cq", cq), ("ck", ck), ("cv", cv)):
                    if cv_.dtype != torch.bfloat16 or not cv_.is_contiguous():
                        st[key] = cv_.to(torch.bfloat16).contiguous()
                v += [S.data_ptr(), st["cq"].data_ptr(), st["ck"].data_ptr(), st["cv"].data_ptr(), 0, 0, 0, 0]
            else:
                ckv, kr = st["c_kv"], st["k_rope"]
                if ckv.dtype != torch.bfloat16 or not ckv.is_contiguous():
                    ckv = ckv.to(torch.bfloat16).contiguous()
                if kr.dtype != torch.bfloat16 or not kr.is_contiguous():
                    kr = kr.to(torch.bfloat16).contiguous()
                L = ckv.shape[0]
                need = L + 1
                if (ckv.storage_offset() == 0 and kr.storage_offset() == 0
                        and ckv.untyped_storage().nbytes() >= need * 1024 and kr.untyped_storage().nbytes() >= need * 128):
                    nckv, nkr, migrate = ckv, kr, 0                       # append in place
                else:                                                      # hand the cache over to a capacity buffer
                    cap = L + 256
                    nckv = torch.empty((cap, 512), dtype=torch.bfloat16, device=dev)
                    nkr = torch.empty((cap, 64), dtype=torch.bfloat16, device=dev)
                    migrate = 1
                v += [0, 0, 0, 0, ckv.data_ptr(), kr.data_ptr(), nckv.data_ptr(), nkr.data_ptr()]
                st["c_kv"] = nckv.as_strided((need, 512), (512, 1))
                st["k_rope"] = nkr.as_strided((need, 64), (64, 1))
        out = torch.empty_like(hidden)
        v += [hidden.data_ptr(), out.data_ptr(), L, migrate]
        ext.run(self._packed["handle"], v)
        return out, state