"""Fused W4A16 Kimi-Linear Hybrid Decode Megakernel.

Decodes batch-1 autoregressively using a single cooperative GPU kernel launch
per decode step. Fuses all 4 layers (3 KDA + 1 MLA), all int4 dequant-GEMVs,
short causal depthwise conv, KDA recurrence, MLA absorbed latent-cache attention,
MoE top-8 routing, expert SwiGLU GEMVs, RMSNorms, and residual additions.
"""
from __future__ import annotations

import math
import os
from dataclasses import dataclass, field

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

# Ensure ninja is in PATH if available in venv
for p in ("/[REDACTED-home]/kb-mega/outputs/runs/20260902_210053_agy_gemini-3.8-flash-high_02_kimi_linear_decode/repo/.venv/bin", "/[REDACTED-home]/kb-mega/.venv/bin"):
    if os.path.exists(p) and p not in os.environ.get("PATH", ""):
        os.environ["PATH"] = f"{p}:{os.environ.get('PATH', '')}"

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


# --------------------------------------------------------------------------- #
# Module Definitions Matching Reference State Dict Layout
# --------------------------------------------------------------------------- #
class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE):
        super().__init__()
        self.in_f, self.out_f, self.group = in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16))


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


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


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


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


class Block(nn.Module):
    def __init__(self, cfg: Config, kind: str):
        super().__init__()
        self.kind = kind
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
        self.moe = MoE(cfg)


