"""Single-launch W4A16 decode kernel for the Kimi-Linear motif.

The CUDA grid is cooperative and persistent.  All projection, attention,
recurrent-state, routing, and expert phases execute inside that one grid; grid
barriers connect phases without returning to Python.  Quantized weights are
unpacked and dequantized in registers and are never materialized.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field

import torch
import torch.nn as nn


EPS = 1.0e-6
GROUP_SIZE = 128


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


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


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
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(in_f // group, 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
        self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))


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


class MLA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, h = cfg.hidden, cfg.mla_heads
        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: 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)


_CPP = r"""
#include <torch/extension.h>
void kimi_launch(std::vector<torch::Tensor>, torch::Tensor, torch::Tensor,
                 std::vector<torch::Tensor>, torch::Tensor, torch::Tensor,
                 torch::Tensor, torch::Tensor, torch::Tensor, int64_t, bool);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("launch", &kimi_launch); }
"""


_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <cfloat>

namespace cg = cooperative_groups;
using B = __nv_bfloat16;

constexpr int D = 2304;
constexpr int H = 32;
constexpr int DK = 128;
constexpr int C = 4096;
constexpr int M = 1024;
constexpr int ACTIVE = 8;
constexpr int SLOTS = 9;
constexpr int E = 64;
constexpr int KV = 512;
constexpr int QR = 64;
constexpr int QN = 128;
constexpr int QDIM = 192;
constexpr int THREADS = 256;

struct QW { const unsigned char* q; const B* s; const B* z; };
struct EW { const unsigned char* q; const B* s; const B* z; };
struct MoEW { const B* router; EW gate, up, down, sg, su, sd; };
struct KL { const B* an; QW q, k, v, g; const B* beta; const B* conv; QW o; const B* mn; MoEW moe; };
struct ML { const B* an; QW q, kva, kvb, o; const B* mn; MoEW moe; };
struct KS { float* S; B* cq; B* ck; B* cv; };
struct P {
    KL kl[3]; ML ml; KS ks[3];
    const B* hidden; B* output;
    const B* cin; const B* kin; B* cout; B* kout;
    float* work; int L; int copy_cache; int nb;
};

__device__ __forceinline__ float lb(const B* p) { return __bfloat162float(*p); }
__device__ __forceinline__ void sb(B* p, float x) { *p = __float2bfloat16_rn(x); }
__device__ __forceinline__ float br(float x) { return __bfloat162float(__float2bfloat16_rn(x)); }
__device__ __forceinline__ float silu(float x) { return x / (1.0f + __expf(-x)); }

__device__ __forceinline__ float iq(unsigned char b, int hi) {
    return float(hi ? (b >> 4) : (b & 15));
}

__device__ __forceinline__ void split_info(int N, int Kpairs, int nb, int bid,
                                            int &n, int &sid, int &ns, int &lo, int &hi) {
    int tiles = (N + THREADS - 1) / THREADS;
    int tile = bid % tiles;
    sid = bid / tiles;
    ns = (nb - 1 - tile) / tiles + 1;
    n = tile * THREADS + threadIdx.x;
    lo = (Kpairs * sid) / ns;
    hi = (Kpairs * (sid + 1)) / ns;
}

__device__ void gemv_part(QW w, const float* x, int K, int N, float* part, int nb) {
    int n, sid, ns, lo, hi;
    split_info(N, K / 2, nb, blockIdx.x, n, sid, ns, lo, hi);
    float acc = 0.0f;
    if (n < N) {
        int g0 = lo >> 6, g1 = (hi + 63) >> 6;
        for (int g = g0; g < g1; ++g) {
            int a = max(lo, g << 6), b = min(hi, (g + 1) << 6);
            float s = lb(w.s + g * N + n), z = lb(w.z + g * N + n);
            float qa = 0.0f, sx = 0.0f;
            for (int p = a; p < b; ++p) {
                unsigned char q = w.q[p * N + n];
                float x0=x[2*p],x1=x[2*p+1];
                qa=fmaf(x0,iq(q,0),qa);qa=fmaf(x1,iq(q,1),qa);sx+=x0+x1;
            }
            acc=fmaf(qa-z*sx,s,acc);
        }
    }
    part[blockIdx.x * THREADS + threadIdx.x] = acc;
}

__device__ void gemv_reduce(float* part, float* out, int N, int nb, bool round_out) {
    int tiles = (N + THREADS - 1) / THREADS;
    int tile = blockIdx.x % tiles, sid = blockIdx.x / tiles;
    int n = tile * THREADS + threadIdx.x;
    if (sid == 0 && n < N) {
        int ns = (nb - 1 - tile) / tiles + 1;
        float v = 0.0f;
        for (int s = 0; s < ns; ++s) v += part[(s * tiles + tile) * THREADS + threadIdx.x];
        out[n] = round_out ? br(v) : v;
    }
}

__device__ void qkvg_part(const KL &l, const float* x, float* p0, float* p1,
                          float* p2, float* p3, int nb) {
    int n, sid, ns, lo, hi;
    split_info(C, D / 2, nb, blockIdx.x, n, sid, ns, lo, hi);
    float aq=0, ak=0, av=0, ag=0;
    if (n < C) {
        for (int gg = lo >> 6; gg < (hi + 63) >> 6; ++gg) {
            int a=max(lo,gg<<6), b=min(hi,(gg+1)<<6);
            float sq=lb(l.q.s+gg*C+n), zq=lb(l.q.z+gg*C+n);
            float sk=lb(l.k.s+gg*C+n), zk=lb(l.k.z+gg*C+n);
            float sv=lb(l.v.s+gg*C+n), zv=lb(l.v.z+gg*C+n);
            float sg=lb(l.g.s+gg*C+n), zg=lb(l.g.z+gg*C+n);
            float tq=0,tk=0,tv=0,tg=0,sx=0;
            for (int p=a; p<b; ++p) {
                float x0=x[2*p], x1=x[2*p+1];
                unsigned char bq=l.q.q[p*C+n], bk=l.k.q[p*C+n];
                unsigned char bv=l.v.q[p*C+n], bg=l.g.q[p*C+n];
                tq=fmaf(x0,iq(bq,0),tq);tq=fmaf(x1,iq(bq,1),tq);
                tk=fmaf(x0,iq(bk,0),tk);tk=fmaf(x1,iq(bk,1),tk);
                tv=fmaf(x0,iq(bv,0),tv);tv=fmaf(x1,iq(bv,1),tv);
                tg=fmaf(x0,iq(bg,0),tg);tg=fmaf(x1,iq(bg,1),tg);sx+=x0+x1;
            }
            aq=fmaf(tq-zq*sx,sq,aq);ak=fmaf(tk-zk*sx,sk,ak);
            av=fmaf(tv-zv*sx,sv,av);ag=fmaf(tg-zg*sx,sg,ag);
        }
    }
    int at=blockIdx.x*THREADS+threadIdx.x;
    p0[at]=aq; p1[at]=ak; p2[at]=av; p3[at]=ag;
}

__device__ void rms(const float* in, const B* w, float* out, float* sh) {
    if (blockIdx.x == 0) {
        float v=0;
        for (int i=threadIdx.x; i<D; i+=THREADS) v=fmaf(in[i],in[i],v);
        sh[threadIdx.x]=v; __syncthreads();
        for(int s=128;s;s>>=1){if(threadIdx.x<s)sh[threadIdx.x]+=sh[threadIdx.x+s];__syncthreads();}
        float rr=rsqrtf(sh[0]/float(D)+1.0e-6f);
        for(int i=threadIdx.x;i<D;i+=THREADS) out[i]=br(in[i]*rr*lb(w+i));
    }
}

__device__ void post_qkvg(const KL &l, KS &st, const float* x,
                           float* q, float* k, float* v, float* decay, float* beta, int nb) {
    int stride=nb*THREADS;
    for(int c=blockIdx.x*THREADS+threadIdx.x;c<C;c+=stride){
        float rq=br(q[c]), rk=br(k[c]), rv=br(v[c]), rg=br(decay[c]);
        B *cq=st.cq, *ck=st.ck, *cv=st.cv;
        const B* cw=l.conv;
        float oq=lb(cq+c)*lb(cw+0*C*4+c*4+0)+lb(cq+C+c)*lb(cw+c*4+1)
                +lb(cq+2*C+c)*lb(cw+c*4+2)+rq*lb(cw+c*4+3);
        float ok=lb(ck+c)*lb(cw+C*4+c*4+0)+lb(ck+C+c)*lb(cw+C*4+c*4+1)
                +lb(ck+2*C+c)*lb(cw+C*4+c*4+2)+rk*lb(cw+C*4+c*4+3);
        float ov=lb(cv+c)*lb(cw+2*C*4+c*4+0)+lb(cv+C+c)*lb(cw+2*C*4+c*4+1)
                +lb(cv+2*C+c)*lb(cw+2*C*4+c*4+2)+rv*lb(cw+2*C*4+c*4+3);
        B aq=cq[C+c], bq=cq[2*C+c], ak=ck[C+c], bk=ck[2*C+c], av=cv[C+c], bv=cv[2*C+c];
        cq[c]=aq; cq[C+c]=bq; sb(cq+2*C+c,rq);
        ck[c]=ak; ck[C+c]=bk; sb(ck+2*C+c,rk);
        cv[c]=av; cv[C+c]=bv; sb(cv+2*C+c,rv);
        q[c]=br(silu(oq)); k[c]=br(silu(ok)); v[c]=br(silu(ov));
        decay[c]=1.0f/(1.0f+__expf(rg));
    }
    if(blockIdx.x==0){
        for(int h=threadIdx.x;h<H;h+=THREADS){
            float a=0; for(int j=0;j<D;++j)a=fmaf(x[j],lb(l.beta+h*D+j),a);
            beta[h]=1.0f/(1.0f+__expf(-a));
        }
    }
}

__device__ void kda_state(KS &st, const float* q, const float* k, const float* v,
                          const float* decay, const float* beta, float* o) {
    if(blockIdx.x < H && threadIdx.x < DK){
        int h=blockIdx.x, d=threadIdx.x; float pred=0;
        for(int j=0;j<DK;++j){
            int ix=(h*DK+j)*DK+d;
            pred=fmaf(st.S[ix]*decay[h*DK+j],k[h*DK+j],pred);
        }
        float acc=0, vd=v[h*DK+d];
        for(int j=0;j<DK;++j){
            int ix=(h*DK+j)*DK+d;
            float ns=st.S[ix]*decay[h*DK+j]+beta[h]*k[h*DK+j]*(vd-pred);
            st.S[ix]=ns; acc=fmaf(ns,q[h*DK+j]*0.08838834764831845f,acc);
        }
        o[h*DK+d]=br(acc);
    }
}

__device__ void residual_attn(float* cur, const float* a, int nb) {
    for(int i=blockIdx.x*THREADS+threadIdx.x;i<D;i+=nb*THREADS) cur[i]=br(cur[i]+br(a[i]));
}

__device__ void route(const MoEW &m, const float* x, float* logits, int* idx, float* rw) {
    if(blockIdx.x==0){
        if(threadIdx.x<E){
            float a=0; for(int j=0;j<D;++j)a=fmaf(x[j],lb(m.router+threadIdx.x*D+j),a);
            logits[threadIdx.x]=a;
        }
        __syncthreads();
        if(threadIdx.x==0){
            float den=0;
            for(int s=0;s<ACTIVE;++s){
                float best=-FLT_MAX; int bi=0;
                for(int e=0;e<E;++e){
                    bool used=false; for(int u=0;u<s;++u)used|=(idx[u]==e);
                    if(!used && logits[e]>best){best=logits[e];bi=e;}
                }
                idx[s]=bi; rw[s]=best;
            }
            float mx=rw[0]; for(int s=1;s<ACTIVE;++s)mx=fmaxf(mx,rw[s]);
            for(int s=0;s<ACTIVE;++s){rw[s]=__expf(rw[s]-mx);den+=rw[s];}
            for(int s=0;s<ACTIVE;++s)rw[s]=rw[s]*(2.446f/den);
        }
    }
}

__device__ void expert_gu_part(const MoEW &m, const float* x, const int* idx,
                               float* p0, float* p1, int nb) {
    int n,sid,ns,lo,hi; split_info(SLOTS*M,D/2,nb,blockIdx.x,n,sid,ns,lo,hi);
    float ag=0,au=0;
    if(n<SLOTS*M){
        int slot=n/M,col=n-slot*M,ex=(slot<ACTIVE?idx[slot]:0);
        EW wg=(slot<ACTIVE?m.gate:m.sg), wu=(slot<ACTIVE?m.up:m.su);
        int eq=ex*(D/2)*M, ep=ex*(D/128)*M;
        for(int gg=lo>>6;gg<(hi+63)>>6;++gg){
            int a=max(lo,gg<<6),b=min(hi,(gg+1)<<6);
            float sg=lb(wg.s+ep+gg*M+col),zg=lb(wg.z+ep+gg*M+col);
            float su=lb(wu.s+ep+gg*M+col),zu=lb(wu.z+ep+gg*M+col);
            float tg=0,tu=0,sx=0;
            for(int p=a;p<b;++p){
                float x0=x[2*p],x1=x[2*p+1];
                unsigned char bg=wg.q[eq+p*M+col],bu=wu.q[eq+p*M+col];
                tg=fmaf(x0,iq(bg,0),tg);tg=fmaf(x1,iq(bg,1),tg);
                tu=fmaf(x0,iq(bu,0),tu);tu=fmaf(x1,iq(bu,1),tu);sx+=x0+x1;
            }
            ag=fmaf(tg-zg*sx,sg,ag);au=fmaf(tu-zu*sx,su,au);
        }
    }
    int at=blockIdx.x*THREADS+threadIdx.x;p0[at]=ag;p1[at]=au;
}

__device__ void expert_gu_reduce(float* p0,float* p1,float* h,int nb){
    int N=SLOTS*M,tiles=(N+THREADS-1)/THREADS,tile=blockIdx.x%tiles,sid=blockIdx.x/tiles;
    int n=tile*THREADS+threadIdx.x;
    if(sid==0&&n<N){int ns=(nb-1-tile)/tiles+1;float g=0,u=0;
        for(int s=0;s<ns;++s){int a=(s*tiles+tile)*THREADS+threadIdx.x;g+=p0[a];u+=p1[a];}
        h[n]=silu(g)*u;}
}

__device__ void expert_down_part(const MoEW &m,const float* h,const int* idx,const float* rw,
                                 float* part,int nb){
    int n,sid,ns,lo,hi;split_info(D,M/2,nb,blockIdx.x,n,sid,ns,lo,hi);float acc=0;
    if(n<D){
        for(int gg=lo>>6;gg<(hi+63)>>6;++gg){
            int a=max(lo,gg<<6),b=min(hi,(gg+1)<<6);
            for(int slot=0;slot<SLOTS;++slot){
                int ex=slot<ACTIVE?idx[slot]:0;float coef=slot<ACTIVE?rw[slot]:1.0f;
                EW wd=slot<ACTIVE?m.down:m.sd;int eq=ex*(M/2)*D,ep=ex*(M/128)*D;
                float s=lb(wd.s+ep+gg*D+n),z=lb(wd.z+ep+gg*D+n);
                float tq=0,sx=0;for(int p=a;p<b;++p){unsigned char bq=wd.q[eq+p*D+n];
                    float x0=h[slot*M+2*p],x1=h[slot*M+2*p+1];
                    tq=fmaf(x0,iq(bq,0),tq);tq=fmaf(x1,iq(bq,1),tq);sx+=x0+x1;}
                acc=fmaf(coef*s, tq-z*sx, acc);
            }
        }
    }
    part[blockIdx.x*THREADS+threadIdx.x]=acc;
}

__device__ void moe_residual(float* part,float* cur,int nb){
    int tiles=(D+THREADS-1)/THREADS,tile=blockIdx.x%tiles,sid=blockIdx.x/tiles;
    int n=tile*THREADS+threadIdx.x;if(sid==0&&n<D){int ns=(nb-1-tile)/tiles+1;float v=0;
        for(int s=0;s<ns;++s)v+=part[(s*tiles+tile)*THREADS+threadIdx.x];
        cur[n]=br(cur[n]+br(v));}
}

__device__ void mla_qkv_part(const ML &l,const float* x,float* part,int nb){
    constexpr int N1=H*QDIM,N2=KV+QR,NT=N1+N2;
    int n,sid,ns,lo,hi;split_info(NT,D/2,nb,blockIdx.x,n,sid,ns,lo,hi);float acc=0;
    if(n<NT){QW w=n<N1?l.q:l.kva;int nn=n<N1?n:n-N1,N=n<N1?N1:N2;
        for(int gg=lo>>6;gg<(hi+63)>>6;++gg){int a=max(lo,gg<<6),b=min(hi,(gg+1)<<6);
            float s=lb(w.s+gg*N+nn),z=lb(w.z+gg*N+nn);
            float tq=0,sx=0;for(int p=a;p<b;++p){unsigned char bq=w.q[p*N+nn];float x0=x[2*p],x1=x[2*p+1];
                tq=fmaf(x0,iq(bq,0),tq);tq=fmaf(x1,iq(bq,1),tq);sx+=x0+x1;}acc=fmaf(tq-z*sx,s,acc);}
    } part[blockIdx.x*THREADS+threadIdx.x]=acc;
}

__device__ void mla_qkv_reduce(float* part,float* out,int nb){
    constexpr int NT=H*QDIM+KV+QR;int tiles=(NT+THREADS-1)/THREADS,tile=blockIdx.x%tiles,sid=blockIdx.x/tiles;
    int n=tile*THREADS+threadIdx.x;if(sid==0&&n<NT){int ns=(nb-1-tile)/tiles+1;float v=0;
        for(int s=0;s<ns;++s)v+=part[(s*tiles+tile)*THREADS+threadIdx.x];out[n]=br(v);}
}

__device__ void mla_rope_cache(float* qkv,B* cout,B* kout,int pos,int nb){
    int stride=nb*THREADS;
    for(int n=blockIdx.x*THREADS+threadIdx.x;n<H*QR/2;n+=stride){
        int h=n/(QR/2),j=n%(QR/2);float ang=pos*__powf(10000.0f,-float(2*j)/float(QR));float co=__cosf(ang),si=__sinf(ang);
        int a=h*QDIM+QN+2*j;float e=qkv[a],o=qkv[a+1];qkv[a]=br(e*co-o*si);qkv[a+1]=br(o*co+e*si);
    }
    for(int r=blockIdx.x*THREADS+threadIdx.x;r<KV;r+=stride)sb(cout+pos*KV+r,qkv[H*QDIM+r]);
    for(int j=blockIdx.x*THREADS+threadIdx.x;j<QR/2;j+=stride){
        float ang=pos*__powf(10000.0f,-float(2*j)/float(QR));float co=__cosf(ang),si=__sinf(ang);
        float e=qkv[H*QDIM+KV+2*j],o=qkv[H*QDIM+KV+2*j+1];
        sb(kout+pos*QR+2*j,e*co-o*si);sb(kout+pos*QR+2*j+1,o*co+e*si);
    }
}

__device__ void q_abs(const ML &l,const float* q,float* qa,int nb){
    int warp=threadIdx.x>>5,lane=threadIdx.x&31,gw=blockIdx.x*8+warp,total=nb*8;
    for(int n=gw;n<KV*H;n+=total){int r=n/H,h=n%H;float a=0;
        for(int d=lane;d<QN;d+=32){int col=h*(QN+DK)+d;unsigned char b=l.kvb.q[(r/2)*(H*(QN+DK))+col];
            float s=lb(l.kvb.s+(r/128)*(H*(QN+DK))+col),z=lb(l.kvb.z+(r/128)*(H*(QN+DK))+col);
            a=fmaf(q[h*QDIM+d],(iq(b,r&1)-z)*s,a);}
        for(int off=16;off;off>>=1)a+=__shfl_down_sync(0xffffffff,a,off);
        if(lane==0)qa[r*H+h]=a;
    }
}

__device__ void mla_scores(const float* qa,const float* q,const B* c,const B* kr,float* scores,int T,int nb){
    int warp=threadIdx.x>>5,h=threadIdx.x&31,gw=blockIdx.x*8+warp,total=nb*8;
    for(int t=gw;t<T;t+=total){float a=0;
        for(int r=0;r<KV;++r)a=fmaf(lb(c+t*KV+r),qa[r*H+h],a);
        for(int d=0;d<QR;++d)a=fmaf(lb(kr+t*QR+d),q[h*QDIM+QN+d],a);
        scores[t*H+h]=a*0.07216878364870322f;
    }
}

__device__ void mla_softmax(const float* scores,float* probs,int T,float* sh){
    if(blockIdx.x<H){int h=blockIdx.x;float mx=-FLT_MAX;
        for(int t=threadIdx.x;t<T;t+=THREADS)mx=fmaxf(mx,scores[t*H+h]);sh[threadIdx.x]=mx;__syncthreads();
        for(int s=128;s;s>>=1){if(threadIdx.x<s)sh[threadIdx.x]=fmaxf(sh[threadIdx.x],sh[threadIdx.x+s]);__syncthreads();}mx=sh[0];
        float sum=0;for(int t=threadIdx.x;t<T;t+=THREADS){float p=__expf(scores[t*H+h]-mx);probs[t*H+h]=p;sum+=p;}
        sh[threadIdx.x]=sum;__syncthreads();for(int s=128;s;s>>=1){if(threadIdx.x<s)sh[threadIdx.x]+=sh[threadIdx.x+s];__syncthreads();}sum=sh[0];
        for(int t=threadIdx.x;t<T;t+=THREADS)probs[t*H+h]/=sum;
    }
}

__device__ void mla_c_part(const float* probs,const B* c,float* part,int T,int nb){
    constexpr int N=H*KV,tiles=N/THREADS;int tile=blockIdx.x%tiles,sid=blockIdx.x/tiles;
    int ns=(nb-1-tile)/tiles+1,n=tile*THREADS+threadIdx.x;
    int lo=(T*sid)/ns,hi=(T*(sid+1))/ns,h=n/KV,r=n%KV;float a=0;
    for(int t=lo;t<hi;++t)a=fmaf(probs[t*H+h],lb(c+t*KV+r),a);
    part[blockIdx.x*THREADS+threadIdx.x]=a;
}

__device__ void mla_c_reduce(float* part,float* cs,int nb){
    constexpr int N=H*KV,tiles=N/THREADS;int tile=blockIdx.x%tiles,sid=blockIdx.x/tiles;
    int n=tile*THREADS+threadIdx.x;if(sid==0){int ns=(nb-1-tile)/tiles+1;float v=0;
        for(int s=0;s<ns;++s)v+=part[(s*tiles+tile)*THREADS+threadIdx.x];cs[n]=v;}
}

__device__ void mla_v_part(const ML &l,const float* cs,float* part,int nb){
    constexpr int NO=H*DK,NW=H*(QN+DK);int n,sid,ns,lo,hi;split_info(NO,KV/2,nb,blockIdx.x,n,sid,ns,lo,hi);float acc=0;
    if(n<NO){int h=n/DK,v=n%DK,col=h*(QN+DK)+QN+v;
        for(int gg=lo>>6;gg<(hi+63)>>6;++gg){int a=max(lo,gg<<6),b=min(hi,(gg+1)<<6);float s=lb(l.kvb.s+gg*NW+col),z=lb(l.kvb.z+gg*NW+col);
            float tq=0,sx=0;for(int p=a;p<b;++p){unsigned char bq=l.kvb.q[p*NW+col];float x0=cs[h*KV+2*p],x1=cs[h*KV+2*p+1];
                tq=fmaf(x0,iq(bq,0),tq);tq=fmaf(x1,iq(bq,1),tq);sx+=x0+x1;}acc=fmaf(tq-z*sx,s,acc);}
    }part[blockIdx.x*THREADS+threadIdx.x]=acc;
}

__device__ void mla_v_reduce(float* part,float* out,int nb){gemv_reduce(part,out,H*DK,nb,true);}

__device__ void run_moe(const MoEW &m,const B* norm,float* cur,float* x,float* p0,float* p1,
                        float* h,float* logits,int* idx,float* rw,float* sh,cg::grid_group &grid,int nb){
    rms(cur,norm,x,sh);grid.sync();route(m,x,logits,idx,rw);grid.sync();
    expert_gu_part(m,x,idx,p0,p1,nb);grid.sync();expert_gu_reduce(p0,p1,h,nb);grid.sync();
    expert_down_part(m,h,idx,rw,p0,nb);grid.sync();moe_residual(p0,cur,nb);grid.sync();
}

__global__ __launch_bounds__(THREADS,1) void kimi_kernel(P p){
    cg::grid_group grid=cg::this_grid();__shared__ float sh[THREADS];int nb=p.nb;
    float* p0=p.work;float* p1=p0+nb*THREADS;float* p2=p1+nb*THREADS;float* p3=p2+nb*THREADS;
    float* x=p3+nb*THREADS;float* q=x+D;float* k=q+C;float* v=k+C;float* aux=v+C;
    float* o=aux+C;float* eh=o+C;float* logits=eh+SLOTS*M;int* idx=(int*)(logits+E);float* rw=(float*)(idx+ACTIVE);
    float* mla=q;float* qa=eh;float* cs=qa+KV*H;float* scores=cs+H*KV;float* probs=scores+(p.L+1)*H;
    int stride=nb*THREADS;
    for(int i=blockIdx.x*THREADS+threadIdx.x;i<D;i+=stride)x[i]=lb(p.hidden+i);
    if(p.copy_cache){
        for(int i=blockIdx.x*THREADS+threadIdx.x;i<p.L*KV;i+=stride)p.cout[i]=p.cin[i];
        for(int i=blockIdx.x*THREADS+threadIdx.x;i<p.L*QR;i+=stride)p.kout[i]=p.kin[i];
    }
    grid.sync();
    for(int z=0;z<3;++z){KL &l=p.kl[z];KS &st=p.ks[z];
        rms(x,l.an,o,sh);grid.sync();
        qkvg_part(l,o,p0,p1,p2,p3,nb);grid.sync();
        gemv_reduce(p0,q,C,nb,false);gemv_reduce(p1,k,C,nb,false);gemv_reduce(p2,v,C,nb,false);gemv_reduce(p3,aux,C,nb,false);grid.sync();
        post_qkvg(l,st,o,q,k,v,aux,logits,nb);grid.sync();
        kda_state(st,q,k,v,aux,logits,o);grid.sync();
        gemv_part(l.o,o,C,D,p0,nb);grid.sync();gemv_reduce(p0,q,D,nb,false);grid.sync();
        residual_attn(x,q,nb);grid.sync();
        run_moe(l.moe,l.mn,x,o,p0,p1,eh,logits,idx,rw,sh,grid,nb);
    }
    ML &l=p.ml;
    rms(x,l.an,o,sh);grid.sync();
    mla_qkv_part(l,o,p0,nb);grid.sync();mla_qkv_reduce(p0,mla,nb);grid.sync();
    mla_rope_cache(mla,p.cout,p.kout,p.L,nb);grid.sync();
    q_abs(l,mla,qa,nb);grid.sync();
    mla_scores(qa,mla,p.cout,p.kout,scores,p.L+1,nb);grid.sync();
    mla_softmax(scores,probs,p.L+1,sh);grid.sync();
    mla_c_part(probs,p.cout,p0,p.L+1,nb);grid.sync();mla_c_reduce(p0,cs,nb);grid.sync();
    mla_v_part(l,cs,p0,nb);grid.sync();mla_v_reduce(p0,o,nb);grid.sync();
    gemv_part(l.o,o,C,D,p0,nb);grid.sync();gemv_reduce(p0,q,D,nb,false);grid.sync();
    residual_attn(x,q,nb);grid.sync();
    run_moe(l.moe,l.mn,x,o,p0,p1,eh,logits,idx,rw,sh,grid,nb);
    for(int i=blockIdx.x*THREADS+threadIdx.x;i<D;i+=stride)sb(p.output+i,x[i]);
}

static B* bp(torch::Tensor &t){return reinterpret_cast<B*>(t.data_ptr<at::BFloat16>());}
static const B* cbp(torch::Tensor &t){return reinterpret_cast<const B*>(t.data_ptr<at::BFloat16>());}
static QW qw(std::vector<torch::Tensor>&w,int&i){QW a{w[i].data_ptr<unsigned char>(),cbp(w[i+1]),cbp(w[i+2])};i+=3;return a;}
static EW ew(std::vector<torch::Tensor>&w,int&i){EW a{w[i].data_ptr<unsigned char>(),cbp(w[i+1]),cbp(w[i+2])};i+=3;return a;}
static MoEW mw(std::vector<torch::Tensor>&w,int&i){MoEW m{};m.router=cbp(w[i++]);m.gate=ew(w,i);m.up=ew(w,i);m.down=ew(w,i);m.sg=ew(w,i);m.su=ew(w,i);m.sd=ew(w,i);return m;}

void kimi_launch(std::vector<torch::Tensor> w,torch::Tensor hidden,torch::Tensor output,
                 std::vector<torch::Tensor> st,torch::Tensor cin,torch::Tensor kin,
                 torch::Tensor cout,torch::Tensor kout,torch::Tensor work,int64_t L,bool copy){
    P p{};int i=0;
    for(int z=0;z<3;++z){KL &l=p.kl[z];l.an=cbp(w[i++]);l.mn=cbp(w[i++]);l.q=qw(w,i);l.k=qw(w,i);l.v=qw(w,i);l.g=qw(w,i);
        l.beta=cbp(w[i++]);l.conv=cbp(w[i++]);l.o=qw(w,i);l.moe=mw(w,i);}
    ML &l=p.ml;l.an=cbp(w[i++]);l.mn=cbp(w[i++]);l.q=qw(w,i);l.kva=qw(w,i);l.kvb=qw(w,i);l.o=qw(w,i);l.moe=mw(w,i);
    for(int z=0,s=0;z<3;++z){p.ks[z].S=st[s++].data_ptr<float>();p.ks[z].cq=bp(st[s++]);p.ks[z].ck=bp(st[s++]);p.ks[z].cv=bp(st[s++]);}
    p.hidden=cbp(hidden);p.output=bp(output);p.cin=cbp(cin);p.kin=cbp(kin);p.cout=bp(cout);p.kout=bp(kout);p.work=work.data_ptr<float>();p.L=int(L);p.copy_cache=copy;
    static int grid=0;if(!grid){int dev,active;cudaGetDevice(&dev);cudaDeviceProp prop;cudaGetDeviceProperties(&prop,dev);
        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active,kimi_kernel,THREADS,0);grid=active*prop.multiProcessorCount;}
    p.nb=grid;void* args[]={&p};cudaError_t e=cudaLaunchCooperativeKernel((void*)kimi_kernel,grid,THREADS,args,0,at::cuda::getCurrentCUDAStream());
    TORCH_CHECK(e==cudaSuccess,"cooperative Kimi kernel launch failed: ",cudaGetErrorString(e));
}
"""


