KernelBench mega · H100
Kimi-Linear Decode GPT-5.6 Sol
4.53×geomean speedup across shapes
manually audited: clean
harnesscodex
Kernel source (redacted)
"""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
20260721_125927_codex_gpt-5.6-sol_02_kimi_linear_decode