# --------------------------------------------------------------------------- #
# CUDA Megakernel Source
# --------------------------------------------------------------------------- #
_CUDA_SRC = r'''
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <math.h>

namespace cg = cooperative_groups;

__inline__ __device__ float silu(float x) {
    return x / (1.0f + expf(-x));
}

__inline__ __device__ float softplus(float x) {
    if (x > 20.0f) return x;
    return log1pf(expf(x));
}

__inline__ __device__ float sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

struct BlockPointers {
    int is_mla;
    int pad;

    const __nv_bfloat16* attn_norm;
    const __nv_bfloat16* moe_norm;

    // KDA weights
    const uint8_t* q_wq;
    const __nv_bfloat16* q_s;
    const __nv_bfloat16* q_z;

    const uint8_t* k_wq;
    const __nv_bfloat16* k_s;
    const __nv_bfloat16* k_z;

    const uint8_t* v_wq;
    const __nv_bfloat16* v_s;
    const __nv_bfloat16* v_z;

    const uint8_t* g_wq;
    const __nv_bfloat16* g_s;
    const __nv_bfloat16* g_z;

    const __nv_bfloat16* beta_w;
    const __nv_bfloat16* conv_w;

    // MLA weights
    const uint8_t* mla_q_wq;
    const __nv_bfloat16* mla_q_s;
    const __nv_bfloat16* mla_q_z;

    const uint8_t* mla_kva_wq;
    const __nv_bfloat16* mla_kva_s;
    const __nv_bfloat16* mla_kva_z;

    const __nv_bfloat16* mla_W_k;
    const __nv_bfloat16* mla_W_v;

    // Output projection
    const uint8_t* o_wq;
    const __nv_bfloat16* o_s;
    const __nv_bfloat16* o_z;

    // MoE weights
    const __nv_bfloat16* router_w;
    const uint8_t* gate_wq;
    const __nv_bfloat16* gate_s;
    const __nv_bfloat16* gate_z;
    const uint8_t* up_wq;
    const __nv_bfloat16* up_s;
    const __nv_bfloat16* up_z;
    const uint8_t* down_wq;
    const __nv_bfloat16* down_s;
    const __nv_bfloat16* down_z;

    // Shared expert
    const uint8_t* s_gate_wq;
    const __nv_bfloat16* s_gate_s;
    const __nv_bfloat16* s_gate_z;
    const uint8_t* s_up_wq;
    const __nv_bfloat16* s_up_s;
    const __nv_bfloat16* s_up_z;
    const uint8_t* s_down_wq;
    const __nv_bfloat16* s_down_s;
    const __nv_bfloat16* s_down_z;
};

struct KernelWorkspace {
    __nv_bfloat16* x_norm;      // [2304]
    __nv_bfloat16* q;           // [6144]
    __nv_bfloat16* k;           // [4096]
    __nv_bfloat16* v;           // [4096]
    __nv_bfloat16* g;           // [4096]
    __nv_bfloat16* beta;        // [32]
    __nv_bfloat16* o;           // [4096]
    __nv_bfloat16* attn_out;    // [2304]
    __nv_bfloat16* kva;         // [576]
    float* q_abs;               // [32, 512]
    float* scores_ws;           // [32, L_max]
    float* o_latent;            // [32, 512]
    int* topk_idx;              // [8]
    float* topk_weights;        // [8]
    float* gate_up;             // [9 * 2048]
    float* h_expert;            // [9 * 1024]
    float* expert_outs;         // [9 * 2304]
    float* moe_out;             // [2304]
};

__inline__ __device__ float compute_gemv_col(
    const __nv_bfloat16* __restrict__ x,
    const uint8_t* __restrict__ w_q,
    const __nv_bfloat16* __restrict__ scales,
    const __nv_bfloat16* __restrict__ zeros,
    int col, int K, int N
) {
    int num_groups = K / 128;
    float total_sum = 0.0f;
    for (int g = 0; g < num_groups; ++g) {
        float scale = __bfloat162float(scales[g * N + col]);
        float zero = __bfloat162float(zeros[g * N + col]);
        float group_acc = 0.0f;
        float group_x_sum = 0.0f;

        
        

        #pragma unroll 8
        for (int p = 0; p < 64; ++p) {
            uint8_t byte = w_q[(g * 64 + p) * N + col];
            
            float u0 = (float)(byte & 0x0F);
            float u1 = (float)(byte >> 4);
            int k0 = g * 128 + 2 * p; float x0 = __bfloat162float(x[k0]);
            float x1 = __bfloat162float(x[k0 + 1]);

            group_acc += x0 * u0 + x1 * u1;
            group_x_sum += x0 + x1;
        }
        total_sum += (group_acc - zero * group_x_sum) * scale;
    }
    return total_sum;
}

__device__ void rmsnorm_block(
    const __nv_bfloat16* __restrict__ x,
    const __nv_bfloat16* __restrict__ w,
    __nv_bfloat16* __restrict__ out,
    int D, cg::grid_group& grid
) {
    if (blockIdx.x == 0) {
        int tid = threadIdx.x;
        int lane = tid % 32;
        int wid = tid / 32;
        __shared__ float smem[32];

        float sum_sq = 0.0f;
        for (int i = tid; i < D; i += blockDim.x) {
            float val = __bfloat162float(x[i]);
            sum_sq += val * val;
        }
        for (int offset = 16; offset > 0; offset /= 2) {
            sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
        }
        if (lane == 0) smem[wid] = sum_sq;
        __syncthreads();

        if (wid == 0) {
            float b_sum = (lane < (blockDim.x / 32)) ? smem[lane] : 0.0f;
            for (int offset = 16; offset > 0; offset /= 2) {
                b_sum += __shfl_down_sync(0xffffffff, b_sum, offset);
            }
            if (lane == 0) {
                float mean = b_sum / (float)D;
                smem[0] = rsqrtf(mean + 1.0e-6f);
            }
        }
        __syncthreads();

        float rscale = smem[0];
        for (int i = tid; i < D; i += blockDim.x) {
            float val = __bfloat162float(x[i]);
            float weight = __bfloat162float(w[i]);
            out[i] = __float2bfloat16(val * rscale * weight);
        }
    }
    grid.sync();
}

__global__ void megakernel(
    __nv_bfloat16* __restrict__ hidden,
    float* S0, __nv_bfloat16* cq0, __nv_bfloat16* ck0, __nv_bfloat16* cv0,
    float* S1, __nv_bfloat16* cq1, __nv_bfloat16* ck1, __nv_bfloat16* cv1,
    float* S2, __nv_bfloat16* cq2, __nv_bfloat16* ck2, __nv_bfloat16* cv2,
    const __nv_bfloat16* __restrict__ old_c_kv,
    __nv_bfloat16* __restrict__ new_c_kv,
    const __nv_bfloat16* __restrict__ old_k_rope,
    __nv_bfloat16* __restrict__ new_k_rope,
    int context_len,
    const BlockPointers* __restrict__ blocks,
    KernelWorkspace ws
) {
    cg::grid_group grid = cg::this_grid();
    int gtid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_threads = gridDim.x * blockDim.x;

    float* S_ptrs[3] = { S0, S1, S2 };
    __nv_bfloat16* cq_ptrs[3] = { cq0, cq1, cq2 };
    __nv_bfloat16* ck_ptrs[3] = { ck0, ck1, ck2 };
    __nv_bfloat16* cv_ptrs[3] = { cv0, cv1, cv2 };

    for (int b = 0; b < 4; ++b) {
        const BlockPointers& bp = blocks[b];

        // 1. Pre-attention RMSNorm
        rmsnorm_block(hidden, bp.attn_norm, ws.x_norm, 2304, grid);

        if (b < 3) {
            // === KDA LAYER ===
            float* cur_S = S_ptrs[b];
            __nv_bfloat16* cur_cq = cq_ptrs[b];
            __nv_bfloat16* cur_ck = ck_ptrs[b];
            __nv_bfloat16* cur_cv = cv_ptrs[b];

            // 1.1 Projections: q, k, v, g, beta (16416 cols)
            for (int col = gtid; col < 16416; col += total_threads) {
                if (col < 4096) {
                    float val = compute_gemv_col(ws.x_norm, bp.q_wq, bp.q_s, bp.q_z, col, 2304, 4096);
                    ws.q[col] = __float2bfloat16(val);
                } else if (col < 8192) {
                    int c = col - 4096;
                    float val = compute_gemv_col(ws.x_norm, bp.k_wq, bp.k_s, bp.k_z, c, 2304, 4096);
                    ws.k[c] = __float2bfloat16(val);
                } else if (col < 12288) {
                    int c = col - 8192;
                    float val = compute_gemv_col(ws.x_norm, bp.v_wq, bp.v_s, bp.v_z, c, 2304, 4096);
                    ws.v[c] = __float2bfloat16(val);
                } else if (col < 16384) {
                    int c = col - 12288;
                    float val = compute_gemv_col(ws.x_norm, bp.g_wq, bp.g_s, bp.g_z, c, 2304, 4096);
                    ws.g[c] = __float2bfloat16(val);
                } else {
                    int h = col - 16384;
                    const __nv_bfloat16* w_row = bp.beta_w + h * 2304;
                    float sum = 0.0f;
                    #pragma unroll 4
                    for (int k = 0; k < 2304; ++k) {
                        sum += __bfloat162float(ws.x_norm[k]) * __bfloat162float(w_row[k]);
                    }
                    ws.beta[h] = __float2bfloat16(sum);
                }
            }
            grid.sync();

            // 1.2 Short Conv on q, k, v (12288 channels)
            for (int c_idx = gtid; c_idx < 12288; c_idx += total_threads) {
                int ch = c_idx % 4096;
                int which = c_idx / 4096;
                __nv_bfloat16* prev = (which == 0) ? cur_cq : ((which == 1) ? cur_ck : cur_cv);
                __nv_bfloat16* cur_val_ptr = (which == 0) ? ws.q : ((which == 1) ? ws.k : ws.v);
                const __nv_bfloat16* conv_w_which = bp.conv_w + which * 4096 * 4;

                float p0 = __bfloat162float(prev[0 * 4096 + ch]);
                float p1 = __bfloat162float(prev[1 * 4096 + ch]);
                float p2 = __bfloat162float(prev[2 * 4096 + ch]);
                float cur = __bfloat162float(cur_val_ptr[ch]);

                float w0 = __bfloat162float(conv_w_which[ch * 4 + 0]);
                float w1 = __bfloat162float(conv_w_which[ch * 4 + 1]);
                float w2 = __bfloat162float(conv_w_which[ch * 4 + 2]);
                float w3 = __bfloat162float(conv_w_which[ch * 4 + 3]);

                float conv_out = p0 * w0 + p1 * w1 + p2 * w2 + cur * w3;
                cur_val_ptr[ch] = __float2bfloat16(silu(conv_out));

                prev[0 * 4096 + ch] = __float2bfloat16(p1);
                prev[1 * 4096 + ch] = __float2bfloat16(p2);
                prev[2 * 4096 + ch] = __float2bfloat16(cur);
            }
            grid.sync();

            // 1.3 KDA Recurrent State Update (Blocks 0..31: 1 block per head)
            if (blockIdx.x < 32) {
                int h = blockIdx.x;
                int tid = threadIdx.x;

                __shared__ float s_q[128];
                __shared__ float s_k[128];
                __shared__ float s_v[128];
                __shared__ float s_g[128];
                __shared__ float s_beta;
                __shared__ float s_diff[128];

                if (tid < 128) {
                    float scale = 0.08838834764831845f;
                    s_q[tid] = __bfloat162float(ws.q[h * 128 + tid]) * scale;
                    s_k[tid] = __bfloat162float(ws.k[h * 128 + tid]);
                    s_v[tid] = __bfloat162float(ws.v[h * 128 + tid]);
                    s_g[tid] = expf(-softplus(__bfloat162float(ws.g[h * 128 + tid])));
                }
                if (tid == 0) {
                    s_beta = sigmoid(__bfloat162float(ws.beta[h]));
                }
                __syncthreads();

                float* S_head = cur_S + h * 128 * 128;

                if (tid < 128) {
                    float g_exp_i = s_g[tid];
                    #pragma unroll 4
                    for (int j = 0; j < 128; ++j) {
                        S_head[tid * 128 + j] *= g_exp_i;
                    }
                }
                __syncthreads();

                if (tid < 128) {
                    float p_j = 0.0f;
                    #pragma unroll 4
                    for (int i = 0; i < 128; ++i) {
                        p_j += s_k[i] * S_head[i * 128 + tid];
                    }
                    s_diff[tid] = s_v[tid] - p_j;
                }
                __syncthreads();

                if (tid < 128) {
                    float beta_k_i = s_beta * s_k[tid];
                    #pragma unroll 4
                    for (int j = 0; j < 128; ++j) {
                        S_head[tid * 128 + j] += beta_k_i * s_diff[j];
                    }
                }
                __syncthreads();

                if (tid < 128) {
                    float o_j = 0.0f;
                    #pragma unroll 4
                    for (int i = 0; i < 128; ++i) {
                        o_j += s_q[i] * S_head[i * 128 + tid];
                    }
                    ws.o[h * 128 + tid] = __float2bfloat16(o_j);
                }
            }
            grid.sync();

            // 1.4 Output Projection: ws.o (4096) -> ws.attn_out (2304)
            for (int col = gtid; col < 2304; col += total_threads) {
                float val = compute_gemv_col(ws.o, bp.o_wq, bp.o_s, bp.o_z, col, 4096, 2304);
                ws.attn_out[col] = __float2bfloat16(val);
            }
            grid.sync();

        } else {
            // === MLA LAYER ===
            // 1.1 Projections: mla_q (6144), kva (576) -> 6720 total cols
            for (int col = gtid; col < 6720; col += total_threads) {
                if (col < 6144) {
                    float val = compute_gemv_col(ws.x_norm, bp.mla_q_wq, bp.mla_q_s, bp.mla_q_z, col, 2304, 6144);
                    ws.q[col] = __float2bfloat16(val);
                } else {
                    int c = col - 6144;
                    float val = compute_gemv_col(ws.x_norm, bp.mla_kva_wq, bp.mla_kva_s, bp.mla_kva_z, c, 2304, 576);
                    ws.kva[c] = __float2bfloat16(val);
                }
            }
            grid.sync();

            int pos = context_len;
            int L = pos + 1;

            // RoPE on q_rope (32 heads x 64 dims = 2048) and k_rope (64 dims)
            for (int idx = gtid; idx < 32 * 64 + 64; idx += total_threads) {
                if (idx < 32 * 64) {
                    int h = idx / 64;
                    int d = idx % 64;
                    if (d % 2 == 0) {
                        int pair = d / 2;
                        float inv = powf(10000.0f, -(2.0f * pair) / 64.0f);
                        float ang = pos * inv;
                        float c_val = cosf(ang);
                        float s_val = sinf(ang);
                        int base_q = h * 192 + 128 + d;
                        float even = __bfloat162float(ws.q[base_q]);
                        float odd = __bfloat162float(ws.q[base_q + 1]);
                        ws.q[base_q] = __float2bfloat16(even * c_val - odd * s_val);
                        ws.q[base_q + 1] = __float2bfloat16(odd * c_val + even * s_val);
                    }
                } else {
                    int d = idx - 32 * 64;
                    if (d % 2 == 0) {
                        int pair = d / 2;
                        float inv = powf(10000.0f, -(2.0f * pair) / 64.0f);
                        float ang = pos * inv;
                        float c_val = cosf(ang);
                        float s_val = sinf(ang);
                        int base_k = 512 + d;
                        float even = __bfloat162float(ws.kva[base_k]);
                        float odd = __bfloat162float(ws.kva[base_k + 1]);
                        ws.kva[base_k] = __float2bfloat16(even * c_val - odd * s_val);
                        ws.kva[base_k + 1] = __float2bfloat16(odd * c_val + even * s_val);
                    }
                }
            }

            // Copy old cache to new cache and append new token
            for (int i = gtid; i < pos * 512; i += total_threads) {
                new_c_kv[i] = old_c_kv[i];
            }
            for (int i = gtid; i < pos * 64; i += total_threads) {
                new_k_rope[i] = old_k_rope[i];
            }
            if (gtid < 512) {
                new_c_kv[pos * 512 + gtid] = ws.kva[gtid];
            }
            if (gtid < 64) {
                new_k_rope[pos * 64 + gtid] = ws.kva[512 + gtid];
            }
            grid.sync();

            // 1.3 MLA Absorbed Attention
            for (int idx = gtid; idx < 32 * 512; idx += total_threads) {
                int h = idx / 512;
                int e = idx % 512;
                const __nv_bfloat16* q_h = ws.q + h * 192;
                const __nv_bfloat16* w_row = bp.mla_W_k + (e * 32 + h) * 128;
                float sum = 0.0f;
                #pragma unroll 4
                for (int d = 0; d < 128; ++d) {
                    sum += __bfloat162float(q_h[d]) * __bfloat162float(w_row[d]);
                }
                ws.q_abs[idx] = sum;
            }
            grid.sync();

            // Head attention (Blocks 0..31: 1 block per head)
            if (blockIdx.x < 32) {
                int h = blockIdx.x;
                int tid = threadIdx.x;
                int wid = tid / 32;
                int lane = tid % 32;

                const float* q_a = ws.q_abs + h * 512;
                const __nv_bfloat16* q_r = ws.q + h * 192 + 128;
                float* my_scores = ws.scores_ws + h * L;
                float mla_scale = 0.07216878364870322f;

                __shared__ float s_qa[512];
                __shared__ float s_qr[64];
                for (int i = tid; i < 512; i += blockDim.x) {
                    s_qa[i] = q_a[i];
                }
                if (tid < 64) {
                    s_qr[tid] = __bfloat162float(q_r[tid]);
                }
                __syncthreads();

                // Warp-parallel score computation across tokens l
                for (int l = wid; l < L; l += 8) {
                    const __nv_bfloat16* c_row = new_c_kv + l * 512;
                    const __nv_bfloat16* kr_row = new_k_rope + l * 64;

                    float dot = 0.0f;
                    #pragma unroll 4
                    for (int p = 0; p < 16; ++p) {
                        int idx = p * 32 + lane;
                        dot += s_qa[idx] * __bfloat162float(c_row[idx]);
                    }
                    dot += s_qr[lane] * __bfloat162float(kr_row[lane]);
                    dot += s_qr[lane + 32] * __bfloat162float(kr_row[lane + 32]);

                    for (int offset = 16; offset > 0; offset /= 2) {
                        dot += __shfl_down_sync(0xffffffff, dot, offset);
                    }
                    if (lane == 0) {
                        my_scores[l] = dot * mla_scale;
                    }
                }
                __syncthreads();

                // Softmax
                float local_max = -1.0e30f;
                for (int l = tid; l < L; l += blockDim.x) {
                    if (my_scores[l] > local_max) local_max = my_scores[l];
                }
                for (int offset = 16; offset > 0; offset /= 2) {
                    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
                }
                __shared__ float s_max[8];
                if (lane == 0) s_max[wid] = local_max;
                __syncthreads();
                if (wid == 0) {
                    float b_max = (lane < 8) ? s_max[lane] : -1.0e30f;
                    for (int offset = 4; offset > 0; offset /= 2) {
                        b_max = fmaxf(b_max, __shfl_down_sync(0xffffffff, b_max, offset));
                    }
                    if (lane == 0) s_max[0] = b_max;
                }
                __syncthreads();
                float global_max = s_max[0];

                float local_sum = 0.0f;
                for (int l = tid; l < L; l += blockDim.x) {
                    float e = expf(my_scores[l] - global_max);
                    my_scores[l] = e;
                    local_sum += e;
                }
                for (int offset = 16; offset > 0; offset /= 2) {
                    local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
                }
                __shared__ float s_sum[8];
                if (lane == 0) s_sum[wid] = local_sum;
                __syncthreads();
                if (wid == 0) {
                    float b_sum = (lane < 8) ? s_sum[lane] : 0.0f;
                    for (int offset = 4; offset > 0; offset /= 2) {
                        b_sum += __shfl_down_sync(0xffffffff, b_sum, offset);
                    }
                    if (lane == 0) s_sum[0] = b_sum;
                }
                __syncthreads();
                float inv_sum = 1.0f / s_sum[0];

                for (int l = tid; l < L; l += blockDim.x) {
                    my_scores[l] *= inv_sum;
                }
                __syncthreads();

                // Compute o_latent[h, e] using linear streaming over l
                float acc0 = 0.0f;
                float acc1 = 0.0f;
                for (int l = 0; l < L; ++l) {
                    float p_l = my_scores[l];
                    const __nv_bfloat16* row = new_c_kv + l * 512;
                    acc0 += p_l * __bfloat162float(row[tid]);
                    acc1 += p_l * __bfloat162float(row[tid + 256]);
                }
                ws.o_latent[h * 512 + tid] = acc0;
                ws.o_latent[h * 512 + tid + 256] = acc1;
            }
            grid.sync();

            // Project o_latent with W_v
            for (int idx = gtid; idx < 32 * 128; idx += total_threads) {
                int h = idx / 128;
                int d = idx % 128;
                const float* ol = ws.o_latent + h * 512;
                float sum = 0.0f;
                #pragma unroll 4
                for (int e = 0; e < 512; ++e) {
                    sum += ol[e] * __bfloat162float(bp.mla_W_v[(e * 32 + h) * 128 + d]);
                }
                ws.o[idx] = __float2bfloat16(sum);
            }
            grid.sync();

            // 1.4 Output Projection
            for (int col = gtid; col < 2304; col += total_threads) {
                float val = compute_gemv_col(ws.o, bp.o_wq, bp.o_s, bp.o_z, col, 4096, 2304);
                ws.attn_out[col] = __float2bfloat16(val);
            }
            grid.sync();
        }

        // Residual Add for Attention
        for (int i = gtid; i < 2304; i += total_threads) {
            float cur_h = __bfloat162float(hidden[i]);
            float att = __bfloat162float(ws.attn_out[i]);
            hidden[i] = __float2bfloat16(cur_h + att);
        }
        grid.sync();

        // === MoE FFN ===
        rmsnorm_block(hidden, bp.moe_norm, ws.x_norm, 2304, grid);

        // Router & Top-8 selection
        if (blockIdx.x == 0) {
            int tid = threadIdx.x;
            __shared__ float s_logits[64];
            __shared__ float s_probs[64];

            if (tid < 64) {
                float dot = 0.0f;
                const __nv_bfloat16* w_row = bp.router_w + tid * 2304;
                #pragma unroll 4
                for (int k = 0; k < 2304; ++k) {
                    dot += __bfloat162float(ws.x_norm[k]) * __bfloat162float(w_row[k]);
                }
                s_logits[tid] = dot;
            }
            __syncthreads();

            if (tid == 0) {
                float max_val = s_logits[0];
                for (int i = 1; i < 64; ++i) {
                    if (s_logits[i] > max_val) max_val = s_logits[i];
                }
                float sum_exp = 0.0f;
                for (int i = 0; i < 64; ++i) {
                    float e = expf(s_logits[i] - max_val);
                    s_probs[i] = e;
                    sum_exp += e;
                }
                for (int i = 0; i < 64; ++i) {
                    s_probs[i] /= sum_exp;
                }

                float sum_top8 = 0.0f;
                for (int rank = 0; rank < 8; ++rank) {
                    float best_p = -1.0f;
                    int best_idx = -1;
                    for (int i = 0; i < 64; ++i) {
                        if (s_probs[i] > best_p) {
                            best_p = s_probs[i];
                            best_idx = i;
                        }
                    }
                    ws.topk_idx[rank] = best_idx;
                    ws.topk_weights[rank] = best_p;
                    sum_top8 += best_p;
                    s_probs[best_idx] = -100.0f;
                }
                float routed_scale = 2.446f / (sum_top8 + 1.0e-9f);
                for (int rank = 0; rank < 8; ++rank) {
                    ws.topk_weights[rank] *= routed_scale;
                }
            }
        }
        grid.sync();

        // 4. MoE Expert GEMVs
        // 4.1 Gate & Up: 18432 columns
        for (int col = gtid; col < 18432; col += total_threads) {
            int j = col / 2048;
            int sub = col % 2048;
            int e = (j < 8) ? ws.topk_idx[j] : 0;
            bool is_gate = (sub < 1024);
            int m = sub % 1024;

            const uint8_t* wq;
            const __nv_bfloat16* s;
            const __nv_bfloat16* z;
            if (j < 8) {
                wq = (is_gate ? bp.gate_wq : bp.up_wq) + e * 1152 * 1024;
                s = (is_gate ? bp.gate_s : bp.up_s) + e * 18 * 1024;
                z = (is_gate ? bp.gate_z : bp.up_z) + e * 18 * 1024;
            } else {
                wq = is_gate ? bp.s_gate_wq : bp.s_up_wq;
                s = is_gate ? bp.s_gate_s : bp.s_up_s;
                z = is_gate ? bp.s_gate_z : bp.s_up_z;
            }
            float val = compute_gemv_col(ws.x_norm, wq, s, z, m, 2304, 1024);
            ws.gate_up[col] = val;
        }
        grid.sync();

        // 4.2 SwiGLU: 9216 elements
        for (int idx = gtid; idx < 9 * 1024; idx += total_threads) {
            int j = idx / 1024;
            int m = idx % 1024;
            float g = ws.gate_up[j * 2048 + m];
            float u = ws.gate_up[j * 2048 + 1024 + m];
            ws.h_expert[idx] = silu(g) * u;
        }
        grid.sync();

        // 4.3 Down GEMV: 20736 columns
        for (int col = gtid; col < 20736; col += total_threads) {
            int j = col / 2304;
            int d = col % 2304;
            int e = (j < 8) ? ws.topk_idx[j] : 0;
            float w = (j < 8) ? ws.topk_weights[j] : 1.0f;

            const uint8_t* wq = (j < 8) ? (bp.down_wq + e * 512 * 2304) : bp.s_down_wq;
            const __nv_bfloat16* s = (j < 8) ? (bp.down_s + e * 8 * 2304) : bp.s_down_s;
            const __nv_bfloat16* z = (j < 8) ? (bp.down_z + e * 8 * 2304) : bp.s_down_z;

            float total_sum = 0.0f;
            for (int g = 0; g < 8; ++g) {
                float scale = __bfloat162float(s[g * 2304 + d]);
                float zero = __bfloat162float(z[g * 2304 + d]);
                float group_acc = 0.0f;
                float group_x_sum = 0.0f;
                #pragma unroll 8
                for (int p = 0; p < 64; ++p) {
                    int k0 = g * 128 + 2 * p;
                    uint8_t byte = wq[(g * 64 + p) * 2304 + d];
                    float u0 = (float)(byte & 0x0F);
                    float u1 = (float)(byte >> 4);
                    float x0 = ws.h_expert[j * 1024 + k0];
                    float x1 = ws.h_expert[j * 1024 + k0 + 1];

                    group_acc += x0 * u0 + x1 * u1;
                    group_x_sum += x0 + x1;
                }
                total_sum += (group_acc - zero * group_x_sum) * scale;
            }
            ws.expert_outs[col] = total_sum * w;
        }
        grid.sync();

        // 4.4 Reduce 9 expert outputs into moe_out
        for (int d = gtid; d < 2304; d += total_threads) {
            float sum = 0.0f;
            #pragma unroll
            for (int j = 0; j < 9; ++j) {
                sum += ws.expert_outs[j * 2304 + d];
            }
            ws.moe_out[d] = sum;
        }
        grid.sync();

        // 5. Residual Add for MoE
        for (int i = gtid; i < 2304; i += total_threads) {
            float cur_h = __bfloat162float(hidden[i]);
            float moe = ws.moe_out[i];
            hidden[i] = __float2bfloat16(cur_h + moe);
        }
        grid.sync();
    }
}

void setup_kda_pointers(
    torch::Tensor blocks_buf, int b,
    torch::Tensor attn_norm, torch::Tensor moe_norm,
    torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
    torch::Tensor k_wq, torch::Tensor k_s, torch::Tensor k_z,
    torch::Tensor v_wq, torch::Tensor v_s, torch::Tensor v_z,
    torch::Tensor g_wq, torch::Tensor g_s, torch::Tensor g_z,
    torch::Tensor beta_w, torch::Tensor conv_w,
    torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
    torch::Tensor router_w,
    torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
    torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
    torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
    torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
    torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
    torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
) {
    BlockPointers* bp = reinterpret_cast<BlockPointers*>(blocks_buf.data_ptr());
    bp[b].is_mla = 0;
    bp[b].attn_norm = reinterpret_cast<const __nv_bfloat16*>(attn_norm.data_ptr());
    bp[b].moe_norm = reinterpret_cast<const __nv_bfloat16*>(moe_norm.data_ptr());

    bp[b].q_wq = reinterpret_cast<const uint8_t*>(q_wq.data_ptr());
    bp[b].q_s = reinterpret_cast<const __nv_bfloat16*>(q_s.data_ptr());
    bp[b].q_z = reinterpret_cast<const __nv_bfloat16*>(q_z.data_ptr());

    bp[b].k_wq = reinterpret_cast<const uint8_t*>(k_wq.data_ptr());
    bp[b].k_s = reinterpret_cast<const __nv_bfloat16*>(k_s.data_ptr());
    bp[b].k_z = reinterpret_cast<const __nv_bfloat16*>(k_z.data_ptr());

    bp[b].v_wq = reinterpret_cast<const uint8_t*>(v_wq.data_ptr());
    bp[b].v_s = reinterpret_cast<const __nv_bfloat16*>(v_s.data_ptr());
    bp[b].v_z = reinterpret_cast<const __nv_bfloat16*>(v_z.data_ptr());

    bp[b].g_wq = reinterpret_cast<const uint8_t*>(g_wq.data_ptr());
    bp[b].g_s = reinterpret_cast<const __nv_bfloat16*>(g_s.data_ptr());
    bp[b].g_z = reinterpret_cast<const __nv_bfloat16*>(g_z.data_ptr());

    bp[b].beta_w = reinterpret_cast<const __nv_bfloat16*>(beta_w.data_ptr());
    bp[b].conv_w = reinterpret_cast<const __nv_bfloat16*>(conv_w.data_ptr());

    bp[b].o_wq = reinterpret_cast<const uint8_t*>(o_wq.data_ptr());
    bp[b].o_s = reinterpret_cast<const __nv_bfloat16*>(o_s.data_ptr());
    bp[b].o_z = reinterpret_cast<const __nv_bfloat16*>(o_z.data_ptr());

    bp[b].router_w = reinterpret_cast<const __nv_bfloat16*>(router_w.data_ptr());
    bp[b].gate_wq = reinterpret_cast<const uint8_t*>(gate_wq.data_ptr());
    bp[b].gate_s = reinterpret_cast<const __nv_bfloat16*>(gate_s.data_ptr());
    bp[b].gate_z = reinterpret_cast<const __nv_bfloat16*>(gate_z.data_ptr());

    bp[b].up_wq = reinterpret_cast<const uint8_t*>(up_wq.data_ptr());
    bp[b].up_s = reinterpret_cast<const __nv_bfloat16*>(up_s.data_ptr());
    bp[b].up_z = reinterpret_cast<const __nv_bfloat16*>(up_z.data_ptr());

    bp[b].down_wq = reinterpret_cast<const uint8_t*>(down_wq.data_ptr());
    bp[b].down_s = reinterpret_cast<const __nv_bfloat16*>(down_s.data_ptr());
    bp[b].down_z = reinterpret_cast<const __nv_bfloat16*>(down_z.data_ptr());

    bp[b].s_gate_wq = reinterpret_cast<const uint8_t*>(s_gate_wq.data_ptr());
    bp[b].s_gate_s = reinterpret_cast<const __nv_bfloat16*>(s_gate_s.data_ptr());
    bp[b].s_gate_z = reinterpret_cast<const __nv_bfloat16*>(s_gate_z.data_ptr());

    bp[b].s_up_wq = reinterpret_cast<const uint8_t*>(s_up_wq.data_ptr());
    bp[b].s_up_s = reinterpret_cast<const __nv_bfloat16*>(s_up_s.data_ptr());
    bp[b].s_up_z = reinterpret_cast<const __nv_bfloat16*>(s_up_z.data_ptr());

    bp[b].s_down_wq = reinterpret_cast<const uint8_t*>(s_down_wq.data_ptr());
    bp[b].s_down_s = reinterpret_cast<const __nv_bfloat16*>(s_down_s.data_ptr());
    bp[b].s_down_z = reinterpret_cast<const __nv_bfloat16*>(s_down_z.data_ptr());
}

void setup_mla_pointers(
    torch::Tensor blocks_buf, int b,
    torch::Tensor attn_norm, torch::Tensor moe_norm,
    torch::Tensor mla_q_wq, torch::Tensor mla_q_s, torch::Tensor mla_q_z,
    torch::Tensor mla_kva_wq, torch::Tensor mla_kva_s, torch::Tensor mla_kva_z,
    torch::Tensor mla_W_k, torch::Tensor mla_W_v,
    torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
    torch::Tensor router_w,
    torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
    torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
    torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
    torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
    torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
    torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
) {
    BlockPointers* bp = reinterpret_cast<BlockPointers*>(blocks_buf.data_ptr());
    bp[b].is_mla = 1;
    bp[b].attn_norm = reinterpret_cast<const __nv_bfloat16*>(attn_norm.data_ptr());
    bp[b].moe_norm = reinterpret_cast<const __nv_bfloat16*>(moe_norm.data_ptr());

    bp[b].mla_q_wq = reinterpret_cast<const uint8_t*>(mla_q_wq.data_ptr());
    bp[b].mla_q_s = reinterpret_cast<const __nv_bfloat16*>(mla_q_s.data_ptr());
    bp[b].mla_q_z = reinterpret_cast<const __nv_bfloat16*>(mla_q_z.data_ptr());

    bp[b].mla_kva_wq = reinterpret_cast<const uint8_t*>(mla_kva_wq.data_ptr());
    bp[b].mla_kva_s = reinterpret_cast<const __nv_bfloat16*>(mla_kva_s.data_ptr());
    bp[b].mla_kva_z = reinterpret_cast<const __nv_bfloat16*>(mla_kva_z.data_ptr());

    bp[b].mla_W_k = reinterpret_cast<const __nv_bfloat16*>(mla_W_k.data_ptr());
    bp[b].mla_W_v = reinterpret_cast<const __nv_bfloat16*>(mla_W_v.data_ptr());

    bp[b].o_wq = reinterpret_cast<const uint8_t*>(o_wq.data_ptr());
    bp[b].o_s = reinterpret_cast<const __nv_bfloat16*>(o_s.data_ptr());
    bp[b].o_z = reinterpret_cast<const __nv_bfloat16*>(o_z.data_ptr());

    bp[b].router_w = reinterpret_cast<const __nv_bfloat16*>(router_w.data_ptr());
    bp[b].gate_wq = reinterpret_cast<const uint8_t*>(gate_wq.data_ptr());
    bp[b].gate_s = reinterpret_cast<const __nv_bfloat16*>(gate_s.data_ptr());
    bp[b].gate_z = reinterpret_cast<const __nv_bfloat16*>(gate_z.data_ptr());

    bp[b].up_wq = reinterpret_cast<const uint8_t*>(up_wq.data_ptr());
    bp[b].up_s = reinterpret_cast<const __nv_bfloat16*>(up_s.data_ptr());
    bp[b].up_z = reinterpret_cast<const __nv_bfloat16*>(up_z.data_ptr());

    bp[b].down_wq = reinterpret_cast<const uint8_t*>(down_wq.data_ptr());
    bp[b].down_s = reinterpret_cast<const __nv_bfloat16*>(down_s.data_ptr());
    bp[b].down_z = reinterpret_cast<const __nv_bfloat16*>(down_z.data_ptr());

    bp[b].s_gate_wq = reinterpret_cast<const uint8_t*>(s_gate_wq.data_ptr());
    bp[b].s_gate_s = reinterpret_cast<const __nv_bfloat16*>(s_gate_s.data_ptr());
    bp[b].s_gate_z = reinterpret_cast<const __nv_bfloat16*>(s_gate_z.data_ptr());

    bp[b].s_up_wq = reinterpret_cast<const uint8_t*>(s_up_wq.data_ptr());
    bp[b].s_up_s = reinterpret_cast<const __nv_bfloat16*>(s_up_s.data_ptr());
    bp[b].s_up_z = reinterpret_cast<const __nv_bfloat16*>(s_up_z.data_ptr());

    bp[b].s_down_wq = reinterpret_cast<const uint8_t*>(s_down_wq.data_ptr());
    bp[b].s_down_s = reinterpret_cast<const __nv_bfloat16*>(s_down_s.data_ptr());
    bp[b].s_down_z = reinterpret_cast<const __nv_bfloat16*>(s_down_z.data_ptr());
}

void launch_megakernel(
    torch::Tensor hidden,
    torch::Tensor S0, torch::Tensor cq0, torch::Tensor ck0, torch::Tensor cv0,
    torch::Tensor S1, torch::Tensor cq1, torch::Tensor ck1, torch::Tensor cv1,
    torch::Tensor S2, torch::Tensor cq2, torch::Tensor ck2, torch::Tensor cv2,
    torch::Tensor old_c_kv, torch::Tensor new_c_kv,
    torch::Tensor old_k_rope, torch::Tensor new_k_rope,
    int context_len,
    torch::Tensor blocks_buffer,
    torch::Tensor ws_buffer
) {
    int dev = 0;
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, dev);
    int num_blocks = prop.multiProcessorCount;
    int num_threads = 256;

    const BlockPointers* bp_d = reinterpret_cast<const BlockPointers*>(blocks_buffer.data_ptr());
    uint8_t* ws_ptr = reinterpret_cast<uint8_t*>(ws_buffer.data_ptr());

    KernelWorkspace ws;
    size_t offset = 0;
    auto alloc_ws = [&](size_t bytes) -> void* {
        void* p = ws_ptr + offset;
        offset = (offset + bytes + 255) & ~255;
        return p;
    };

    ws.x_norm = (__nv_bfloat16*)alloc_ws(2304 * sizeof(__nv_bfloat16));
    ws.q = (__nv_bfloat16*)alloc_ws(6144 * sizeof(__nv_bfloat16));
    ws.k = (__nv_bfloat16*)alloc_ws(4096 * sizeof(__nv_bfloat16));
    ws.v = (__nv_bfloat16*)alloc_ws(4096 * sizeof(__nv_bfloat16));
    ws.g = (__nv_bfloat16*)alloc_ws(4096 * sizeof(__nv_bfloat16));
    ws.beta = (__nv_bfloat16*)alloc_ws(32 * sizeof(__nv_bfloat16));
    ws.o = (__nv_bfloat16*)alloc_ws(4096 * sizeof(__nv_bfloat16));
    ws.attn_out = (__nv_bfloat16*)alloc_ws(2304 * sizeof(__nv_bfloat16));
    ws.kva = (__nv_bfloat16*)alloc_ws(576 * sizeof(__nv_bfloat16));
    ws.q_abs = (float*)alloc_ws(32 * 512 * sizeof(float));
    ws.scores_ws = (float*)alloc_ws(32 * 32768 * sizeof(float));
    ws.o_latent = (float*)alloc_ws(32 * 512 * sizeof(float));
    ws.topk_idx = (int*)alloc_ws(8 * sizeof(int));
    ws.topk_weights = (float*)alloc_ws(8 * sizeof(float));
    ws.gate_up = (float*)alloc_ws(9 * 2048 * sizeof(float));
    ws.h_expert = (float*)alloc_ws(9 * 1024 * sizeof(float));
    ws.expert_outs = (float*)alloc_ws(9 * 2304 * sizeof(float));
    ws.moe_out = (float*)alloc_ws(2304 * sizeof(float));

    __nv_bfloat16* h_ptr = reinterpret_cast<__nv_bfloat16*>(hidden.data_ptr());
    float* s0_ptr = S0.data_ptr<float>();
    __nv_bfloat16* cq0_ptr = reinterpret_cast<__nv_bfloat16*>(cq0.data_ptr());
    __nv_bfloat16* ck0_ptr = reinterpret_cast<__nv_bfloat16*>(ck0.data_ptr());
    __nv_bfloat16* cv0_ptr = reinterpret_cast<__nv_bfloat16*>(cv0.data_ptr());

    float* s1_ptr = S1.data_ptr<float>();
    __nv_bfloat16* cq1_ptr = reinterpret_cast<__nv_bfloat16*>(cq1.data_ptr());
    __nv_bfloat16* ck1_ptr = reinterpret_cast<__nv_bfloat16*>(ck1.data_ptr());
    __nv_bfloat16* cv1_ptr = reinterpret_cast<__nv_bfloat16*>(cv1.data_ptr());

    float* s2_ptr = S2.data_ptr<float>();
    __nv_bfloat16* cq2_ptr = reinterpret_cast<__nv_bfloat16*>(cq2.data_ptr());
    __nv_bfloat16* ck2_ptr = reinterpret_cast<__nv_bfloat16*>(ck2.data_ptr());
    __nv_bfloat16* cv2_ptr = reinterpret_cast<__nv_bfloat16*>(cv2.data_ptr());

    const __nv_bfloat16* old_ckv_ptr = reinterpret_cast<const __nv_bfloat16*>(old_c_kv.data_ptr());
    __nv_bfloat16* new_ckv_ptr = reinterpret_cast<__nv_bfloat16*>(new_c_kv.data_ptr());
    const __nv_bfloat16* old_kr_ptr = reinterpret_cast<const __nv_bfloat16*>(old_k_rope.data_ptr());
    __nv_bfloat16* new_kr_ptr = reinterpret_cast<__nv_bfloat16*>(new_k_rope.data_ptr());

    void* kernel_args[] = {
        &h_ptr,
        &s0_ptr, &cq0_ptr, &ck0_ptr, &cv0_ptr,
        &s1_ptr, &cq1_ptr, &ck1_ptr, &cv1_ptr,
        &s2_ptr, &cq2_ptr, &ck2_ptr, &cv2_ptr,
        &old_ckv_ptr, &new_ckv_ptr, &old_kr_ptr, &new_kr_ptr,
        &context_len,
        &bp_d,
        &ws
    };

    cudaLaunchCooperativeKernel((void*)megakernel, num_blocks, num_threads, kernel_args, 0, 0);
}
'''