_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        # The runner's nvcc shim is outside CUDA_HOME; point the standard PyTorch
        # extension builder at the installed toolkit before importing it.
        os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
        from torch.utils.cpp_extension import load_inline

        _extension = load_inline(
            name="kimi_linear_single_kernel_v7",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=False,
        )
    return _extension


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._weight_args = None
        self._workspace = None

    @staticmethod
    def _add_q(dst, q):
        dst.extend((q.w_q, q.scales, q.zeros))

    @staticmethod
    def _add_e(dst, q):
        dst.extend((q.w_q, q.scales, q.zeros))

    def _weights(self):
        if self._weight_args is not None:
            return self._weight_args
        out = []
        for b in self.blocks:
            out.extend((b.attn_norm, b.moe_norm))
            if b.kind == "K":
                for q in (b.attn.q_proj, b.attn.k_proj, b.attn.v_proj, b.attn.g_proj):
                    self._add_q(out, q)
                out.extend((b.attn.beta_proj.weight, b.attn.conv_w))
                self._add_q(out, b.attn.o_proj)
            else:
                for q in (b.attn.q_proj, b.attn.kv_a, b.attn.kv_b, b.attn.o_proj):
                    self._add_q(out, q)
            out.append(b.moe.router.weight)
            for q in (b.moe.gate, b.moe.up, b.moe.down, b.moe.s_gate, b.moe.s_up, b.moe.s_down):
                self._add_e(out, q)
        self._weight_args = out
        return out

    @staticmethod
    def _extended(cache, rows, width):
        need = rows * width * cache.element_size()
        available = cache.untyped_storage().nbytes() - cache.storage_offset() * cache.element_size()
        if available >= need:
            return cache.as_strided((rows, width), (width, 1)), False
        capacity = max(rows + 63, ((rows + 63) // 64) * 64)
        buf = torch.empty((capacity, width), dtype=cache.dtype, device=cache.device)
        return buf[:rows], True

    def step(self, hidden, state):
        mla_state = state[3]
        old_c, old_k = mla_state["c_kv"], mla_state["k_rope"]
        length = old_c.shape[0]
        new_c, copy_c = self._extended(old_c, length + 1, 512)
        new_k, copy_k = self._extended(old_k, length + 1, 64)
        copy_cache = copy_c or copy_k
        # Both reference-created cache tensors have exact-sized storage; after
        # the first step both views have spare storage and advance together.
        if copy_c != copy_k:
            new_c = torch.empty((length + 64, 512), dtype=old_c.dtype, device=old_c.device)[: length + 1]
            new_k = torch.empty((length + 64, 64), dtype=old_k.dtype, device=old_k.device)[: length + 1]
            copy_cache = True

        needed = 4 * 256 * 256 + 70000 + 2 * (length + 1) * 32
        if self._workspace is None or self._workspace.numel() < needed:
            self._workspace = torch.empty(needed + 4096, dtype=torch.float32, device=hidden.device)
        output = torch.empty_like(hidden)
        st_args = []
        for s in state[:3]:
            st_args.extend((s["S"], s["cq"], s["ck"], s["cv"]))
        _get_extension().launch(
            self._weights(), hidden, output, st_args, old_c, old_k, new_c, new_k,
            self._workspace, length, copy_cache,
        )
        mla_state["c_kv"] = new_c
        mla_state["k_rope"] = new_k
        return output, state