"""Single launch CUDA decode of the four layer W4A16 Kimi Linear motif."""
from __future__ import annotations

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


class QuantLinear(nn.Module):
    def __init__(self, k, n):
        super().__init__()
        self.register_buffer("w_q", torch.zeros(k // 2, n, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(k // 128, n, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(k // 128, n, dtype=torch.bfloat16))


class QuantExperts(nn.Module):
    def __init__(self, e, k, n):
        super().__init__()
        self.register_buffer("w_q", torch.zeros(e, k // 2, n, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(e, k // 128, n, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(e, k // 128, n, dtype=torch.bfloat16))


class KDA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, c = cfg.hidden, cfg.kda_heads * cfg.kda_head_dim
        self.q_proj = QuantLinear(d, c)
        self.k_proj = QuantLinear(d, c)
        self.v_proj = QuantLinear(d, c)
        self.g_proj = QuantLinear(d, c)
        self.beta_proj = nn.Linear(d, cfg.kda_heads, 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)


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))
        self.kv_a = QuantLinear(d, cfg.kv_lora + cfg.qk_rope)
        self.kv_b = QuantLinear(cfg.kv_lora, h * (cfg.qk_nope + cfg.v_head))
        self.o_proj = QuantLinear(h * cfg.v_head, d)


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


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


CPP_SRC = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <vector>
using bf = __nv_bfloat16;
struct W { const unsigned char* q; const bf* s; const bf* z; int K; int N; };
struct A {
  W w[44];
  const bf *norm[4], *mn[4], *conv[3], *beta[3], *router[4];
  float *S[3]; bf *cq[3], *ck[3], *cv[3];
  const bf *oldc, *oldr; bf *newc, *newr;
  const bf *input; bf *output; float *t; int L;
};

void launch(A, cudaStream_t);
void run(std::vector<torch::Tensor> ts, int L) {
  A a{};a.input=(const bf*)ts[0].data_ptr();a.output=(bf*)ts[1].data_ptr();a.t=(float*)ts[2].data_ptr();a.L=L;
  for(int i=0;i<3;i++) {int p=3+i*4;a.S[i]=(float*)ts[p].data_ptr();a.cq[i]=(bf*)ts[p+1].data_ptr();a.ck[i]=(bf*)ts[p+2].data_ptr();a.cv[i]=(bf*)ts[p+3].data_ptr();}
  a.oldc=(const bf*)ts[15].data_ptr();a.oldr=(const bf*)ts[16].data_ptr();a.newc=(bf*)ts[17].data_ptr();a.newr=(bf*)ts[18].data_ptr();
  for(int i=0;i<4;i++) {int p=19+i*5;a.norm[i]=(const bf*)ts[p].data_ptr();a.mn[i]=(const bf*)ts[p+1].data_ptr();if(i<3) {a.conv[i]=(const bf*)ts[p+2].data_ptr();a.beta[i]=(const bf*)ts[p+3].data_ptr();}a.router[i]=(const bf*)ts[p+4].data_ptr();}
  for(int i=0;i<44;i++) {int p=39+i*3;a.w[i].q=(const unsigned char*)ts[p].data_ptr();a.w[i].s=(const bf*)ts[p+1].data_ptr();a.w[i].z=(const bf*)ts[p+2].data_ptr();a.w[i].K=ts[p].size(ts[p].dim()-2)*2;a.w[i].N=ts[p].size(ts[p].dim()-1);}
  launch(a, at::cuda::getCurrentCUDAStream());
}'''

CUDA_SRC = r'''
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cooperative_groups.h>
#define GID (blockIdx.x*blockDim.x+threadIdx.x)
#define GSTRIDE (blockDim.x*gridDim.x)
using bf = __nv_bfloat16;
struct W { const unsigned char* q; const bf* s; const bf* z; int K; int N; };
struct A {
  W w[44];
  const bf *norm[4], *mn[4], *conv[3], *beta[3], *router[4];
  float *S[3]; bf *cq[3], *ck[3], *cv[3];
  const bf *oldc, *oldr; bf *newc, *newr;
  const bf *input; bf *output; float *t; int L;
};

__device__ __forceinline__ float B(const bf v) { return __bfloat162float(v); }
__device__ __forceinline__ float R(float v) { return B(__float2bfloat16_rn(v)); }
__device__ __forceinline__ float sigmoid(float v) { return 1.f / (1.f + expf(-v)); }
__device__ __forceinline__ float silu(float v) { return v * sigmoid(v); }

__device__ __noinline__ void gemv(const A& a, int wi, const float* x, float* y, int e, bool round_out) {
  const W& w = a.w[wi];
  int K=w.K,N=w.N,eo=e<0?0:e;
  int tile_width=N<=1024?8:16;
  int ksplit=128/tile_width;
  int lane_n=threadIdx.x&(tile_width-1),lane_k=threadIdx.x/(tile_width);
  __shared__ float part[128];
  int tiles=(N+tile_width-1)/tile_width;
  for(int tile=blockIdx.x;tile<tiles;tile+=gridDim.x) {
    int n=tile*tile_width+lane_n;float sum=0.f;
    if(n<N) {
      for(int g=lane_k;g<K/128;g+=ksplit) {
        float scale=B(w.s[(eo*(K/128)+g)*N+n]);
        float zero=B(w.z[(eo*(K/128)+g)*N+n]);
        for(int p=0;p<64;p++) {
          int k=g*128+2*p;
          unsigned char q=w.q[(eo*(K/2)+k/2)*N+n];
          sum=fmaf(x[k],(float(q&15)-zero)*scale,sum);
          sum=fmaf(x[k+1],(float(q>>4)-zero)*scale,sum);
        }
      }
    }
    part[threadIdx.x]=sum;
    __syncthreads();
    if(lane_k==0 && n<N) {
      float v=0.f;for(int i=0;i<ksplit;i++)v+=part[i*tile_width+lane_n];
      y[n]=round_out?R(v):v;
    }
    __syncthreads();
  }
  cooperative_groups::this_grid().sync();
}

__device__ __noinline__ void norm(const A& a, const float* x, const bf* w, float* y) {
  int tid=GID; float p=0.f;
  for(int i=tid;i<2304;i+=GSTRIDE) p+=x[i]*x[i];
  a.t[920000+tid]=p; cooperative_groups::this_grid().sync();
  if(tid==0) { float s=0.f; for(int i=0;i<GSTRIDE;i++) s+=a.t[920000+i]; a.t[920300]=rsqrtf(s/2304.f+1.e-6f); }
  cooperative_groups::this_grid().sync(); float scale=a.t[920300];
  for(int i=tid;i<2304;i+=GSTRIDE) y[i]=R(x[i]*scale*B(w[i]));
  cooperative_groups::this_grid().sync();
}

__device__ __noinline__ void small(const bf* w,const float* x,float* y,int N,int K) {
  int tid=GID;
  for(int n=tid;n<N;n+=GSTRIDE) { float v=0.f; for(int k=0;k<K;k++) v=fmaf(x[k],B(w[n*K+k]),v); y[n]=R(v); }
  cooperative_groups::this_grid().sync();
}

__device__ __noinline__ void conv(const A& a, int layer, int which, float* y) {
  bf* c = which==0 ? a.cq[layer] : (which==1 ? a.ck[layer] : a.cv[layer]);
  const bf* w=a.conv[layer]+which*4096*4;
  for(int n=GID;n<4096;n+=GSTRIDE) {
    float v=B(c[n])*B(w[n*4]) + B(c[4096+n])*B(w[n*4+1]);
    v+=B(c[8192+n])*B(w[n*4+2]) + y[n]*B(w[n*4+3]);
    c[n]=c[4096+n]; c[4096+n]=c[8192+n]; c[8192+n]=__float2bfloat16_rn(y[n]);
    y[n]=R(silu(v));
  }
  cooperative_groups::this_grid().sync();
}

__device__ __noinline__ void kda_state(const A& a,int layer,const float* q,const float* k,const float* v,const float* g,const float* beta,float* out) {
  int j=threadIdx.x;
  for(int h=blockIdx.x;h<32;h+=gridDim.x) {
    if(j<128) {
      float pred=0.f;
      for(int i=0;i<128;i++) {
        int ix=(h*128+i)*128+j;
        float dec=sigmoid(-g[h*128+i]);
        pred=fmaf(a.S[layer][ix]*dec,k[h*128+i],pred);
      }
      float b=sigmoid(beta[h]);
      float delta=b*(v[h*128+j]-pred), o=0.f;
      for(int i=0;i<128;i++) {
        int ix=(h*128+i)*128+j;
        float updated=a.S[layer][ix]*sigmoid(-g[h*128+i])+k[h*128+i]*delta;
        a.S[layer][ix]=updated;
        o=fmaf(updated,q[h*128+i]*0.08838834764831845f,o);
      }
      out[h*128+j]=R(o);
    }
  }
  cooperative_groups::this_grid().sync();
}

__device__ __noinline__ void moe(const A& a,int layer,const float* x,float* out) {
  float* t=a.t; const int base=layer*11;
  small(a.router[layer],x,t+88000,64,2304);
  if(GID==0) {
    bool chosen[64]={false}; float val[8]; int ix[8];
    for(int j=0;j<8;j++) {
      float best=-INFINITY; int bi=0;
      for(int e=0;e<64;e++) if(!chosen[e] && t[88000+e]>best) { best=t[88000+e]; bi=e; }
      chosen[bi]=true; ix[j]=bi; val[j]=best;
    }
    float mx=val[0],den=0.f; for(int j=0;j<8;j++) den+=expf(val[j]-mx);
    for(int j=0;j<8;j++) {t[900200+j]=float(ix[j]);t[900208+j]=2.446f*expf(val[j]-mx)/den;}
  }
  cooperative_groups::this_grid().sync();
  for(int n=GID;n<2304;n+=GSTRIDE) t[65536+n]=0.f;
  cooperative_groups::this_grid().sync();
  for(int j=0;j<9;j++) {
    int e=j<8 ? int(t[900200+j]) : 0;
    int wi=base+(j<8 ? 5 : 8);
    gemv(a,wi,x,t+73728,e,false);
    gemv(a,wi+1,x,t+76800,e,false);
    for(int n=GID;n<1024;n+=GSTRIDE) t[79872+n]=silu(t[73728+n])*t[76800+n];
    cooperative_groups::this_grid().sync();
    gemv(a,wi+2,t+79872,t+82944,e,false);
    float weight=j<8 ? t[900208+j] : 1.f;
    for(int n=GID;n<2304;n+=GSTRIDE) t[65536+n]+=weight*t[82944+n];
    cooperative_groups::this_grid().sync();
  }
  for(int n=GID;n<2304;n+=GSTRIDE) out[n]=R(out[n]+R(t[65536+n]));
  cooperative_groups::this_grid().sync();
}

__device__ __forceinline__ float deq(const W& w,int k,int n) {
  unsigned char q=w.q[(k/2)*w.N+n];
  float z=B(w.z[(k/128)*w.N+n]),s=B(w.s[(k/128)*w.N+n]);
  return R(R(float((k&1)?q>>4:q&15)-z)*s);
}

__device__ __noinline__ void mla(const A& a,int layer,const float* x,float* attn_out) {
  float* t=a.t; int wi=layer*11;
  gemv(a,wi+0,x,t+16384,-1,true);
  gemv(a,wi+1,x,t+24576,-1,true);
  for(int j=GID;j<32;j+=GSTRIDE) {
    float ang=float(a.L)*powf(10000.f,-float(2*j)/64.f), co=cosf(ang),si=sinf(ang);
    float e=t[24576+512+2*j],o=t[24576+512+2*j+1];
    a.newr[a.L*64+2*j]=__float2bfloat16_rn(e*co-o*si);
    a.newr[a.L*64+2*j+1]=__float2bfloat16_rn(o*co+e*si);
    for(int h=0;h<32;h++) {
      int ix=16384+h*192+128+2*j;
      e=t[ix];o=t[ix+1];t[ix]=R(e*co-o*si);t[ix+1]=R(o*co+e*si);
    }
  }
  for(int j=GID;j<512;j+=GSTRIDE) a.newc[a.L*512+j]=__float2bfloat16_rn(t[24576+j]);
  for(int j=GID;j<a.L*512;j+=GSTRIDE) a.newc[j]=a.oldc[j];
  for(int j=GID;j<a.L*64;j+=GSTRIDE) a.newr[j]=a.oldr[j];
  cooperative_groups::this_grid().sync();
  const W& wb=a.w[wi+2];
  for(int hc=GID;hc<32*512;hc+=GSTRIDE) {
    int h=hc/512,c=hc%512;float v=0.f;
    for(int j=0;j<128;j++) v=fmaf(t[16384+h*192+j],deq(wb,c,h*256+j),v);
    t[800000+hc]=v;
  }
  cooperative_groups::this_grid().sync();
  for(int p=GID;p<(a.L+1)*32;p+=GSTRIDE) {
    int l=p/32,h=p%32;float v=0.f;
    for(int c=0;c<512;c++) v=fmaf(t[800000+h*512+c],B(a.newc[l*512+c]),v);
    for(int j=0;j<64;j++) v=fmaf(t[16384+h*192+128+j],B(a.newr[l*64+j]),v);
    t[100000+p]=v*0.07216878364870322f;
  }
  cooperative_groups::this_grid().sync();
  if(GID<32) {
    int h=GID;float mx=-INFINITY;
    for(int l=0;l<=a.L;l++) mx=fmaxf(mx,t[100000+l*32+h]);
    float den=0.f;for(int l=0;l<=a.L;l++) den+=expf(t[100000+l*32+h]-mx);
    t[900000+h]=mx;t[900100+h]=1.f/den;
  }
  cooperative_groups::this_grid().sync();
  for(int p=GID;p<(a.L+1)*32;p+=GSTRIDE) {
    int h=p%32;
    t[100000+p]=expf(t[100000+p]-t[900000+h])*t[900100+h];
  }
  cooperative_groups::this_grid().sync();
  for(int hc=GID;hc<32*512;hc+=GSTRIDE) {
    int h=hc/512,c=hc%512;float sum=0.f;
    for(int l=0;l<=a.L;l++) sum=fmaf(t[100000+l*32+h],B(a.newc[l*512+c]),sum);
    t[820000+hc]=sum;
  }
  cooperative_groups::this_grid().sync();
  for(int j=GID;j<4096;j+=GSTRIDE) {
    int h=j/128,n=h*256+128+j%128;float v=0.f;
    for(int c=0;c<512;c++) v=fmaf(t[820000+h*512+c],deq(wb,c,n),v);
    t[49152+j]=R(v);
  }
  cooperative_groups::this_grid().sync();
  gemv(a,wi+3,t+49152,attn_out,-1,true);
}

__global__ void decode(A a) {
  float* t=a.t;int tid=GID;
  for(int j=tid;j<2304;j+=GSTRIDE)t[j]=B(a.input[j]);
  cooperative_groups::this_grid().sync();
  for(int layer=0;layer<4;layer++) {
    int wi=layer*11;
    norm(a,t,a.norm[layer],t+8192);
    if(layer<3) {
      gemv(a,wi+0,t+8192,t+16384,-1,true);
      gemv(a,wi+1,t+8192,t+24576,-1,true);
      gemv(a,wi+2,t+8192,t+32768,-1,true);
      gemv(a,wi+3,t+8192,t+40960,-1,true);
      small(a.beta[layer],t+8192,t+88000,32,2304);
      conv(a,layer,0,t+16384);conv(a,layer,1,t+24576);conv(a,layer,2,t+32768);
      kda_state(a,layer,t+16384,t+24576,t+32768,t+40960,t+88000,t+49152);
      gemv(a,wi+4,t+49152,t+57344,-1,true);
    } else mla(a,layer,t+8192,t+57344);
    for(int j=tid;j<2304;j+=GSTRIDE)t[57344+j]=R(t[j]+t[57344+j]);
    cooperative_groups::this_grid().sync();
    norm(a,t+57344,a.mn[layer],t+8192);
    moe(a,layer,t+8192,t+57344);
    for(int j=tid;j<2304;j+=GSTRIDE)t[j]=t[57344+j];
    cooperative_groups::this_grid().sync();
  }
  for(int j=tid;j<2304;j+=GSTRIDE)a.output[j]=__float2bfloat16_rn(t[j]);
}

void launch(A a, cudaStream_t stream) { void* params[] = {&a}; cudaLaunchCooperativeKernel((void*)decode, dim3(128), dim3(128), params, 0, stream); }
'''


_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="kimi_linear_single_launch_v1",
            cpp_sources=CPP_SRC,
            cuda_sources=CUDA_SRC,
            functions=["run"],
            extra_cuda_cflags=["-O3", "-std=c++17", "--use_fast_math", "-ccbin=/usr/bin/g++-14", "-gencode=arch=compute_120,code=sm_120"],
            extra_cflags=["-O2"],
            with_cuda=True,
            no_implicit_headers=True,
            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._scratch = None

    def step(self, hidden, state):
        if self._scratch is None or self._scratch.device != hidden.device:
            self._scratch = torch.empty(1000000, device=hidden.device, dtype=torch.float32)
        output = torch.empty_like(hidden)
        L = state[3]["c_kv"].shape[0]
        oldc, oldr = state[3]["c_kv"], state[3]["k_rope"]
        newc = torch.empty((L + 1, 512), device=hidden.device, dtype=hidden.dtype)
        newr = torch.empty((L + 1, 64), device=hidden.device, dtype=hidden.dtype)
        ts = [hidden, output, self._scratch]
        for st in state[:3]:
            ts.extend((st["S"], st["cq"], st["ck"], st["cv"]))
        ts.extend((oldc, oldr, newc, newr))
        for blk in self.blocks:
            conv = blk.attn.conv_w if blk.kind == "K" else blk.attn_norm
            beta = blk.attn.beta_proj.weight if blk.kind == "K" else blk.attn_norm
            ts.extend((blk.attn_norm, blk.moe_norm, conv, beta, blk.moe.router.weight))
        for blk in self.blocks:
            if blk.kind == "K":
                linear = [blk.attn.q_proj, blk.attn.k_proj, blk.attn.v_proj, blk.attn.g_proj, blk.attn.o_proj]
            else:
                linear = [blk.attn.q_proj, blk.attn.kv_a, blk.attn.kv_b, blk.attn.o_proj, blk.attn.o_proj]
            linear += [blk.moe.gate, blk.moe.up, blk.moe.down,
                       blk.moe.s_gate, blk.moe.s_up, blk.moe.s_down]
            for mod in linear:
                ts.extend((mod.w_q, mod.scales, mod.zeros))
        _get_extension().run(ts, L)
        state[3]["c_kv"] = newc
        state[3]["k_rope"] = newr
        return output, state