pub const QWEN38_EXTRA_KERNELS: &str = r#"
// A plain argmax that does NOT advance the position — the MTP draft is a guess about the future, not a step
// through it, so it must not touch the decode state it is speculating on.
__global__ void gk_argmax_only(const float* logits, int vocab, int* out){
__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) *out=bi[0];
}
// Concatenate the two normed halves the nextn adapter consumes: [ rmsnorm(embed(t),enorm) | rmsnorm(h,hnorm) ].
__global__ void gk_mtp_cat(float* __restrict__ out, const float* __restrict__ e, const float* __restrict__ h,
const float* __restrict__ we, const float* __restrict__ wh, int n, float eps){
__shared__ float sh[256];
const float* src = blockIdx.x == 0 ? e : h;
const float* w = blockIdx.x == 0 ? we : wh;
float* dst = out + (long)blockIdx.x * n;
float s=0.f; for(int i=threadIdx.x;i<n;i+=blockDim.x) s+=src[i]*src[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) dst[i]=src[i]*sc*w[i];
}
// ══ SPLIT-K FOR THE NARROW PROJECTIONS ══════════════════════════════════════════════════════════
// One warp per output row means the o-projection and the down-projection (both N=5120) launch only 640
// blocks — 38.8 of an H100's 64 warps per SM, against 100% for the wide matrices. Those two are 31.7% of the
// model's weights, so a third of every token runs at 61% occupancy. Splitting K into S chunks multiplies the
// block count by S and fills the machine.
//
// The combine is a SEPARATE kernel summing the partials in FIXED ascending order, not an atomicAdd: float
// addition is not associative, so an atomic would make the result depend on scheduler order and the same
// prompt could produce different tokens on different runs. Determinism is not negotiable here.
//
// This is NOT bit-identical to the rolled form — chunking K re-associates the sum — so it is gated on the
// golden token stream and on distance from the f64 oracle, never adopted on throughput alone.
__global__ void gk_gemv_hl_sk(float* __restrict__ part, const signed char* __restrict__ wq,
const __half* __restrict__ ws, const signed char* __restrict__ xh,
const signed char* __restrict__ xl, const float* __restrict__ xs,
int K, int N, int S){
int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return;
int sl=blockIdx.y, lane=threadIdx.x&31;
const int4* w16=(const int4*)(wq+(long)o*K);
const int4* h16=(const int4*)xh; const int4* l16=(const int4*)xl;
const __half* wsr=ws+(long)o*(K/32);
int nq=K/512, per=(nq+S-1)/S, b0=sl*per, be=b0+per<nq?b0+per:nq;
float acc=0.f;
for(int b=b0;b<be;b++){
int4 w=w16[b*32+lane], hh=h16[b*32+lane], ll=l16[b*32+lane];
int blk=b*16+lane/2;
int dh=__dp4a(w.x,hh.x,0)+__dp4a(w.y,hh.y,0)+__dp4a(w.z,hh.z,0)+__dp4a(w.w,hh.w,0);
int dl=__dp4a(w.x,ll.x,0)+__dp4a(w.y,ll.y,0)+__dp4a(w.z,ll.z,0)+__dp4a(w.w,ll.w,0);
acc += ((float)dh + (float)dl*(1.0f/127.0f)) * (__half2float(wsr[blk])*xs[blk]);
}
if(sl==S-1){ int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=(float)xh[j] + (float)xl[j]*(1.0f/127.0f);
acc += a*(float)wq[(long)o*K+j]*(__half2float(wsr[blk])*xs[blk]); } }
for(int t=16;t>0;t>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,t);
if(lane==0) part[(long)sl*N+o]=acc;
}
// Fixed ascending order over the K-chunks: same inputs always give the same bits.
__global__ void gk_sk_combine(float* __restrict__ y, const float* __restrict__ part, int N, int S){
int o=blockIdx.x*blockDim.x+threadIdx.x; if(o>=N) return;
float a=0.f;
for(int s=0;s<S;++s) a+=part[(long)s*N+o];
y[o]=a;
}
// Four K-blocks in flight. Two measured +5.2% on the two-pass lane (38.30 -> 40.29); this is the same idea
// pushed one step further, and stops there because the tail handling and register pressure both start to
// cost more than the extra loads return. Accumulation order is unchanged, so it is still bit-identical.
__global__ void gk_gemv_hlu4(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xh, const signed char* __restrict__ xl,
const float* __restrict__ 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* h16=(const int4*)xh; const int4* l16=(const int4*)xl;
const __half* wsr=ws+(long)o*(K/32);
float acc=0.f; int nq=K/512; int b=0;
for(; b+3<nq; b+=4){
int4 w0=w16[b*32+lane], h0=h16[b*32+lane], l0=l16[b*32+lane];
int4 w1=w16[(b+1)*32+lane], h1=h16[(b+1)*32+lane], l1=l16[(b+1)*32+lane];
int4 w2=w16[(b+2)*32+lane], h2=h16[(b+2)*32+lane], l2=l16[(b+2)*32+lane];
int4 w3=w16[(b+3)*32+lane], h3=h16[(b+3)*32+lane], l3=l16[(b+3)*32+lane];
int k0=b*16+lane/2, k1=(b+1)*16+lane/2, k2=(b+2)*16+lane/2, k3=(b+3)*16+lane/2;
int a0=__dp4a(w0.x,h0.x,0)+__dp4a(w0.y,h0.y,0)+__dp4a(w0.z,h0.z,0)+__dp4a(w0.w,h0.w,0);
int c0=__dp4a(w0.x,l0.x,0)+__dp4a(w0.y,l0.y,0)+__dp4a(w0.z,l0.z,0)+__dp4a(w0.w,l0.w,0);
int a1=__dp4a(w1.x,h1.x,0)+__dp4a(w1.y,h1.y,0)+__dp4a(w1.z,h1.z,0)+__dp4a(w1.w,h1.w,0);
int c1=__dp4a(w1.x,l1.x,0)+__dp4a(w1.y,l1.y,0)+__dp4a(w1.z,l1.z,0)+__dp4a(w1.w,l1.w,0);
int a2=__dp4a(w2.x,h2.x,0)+__dp4a(w2.y,h2.y,0)+__dp4a(w2.z,h2.z,0)+__dp4a(w2.w,h2.w,0);
int c2=__dp4a(w2.x,l2.x,0)+__dp4a(w2.y,l2.y,0)+__dp4a(w2.z,l2.z,0)+__dp4a(w2.w,l2.w,0);
int a3=__dp4a(w3.x,h3.x,0)+__dp4a(w3.y,h3.y,0)+__dp4a(w3.z,h3.z,0)+__dp4a(w3.w,h3.w,0);
int c3=__dp4a(w3.x,l3.x,0)+__dp4a(w3.y,l3.y,0)+__dp4a(w3.z,l3.z,0)+__dp4a(w3.w,l3.w,0);
acc += ((float)a0 + (float)c0*(1.0f/127.0f)) * (__half2float(wsr[k0])*xs[k0]);
acc += ((float)a1 + (float)c1*(1.0f/127.0f)) * (__half2float(wsr[k1])*xs[k1]);
acc += ((float)a2 + (float)c2*(1.0f/127.0f)) * (__half2float(wsr[k2])*xs[k2]);
acc += ((float)a3 + (float)c3*(1.0f/127.0f)) * (__half2float(wsr[k3])*xs[k3]);
}
for(; b<nq; b++){
int4 w=w16[b*32+lane], hh=h16[b*32+lane], ll=l16[b*32+lane]; int blk=b*16+lane/2;
int dh=__dp4a(w.x,hh.x,0)+__dp4a(w.y,hh.y,0)+__dp4a(w.z,hh.z,0)+__dp4a(w.w,hh.w,0);
int dl=__dp4a(w.x,ll.x,0)+__dp4a(w.y,ll.y,0)+__dp4a(w.z,ll.z,0)+__dp4a(w.w,ll.w,0);
acc += ((float)dh + (float)dl*(1.0f/127.0f)) * (__half2float(wsr[blk])*xs[blk]);
}
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=(float)xh[j] + (float)xl[j]*(1.0f/127.0f);
acc += a*(float)wq[(long)o*K+j]*(__half2float(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;
}
// ══ TWO K-BLOCKS IN FLIGHT ══════════════════════════════════════════════════════════════════════
// Widening in N (two output rows per warp) was MEASURED slower on A100 -- 41.10 against 42.28 -- because it
// halves the warp count, and this GEMV needs the parallelism more than it needs the reuse. Unrolling in K
// costs no warps: the same warp issues the loads for two consecutive 512-weight blocks before consuming
// either, so there are twice as many loads in flight per warp with the same occupancy.
//
// The accumulation ORDER is deliberately unchanged -- one accumulator, blocks folded in ascending order --
// so this is bit-identical to the rolled form and cannot move a token.
__global__ void gk_gemv_hu(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xq, const float* __restrict__ 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 __half* wsr=ws+(long)o*(K/32);
float acc=0.f; int nq=K/512; int b=0;
for(; b+1<nq; b+=2){
int4 wa=w16[b*32+lane], xa=x16[b*32+lane];
int4 wb=w16[(b+1)*32+lane], xb=x16[(b+1)*32+lane];
int ba=b*16+lane/2, bb=(b+1)*16+lane/2;
int da=__dp4a(wa.x,xa.x,0)+__dp4a(wa.y,xa.y,0)+__dp4a(wa.z,xa.z,0)+__dp4a(wa.w,xa.w,0);
int db=__dp4a(wb.x,xb.x,0)+__dp4a(wb.y,xb.y,0)+__dp4a(wb.z,xb.z,0)+__dp4a(wb.w,xb.w,0);
acc += (float)da * (__half2float(wsr[ba])*xs[ba]);
acc += (float)db * (__half2float(wsr[bb])*xs[bb]);
}
for(; 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 * (__half2float(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])*(__half2float(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;
}
// The same unrolling for the two-pass lane, where each iteration already has three streams to keep in flight.
__global__ void gk_gemv_hlu(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xh, const signed char* __restrict__ xl,
const float* __restrict__ 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* h16=(const int4*)xh; const int4* l16=(const int4*)xl;
const __half* wsr=ws+(long)o*(K/32);
float acc=0.f; int nq=K/512; int b=0;
for(; b+1<nq; b+=2){
int4 wa=w16[b*32+lane], ha=h16[b*32+lane], la=l16[b*32+lane];
int4 wb=w16[(b+1)*32+lane], hb=h16[(b+1)*32+lane], lb=l16[(b+1)*32+lane];
int ba=b*16+lane/2, bb=(b+1)*16+lane/2;
int dha=__dp4a(wa.x,ha.x,0)+__dp4a(wa.y,ha.y,0)+__dp4a(wa.z,ha.z,0)+__dp4a(wa.w,ha.w,0);
int dla=__dp4a(wa.x,la.x,0)+__dp4a(wa.y,la.y,0)+__dp4a(wa.z,la.z,0)+__dp4a(wa.w,la.w,0);
int dhb=__dp4a(wb.x,hb.x,0)+__dp4a(wb.y,hb.y,0)+__dp4a(wb.z,hb.z,0)+__dp4a(wb.w,hb.w,0);
int dlb=__dp4a(wb.x,lb.x,0)+__dp4a(wb.y,lb.y,0)+__dp4a(wb.z,lb.z,0)+__dp4a(wb.w,lb.w,0);
acc += ((float)dha + (float)dla*(1.0f/127.0f)) * (__half2float(wsr[ba])*xs[ba]);
acc += ((float)dhb + (float)dlb*(1.0f/127.0f)) * (__half2float(wsr[bb])*xs[bb]);
}
for(; b<nq; b++){
int4 w=w16[b*32+lane], hh=h16[b*32+lane], ll=l16[b*32+lane]; int blk=b*16+lane/2;
int dh=__dp4a(w.x,hh.x,0)+__dp4a(w.y,hh.y,0)+__dp4a(w.z,hh.z,0)+__dp4a(w.w,hh.w,0);
int dl=__dp4a(w.x,ll.x,0)+__dp4a(w.y,ll.y,0)+__dp4a(w.z,ll.z,0)+__dp4a(w.w,ll.w,0);
acc += ((float)dh + (float)dl*(1.0f/127.0f)) * (__half2float(wsr[blk])*xs[blk]);
}
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=(float)xh[j] + (float)xl[j]*(1.0f/127.0f);
acc += a*(float)wq[(long)o*K+j]*(__half2float(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;
}
// ══ TWO OUTPUT ROWS PER WARP ════════════════════════════════════════════════════════════════════
// A batch-1 GEMV at these shapes is not bandwidth-saturated -- it is LATENCY bound: one row per warp gives
// each warp a single outstanding weight load per iteration, and the measured 1050 GB/s against an A100 that
// delivers ~1500 is the gap that leaves. Giving a warp two adjacent rows doubles the loads in flight, and
// the activation word (and its scale) is fetched once and used twice, which halves the L2 traffic the
// activation generates across the 7760 warps that read it.
__global__ void gk_gemv_h2(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xq, const float* __restrict__ xs, int K, int N){
int r=((blockIdx.x*blockDim.x+threadIdx.x)>>5)*2; if(r>=N) return; int lane=threadIdx.x&31;
bool r1 = (r+1)<N;
const int4* x16=(const int4*)xq;
const int4* wa=(const int4*)(wq+(long)r*K);
const int4* wb=(const int4*)(wq+(long)(r+1)*K);
const __half* sa=ws+(long)r*(K/32); const __half* sb=ws+(long)(r+1)*(K/32);
float acc0=0.f, acc1=0.f; int nq=K/512;
for(int b=0;b<nq;b++){
int4 xx=x16[b*32+lane]; int blk=b*16+lane/2; float xsv=xs[blk];
int4 w0=wa[b*32+lane];
int d0=__dp4a(w0.x,xx.x,0)+__dp4a(w0.y,xx.y,0)+__dp4a(w0.z,xx.z,0)+__dp4a(w0.w,xx.w,0);
acc0 += (float)d0 * (__half2float(sa[blk])*xsv);
if(r1){ int4 w1=wb[b*32+lane];
int d1=__dp4a(w1.x,xx.x,0)+__dp4a(w1.y,xx.y,0)+__dp4a(w1.z,xx.z,0)+__dp4a(w1.w,xx.w,0);
acc1 += (float)d1 * (__half2float(sb[blk])*xsv); }
}
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32; float xv=(float)xq[j]*xs[blk];
acc0 += xv*(float)wq[(long)r*K+j]*__half2float(sa[blk]);
if(r1) acc1 += xv*(float)wq[(long)(r+1)*K+j]*__half2float(sb[blk]); }
for(int s=16;s>0;s>>=1){ acc0+=__shfl_down_sync(0xffffffffu,acc0,s); acc1+=__shfl_down_sync(0xffffffffu,acc1,s); }
if(lane==0){ y[r]=acc0; if(r1) y[r+1]=acc1; }
}
// The same widening for the two-pass lane, where it pays double: the activation's HIGH and LOW words are both
// fetched once per iteration and used for both rows.
__global__ void gk_gemv_hl2(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xh, const signed char* __restrict__ xl,
const float* __restrict__ xs, int K, int N){
int r=((blockIdx.x*blockDim.x+threadIdx.x)>>5)*2; if(r>=N) return; int lane=threadIdx.x&31;
bool r1 = (r+1)<N;
const int4* h16=(const int4*)xh; const int4* l16=(const int4*)xl;
const int4* wa=(const int4*)(wq+(long)r*K);
const int4* wb=(const int4*)(wq+(long)(r+1)*K);
const __half* sa=ws+(long)r*(K/32); const __half* sb=ws+(long)(r+1)*(K/32);
float acc0=0.f, acc1=0.f; int nq=K/512;
for(int b=0;b<nq;b++){
int4 hh=h16[b*32+lane], ll=l16[b*32+lane]; int blk=b*16+lane/2; float xsv=xs[blk];
int4 w0=wa[b*32+lane];
int dh0=__dp4a(w0.x,hh.x,0)+__dp4a(w0.y,hh.y,0)+__dp4a(w0.z,hh.z,0)+__dp4a(w0.w,hh.w,0);
int dl0=__dp4a(w0.x,ll.x,0)+__dp4a(w0.y,ll.y,0)+__dp4a(w0.z,ll.z,0)+__dp4a(w0.w,ll.w,0);
acc0 += ((float)dh0 + (float)dl0*(1.0f/127.0f)) * (__half2float(sa[blk])*xsv);
if(r1){ int4 w1=wb[b*32+lane];
int dh1=__dp4a(w1.x,hh.x,0)+__dp4a(w1.y,hh.y,0)+__dp4a(w1.z,hh.z,0)+__dp4a(w1.w,hh.w,0);
int dl1=__dp4a(w1.x,ll.x,0)+__dp4a(w1.y,ll.y,0)+__dp4a(w1.z,ll.z,0)+__dp4a(w1.w,ll.w,0);
acc1 += ((float)dh1 + (float)dl1*(1.0f/127.0f)) * (__half2float(sb[blk])*xsv); }
}
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=((float)xh[j] + (float)xl[j]*(1.0f/127.0f))*xs[blk];
acc0 += a*(float)wq[(long)r*K+j]*__half2float(sa[blk]);
if(r1) acc1 += a*(float)wq[(long)(r+1)*K+j]*__half2float(sb[blk]); }
for(int s=16;s>0;s>>=1){ acc0+=__shfl_down_sync(0xffffffffu,acc0,s); acc1+=__shfl_down_sync(0xffffffffu,acc1,s); }
if(lane==0){ y[r]=acc0; if(r1) y[r+1]=acc1; }
}
// ══ TWO-PASS INT8 ACTIVATIONS ═══════════════════════════════════════════════════════════════════
// A batch-1 GEMV reads a whole weight matrix and ONE activation vector, so the activation costs no
// measurable bandwidth -- but rounding it onto a single int8 grid costs accuracy at every one of the 300+
// projections in a token. Splitting it into a high and a low int8 word represents it as
// a ~= s * (qh + ql/127), | a/s - (qh + ql/127) | <= 0.5/127
// which is 254x tighter, and keeps `dp4a`. That last part is the point: the obvious alternative -- feeding
// the GEMV f32 activations -- was MEASURED at 32.27 tok/s against 42.28, because one int-to-float convert
// per weight is quarter rate on this chip. Two words of int8 double the dp4a count instead, on a kernel
// that is memory bound, and the weight matrix is still read exactly once.
__global__ void gk_gemv_hl(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xh, const signed char* __restrict__ xl,
const float* __restrict__ 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* h16=(const int4*)xh; const int4* l16=(const int4*)xl;
const __half* 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], hh=h16[b*32+lane], ll=l16[b*32+lane];
int blk=b*16+lane/2;
int dh=__dp4a(w.x,hh.x,0)+__dp4a(w.y,hh.y,0)+__dp4a(w.z,hh.z,0)+__dp4a(w.w,hh.w,0);
int dl=__dp4a(w.x,ll.x,0)+__dp4a(w.y,ll.y,0)+__dp4a(w.z,ll.z,0)+__dp4a(w.w,ll.w,0);
acc += ((float)dh + (float)dl*(1.0f/127.0f)) * (__half2float(wsr[blk])*xs[blk]);
}
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=(float)xh[j] + (float)xl[j]*(1.0f/127.0f);
acc += a*(float)wq[(long)o*K+j]*(__half2float(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;
}
// The producers, each the single-word kernel with its final rounding replaced by the two-word split. The
// scale and every float value are computed exactly as before, so the high word is bit-identical to what the
// single-word lane produced and the low word is pure addition.
__global__ void gk_quant_act_hl(const float* x, signed char* xq, signed char* xl, 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;
lgx_q8hl(v/sc, xq, xl, (long)blk*32+lane);
}
__global__ void gk_rmsnorm_q8f_hl(const float* __restrict__ x, signed char* __restrict__ xq,
signed char* __restrict__ xl, float* __restrict__ xs,
float* __restrict__ xf, const float* __restrict__ w, int n, float eps){
__shared__ float sh[256];
float s=0.f; 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); __syncthreads();
for(int i=threadIdx.x;i<n;i+=blockDim.x) xf[i]=x[i]*sc*w[i];
__syncthreads();
int nb=n/32;
for(int b=threadIdx.x;b<nb;b+=blockDim.x){
float amax=1e-30f;
for(int j=0;j<32;++j){ float a=fabsf(xf[b*32+j]); if(a>amax) amax=a; }
float q=amax/127.0f; xs[b]=q;
for(int j=0;j<32;++j) lgx_q8hl(xf[b*32+j]/q, xq, xl, (long)b*32+j);
}
}
__global__ void gk_add_rmsnorm_q8f_hl(float* __restrict__ x, const float* __restrict__ y,
signed char* __restrict__ xq, signed char* __restrict__ xl,
float* __restrict__ xs, float* __restrict__ xf,
const float* __restrict__ w, int n, float eps){
__shared__ float sh[256];
float s=0.f;
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); __syncthreads();
for(int i=threadIdx.x;i<n;i+=blockDim.x) xf[i]=x[i]*sc*w[i];
__syncthreads();
int nb=n/32;
for(int b=threadIdx.x;b<nb;b+=blockDim.x){
float amax=1e-30f;
for(int j=0;j<32;++j){ float a=fabsf(xf[b*32+j]); if(a>amax) amax=a; }
float q=amax/127.0f; xs[b]=q;
for(int j=0;j<32;++j) lgx_q8hl(xf[b*32+j]/q, xq, xl, (long)b*32+j);
}
}
__global__ void gk_add_rmsnorm_q8_hl(float* x, const float* y, float* out, const float* w,
signed char* oq, signed char* ol, 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; lgx_q8hl(v/qs, oq, ol, (long)b*32+lane); }
}
__global__ void gk_silu_had_q8_hl(const float* __restrict__ g, const float* __restrict__ u,
signed char* __restrict__ hq, signed char* __restrict__ hl_,
float* __restrict__ hs, int n){
int blk = blockIdx.x * blockDim.x + threadIdx.x; int nb = n / 32; if(blk >= nb) return;
float amax = 1e-30f; float t[32];
for(int j = 0; j < 32; ++j){ float gv = g[blk*32+j]; float h = (gv / (1.0f + __expf(-gv))) * u[blk*32+j];
t[j] = h; float a = fabsf(h); if(a > amax) amax = a; }
float sc = amax / 127.0f; hs[blk] = sc;
for(int j = 0; j < 32; ++j) lgx_q8hl(t[j]/sc, hq, hl_, (long)blk*32+j);
}
// The Q8_0 GEMV that keeps its ACTIVATION in f32: the weights are still int8 with a __half block scale, so
// the memory traffic — the thing that decides batch-1 decode time — is byte for byte the same as the dp4a
// form. Only the arithmetic changes, from a quarter of a dp4a per weight to about three float ops, which
// leaves the kernel comfortably memory bound. What it removes is the activation quantization error, which at
// batch 1 was never buying anything: one activation vector against a whole weight matrix is not where the
// bandwidth goes.
__global__ void gk_gemv_deq_h(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const float* __restrict__ 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 __half* wsr=ws+(long)o*(K/32);
const float4* x4=(const float4*)x;
float acc=0.f; int nq=K/512;
// The four weight words stay in REGISTERS and are unpacked by shifting. Reading them back through a
// `(signed char*)&wv` alias instead puts the int4 in local memory, and the resulting spill cost 12% of
// throughput -- on a kernel whose whole premise was that it would be free.
#define LGX_W(W,S) ((float)(int)(signed char)(((W)>>(S))&0xff))
for(int b=0;b<nq;b++){
int4 wv=w16[b*32+lane];
int base4=b*128+lane*4;
float4 a0=x4[base4], a1=x4[base4+1], a2=x4[base4+2], a3=x4[base4+3];
float p = LGX_W(wv.x,0)*a0.x + LGX_W(wv.x,8)*a0.y + LGX_W(wv.x,16)*a0.z + LGX_W(wv.x,24)*a0.w
+ LGX_W(wv.y,0)*a1.x + LGX_W(wv.y,8)*a1.y + LGX_W(wv.y,16)*a1.z + LGX_W(wv.y,24)*a1.w
+ LGX_W(wv.z,0)*a2.x + LGX_W(wv.z,8)*a2.y + LGX_W(wv.z,16)*a2.z + LGX_W(wv.z,24)*a2.w
+ LGX_W(wv.w,0)*a3.x + LGX_W(wv.w,8)*a3.y + LGX_W(wv.w,16)*a3.z + LGX_W(wv.w,24)*a3.w;
acc += p*__half2float(wsr[b*16+lane/2]);
}
#undef LGX_W
int done=nq*512;
for(int j=done+lane;j<K;j+=32){ int blk=j/32; acc += (float)wq[(long)o*K+j]*__half2float(wsr[blk])*x[j]; }
for(int s=16;s>0;s>>=1) acc+=__shfl_down_sync(0xffffffffu,acc,s); if(lane==0) y[o]=acc;
}
// The Q8_0 GEMV with the block scale read as __half — identical arithmetic, 2 bytes per 32 weights instead
// of 4. At batch 1 the whole decode is a stream of these, so the 5.6% of weight traffic this returns is
// 5.6% off the token time. The scale was an f16 in the checkpoint; widening it on the way to the GPU only
// ever cost bandwidth.
__global__ void gk_gemv_h(float* y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ xq, const float* __restrict__ 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 __half* 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 * (__half2float(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])*(__half2float(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_embed_h(float* x, const signed char* wq, const __half* 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]*__half2float(ws[tok*(E/32)+t/32]);
}
// RMSNorm that emits BOTH the int8-quantized activation (for the big Q8_0 GEMVs) and the f32 normed vector
// (for GDN's two small f32-weight gate projections). Emitting both from one pass costs a 20 KB store and
// saves re-reading and re-normalizing the hidden state.
// The same, with the RESIDUAL ADD fused in. Every layer but the first is entered from a residual add, and
// the add is a full pass over the hidden state that the norm is about to read again — so the layer entry
// folds `x += y` into the reduction it already does. Removes one launch and one round-trip per layer.
__global__ void gk_add_rmsnorm_q8f(float* __restrict__ x, const float* __restrict__ y,
signed char* __restrict__ xq, float* __restrict__ xs,
float* __restrict__ xf, const float* __restrict__ w, int n, float eps){
__shared__ float sh[256];
float s = 0.f;
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); __syncthreads();
for(int i = threadIdx.x; i < n; i += blockDim.x) xf[i] = x[i] * sc * w[i];
__syncthreads();
int nb = n / 32;
for(int b = threadIdx.x; b < nb; b += blockDim.x){
float amax = 1e-30f;
for(int j = 0; j < 32; ++j){ float a = fabsf(xf[b*32+j]); if(a > amax) amax = a; }
float q = amax / 127.0f; xs[b] = q;
for(int j = 0; j < 32; ++j) xq[b*32+j] = (signed char)__float2int_rn(xf[b*32+j] / q);
}
}
__global__ void gk_rmsnorm_q8f(const float* __restrict__ x, signed char* __restrict__ xq, float* __restrict__ xs,
float* __restrict__ xf, const float* __restrict__ w, int n, float eps){
__shared__ float sh[256];
float s = 0.f; 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); __syncthreads();
for(int i = threadIdx.x; i < n; i += blockDim.x) xf[i] = x[i] * sc * w[i];
__syncthreads();
int nb = n / 32;
for(int b = threadIdx.x; b < nb; b += blockDim.x){
float amax = 1e-30f;
for(int j = 0; j < 32; ++j){ float a = fabsf(xf[b*32+j]); if(a > amax) amax = a; }
float q = amax / 127.0f; xs[b] = q;
for(int j = 0; j < 32; ++j) xq[b*32+j] = (signed char)__float2int_rn(xf[b*32+j] / q);
}
}
// Diagnostic: L2 norm of a vector, so a debug build can report where the hidden state collapses or explodes.
__global__ void gk_dbg3(const float* __restrict__ v, float* __restrict__ out){
if(threadIdx.x < 3) out[threadIdx.x] = v[threadIdx.x];
}
__global__ void gk_l2probe(const float* __restrict__ v, int n, float* __restrict__ out){
__shared__ float sh[256]; float s = 0.f;
for(int i = threadIdx.x; i < n; i += blockDim.x) s += v[i] * v[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(); }
if(threadIdx.x == 0) *out = sqrtf(sh[0]);
}
__global__ void gk_gemv_f32w(float* __restrict__ y, const float* __restrict__ w,
const float* __restrict__ x, int K, int N){
int o = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; if(o >= N) return;
int lane = threadIdx.x & 31;
const float* wr = w + (long)o * K;
float acc = 0.f;
for(int j = lane; j < K; j += 32) acc += wr[j] * x[j];
for(int s = 16; s > 0; s >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, s);
if(lane == 0) y[o] = acc;
}
// Split the fused projection's k|v halves into their caches at the current position.
__global__ void gk_kv_store(float* __restrict__ kc, float* __restrict__ vc, const int* __restrict__ d_pos,
int KD, const float* __restrict__ kv){
int i = blockIdx.x * blockDim.x + threadIdx.x; if(i >= 2 * KD) return;
long base = (long)(*d_pos) * KD;
if(i < KD) kc[base + i] = kv[i]; else vc[base + (i - KD)] = kv[i];
}
// silu(gate) * up, with the int8 re-quantization fused into the epilogue.
// The same SwiGLU epilogue, but leaving the result in f32 instead of an int8 grid. Used by the deq lane,
// where the projection that consumes it reads f32 activations.
__global__ void gk_silu_had_f(const float* __restrict__ g, const float* __restrict__ u,
float* __restrict__ out, int n){
int i = blockIdx.x * blockDim.x + threadIdx.x; if(i >= n) return;
float gv = g[i];
out[i] = (gv / (1.0f + __expf(-gv))) * u[i];
}
__global__ void gk_silu_had_q8(const float* __restrict__ g, const float* __restrict__ u,
signed char* __restrict__ hq, float* __restrict__ hs, int n){
int blk = blockIdx.x * blockDim.x + threadIdx.x; int nb = n / 32; if(blk >= nb) return;
float amax = 1e-30f; float t[32];
for(int j = 0; j < 32; ++j){ float gv = g[blk*32+j]; float h = (gv / (1.0f + __expf(-gv))) * u[blk*32+j];
t[j] = h; float a = fabsf(h); if(a > amax) amax = a; }
float sc = amax / 127.0f; hs[blk] = sc;
for(int j = 0; j < 32; ++j) hq[blk*32+j] = (signed char)__float2int_rn(t[j] / sc);
}
"#;Expand description
The last two kernels the hybrid needs: an f32-weight GEMV for GDN’s tiny gate projections (quantizing a
[d, 48] matrix would risk the gates for a rounding error that buys 0.26% of bandwidth), a KV-cache store
that splits the fused projection’s k and v halves into their caches, and the SwiGLU epilogue fused with the
activation re-quantization so the FFN’s middle never round-trips through global memory as f32.