Skip to main content

GEN_DEVICE_KERNELS

Constant GEN_DEVICE_KERNELS 

Source
pub const GEN_DEVICE_KERNELS: &str = r#"
__global__ void gk_embed(float* x, const signed char* wq, const float* ws, const int* d_tok, int E){
  int t=blockIdx.x*blockDim.x+threadIdx.x; if(t>=E) return;
  long tok=*d_tok; x[t]=(float)wq[tok*E+t]*ws[tok*(E/32)+t/32];
}
__global__ void gk_rmsnorm(float* out, const float* x, const float* w, int n, float eps){
  __shared__ float sh[256]; float s=0; for(int i=threadIdx.x;i<n;i+=blockDim.x) s+=x[i]*x[i];
  sh[threadIdx.x]=s; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
  float sc=rsqrtf(sh[0]/n+eps); for(int i=threadIdx.x;i<n;i+=blockDim.x) out[i]=x[i]*sc*w[i];
}
// FUSED residual-add + RMSNorm: `x += y` (the residual, kept for the next block) AND `out = rmsnorm(x)*w`, in
// ONE kernel launch and ONE pass over x. The reduction is the SAME 256-thread tree as gk_rmsnorm over the SAME
// summed values (v = x[i]+y[i]), so the float sum is bit-identical ⇒ token-exact. Removes a kernel node and the
// residual write/re-read per layer — the win in the latency-bound batch-1 H100 decode.
__global__ void gk_add_rmsnorm(float* x, const float* y, float* out, const float* w, int n, float eps){
  __shared__ float sh[256]; float s=0; for(int i=threadIdx.x;i<n;i+=blockDim.x){ float v=x[i]+y[i]; x[i]=v; s+=v*v; }
  sh[threadIdx.x]=s; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
  float sc=rsqrtf(sh[0]/n+eps); for(int i=threadIdx.x;i<n;i+=blockDim.x) out[i]=x[i]*sc*w[i];
}
// FUSED residual-add + RMSNorm + INT8 QUANT: like gk_add_rmsnorm, but ALSO emits `out` quantized to int8 per
// 32-block (oq + per-block scale os_) in the SAME kernel — folding the standalone `moe_qact(x2n)` away. `out`
// (float) still feeds the router; `oq`/`os_` (int8) feed the experts. The quant replicates moe_qact's per-32
// amax/127 scheme EXACTLY (one warp per 32-block, shfl amax) ⇒ the experts see bit-identical int8 (token-exact).
__global__ void gk_add_rmsnorm_q8(float* x, const float* y, float* out, const float* w, signed char* oq, float* os_, int n, float eps){
  __shared__ float sh[256]; float s=0; for(int i=threadIdx.x;i<n;i+=blockDim.x){ float v=x[i]+y[i]; x[i]=v; s+=v*v; }
  sh[threadIdx.x]=s; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
  float scale=rsqrtf(sh[0]/n+eps); for(int i=threadIdx.x;i<n;i+=blockDim.x) out[i]=x[i]*scale*w[i];
  __syncthreads();
  int warp=threadIdx.x>>5, lane=threadIdx.x&31, nb=n/32;
  for(int b=warp;b<nb;b+=(blockDim.x>>5)){ float v=out[b*32+lane], av=fabsf(v);
    for(int t=16;t>0;t>>=1) av=fmaxf(av,__shfl_down_sync(0xffffffffu,av,t));
    av=__shfl_sync(0xffffffffu,av,0); float amax=av>1e-30f?av:1e-30f; float qs=amax/127.0f;
    if(lane==0) os_[b]=qs; int q=(int)rintf(v/qs); if(q>127)q=127; if(q<-128)q=-128; oq[b*32+lane]=(signed char)q; }
}
__global__ void gk_rmsnorm_q8(const float* x, signed char* xq, float* xs, const float* w, int n, float eps){
  __shared__ float red[256]; float s=0; for(int i=threadIdx.x;i<n;i+=blockDim.x) s+=x[i]*x[i];
  red[threadIdx.x]=s; __syncthreads();
  for(int k=128;k>0;k>>=1){ if(threadIdx.x<k) red[threadIdx.x]+=red[threadIdx.x+k]; __syncthreads(); }
  float rms=rsqrtf(red[0]/n+eps); int warp=threadIdx.x>>5, lane=threadIdx.x&31, nb=n/32;
  for(int b=warp;b<nb;b+=8){ float y=x[b*32+lane]*rms*w[b*32+lane], a=fabsf(y);
    for(int sh=16;sh>0;sh>>=1) a=fmaxf(a,__shfl_down_sync(0xffffffffu,a,sh));
    a=__shfl_sync(0xffffffffu,a,0); float sc=a>0?a/127.0f:1.0f;
    if(lane==0) xs[b]=sc; xq[b*32+lane]=(signed char)__float2int_rn(y/sc); }
}
// FUSED residual-add + RMSNorm + int8 quant for the ATTENTION norm: `x += y` (residual kept) AND quantize
// rmsnorm(x)*w to int8 (xq/xs) — folds the post-MoE `gk_add` into the NEXT layer's attn norm+quant. Identical
// 128-thread reduction + per-32 quant as gk_rmsnorm_q8 over v=x[i]+y[i] ⇒ bit-identical ⇒ token-exact.
__global__ void gk_add_rmsnorm_q8a(float* x, const float* y, signed char* xq, float* xs, const float* w, int n, float eps){
  __shared__ float red[256]; float s=0; for(int i=threadIdx.x;i<n;i+=blockDim.x){ float v=x[i]+y[i]; x[i]=v; s+=v*v; }
  red[threadIdx.x]=s; __syncthreads();
  for(int k=128;k>0;k>>=1){ if(threadIdx.x<k) red[threadIdx.x]+=red[threadIdx.x+k]; __syncthreads(); }
  float rms=rsqrtf(red[0]/n+eps); int warp=threadIdx.x>>5, lane=threadIdx.x&31, nb=n/32;
  for(int b=warp;b<nb;b+=8){ float y2=x[b*32+lane]*rms*w[b*32+lane], a=fabsf(y2);
    for(int sh=16;sh>0;sh>>=1) a=fmaxf(a,__shfl_down_sync(0xffffffffu,a,sh));
    a=__shfl_sync(0xffffffffu,a,0); float sc=a>0?a/127.0f:1.0f;
    if(lane==0) xs[b]=sc; xq[b*32+lane]=(signed char)__float2int_rn(y2/sc); }
}
__global__ void gk_quant_act(const float* x, signed char* xq, float* xs, int K){
  int blk=blockIdx.x*(blockDim.x/32)+threadIdx.x/32; if((long)blk*32>=K) return; int lane=threadIdx.x&31;
  float v=x[blk*32+lane], a=fabsf(v);
  for(int s=16;s>0;s>>=1) a=fmaxf(a,__shfl_down_sync(0xffffffffu,a,s)); a=__shfl_sync(0xffffffffu,a,0);
  float sc=a>0?a/127.0f:1.0f; if(lane==0) xs[blk]=sc; xq[blk*32+lane]=(signed char)__float2int_rn(v/sc);
}
__global__ void gk_gemv(float* y, const signed char* wq, const float* ws, const signed char* xq, const float* xs, int K, int N){
  int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return; int lane=threadIdx.x&31;
  const int4* w16=(const int4*)(wq+(long)o*K); const int4* x16=(const int4*)xq; const float* wsr=ws+(long)o*(K/32);
  float acc=0.f; int nq=K/512;
  for(int b=0;b<nq;b++){ int4 w=w16[b*32+lane], xx=x16[b*32+lane]; int blk=b*16+lane/2;
    int d=__dp4a(w.x,xx.x,0)+__dp4a(w.y,xx.y,0)+__dp4a(w.z,xx.z,0)+__dp4a(w.w,xx.w,0);
    acc += (float)d * (wsr[blk]*xs[blk]); }
  int done=nq*512; for(int j=done+lane;j<K;j+=32){ int blk=j/32; acc += (float)((int)xq[j]*(int)wq[(long)o*K+j])*(wsr[blk]*xs[blk]); }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) y[o]=acc;
}
__global__ void gk_gemv_kv(float* cache, const int* d_pos, int KD, const signed char* wq, const float* ws, const signed char* xq, const float* xs, int K, int N){
  int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return; int lane=threadIdx.x&31;
  const int4* w16=(const int4*)(wq+(long)o*K); const int4* x16=(const int4*)xq; const float* wsr=ws+(long)o*(K/32);
  float acc=0.f; int nq=K/512;
  for(int b=0;b<nq;b++){ int4 w=w16[b*32+lane], xx=x16[b*32+lane]; int blk=b*16+lane/2;
    int d=__dp4a(w.x,xx.x,0)+__dp4a(w.y,xx.y,0)+__dp4a(w.z,xx.z,0)+__dp4a(w.w,xx.w,0);
    acc += (float)d * (wsr[blk]*xs[blk]); }
  int done=nq*512; for(int j=done+lane;j<K;j+=32){ int blk=j/32; acc += (float)((int)xq[j]*(int)wq[(long)o*K+j])*(wsr[blk]*xs[blk]); }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) cache[(long)(*d_pos)*KD+o]=acc;
}
// FUSED Q/K/V projection: one launch over QD+2·KD output warps, each picking its weight matrix by output range
// (Q→qbuf, K→cache[pos], V→cache[pos]) and reading the SHARED int8 activation once. Replaces three separate
// launch-bound gemvs with one — measured the single largest decode-throughput win (A100: +7 tok/s), because
// batch-1 decode is dominated by per-kernel launch/latency, not weight bandwidth.
__global__ void gk_gemv_qkv(float* qbuf,float* kc,float* vc,const int* d_pos,int KD,
    const signed char* wqq,const float* wqs,const signed char* wkq,const float* wks,const signed char* wvq,const float* wvs,
    const signed char* xq,const float* xs,int K,int QD){
  int og=(blockIdx.x*blockDim.x+threadIdx.x)>>5; int lane=threadIdx.x&31; int Ntot=QD+2*KD; if(og>=Ntot) return;
  const signed char* wq; const float* ws; int o; float* dst; long dsti;
  if(og<QD){ wq=wqq; ws=wqs; o=og; dst=qbuf; dsti=o; }
  else if(og<QD+KD){ wq=wkq; ws=wks; o=og-QD; dst=kc; dsti=(long)(*d_pos)*KD+o; }
  else { wq=wvq; ws=wvs; o=og-QD-KD; dst=vc; dsti=(long)(*d_pos)*KD+o; }
  const int4* w16=(const int4*)(wq+(long)o*K); const int4* x16=(const int4*)xq; const float* wsr=ws+(long)o*(K/32);
  float acc=0.f; int nq=K/512;
  for(int b=0;b<nq;b++){ int4 w=w16[b*32+lane], xx=x16[b*32+lane]; int blk=b*16+lane/2;
    int d=__dp4a(w.x,xx.x,0)+__dp4a(w.y,xx.y,0)+__dp4a(w.z,xx.z,0)+__dp4a(w.w,xx.w,0);
    acc += (float)d * (wsr[blk]*xs[blk]); }
  int done=nq*512; for(int j=done+lane;j<K;j+=32){ int blk=j/32; acc += (float)((int)xq[j]*(int)wq[(long)o*K+j])*(wsr[blk]*xs[blk]); }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) dst[dsti]=acc;
}
__global__ void gk_qk_norm_rope(float* base, const int* d_pos, int stride, const float* nw, int n_heads, int hd, float rbase, float eps){
  int head=blockIdx.x; if(head>=n_heads) return; float* h=base+(long)(*d_pos)*stride+head*hd;
  __shared__ float sh[128]; float s=0; for(int i=threadIdx.x;i<hd;i+=blockDim.x) s+=h[i]*h[i]; sh[threadIdx.x]=s; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
  float sc=rsqrtf(sh[0]/hd+eps); __syncthreads();
  for(int i=threadIdx.x;i<hd;i+=blockDim.x) h[i]=h[i]*sc*nw[i]; __syncthreads();
  int half=hd/2; int pos=*d_pos;
  for(int i=threadIdx.x;i<half;i+=blockDim.x){ float f=powf(rbase,-2.0f*i/hd); float a=pos*f,c=cosf(a),si=sinf(a);
    float x0=h[i],x1=h[i+half]; h[i]=x0*c-x1*si; h[i+half]=x0*si+x1*c; }
}
// FUSED Q+K per-head RMSNorm+RoPE: one launch over NH+NKV head-blocks, picking Q params (block<NH) or K params
// (block≥NH). Replaces two underfilled launches (NH and NKV blocks) with one — the Q and K norm-rope are the
// same body over different (base,stride,weight), so a branch on blockIdx fuses them (A100: +4 tok/s).
__global__ void gk_qk_norm_rope2(float* qbase, float* kbase, const int* d_pos, int kstride, const float* qnw, const float* knw, int nh, int nkv, int hd, float rbase, float eps){
  int b=blockIdx.x; float* base; const float* nw; int stride; int head;
  if(b<nh){ base=qbase; nw=qnw; stride=0; head=b; } else { base=kbase; nw=knw; stride=kstride; head=b-nh; }
  float* h=base+(long)(*d_pos)*stride+head*hd;
  __shared__ float sh[128]; float s=0; for(int i=threadIdx.x;i<hd;i+=blockDim.x) s+=h[i]*h[i]; sh[threadIdx.x]=s; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
  float sc=rsqrtf(sh[0]/hd+eps); __syncthreads();
  for(int i=threadIdx.x;i<hd;i+=blockDim.x) h[i]=h[i]*sc*nw[i]; __syncthreads();
  int half=hd/2; int pos=*d_pos;
  for(int i=threadIdx.x;i<half;i+=blockDim.x){ float f=powf(rbase,-2.0f*i/hd); float a=pos*f,c=cosf(a),si=sinf(a);
    float x0=h[i],x1=h[i+half]; h[i]=x0*c-x1*si; h[i+half]=x0*si+x1*c; }
}
// Flash-attention-style single-query decode: one block per query head, a WARP per cached key (parallel qh·kj
// dot, warp-reduced), a parallel block-softmax (shuffle max/sum + shared cross-warp reduce), then a parallel
// weighted sum over V. Q is normed+roped by the preceding gk_qk_norm_rope; fusing it in here perturbs the float
// order enough to flip a precise near-tie AND doesn't cut graph time (the work just moves) — so kept separate.
__global__ void gk_attn(float* out, const float* q, const float* kc, const float* vc, const int* d_pos, int n_head, int n_kv, int hd, float scale){
  int head=blockIdx.x; if(head>=n_head) return; int kvh=head/(n_head/n_kv); int L=*d_pos+1; int KD=n_kv*hd;
  const float* qh=q+(long)head*hd; extern __shared__ float sc[]; __shared__ float red[32];
  int lane=threadIdx.x&31, wp=threadIdx.x>>5, nw=(blockDim.x+31)>>5;
  for(int j=wp;j<L;j+=nw){ const float* kj=kc+(long)j*KD+kvh*hd; float s=0.f; for(int e=lane;e<hd;e+=32) s+=qh[e]*kj[e];
    for(int r=16;r>0;r>>=1) s+=__shfl_down_sync(0xffffffffu,s,r); if(lane==0) sc[j]=s*scale; }
  __syncthreads();
  float m=-1e30f; for(int j=threadIdx.x;j<L;j+=blockDim.x) m=fmaxf(m,sc[j]);
  for(int r=16;r>0;r>>=1) m=fmaxf(m,__shfl_down_sync(0xffffffffu,m,r)); if(lane==0) red[wp]=m; __syncthreads();
  if(threadIdx.x==0){ float mm=red[0]; for(int i=1;i<nw;i++) mm=fmaxf(mm,red[i]); red[0]=mm; } __syncthreads(); float mx=red[0];
  float ss=0.f; for(int j=threadIdx.x;j<L;j+=blockDim.x){ float e=expf(sc[j]-mx); sc[j]=e; ss+=e; }
  for(int r=16;r>0;r>>=1) ss+=__shfl_down_sync(0xffffffffu,ss,r); if(lane==0) red[wp]=ss; __syncthreads();
  if(threadIdx.x==0){ float su=0.f; for(int i=0;i<nw;i++) su+=red[i]; red[0]=su; } __syncthreads(); float den=red[0];
  for(int e=threadIdx.x;e<hd;e+=blockDim.x){ float a=0.f; for(int j=0;j<L;j++) a+=sc[j]*vc[(long)j*KD+kvh*hd+e]; out[(long)head*hd+e]=a/den; }
}
// ONLINE-SOFTMAX flash-attention decode (llama's fattn-vec shape): one block per query head, warps stream
// DISJOINT slabs of the KV cache, each keeping a RUNNING (max, sum) + weighted-V accumulator in REGISTERS —
// a single pass, warp-shuffle score reduce (no per-key __syncthreads, no shared score array), then ONE
// cross-warp flash-combine. Q lives in registers. Replaces the multi-pass score→sync→max→sync→exp→sync→V form.
__global__ void gk_attn_online(float* out, const float* q, const float* kc, const float* vc, const int* d_pos, int n_head, int n_kv, int hd, float scale){
  int head=blockIdx.x; if(head>=n_head) return; int kvh=head/(n_head/n_kv); int L=*d_pos+1; int KD=n_kv*hd;
  const float* qh=q+(long)head*hd;
  int lane=threadIdx.x&31, wp=threadIdx.x>>5, nw=blockDim.x>>5; int dpl=hd>>5;
  float qreg[8], acc[8]; for(int i=0;i<dpl;i++){ qreg[i]=qh[lane+(i<<5)]; acc[i]=0.f; }
  float m=-1e30f, dsum=0.f;
  for(int j=wp;j<L;j+=nw){
    const float* kj=kc+(long)j*KD+kvh*hd; float part=0.f;
    for(int i=0;i<dpl;i++) part+=qreg[i]*kj[lane+(i<<5)];
    for(int r=16;r>0;r>>=1) part+=__shfl_xor_sync(0xffffffffu,part,r);
    float sj=part*scale; float mn=fmaxf(m,sj), c=__expf(m-mn), p=__expf(sj-mn); m=mn; dsum=dsum*c+p;
    const float* vj=vc+(long)j*KD+kvh*hd;
    for(int i=0;i<dpl;i++) acc[i]=acc[i]*c + p*vj[lane+(i<<5)];
  }
  __shared__ float sm[8], sd[8], sa[8][128];
  if(lane==0){ sm[wp]=m; sd[wp]=dsum; }
  for(int i=0;i<dpl;i++) sa[wp][lane+(i<<5)]=acc[i];
  __syncthreads();
  float M=-1e30f; for(int w=0;w<nw;w++) M=fmaxf(M,sm[w]);
  float D=0.f; for(int w=0;w<nw;w++) D+=sd[w]*__expf(sm[w]-M);
  for(int d=threadIdx.x; d<hd; d+=blockDim.x){ float a=0.f; for(int w=0;w<nw;w++) a+=sa[w][d]*__expf(sm[w]-M); out[(long)head*hd+d]=a/D; }
}
// SPLIT-K flash-attention (llama's fattn-vec `parallel_blocks`): P blocks PER query head each stream a disjoint
// slab of the KV cache with a register online-softmax → a partial (local max, denom, un-normalized weighted-V)
// in scratch; grid = n_head*P for FULL SM occupancy at decode (32 heads × 4 = 128 blocks vs 32). `gk_attn_comb`
// then flash-merges the P partials per head. This is the measured +14 tok/s edge llama gets from `-fa`.
__global__ void gk_attn_part(const float* q, const float* kc, const float* vc, const int* d_pos, int n_head, int n_kv, int hd, float scale, int P, float* acc_part, float* m_part, float* d_part){
  int b=blockIdx.x; int head=b/P, p=b%P; if(head>=n_head) return;
  int kvh=head/(n_head/n_kv); int L=*d_pos+1; int KD=n_kv*hd; const float* qh=q+(long)head*hd;
  int lane=threadIdx.x&31, wp=threadIdx.x>>5, nw=blockDim.x>>5; int dpl=hd>>5;
  float qreg[8], acc[8]; for(int i=0;i<dpl;i++){ qreg[i]=qh[lane+(i<<5)]; acc[i]=0.f; }
  float m=-1e30f, dsum=0.f;
  for(int idx=wp; ; idx+=nw){ int j=p+idx*P; if(j>=L) break;
    const float* kj=kc+(long)j*KD+kvh*hd; float pt=0.f;
    for(int i=0;i<dpl;i++) pt+=qreg[i]*kj[lane+(i<<5)];
    for(int r=16;r>0;r>>=1) pt+=__shfl_xor_sync(0xffffffffu,pt,r);
    float sj=pt*scale; float mn=fmaxf(m,sj), c=__expf(m-mn), pp=__expf(sj-mn); m=mn; dsum=dsum*c+pp;
    const float* vj=vc+(long)j*KD+kvh*hd; for(int i=0;i<dpl;i++) acc[i]=acc[i]*c + pp*vj[lane+(i<<5)];
  }
  __shared__ float sm[8], sd[8], sa[8][128];
  if(lane==0){ sm[wp]=m; sd[wp]=dsum; }
  for(int i=0;i<dpl;i++) sa[wp][lane+(i<<5)]=acc[i];
  __syncthreads();
  float M=-1e30f; for(int w=0;w<nw;w++) M=fmaxf(M,sm[w]);
  float D=0.f; for(int w=0;w<nw;w++) D+=sd[w]*__expf(sm[w]-M);
  if(threadIdx.x==0){ m_part[b]=M; d_part[b]=D; }
  for(int d=threadIdx.x; d<hd; d+=blockDim.x){ float a=0.f; for(int w=0;w<nw;w++) a+=sa[w][d]*__expf(sm[w]-M); acc_part[(long)b*hd+d]=a; }
}
__global__ void gk_attn_comb(float* out, int n_head, int hd, int P, const float* acc_part, const float* m_part, const float* d_part){
  int head=blockIdx.x; if(head>=n_head) return; __shared__ float GM, GD;
  if(threadIdx.x==0){ float M=-1e30f; for(int p=0;p<P;p++) M=fmaxf(M,m_part[head*P+p]); GM=M;
    float D=0.f; for(int p=0;p<P;p++) D+=d_part[head*P+p]*__expf(m_part[head*P+p]-M); GD=D; }
  __syncthreads(); float M=GM, D=GD;
  for(int d=threadIdx.x; d<hd; d+=blockDim.x){ float a=0.f;
    for(int p=0;p<P;p++) a += acc_part[(long)(head*P+p)*hd+d]*__expf(m_part[head*P+p]-M);
    out[(long)head*hd+d]=a/D; }
}
// FUSED split-K combine + activation-quant: the combine already lands each head's hd attention outputs in one
// block, so it quantizes them into int8 32-blocks in its epilogue (per-32-block warp max-abs) instead of a
// separate gk_quant_act launch feeding the O-projection. One fewer kernel on the decode critical path (+~1).
__global__ void gk_attn_comb_q(signed char* aq, float* as_, int n_head, int hd, int P, const float* acc_part, const float* m_part, const float* d_part){
  int head=blockIdx.x; if(head>=n_head) return; __shared__ float GM,GD; __shared__ float so[128];
  if(threadIdx.x==0){ float M=-1e30f; for(int p=0;p<P;p++) M=fmaxf(M,m_part[head*P+p]); GM=M;
    float D=0.f; for(int p=0;p<P;p++) D+=d_part[head*P+p]*__expf(m_part[head*P+p]-M); GD=D; }
  __syncthreads(); float M=GM,D=GD;
  for(int d=threadIdx.x; d<hd; d+=blockDim.x){ float a=0.f; for(int p=0;p<P;p++) a+=acc_part[(long)(head*P+p)*hd+d]*__expf(m_part[head*P+p]-M); so[d]=a/D; }
  __syncthreads();
  int nb=hd/32, warp=threadIdx.x>>5, lane=threadIdx.x&31, nw=blockDim.x>>5;
  for(int blk=warp; blk<nb; blk+=nw){ int gi=head*nb+blk; float v=so[blk*32+lane], av=fabsf(v);
    for(int s2=16;s2>0;s2>>=1) av=fmaxf(av,__shfl_down_sync(0xffffffffu,av,s2));
    av=__shfl_sync(0xffffffffu,av,0); float amax=av>1e-30f?av:1e-30f; float qsc=amax/127.0f;
    if(lane==0) as_[gi]=qsc; int q=(int)rintf(v/qsc); if(q>127)q=127; if(q<-128)q=-128; aq[gi*32+lane]=(signed char)q; }
}
__global__ void gk_add(float* x, const float* y, int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n) x[i]+=y[i]; }
__global__ void gk_argmax_adv(const float* logits, int vocab, int* d_tok, int* d_pos){
  __shared__ float bv[256]; __shared__ int bi[256]; float bestv=-1e30f; int besti=0;
  for(int i=threadIdx.x;i<vocab;i+=blockDim.x){ float v=logits[i]; if(v>bestv){bestv=v;besti=i;} }
  bv[threadIdx.x]=bestv; bi[threadIdx.x]=besti; __syncthreads();
  for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k && bv[threadIdx.x+k]>bv[threadIdx.x]){ bv[threadIdx.x]=bv[threadIdx.x+k]; bi[threadIdx.x]=bi[threadIdx.x+k]; } __syncthreads(); }
  if(threadIdx.x==0){ *d_tok=bi[0]; *d_pos=*d_pos+1; }
}
// ── PRECISE (dequant) lane — llama's batch-1 recipe: f32 activation × int8 weight (no activation quantization),
// so the logits have clear margins (no int8-act near-ties) and the attention can be FUSED without flipping a
// token. Memory-bound on the same int8 weights ⇒ ~same bandwidth as the dp4a lane, but robustly llama-exact. ──
__global__ void gk_gemv_deq(float* y, const signed char* wq, const float* ws, const float* x, int K, int N){
  int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return; int lane=threadIdx.x&31;
  const int4* w16=(const int4*)(wq+(long)o*K); const float* wsr=ws+(long)o*(K/32); float acc=0.f; int nq=K/512;
  for(int b=0;b<nq;b++){ int4 wv=w16[b*32+lane]; int blk=b*16+lane/2; float sc=wsr[blk];
    const signed char* wb=(const signed char*)&wv; const float* xt=x+b*512+lane*16;
    float ba=0.f; for(int i=0;i<16;i++) ba += (float)wb[i]*xt[i]; acc += sc*ba; }
  for(int t=nq*512+lane;t<K;t+=32){ int blk=t/32; acc += wsr[blk]*(float)wq[(long)o*K+t]*x[t]; }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) y[o]=acc;
}
__global__ void gk_gemv_kv_deq(float* cache, const int* d_pos, int KD, const signed char* wq, const float* ws, const float* x, int K, int N){
  int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return; int lane=threadIdx.x&31;
  const int4* w16=(const int4*)(wq+(long)o*K); const float* wsr=ws+(long)o*(K/32); float acc=0.f; int nq=K/512;
  for(int b=0;b<nq;b++){ int4 wv=w16[b*32+lane]; int blk=b*16+lane/2; float sc=wsr[blk];
    const signed char* wb=(const signed char*)&wv; const float* xt=x+b*512+lane*16;
    float ba=0.f; for(int i=0;i<16;i++) ba += (float)wb[i]*xt[i]; acc += sc*ba; }
  for(int t=nq*512+lane;t<K;t+=32){ int blk=t/32; acc += wsr[blk]*(float)wq[(long)o*K+t]*x[t]; }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) cache[(long)(*d_pos)*KD+o]=acc;
}
// FUSED flash-attention: Q per-head QK-RMSNorm + RoPE + scores + block-softmax + weighted-V, all in one kernel.
__global__ void gk_fused_attn(float* out, const float* q, const float* kc, const float* vc, const int* d_pos, int n_head, int n_kv, int hd, float scale, const float* qnw, float rbase, float eps){
  int head=blockIdx.x; if(head>=n_head) return; int kvh=head/(n_head/n_kv); int pos=*d_pos; int L=pos+1; int KD=n_kv*hd;
  extern __shared__ float sc[]; __shared__ float red[32]; __shared__ float qh[256];
  int lane=threadIdx.x&31, wp=threadIdx.x>>5, nw=(blockDim.x+31)>>5;
  const float* qsrc=q+(long)head*hd;
  float qsum=0.f; for(int i=threadIdx.x;i<hd;i+=blockDim.x) qsum+=qsrc[i]*qsrc[i];
  for(int r=16;r>0;r>>=1) qsum+=__shfl_down_sync(0xffffffffu,qsum,r); if(lane==0) red[wp]=qsum; __syncthreads();
  if(threadIdx.x==0){ float s=0.f; for(int i=0;i<nw;i++) s+=red[i]; red[0]=s; } __syncthreads();
  float qrms=rsqrtf(red[0]/hd+eps);
  for(int i=threadIdx.x;i<hd;i+=blockDim.x) qh[i]=qsrc[i]*qrms*qnw[i]; __syncthreads();
  int half=hd/2;
  for(int i=threadIdx.x;i<half;i+=blockDim.x){ float f=powf(rbase,-2.0f*i/hd); float a=pos*f,c=cosf(a),si=sinf(a); float x0=qh[i],x1=qh[i+half]; qh[i]=x0*c-x1*si; qh[i+half]=x0*si+x1*c; }
  __syncthreads();
  for(int j=wp;j<L;j+=nw){ const float* kj=kc+(long)j*KD+kvh*hd; float s=0.f; for(int e=lane;e<hd;e+=32) s+=qh[e]*kj[e];
    for(int r=16;r>0;r>>=1) s+=__shfl_down_sync(0xffffffffu,s,r); if(lane==0) sc[j]=s*scale; }
  __syncthreads();
  float m=-1e30f; for(int j=threadIdx.x;j<L;j+=blockDim.x) m=fmaxf(m,sc[j]);
  for(int r=16;r>0;r>>=1) m=fmaxf(m,__shfl_down_sync(0xffffffffu,m,r)); if(lane==0) red[wp]=m; __syncthreads();
  if(threadIdx.x==0){ float mm=red[0]; for(int i=1;i<nw;i++) mm=fmaxf(mm,red[i]); red[0]=mm; } __syncthreads(); float mx=red[0];
  float ss=0.f; for(int j=threadIdx.x;j<L;j+=blockDim.x){ float e=expf(sc[j]-mx); sc[j]=e; ss+=e; }
  for(int r=16;r>0;r>>=1) ss+=__shfl_down_sync(0xffffffffu,ss,r); if(lane==0) red[wp]=ss; __syncthreads();
  if(threadIdx.x==0){ float su=0.f; for(int i=0;i<nw;i++) su+=red[i]; red[0]=su; } __syncthreads(); float den=red[0];
  for(int e=threadIdx.x;e<hd;e+=blockDim.x){ float a=0.f; for(int j=0;j<L;j++) a+=sc[j]*vc[(long)j*KD+kvh*hd+e]; out[(long)head*hd+e]=a/den; }
}
// Dequant MoE (f32 activation): gate/up SwiGLU + down-combine over the top-K experts, no activation quant.
__global__ void gk_moe_gate_up_deq(const float* x, const signed char* wg_q, const float* wg_s, const signed char* wu_q, const float* wu_s, const int* idx, float* h, int K, int N, int Kd){
  long row=(blockIdx.x*(long)blockDim.x+threadIdx.x)>>5; if(row>=(long)K*N) return; int lane=threadIdx.x&31;
  int kk=row/N, j=row%N; int e=idx[kk]; int nb=Kd/32, nq=Kd/512;
  const signed char* gwr=wg_q+((long)e*N+j)*Kd; const float* gsr=wg_s+((long)e*N+j)*nb;
  const signed char* uwr=wu_q+((long)e*N+j)*Kd; const float* usr=wu_s+((long)e*N+j)*nb;
  const int4* gw16=(const int4*)gwr; const int4* uw16=(const int4*)uwr; float ga=0.f, ua=0.f;
  for(int b=0;b<nq;b++){ int blk=b*16+lane/2; int4 gv=gw16[b*32+lane], uv=uw16[b*32+lane];
    const signed char* gb=(const signed char*)&gv; const signed char* ub=(const signed char*)&uv; const float* xt=x+b*512+lane*16;
    float gsc=gsr[blk], usc=usr[blk]; float gba=0.f, uba=0.f;
    for(int i=0;i<16;i++){ float xi=xt[i]; gba+=(float)gb[i]*xi; uba+=(float)ub[i]*xi; } ga+=gsc*gba; ua+=usc*uba; }
  for(int t=nq*512+lane;t<Kd;t+=32){ int blk=t/32; float xi=x[t]; ga+=gsr[blk]*(float)gwr[t]*xi; ua+=usr[blk]*(float)uwr[t]*xi; }
  for(int s=16;s>0;s>>=1){ ga+=__shfl_down_sync(0xffffffffu,ga,s); ua+=__shfl_down_sync(0xffffffffu,ua,s); }
  if(lane==0) h[row]=(ga/(1.0f+expf(-ga)))*ua;
}
__global__ void gk_moe_down_deq(const float* hbuf, const signed char* wd_q, const float* wd_s, const int* idx, const float* wc, float* H, int K, int N, int Kd){
  int o=blockIdx.x; int warp=threadIdx.x>>5, lane=threadIdx.x&31; int e=idx[warp]; int nb=Kd/32, nq=Kd/512;
  const float* hb=hbuf+(long)warp*Kd; const signed char* wr=wd_q+((long)e*N+o)*Kd; const float* wsr=wd_s+((long)e*N+o)*nb;
  const int4* w16=(const int4*)wr; float acc=0.f;
  for(int b=0;b<nq;b++){ int blk=b*16+lane/2; int4 wv=w16[b*32+lane]; const signed char* wb=(const signed char*)&wv; const float* ht=hb+b*512+lane*16;
    float sc=wsr[blk]; float ba=0.f; for(int i=0;i<16;i++) ba += (float)wb[i]*ht[i]; acc += sc*ba; }
  for(int t=nq*512+lane;t<Kd;t+=32){ int blk=t/32; acc += wsr[blk]*(float)wr[t]*hb[t]; }
  for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s);
  __shared__ float part[64]; if(lane==0) part[warp]=wc[warp]*acc; __syncthreads();
  if(threadIdx.x==0){ float sum=0.f; for(int kk=0;kk<K;kk++) sum+=part[kk]; H[o]=sum; }
}
"#;