_CPP_SRC = '''
void setup_kda_pointers(
    torch::Tensor blocks_buf, int b,
    torch::Tensor attn_norm, torch::Tensor moe_norm,
    torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
    torch::Tensor k_wq, torch::Tensor k_s, torch::Tensor k_z,
    torch::Tensor v_wq, torch::Tensor v_s, torch::Tensor v_z,
    torch::Tensor g_wq, torch::Tensor g_s, torch::Tensor g_z,
    torch::Tensor beta_w, torch::Tensor conv_w,
    torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
    torch::Tensor router_w,
    torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
    torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
    torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
    torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
    torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
    torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
);

void setup_mla_pointers(
    torch::Tensor blocks_buf, int b,
    torch::Tensor attn_norm, torch::Tensor moe_norm,
    torch::Tensor mla_q_wq, torch::Tensor mla_q_s, torch::Tensor mla_q_z,
    torch::Tensor mla_kva_wq, torch::Tensor mla_kva_s, torch::Tensor mla_kva_z,
    torch::Tensor mla_W_k, torch::Tensor mla_W_v,
    torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
    torch::Tensor router_w,
    torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
    torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
    torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
    torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
    torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
    torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
);

void launch_megakernel(
    torch::Tensor hidden,
    torch::Tensor S0, torch::Tensor cq0, torch::Tensor ck0, torch::Tensor cv0,
    torch::Tensor S1, torch::Tensor cq1, torch::Tensor ck1, torch::Tensor cv1,
    torch::Tensor S2, torch::Tensor cq2, torch::Tensor ck2, torch::Tensor cv2,
    torch::Tensor old_c_kv, torch::Tensor new_c_kv,
    torch::Tensor old_k_rope, torch::Tensor new_k_rope,
    int context_len,
    torch::Tensor blocks_buffer,
    torch::Tensor ws_buffer
);
'''

_MEGA_MOD = None


def _get_megakernel_mod():
    global _MEGA_MOD
    if _MEGA_MOD is None:
        _MEGA_MOD = load_inline(
            "kimi_megakernel_ext",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC,
            functions=["setup_kda_pointers", "setup_mla_pointers", "launch_megakernel"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
        )
    return _MEGA_MOD


# --------------------------------------------------------------------------- #
# Main Model Class Exposing step(hidden, state) -> (hidden, state)
# --------------------------------------------------------------------------- #
class Model(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self.blocks_buf_cpu = torch.zeros(4 * 512, dtype=torch.uint8, device="cpu")
        self.blocks_buf_cuda = None
        self.ws_buffer = None
        self.W_k = None
        self.W_v = None
        self._pointers_ready = False

    def _setup_pointers(self):
        mod = _get_megakernel_mod()
        if self.blocks_buf_cuda is None:
            self.blocks_buf_cuda = torch.empty(4 * 512, dtype=torch.uint8, device="cuda:0")
            self.ws_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")

        # Dequantize MLA kv_b for absorb attention
        mla_mod = self.blocks[3].attn
        w_kv_b = self._dequant_kvb(mla_mod.kv_b)
        w_kv_b = w_kv_b.view(self.cfg.kv_lora, self.cfg.mla_heads, self.cfg.qk_nope + self.cfg.v_head)
        self.W_k = w_kv_b[:, :, : self.cfg.qk_nope].contiguous()
        self.W_v = w_kv_b[:, :, self.cfg.qk_nope :].contiguous()

        for b in range(3):
            blk = self.blocks[b]
            mod.setup_kda_pointers(
                self.blocks_buf_cpu, b,
                blk.attn_norm, blk.moe_norm,
                blk.attn.q_proj.w_q, blk.attn.q_proj.scales, blk.attn.q_proj.zeros,
                blk.attn.k_proj.w_q, blk.attn.k_proj.scales, blk.attn.k_proj.zeros,
                blk.attn.v_proj.w_q, blk.attn.v_proj.scales, blk.attn.v_proj.zeros,
                blk.attn.g_proj.w_q, blk.attn.g_proj.scales, blk.attn.g_proj.zeros,
                blk.attn.beta_proj.weight, blk.attn.conv_w,
                blk.attn.o_proj.w_q, blk.attn.o_proj.scales, blk.attn.o_proj.zeros,
                blk.moe.router.weight,
                blk.moe.gate.w_q, blk.moe.gate.scales, blk.moe.gate.zeros,
                blk.moe.up.w_q, blk.moe.up.scales, blk.moe.up.zeros,
                blk.moe.down.w_q, blk.moe.down.scales, blk.moe.down.zeros,
                blk.moe.s_gate.w_q, blk.moe.s_gate.scales, blk.moe.s_gate.zeros,
                blk.moe.s_up.w_q, blk.moe.s_up.scales, blk.moe.s_up.zeros,
                blk.moe.s_down.w_q, blk.moe.s_down.scales, blk.moe.s_down.zeros
            )

        blk3 = self.blocks[3]
        mod.setup_mla_pointers(
            self.blocks_buf_cpu, 3,
            blk3.attn_norm, blk3.moe_norm,
            blk3.attn.q_proj.w_q, blk3.attn.q_proj.scales, blk3.attn.q_proj.zeros,
            blk3.attn.kv_a.w_q, blk3.attn.kv_a.scales, blk3.attn.kv_a.zeros,
            self.W_k, self.W_v,
            blk3.attn.o_proj.w_q, blk3.attn.o_proj.scales, blk3.attn.o_proj.zeros,
            blk3.moe.router.weight,
            blk3.moe.gate.w_q, blk3.moe.gate.scales, blk3.moe.gate.zeros,
            blk3.moe.up.w_q, blk3.moe.up.scales, blk3.moe.up.zeros,
            blk3.moe.down.w_q, blk3.moe.down.scales, blk3.moe.down.zeros,
            blk3.moe.s_gate.w_q, blk3.moe.s_gate.scales, blk3.moe.s_gate.zeros,
            blk3.moe.s_up.w_q, blk3.moe.s_up.scales, blk3.moe.s_up.zeros,
            blk3.moe.s_down.w_q, blk3.moe.s_down.scales, blk3.moe.s_down.zeros
        )
        self.blocks_buf_cuda.copy_(self.blocks_buf_cpu)
        self._pointers_ready = True

    @staticmethod
    def _dequant_kvb(ql: QuantLinear) -> torch.Tensor:
        K = ql.in_f
        wu = torch.empty((K, ql.out_f), dtype=torch.uint8, device=ql.w_q.device)
        wu[[REDACTED: IP]] = ql.w_q & 0xF
        wu[[REDACTED: IP]] = (ql.w_q >> 4) & 0xF
        s = ql.scales.repeat_interleave(ql.group, dim=0)
        z = ql.zeros.repeat_interleave(ql.group, dim=0)
        return (wu.to(torch.bfloat16) - z) * s

    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        res = super().load_state_dict(state_dict, strict=strict, assign=assign)
        self._pointers_ready = False
        return res

    def step(self, hidden: torch.Tensor, state: list[dict[str, torch.Tensor]]) -> tuple[torch.Tensor, list[dict[str, torch.Tensor]]]:
        if not self._pointers_ready:
            self._setup_pointers()

        mod = _get_megakernel_mod()
        pos = state[3]["c_kv"].shape[0]

        # Allocate new cache tensors via CPU metadata (zero GPU kernel launches)
        new_c_kv = torch.empty((pos + 1, self.cfg.kv_lora), dtype=hidden.dtype, device=hidden.device)
        new_k_rope = torch.empty((pos + 1, self.cfg.qk_rope), dtype=hidden.dtype, device=hidden.device)

        # Single Megakernel Launch (invoked exactly once per step)
        mod.launch_megakernel(
            hidden,
            state[0]["S"], state[0]["cq"], state[0]["ck"], state[0]["cv"],
            state[1]["S"], state[1]["cq"], state[1]["ck"], state[1]["cv"],
            state[2]["S"], state[2]["cq"], state[2]["ck"], state[2]["cv"],
            state[3]["c_kv"], new_c_kv, state[3]["k_rope"], new_k_rope,
            pos, self.blocks_buf_cuda, self.ws_buffer
        )

        state[3]["c_kv"] = new_c_kv
        state[3]["k_rope"] = new_k_rope
        return hidden, state