pub const QWEN38_BATCHED_KERNELS: &str = r#"
// `BMAX` bounds the per-thread accumulator array in the staged GEMVs. The emitted program and the
// equivalence tests both define it ahead of this text, but the kernel set must still compile when a
// caller does not — a shipped kernel string that fails on a missing macro is a trap for the next
// harness that includes it. The guard is inert whenever a real batch bound was supplied.
// `BTG` (sequences per warp-tile) and `MAXSEQ` (KV-cache depth) were referenced here without
// defaults long before the MMQ work — the self-containment test found them the moment it existed.
// Same class, same fix.
#ifndef BTG
#define BTG 8
#endif
#ifndef MAXSEQ
#define MAXSEQ 64
#endif
#ifndef BMAX
#define BMAX 32
#endif
// The MMQ tile's shape constants. Same rule as BMAX and for the same reason: the emitted program and
// the equivalence tests define these ahead of this text, but the kernel set must still compile when a
// caller does not. Adding them WITHOUT guards broke nine tests an hour after BMAX was fixed — the
// identical failure, reintroduced. `the_kernel_set_compiles_with_no_defines_supplied` is the standing
// check that stops a third instance.
#ifndef MMQ_TJ
#define MMQ_TJ 32
#endif
#ifndef MMQ_RI
#define MMQ_RI 2
#endif
#ifndef MMQ_RJ
#define MMQ_RJ 2
#endif
#ifndef MMQ_KC
#define MMQ_KC 256
#endif
#ifndef MMQ_ACC
#define MMQ_ACC double
#endif
// The tensor-core tile's shape constants, same guard rule and same reason as the MMQ set above.
// TC_TJ must be 8, 16, 32 or 64: the warp -> sub-tile map needs TC_TJ/8 to divide the 8 warps.
#ifndef TC_TJ
#define TC_TJ 32
#endif
// How many 256-thread tile blocks ptxas must keep resident per SM. MEASURED on an H100 NVL
// (2026-09-20): the tile compiles to 102 registers, so 256 threads x 102 = 26,112 of the SM's 65,536
// and only TWO blocks fit -- 25% occupancy. nsys puts the tile at 70.9% of the step while it uses
// 9.9% of tensor-core peak and 24% of achievable bandwidth: it is neither math- nor bandwidth-bound,
// it is OCCUPANCY-bound, which is why halving its arithmetic bought only 6%.
// 1 lets ptxas take every register it wants — the behaviour this replaces.
// The analogous cap was worth +18.9% on the 30B MoE trio, where the optimum was an INTERIOR point
// (mb=5 best, mb=6 spilled), so this is a knob to sweep and never a constant to guess.
#ifndef TC_DYNSMEM
#define TC_DYNSMEM 0
#endif
// ── IN-KERNEL STALL ATTRIBUTION ──────────────────────────────────────────────────────────────
// `ncu` is permission-blocked in RunPod containers, which is why every attribution in this campaign
// came from ablation and roofline arithmetic. `clock64()` needs no privileges and answers the one
// question those cannot: where does a warp actually WAIT? One lane per warp accumulates, so the
// counters cost a handful of atomics per chunk and the shipped path compiles them away entirely.
#ifndef TC_PROBE
#define TC_PROBE 0
#endif
#ifndef TC_TILED_W
#define TC_TILED_W 0
#endif
#if TC_PROBE
__device__ unsigned long long bq_tc_clk[4];
#define BQ_CLK(v) unsigned long long v = clock64();
#define BQ_ACC(i,a,b) if((threadIdx.x&31)==0) atomicAdd(&bq_tc_clk[i],(unsigned long long)((b)-(a)));
#else
#define BQ_CLK(v)
#define BQ_ACC(i,a,b)
#endif
// ★ THE FALLBACK MUST BE THE SHIPPED VALUE. It used to be 1, which is a residency the program never
// asks for: the emitter always writes `tc_occ_def` (3 on Hopper, 2 below) ahead of this guard. Anything
// that compiles these kernels WITHOUT the emitter -- the decode bandwidth bench, most obviously -- then
// silently measured one block per SM. That is 167 GB/s against the shipped 199 at the decode shape on a
// 3060: the instrument was slandering the tile by 19%, which is the same defect as the bench that once
// ran SPLIT=1. A default that no shipped configuration uses is a trap, not a default.
#ifndef TC_MINBLOCKS
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
#define TC_MINBLOCKS 3
#else
#define TC_MINBLOCKS 2
#endif
#endif
// How many K-chunks the tile keeps staged at once. 1 is the shipped, measured behaviour: a ONE-deep
// prefetch held in REGISTERS (`pw`), which costs ~32 registers at B=8 and is why the tile compiles to
// 181 there (one block per SM). A deeper pipeline has to live in shared memory instead, which is what
// `cp.async` and Hopper's TMA both exist to fill -- and the measured need is large: the tile moves
// 24.33 GB at 728 GB/s while `bq_gdn_step` moves 9.6 GB at 2.0 TB/s in the same step on the same chip.
// This dimension is the structural prerequisite for either. At TC_STAGES=1 every stage index folds to
// 0 and the emitted code is the one that produced every number on record.
#ifndef TC_STAGES
#define TC_STAGES 1
#endif
#ifndef TC_WM
#define TC_WM 4
#endif
#ifndef TC_KC
#define TC_KC 128
#endif
// The row-parallel GEMV's shape. MV_TJ sequences live in registers, MV_RB output rows share each
// activation load, and MV_NW warps cooperate on one row-group's K reduction. Same guard rule as above.
#ifndef MV_TJ
#define MV_TJ 8
#endif
#ifndef MV_RB
#define MV_RB 2
#endif
#ifndef MV_NW
#define MV_NW 1
#endif
// The residual-norm-quantise kernels reduce over the whole model dimension, so they take ONE BLOCK
// PER SEQUENCE — at two sequences that is two blocks of a 132-SM part, and the block width is the
// only knob left that buys memory-level parallelism. It was a literal 256; measured on an A100 at
// B=2, widening to 1024 took the step from 31.05 ms to 29.51 ms (+5%) for no other change. Sizing a
// block by a constant when its grid cannot fill the machine is the bug; the constant is named here
// so the next person sees the choice instead of the number.
#ifndef RMSN_NT
#define RMSN_NT 1024
#endif
// How many chunks the split residual norm cuts the model dimension into. Only read on the fast lane.
#ifndef NORM_G
#define NORM_G 20
#endif
// ══ THE BATCHED GEMV, STAGED ════════════════════════════════════════════════════════════════════
// Corrected target. llama.cpp does NOT switch to tensor cores at these batch sizes — mmq.cu:327 reads
// `!fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE` with that constant set to 64, so for
// every batch below 64 it runs a DP4A kernel. Its advantage over ours was never the instruction; it was the
// STAGING. Our register-tiled form re-reads each sequence's activation slab from L2 for every output row,
// which at B=32 is ~16x more activation traffic than weight traffic and is what pinned us at 128 GB/s of an
// available 3000.
//
// Here the activation tile for ALL sequences is staged once per 512-wide K block in shared memory, and every
// warp in the block reads it from there. Weights still stream from global — they are touched once and never
// reused, so staging them would only add a copy.
//
// Structure follows llama.cpp's MMQ (ggml/src/ggml-cuda/mmq.cuh, MIT, (c) 2023-2024 The ggml authors);
// the code is ours. Crucially the ACCUMULATION ORDER is unchanged from `bq_gemv_h` — same ascending K
// blocks, same per-block scale product — so this is bit-identical, not merely close.
__global__ void bq_gemv_h_smem(float* __restrict__ Y, const signed char* __restrict__ wq,
const __half* __restrict__ ws, const signed char* __restrict__ Xq,
const float* __restrict__ Xs, int K, int N, int Bn, int YS){
int warps = blockDim.x >> 5;
int o = blockIdx.x * warps + (threadIdx.x >> 5);
int lane = threadIdx.x & 31, tid = threadIdx.x;
int KB = K / 32;
extern __shared__ char gsmem[];
signed char* Asm = (signed char*)gsmem; // Bn x 512 staged activations
float* Ass = (float*)(gsmem + Bn * 512); // Bn x 16 staged scales
float acc[BMAX];
#pragma unroll
for(int t=0;t<BMAX;t++) acc[t]=0.f;
const signed char* wrow = wq + (long)(o < N ? o : 0) * K;
const __half* wsr = ws + (long)(o < N ? o : 0) * KB;
for(int k0=0;k0<K;k0+=512){
for(int i=tid;i<Bn*512;i+=blockDim.x){ int b=i>>9, c=i&511;
Asm[i] = (k0+c<K) ? Xq[(long)b*K+k0+c] : (signed char)0; }
for(int i=tid;i<Bn*16;i+=blockDim.x){ int b=i>>4, c=i&15;
Ass[i] = (k0/32+c<KB) ? Xs[(long)b*KB+k0/32+c] : 0.f; }
__syncthreads();
if(o < N && k0+lane*16 < K){
int4 w = *(const int4*)(wrow + k0 + lane*16);
int blk = lane >> 1; // which 32-block of this 512 the lane covers
float wsb = __half2float(wsr[k0/32 + blk]);
#pragma unroll
for(int t=0;t<BMAX;t++){
if(t>=Bn) break;
int4 xx = *(const int4*)(Asm + t*512 + lane*16);
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[t] += (float)d * (wsb * Ass[t*16 + blk]);
}
}
__syncthreads();
}
#pragma unroll
for(int t=0;t<BMAX;t++){
if(t>=Bn) break;
float a=acc[t];
for(int r=16;r>0;r>>=1) a+=__shfl_down_sync(0xffffffffu,a,r);
if(lane==0 && o<N) Y[(long)t*YS+o]=a;
}
}
// The TWO-PASS activation form, staged. This is the lane the accuracy result rests on — `bq_gemv_hl`
// is what runs when the emitter is asked for the golden/oracle-closest decode — and it was the one
// lane the staging fix did not cover, so it kept paying the L2 re-read that pinned the batch at
// 128 GB/s of an available 3000.
//
// Both activation words are staged, so the shared-memory footprint is 2*Bn*512 + Bn*64 bytes
// (34 KB at Bn=32, inside the 48 KB default block limit). Weights still stream from global: each
// weight byte is read once per output row and never reused, so staging them would only add a copy.
//
// ACCUMULATION ORDER IS UNCHANGED from `bq_gemv_hl`: the same ascending 512-wide K chunks, the same
// `blk = k0/32 + lane/2` block-scale index, and the same
// `((float)dh + (float)dl/127) * (wsb * xs)` grouping per chunk. That is what makes this a free
// substitution rather than a numerically different kernel — the golden token stream cannot move.
// Every K in this model is a multiple of 512 (5120, 17408, 2048, 6144, 1024, 248320), so the
// zero-padded staging covers the whole reduction and `bq_gemv_hl`'s scalar tail loop is empty.
__global__ void bq_gemv_hl_smem(float* __restrict__ 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 Bn, int YS){
int warps = blockDim.x >> 5;
int o = blockIdx.x * warps + (threadIdx.x >> 5);
int lane = threadIdx.x & 31, tid = threadIdx.x;
int KB = K / 32;
extern __shared__ char gsmem[];
signed char* Hsm = (signed char*)gsmem; // Bn x 512 staged hi words
signed char* Lsm = (signed char*)(gsmem + Bn * 512); // Bn x 512 staged lo words
float* Ass = (float*)(gsmem + 2 * Bn * 512); // Bn x 16 staged block scales
float acc[BMAX];
#pragma unroll
for(int t=0;t<BMAX;t++) acc[t]=0.f;
const signed char* wrow = wq + (long)(o < N ? o : 0) * K;
const __half* wsr = ws + (long)(o < N ? o : 0) * KB;
for(int k0=0;k0<K;k0+=512){
for(int i=tid;i<Bn*512;i+=blockDim.x){ int b=i>>9, c=i&511;
int src = k0+c;
Hsm[i] = (src<K) ? Xh[(long)b*K+src] : (signed char)0;
Lsm[i] = (src<K) ? Xl[(long)b*K+src] : (signed char)0; }
for(int i=tid;i<Bn*16;i+=blockDim.x){ int b=i>>4, c=i&15;
Ass[i] = (k0/32+c<KB) ? Xs[(long)b*KB+k0/32+c] : 0.f; }
__syncthreads();
if(o < N && k0+lane*16 < K){
int4 w = *(const int4*)(wrow + k0 + lane*16);
int blk = lane >> 1;
float wsb = __half2float(wsr[k0/32 + blk]);
#pragma unroll
for(int t=0;t<BMAX;t++){
if(t>=Bn) break;
int4 hh = *(const int4*)(Hsm + t*512 + lane*16);
int4 ll = *(const int4*)(Lsm + t*512 + lane*16);
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[t] += ((float)dh + (float)dl*(1.0f/127.0f)) * (wsb * Ass[t*16 + blk]);
}
}
__syncthreads();
}
#pragma unroll
for(int t=0;t<BMAX;t++){
if(t>=Bn) break;
float a=acc[t];
for(int r=16;r>0;r>>=1) a+=__shfl_down_sync(0xffffffffu,a,r);
if(lane==0 && o<N) Y[(long)t*YS+o]=a;
}
}
// ══ THE MMQ-SHAPED TILE — the lowering that actually wins at batch ══════════════════════════════
//
// MEASURED, H100 NVL, isolated projection N=17408 K=5120, double accumulators, KC=256:
// B=16 0.185 -> 0.194 ms (0.95x, LOSS)
// B=32 0.347 -> 0.285 ms (1.22x)
// B=64 0.673 -> 0.484 ms (1.39x)
// The win grows with the batch, so this is a shape rule with a threshold, not a replacement.
//
// WHY the register-tiled form runs out of road, which is the whole reason this kernel exists. It
// reuses each WEIGHT element across all B sequences but each ACTIVATION element across exactly one
// output row, so at B=32 it issues 64 activation loads per weight load. Neither escape works inside a
// warp-register tile: BTG must cover the batch or the weight matrix is re-streamed (4x at B=32), and
// rows-per-warp needs RI*BTG accumulators — registers measured 40 -> 63 -> 96 -> 168 as RI went 1 to
// 8. Shared memory breaks the coupling: the tile size stops being a register-count problem, and each
// thread keeps only its own RI x RJ accumulators.
//
// Two properties are load-bearing and were both learned the expensive way:
// * The thread -> output map must be a BIJECTION onto the tile. With (TI/RI)*(TJ/RJ) < NT the
// surplus threads wrap and recompute outputs another thread already produced: silently correct,
// 4x the work. The static_assert below is that bug, caught.
// * The shared rows must be PADDED. Unpadded rows are KC bytes apart and shared memory repeats
// every 128 B, so every thread of a warp — which differ by row — hits the same bank: a 32-way
// conflict on every read of the compute loop. One int of padding took the kernel from 0.704 ms
// to 0.450 ms on the development card. It is the difference between losing and winning.
//
// Accumulation is in double: the per-32-block int32 dot is exact, so the only rounding is the sum of
// K/32 float terms, and fp64 is half-rate on a datacentre part (measured free here: 1.22x with double
// against 1.25x with float) while being 10x more accurate than the float sum. On consumer cards fp64
// is 1/32 rate and this choice would be wrong — which is why it is a knob, not a constant.
//
// Tile shape and K-step follow llama.cpp's MMQ (ggml/src/ggml-cuda/mmq.cuh, MIT, (c) 2023-2024 The
// ggml authors): 8 warps, 256 K-elements staged per barrier round. A sweep of KC over 128/256/512
// independently picked 256, which is their MMQ_ITER_K. The DESIGN PARAMETERS are theirs; the code is
// ours. Attribution recorded in benchmarks/results/qwen38_27b.json.
//
// This re-associates the K reduction, so unlike the staged GEMVs it is NOT bit-identical to the
// register-tiled form; it is held to a tolerance by `batched_equivalence_lock`, exactly as the
// tensor-core lowering is, and the golden token stream is the acceptance test.
__global__ void __launch_bounds__(256) bq_gemm_mmq(
float* __restrict__ 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 B, int YS){
constexpr int NCG = MMQ_TJ / MMQ_RJ; // column groups fixed by the sequence tile
constexpr int NRG = 256 / NCG; // row groups take the rest of the block
constexpr int TI = NRG * MMQ_RI; // ...which DERIVES the output-row tile
constexpr int SW = MMQ_KC/4 + 1; // ints per shared row; +1 puts rows on consecutive banks
static_assert(MMQ_TJ % MMQ_RJ == 0 && NCG * NRG == 256,
"thread -> output map must be a bijection onto the tile");
__shared__ int sW[TI][SW];
__shared__ int sH[MMQ_TJ][SW];
__shared__ int sL[MMQ_TJ][SW];
__shared__ __half sWs[TI][MMQ_KC/32];
__shared__ float sXs[MMQ_TJ][MMQ_KC/32];
const int tid = threadIdx.x;
const int i0 = blockIdx.x * TI;
const int KB = K/32;
const int rg = tid % NRG, cg = tid / NRG;
const int ir0 = rg*MMQ_RI, jr0 = cg*MMQ_RJ;
MMQ_ACC acc[MMQ_RI][MMQ_RJ];
#pragma unroll
for(int a=0;a<MMQ_RI;a++)
#pragma unroll
for(int b=0;b<MMQ_RJ;b++) acc[a][b]=(MMQ_ACC)0;
for(int k0=0;k0<K;k0+=MMQ_KC){
for(int idx=tid; idx<TI*(MMQ_KC/16); idx+=256){
int r=idx/(MMQ_KC/16), c=idx%(MMQ_KC/16); int o=i0+r;
const int4 v = (o<N && k0+c*16<K) ? ((const int4*)(wq+(long)o*K+k0))[c] : make_int4(0,0,0,0);
sW[r][c*4+0]=v.x; sW[r][c*4+1]=v.y; sW[r][c*4+2]=v.z; sW[r][c*4+3]=v.w;
}
for(int idx=tid; idx<MMQ_TJ*(MMQ_KC/16); idx+=256){
int r=idx/(MMQ_KC/16), c=idx%(MMQ_KC/16);
const int4 h = (r<B && k0+c*16<K) ? ((const int4*)(Xh+(long)r*K+k0))[c] : make_int4(0,0,0,0);
const int4 l = (r<B && k0+c*16<K) ? ((const int4*)(Xl+(long)r*K+k0))[c] : make_int4(0,0,0,0);
sH[r][c*4+0]=h.x; sH[r][c*4+1]=h.y; sH[r][c*4+2]=h.z; sH[r][c*4+3]=h.w;
sL[r][c*4+0]=l.x; sL[r][c*4+1]=l.y; sL[r][c*4+2]=l.z; sL[r][c*4+3]=l.w;
}
for(int idx=tid; idx<TI*(MMQ_KC/32); idx+=256){
int r=idx/(MMQ_KC/32), c=idx%(MMQ_KC/32); int o=i0+r;
sWs[r][c] = (o<N && k0/32+c<KB) ? ws[(long)o*KB + k0/32 + c] : __float2half(0.f);
}
for(int idx=tid; idx<MMQ_TJ*(MMQ_KC/32); idx+=256){
int r=idx/(MMQ_KC/32), c=idx%(MMQ_KC/32);
sXs[r][c] = (r<B && k0/32+c<KB) ? Xs[(long)r*KB + k0/32 + c] : 0.f;
}
__syncthreads();
#pragma unroll
for(int blk=0; blk<MMQ_KC/32; ++blk){
#pragma unroll
for(int a=0;a<MMQ_RI;a++){
float wsb = __half2float(sWs[ir0+a][blk]);
#pragma unroll
for(int b=0;b<MMQ_RJ;b++){
int dh=0, dl=0;
#pragma unroll
for(int q=0;q<8;q++){
int w = sW[ir0+a][blk*8+q];
dh = __dp4a(w, sH[jr0+b][blk*8+q], dh);
dl = __dp4a(w, sL[jr0+b][blk*8+q], dl);
}
acc[a][b] += (MMQ_ACC)(((float)dh + (float)dl*(1.0f/127.0f)) * (wsb * sXs[jr0+b][blk]));
}
}
}
__syncthreads();
}
#pragma unroll
for(int a=0;a<MMQ_RI;a++){ int o=i0+ir0+a; if(o>=N) continue;
#pragma unroll
for(int b=0;b<MMQ_RJ;b++){ int j=jr0+b; if(j>=B) continue; Y[(long)j*YS+o]=(float)acc[a][b]; } }
}
// ══ THE SAME matmul, LOWERED AS A GEMM ══════════════════════════════════════════════════════════
// At batch 1 a projection is a GEMV and one warp per output row is the right shape. At batch B it is a GEMM,
// and that shape becomes wrong: the register-tiled form issues ~64 dp4a per 512-weight block and re-reads
// every sequence's activation slab through L2, which measured 128 GB/s of an available 3000 at B=32 — compute
// bound, tensor cores idle, and 3.76x behind llama.cpp, which switches to its MMQ kernels for exactly this.
//
// Nothing about the model changes. The `.lg` still says `matmul`; the COMPILER picks the lowering from the
// shape it is given. That is the same decision llama.cpp makes between mmvq and mmq, made in the emitter
// instead of by hand.
//
// Golden survives because the int8 dot product accumulates in int32 — exact, no rounding — and the per-32
// block scales are applied in registers in the SAME float grouping the dp4a path uses
// (`(float)c * (ws * xs)` per k-block, summed in k order). The arithmetic that decides a token is unchanged.
// Tile shape and K-step follow llama.cpp's MMQ configuration (ggml/src/ggml-cuda/mmq.cuh,
// mmq-config-ampere.cuh): MMQ_ITER_K = 256, 8 warps, shared-memory tiles padded to dodge bank conflicts.
// llama.cpp is MIT-licensed, Copyright (c) 2023-2024 The ggml authors; the DESIGN PARAMETERS below are taken
// from it, the code is ours. Attribution recorded in benchmarks/results/qwen38_27b.json.
//
// Why it matters, measured: our first version staged only 32 K-elements per __syncthreads round, so a
// K=5120 projection paid 320 barriers per output tile where llama.cpp pays 40. Same tensor-core instruction,
// 8x the synchronization. Staging 256 K per round is the single structural difference.
#define TBM 16
#define TBN 64
#define KSTEP 256
__global__ void bq_gemm_mma(float* __restrict__ Y, const signed char* __restrict__ Wq,
const __half* __restrict__ Ws, const signed char* __restrict__ Xq,
const float* __restrict__ Xs, int K, int N, int M, int YS){
int KB=K/32; int mblk=blockIdx.y*TBM, nblk=blockIdx.x*TBN;
int warp=threadIdx.x>>5, lane=threadIdx.x&31, tid=threadIdx.x; int gid=lane>>2, tq=lane&3;
// +4 bytes of row padding: 32-byte rows make every lane of a quarter-warp hit the same shared-memory bank.
const int XP = KSTEP + 4, WP = KSTEP + 4;
extern __shared__ signed char smem[];
signed char* Xsm = smem; // TBM rows x XP
signed char* Wsm = smem + TBM*XP; // TBN rows x WP
float d0=0.f,d1=0.f,d2=0.f,d3=0.f;
for(int k0=0;k0<K;k0+=KSTEP){
// Stage KSTEP of K at once — one barrier pair per 256 elements instead of per 32.
for(int i=tid;i<TBM*KSTEP;i+=256){ int r=i/KSTEP,c=i%KSTEP;
Xsm[r*XP+c] = (mblk+r<M && k0+c<K) ? Xq[(long)(mblk+r)*K+k0+c] : (signed char)0; }
for(int i=tid;i<TBN*KSTEP;i+=256){ int r=i/KSTEP,c=i%KSTEP;
Wsm[r*WP+c] = (nblk+r<N && k0+c<K) ? Wq[(long)(nblk+r)*K+k0+c] : (signed char)0; }
__syncthreads();
// Then run the mma chain over the staged block with no further barriers.
for(int ks=0; ks<KSTEP && k0+ks<K; ks+=32){
int kb=(k0+ks)>>5;
int a0=*(const int*)(Xsm+gid*XP+ks+tq*4);
int a1=*(const int*)(Xsm+(gid+8)*XP+ks+tq*4);
int a2=*(const int*)(Xsm+gid*XP+ks+16+tq*4);
int a3=*(const int*)(Xsm+(gid+8)*XP+ks+16+tq*4);
int b0=*(const int*)(Wsm+(warp*8+gid)*WP+ks+tq*4);
int b1=*(const int*)(Wsm+(warp*8+gid)*WP+ks+16+tq*4);
int c0=0,c1=0,c2=0,c3=0;
asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
: "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3) : "r"(a0),"r"(a1),"r"(a2),"r"(a3), "r"(b0),"r"(b1));
int r0=mblk+gid, r1=mblk+gid+8, cbase=nblk+warp*8;
float xs0 = (r0<M) ? Xs[(long)r0*KB+kb] : 0.f;
float xs1 = (r1<M) ? Xs[(long)r1*KB+kb] : 0.f;
float ws0 = (cbase+tq*2 <N) ? __half2float(Ws[(long)(cbase+tq*2)*KB+kb]) : 0.f;
float ws1 = (cbase+tq*2+1 <N) ? __half2float(Ws[(long)(cbase+tq*2+1)*KB+kb]) : 0.f;
d0+=(float)c0*(ws0*xs0); d1+=(float)c1*(ws1*xs0);
d2+=(float)c2*(ws0*xs1); d3+=(float)c3*(ws1*xs1);
}
__syncthreads();
}
int r0=mblk+gid, r1=mblk+gid+8, cc0=nblk+warp*8+tq*2, cc1=cc0+1;
if(r0<M){ if(cc0<N) Y[(long)r0*YS+cc0]=d0; if(cc1<N) Y[(long)r0*YS+cc1]=d1; }
if(r1<M){ if(cc0<N) Y[(long)r1*YS+cc0]=d2; if(cc1<N) Y[(long)r1*YS+cc1]=d3; }
}
// ── THE KERNEL THAT MAKES BATCHING WORTH IT ─────────────────────────────────────────────────────
// One warp owns one output row. It walks the row's weights ONCE, and for each 512-weight block it reuses the
// loaded `int4` across BTG sequences held in registers. Weight traffic is therefore independent of the batch:
// the same bytes that produced one token now produce BTG of them.
__global__ void bq_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 B, int YS){
int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return;
int lane=threadIdx.x&31; int b0=blockIdx.y*BTG;
const int4* w16=(const int4*)(wq+(long)o*K); const __half* wsr=ws+(long)o*(K/32);
int nq=K/512, KB=K/32;
float acc[BTG];
#pragma unroll
for(int t=0;t<BTG;t++) acc[t]=0.f;
for(int c=0;c<nq;c++){
int4 w=w16[c*32+lane]; int blk=c*16+lane/2; float wsb=__half2float(wsr[blk]);
#pragma unroll
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break;
int4 xx=((const int4*)(Xq+(long)bb*K))[c*32+lane];
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[t] += (float)d * (wsb*Xs[(long)bb*KB+blk]);
}
}
int done=nq*512;
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
acc[t] += (float)((int)Xq[(long)bb*K+j]*(int)wq[(long)o*K+j])
* (__half2float(wsr[blk])*Xs[(long)bb*KB+blk]); }
}
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break; float a=acc[t];
for(int s=16;s>0;s>>=1) a+=__shfl_down_sync(0xffffffffu,a,s);
if(lane==0) Y[(long)bb*YS+o]=a;
}
}
// The two-pass activation form, batched the same way: the weight is still read once per row for the whole
// batch, and each sequence contributes two dp4a instead of one.
__global__ void bq_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 B, int YS){
int o=(blockIdx.x*blockDim.x+threadIdx.x)>>5; if(o>=N) return;
int lane=threadIdx.x&31; int b0=blockIdx.y*BTG;
const int4* w16=(const int4*)(wq+(long)o*K); const __half* wsr=ws+(long)o*(K/32);
int nq=K/512, KB=K/32;
float acc[BTG];
#pragma unroll
for(int t=0;t<BTG;t++) acc[t]=0.f;
for(int c=0;c<nq;c++){
int4 w=w16[c*32+lane]; int blk=c*16+lane/2; float wsb=__half2float(wsr[blk]);
#pragma unroll
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break;
int4 hh=((const int4*)(Xh+(long)bb*K))[c*32+lane];
int4 ll=((const int4*)(Xl+(long)bb*K))[c*32+lane];
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[t] += ((float)dh + (float)dl*(1.0f/127.0f)) * (wsb*Xs[(long)bb*KB+blk]);
}
}
int done=nq*512;
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break;
for(int j=done+lane;j<K;j+=32){ int blk=j/32;
float a=(float)Xh[(long)bb*K+j] + (float)Xl[(long)bb*K+j]*(1.0f/127.0f);
acc[t] += a*(float)wq[(long)o*K+j]*(__half2float(wsr[blk])*Xs[(long)bb*KB+blk]); }
}
for(int t=0;t<BTG;t++){
int bb=b0+t; if(bb>=B) break; float a=acc[t];
for(int s=16;s>0;s>>=1) a+=__shfl_down_sync(0xffffffffu,a,s);
if(lane==0) Y[(long)bb*YS+o]=a;
}
}
// The tiny f32-weight projection GDN needs for its gates. Same shape, no quantization.
// ── THE GDN's SMALL f32 PROJECTION, ROW-PARALLEL ────────────────────────────────────────────────
// Third instance of the same shape mistake, and the clearest: one warp per output row, with only
// N = 2*v_heads = 96 rows. That is 96 warps per sequence — 192 warps, ~6k threads, on a 132-SM part
// — each grinding a K=5120 reduction 32 lanes at a time. Measured 0.45 ms over 48 launches.
//
// F32W_NW warps now share each row's reduction and rejoin through shared memory, the same fix as
// `bq_gemv_wide` and for the same reason. Unlike that one it is NOT bit-identical: splitting a
// reduction across warps re-associates it. The GDN gates it feeds are a sigmoid and a softplus decay,
// both smooth and both far from a decision boundary, so a one-ulp move there cannot change a token —
// but it is a different number, so it is asserted against the original within a bound rather than
// claimed equal.
#ifndef F32W_NW
#define F32W_NW 8
#endif
__global__ void __launch_bounds__(256) bq_gemv_f32w(float* Y, const float* __restrict__ w,
const float* __restrict__ X, int K, int N, int B){
const int wpr=F32W_NW; // warps cooperating on one row
const int rows_per_block=(256/32)/wpr; // rows a block covers
const int lane=threadIdx.x&31, warp=threadIdx.x>>5;
const int o=blockIdx.x*rows_per_block + warp/wpr;
const int sub=warp%wpr;
const int bb=blockIdx.y; if(bb>=B) return;
__shared__ float red[256/32];
float a=0.f;
if(o<N){ const float* wr=w+(long)o*K; const float* x=X+(long)bb*K;
for(int j=lane+sub*32;j<K;j+=32*wpr) a+=wr[j]*x[j]; }
for(int s=16;s>0;s>>=1) a+=__shfl_down_sync(0xffffffffu,a,s);
if(lane==0) red[warp]=a;
__syncthreads();
if(o<N && sub==0 && lane==0){
float t=0.f;
for(int q=0;q<wpr;q++) t+=red[(warp/wpr)*wpr+q]; // fixed ascending order: deterministic
Y[(long)bb*N+o]=t;
}
}
// ── the per-sequence pointwise and reduction kernels: one block per sequence ─────────────────────
__global__ void bq_embed(float* X, const signed char* wq, const __half* ws, const int* D_tok, int E){
int bb=blockIdx.y; long tok=D_tok[bb];
for(int t=blockIdx.x*blockDim.x+threadIdx.x; t<E; t+=gridDim.x*blockDim.x)
X[(long)bb*E+t]=(float)wq[tok*E+t]*__half2float(ws[tok*(E/32)+t/32]);
}
__global__ void bq_rmsnorm_q8f_hl(const float* X, signed char* Xq, signed char* Xl, float* Xs,
float* Xf, const float* w, int n, float eps){
int bb=blockIdx.x; const float* x=X+(long)bb*n; float* xf=Xf+(long)bb*n;
__shared__ float sh[RMSN_NT];
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[(long)bb*nb+b]=q;
for(int j=0;j<32;++j) lgx_q8hl(xf[b*32+j]/q, Xq+(long)bb*n, Xl+(long)bb*n, (long)b*32+j);
}
}
// The residual-add + RMSNorm + requantise that follows a projection. `SK` folds the split-K
// reduction in: instead of reading a GEMM output some earlier kernel summed, it sums the S partials
// itself, in the SAME order `bq_sk_reduce` used (0.f, then s ascending), so the float result is
// bit-identical -- and the separate launch, plus a full write and read of Y, disappear.
template<bool SK>
__device__ __forceinline__ void bq_add_rmsnorm_q8f_hl_body(float* X, const float* Y, int S,
signed char* Xq, signed char* Xl,
float* Xs, float* Xf, const float* w, int n, float eps){
int bb=blockIdx.x; float* x=X+(long)bb*n; const float* y=Y+(long)bb*n; float* xf=Xf+(long)bb*n;
__shared__ float sh[RMSN_NT];
const long skst=(long)n*gridDim.x;
float s=0.f; for(int i=threadIdx.x;i<n;i+=blockDim.x){
float yv; if(SK){ yv=0.f; for(int q=0;q<S;q++) yv+=y[(long)q*skst+i]; } else yv=y[i];
float v=x[i]+yv; 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[(long)bb*nb+b]=q;
for(int j=0;j<32;++j) lgx_q8hl(xf[b*32+j]/q, Xq+(long)bb*n, Xl+(long)bb*n, (long)b*32+j);
}
}
__global__ void bq_add_rmsnorm_q8f_hl(float* X, const float* Y, signed char* Xq, signed char* Xl,
float* Xs, float* Xf, const float* w, int n, float eps){
bq_add_rmsnorm_q8f_hl_body<false>(X,Y,1,Xq,Xl,Xs,Xf,w,n,eps);
}
__global__ void bq_add_rmsnorm_q8f_hl_sk(float* X, const float* P, int S, signed char* Xq, signed char* Xl,
float* Xs, float* Xf, const float* w, int n, float eps){
bq_add_rmsnorm_q8f_hl_body<true>(X,P,S,Xq,Xl,Xs,Xf,w,n,eps);
}
__global__ void bq_add(float* X, const float* Y, int n){
int bb=blockIdx.y; int i=blockIdx.x*blockDim.x+threadIdx.x;
if(i<n) X[(long)bb*n+i]+=Y[(long)bb*n+i];
}
__global__ void bq_gate_mul(float* A, int as_, const float* G, int gs, int n){
int bb=blockIdx.y; int i=blockIdx.x*blockDim.x+threadIdx.x;
if(i<n){ float g=G[(long)bb*gs+i]; long o=(long)bb*as_+i; A[o]*= 1.0f/(1.0f+__expf(-g)); }
}
// The attention tail: gate, then requantise. Two launches over the SAME grid, the second reading back
// what the first wrote to global memory one instruction earlier. Fused, the gated value never leaves a
// register, and the amax reduction the quantiser already does covers it. Bit-identical: the expression
// producing `v` is unchanged, and `A` is still written because the next step's oracle reads it.
__global__ void bq_gate_mul_q8_hl(float* A, int as_, const float* G, int gs,
signed char* Xq, signed char* Xl, float* Xs, int K){
int bb=blockIdx.y;
int blk=blockIdx.x*(blockDim.x/32)+threadIdx.x/32; if((long)blk*32>=K) return; int lane=threadIdx.x&31;
const long i=(long)blk*32+lane;
const float g=G[(long)bb*gs+i];
const long o=(long)bb*as_+i;
const float v=A[o]*(1.0f/(1.0f+__expf(-g)));
A[o]=v;
float 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[(long)bb*(K/32)+blk]=sc;
lgx_q8hl(v/sc, Xq+(long)bb*K, Xl+(long)bb*K, i);
}
__global__ void bq_quant_act_hl(const float* X, signed char* Xq, signed char* Xl, float* Xs, int K){
int bb=blockIdx.y;
int blk=blockIdx.x*(blockDim.x/32)+threadIdx.x/32; if((long)blk*32>=K) return; int lane=threadIdx.x&31;
const float* x=X+(long)bb*K;
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[(long)bb*(K/32)+blk]=sc;
lgx_q8hl(v/sc, Xq+(long)bb*K, Xl+(long)bb*K, (long)blk*32+lane);
}
// ── THE SwiGLU EPILOGUE, PARALLELISED BY THE RIGHT UNIT ─────────────────────────────────────────
// This measured 1.14 ms of a 20.3 ms step at B=2 — 64 launches at 17.8 us — and the reason was a
// grid sized by the wrong quantity. The work unit is a 32-wide quantisation block, and the kernel
// gave one THREAD to each: `Gs(dff)` launches dff/256 blocks of 256 threads, so 17408 threads start
// and `if(blk>=nb) return` sends all but dff/32 = 544 of them home. 1.6% of the threads did the work,
// each looping 32 elements through a `float t[32]` register array.
//
// One WARP per 32-block instead of one thread fixes it: every lane owns one element, the block's
// amax comes from a shuffle reduction, and the same launch geometry now has every thread working.
//
// It is BIT-IDENTICAL, not a tolerance. `fmaxf` is associative and exact, so the shuffle reduction
// produces the same `amax` the serial loop did; `sc` is therefore the same float, and each lane then
// runs the identical quantisation on the identical value. Only which thread computes what has moved.
template<bool SK>
__device__ __forceinline__ void bq_silu_had_q8_hl_body(const float* F, int S, signed char* Hq, signed char* Hl, float* Hs, int n){
const int bb=blockIdx.y, nb=n/32;
const int warp=(blockIdx.x*blockDim.x+threadIdx.x)>>5, lane=threadIdx.x&31;
if(warp>=nb) return;
const float* g=F+(long)bb*2*n; const float* u=g+n;
const long idx=(long)warp*32+lane;
const long skst=(long)2*n*gridDim.y;
float gv,uv;
if(SK){ gv=0.f; uv=0.f; for(int q=0;q<S;q++){ gv+=g[(long)q*skst+idx]; uv+=u[(long)q*skst+idx]; } }
else { gv=g[idx]; uv=u[idx]; }
const float h=(gv/(1.0f+__expf(-gv)))*uv;
float a=fabsf(h);
for(int s=16;s;s>>=1) a=fmaxf(a,__shfl_xor_sync(0xffffffffu,a,s));
const float sc=fmaxf(a,1e-30f)/127.0f; // the serial form seeded amax at 1e-30
if(lane==0) Hs[(long)bb*nb+warp]=sc;
lgx_q8hl(h/sc, Hq+(long)bb*n, Hl+(long)bb*n, idx);
}
__global__ void bq_silu_had_q8_hl(const float* F, signed char* Hq, signed char* Hl, float* Hs, int n){
bq_silu_had_q8_hl_body<false>(F,1,Hq,Hl,Hs,n);
}
__global__ void bq_silu_had_q8_hl_sk(const float* P, int S, signed char* Hq, signed char* Hl, float* Hs, int n){
bq_silu_had_q8_hl_body<true>(P,S,Hq,Hl,Hs,n);
}
// Greedy pick over R rows of logits, WITHOUT advancing any position. Decode's `bq_argmax_adv` folds
// the position bump into the pick because a step always consumes exactly one token; a speculative step
// does not know how many it consumed until the drafts have been checked against these very rows, so
// the two have to come apart. Same reduction tree and same tie-break (first index wins on >), so a row
// scored here picks the token decode would have picked.
// Lay the speculative chunk out the way a T>1 forward expects it: row t*B+b. Row 0 is the token the
// model actually produced last; row 1 is the draft being checked. Nothing else in the verify pass
// needs to know it is speculative -- it is an ordinary two-token chunk.
// The draft head conditions on the embedding of the token just produced AND the hidden state of the
// token BEFORE it. After a verify pass `x` holds both rows -- position p in row 0, p+1 in row 1 --
// and which is "the one before" depends on whether the draft was accepted. Leaving it at row 0 makes
// the draft condition on a state one token stale exactly when it was RIGHT. That diverges the stream
// from greedy at the second token, and every throughput measurement was blind to it.
__global__ void bq_mtp_pick_hidden(float* __restrict__ X, const int* __restrict__ Acc, int B, int E){
int b=blockIdx.y; if(!Acc[b]) return;
for(int i=blockIdx.x*blockDim.x+threadIdx.x; i<E; i+=gridDim.x*blockDim.x)
X[(long)b*E+i] = X[(long)(B+b)*E+i];
}
__global__ void bq_mtp_pack(int* __restrict__ D_tok, const int* __restrict__ D_draft, int B){
int b=blockIdx.x*blockDim.x+threadIdx.x; if(b>=B) return;
D_tok[B+b]=D_draft[b];
}
// Accept per sequence: the draft stands only if the model's own pick at that position agrees.
// `V` holds the verify pass's greedy picks in the same row-major order, so V[b] is the TRUE token at
// position p+1 and V[B+b] is the token at p+2 that is only real if the draft was.
__global__ void bq_mtp_accept(const int* __restrict__ V, const int* __restrict__ D_draft,
int* __restrict__ D_tok, int* __restrict__ D_pos,
int* __restrict__ Acc, int B){
int b=blockIdx.x*blockDim.x+threadIdx.x; if(b>=B) return;
int ok = (V[b]==D_draft[b]) ? 1 : 0;
Acc[b] = ok;
if(ok) atomicAdd(Acc+B, 1); // Acc[B] is the running accept total
D_tok[b]= ok ? V[B+b] : V[b]; // next step consumes p+2 on accept, p+1 on reject
D_pos[b]= D_pos[b] + 1 + ok;
}
// Roll the recurrence back for the sequences whose draft was wrong. Accepted sequences keep the
// state the verify pass left (it consumed a token that turned out to be real); rejected ones take
// the checkpoint written after token 0, which the lock proves is exactly what a one-token chunk
// leaves. Copying only the rejected slices is why this is cheap at a high accept rate.
// ★ THE RECURRENT STATE'S STORAGE WIDTH. It is 151 MB per sequence — 9.66 GB of read-modify-write per
// decode step at B=32, and at 4.8 GB resident it is the largest term in the step that cannot hide in
// L2. Halving it halves that traffic AND doubles the concurrency ceiling, which is a serving metric of
// its own: vLLM carries the identical state and its baseline here died on "max_num_seqs exceeds
// available Mamba cache blocks".
//
// Ten mantissa bits, not eight. Measured against an f64 reference the output drifts ~1e-3 with f16 and
// 1.7e-2 to 2.6e-1 with bf16 — and the drift does NOT grow with sequence length, because the decay
// multiplies an error injected at step t by alpha^(T-t) before it reaches the output, so old error dies
// instead of accumulating. That is what licenses a half-width state at all; it is asserted by
// `the_gated_recurrence_tolerates_a_half_width_state_but_only_with_ten_mantissa_bits`.
//
// The body touches the state EXACTLY TWICE — one load into registers, one store at the end — so the
// width is a typedef and nothing else in the recurrence changes. Compute stays f32 throughout.
// ★ THE KV CACHE'S STORAGE WIDTH. The state lane alone does NOT double the concurrency ceiling, which
// is a claim this ledger had to retract: the GDN state is one of three per-sequence terms and at 2K
// context the KV cache is the LARGER of them, because it is kept as `float`. Per sequence at 2K:
// 151 MB state + 268 MB KV + 8 MB conv. Halving the state alone is 427 -> 352 MB (1.21x more
// sequences); halving both is 427 -> 218 MB (1.97x). Concurrency is a serving metric and vLLM carries
// the identical recurrent state, so this is the half that makes the pair worth having.
//
// f16 is conservative here — attention is a weighted average over many keys, production engines ship
// fp8 KV, and this keeps ten mantissa bits. Compute stays f32: only the cached representation narrows.
#ifdef BQ_KV16
typedef __half bq_kv_t;
__device__ __forceinline__ float bq_kv_ld(__half v){ return __half2float(v); }
__device__ __forceinline__ __half bq_kv_cv(float v){ return __float2half(v); }
#else
typedef float bq_kv_t;
__device__ __forceinline__ float bq_kv_ld(float v){ return v; }
__device__ __forceinline__ float bq_kv_cv(float v){ return v; }
#endif
#ifdef BQ_STATE16
typedef __half bq_state_t;
__device__ __forceinline__ float bq_st_ld(__half v){ return __half2float(v); }
__device__ __forceinline__ __half bq_st_cv(float v){ return __float2half(v); }
#else
typedef float bq_state_t;
__device__ __forceinline__ float bq_st_ld(float v){ return v; }
__device__ __forceinline__ float bq_st_cv(float v){ return v; }
#endif
__global__ void bq_mtp_rollback(bq_state_t* __restrict__ S, const bq_state_t* __restrict__ SM,
const int* __restrict__ Acc, long per_seq){
int b=blockIdx.y; if(Acc[b]) return;
long i=(long)blockIdx.x*blockDim.x+threadIdx.x, stride=(long)gridDim.x*blockDim.x;
bq_state_t* d=S+(long)b*per_seq; const bq_state_t* src=SM+(long)b*per_seq;
for(; i<per_seq; i+=stride) d[i]=src[i];
}
// ── PRE-TILED WEIGHTS ────────────────────────────────────────────────────────────────────────────
// The decode GEMM waits 67% of its time for data (measured in-kernel with clock64(), since ncu is
// permission-blocked in RunPod containers) while sitting at 42% of a 3180 GB/s pure-read roofline.
// Neither depth, occupancy, run length, tile shape nor the instruction moved it, because none of them
// change the ACCESS PATTERN: row-major weights make a block read TI separate KC-byte streams, and a
// reader that streams contiguously gets 95% of the roofline.
//
// Storing each (n-tile, k-chunk) contiguously turns that into ONE streaming read of TI*KC bytes. The
// transform runs once at load: it is a pure permutation of the same bytes, so the arithmetic is
// untouched and the equivalence lock must still see bit-identity.
__global__ void bq_tile_w(signed char* __restrict__ dst, const signed char* __restrict__ src,
int N, int K, int TI, int KC){
long tiles_k = (long)((K + KC - 1) / KC);
long i = (long)blockIdx.x * blockDim.x + threadIdx.x, stride = (long)gridDim.x * blockDim.x;
long total = (long)((N + TI - 1) / TI) * tiles_k * TI * KC;
for(; i < total; i += stride){
long within = i % ((long)TI * KC);
long tile = i / ((long)TI * KC);
int r = (int)(within / KC), c = (int)(within % KC);
int nt = (int)(tile / tiles_k), kt = (int)(tile % tiles_k);
int n = nt * TI + r, k = kt * KC + c;
dst[i] = (n < N && k < K) ? src[(long)n * K + k] : (signed char)0;
}
}
// The same permutation for the per-32 scales, which are read on the same footing.
__global__ void bq_tile_ws(__half* __restrict__ dst, const __half* __restrict__ src,
int N, int K, int TI, int KC){
int NB = K / 32, NBC = KC / 32;
long tiles_k = (long)((K + KC - 1) / KC);
long i = (long)blockIdx.x * blockDim.x + threadIdx.x, stride = (long)gridDim.x * blockDim.x;
long total = (long)((N + TI - 1) / TI) * tiles_k * TI * NBC;
for(; i < total; i += stride){
long within = i % ((long)TI * NBC);
long tile = i / ((long)TI * NBC);
int r = (int)(within / NBC), c = (int)(within % NBC);
int nt = (int)(tile / tiles_k), kt = (int)(tile % tiles_k);
int n = nt * TI + r, b = kt * NBC + c;
dst[i] = (n < N && b < NB) ? src[(long)n * NB + b] : __float2half(0.f);
}
}
// ── THE fp8 LANE ─────────────────────────────────────────────────────────────────────────────────
// Q8_0 carries a scale every 32 elements, so the accumulator must be read out, converted and rescaled
// FOUR times per mma issue. fp8 e4m3 accumulates in fp32 natively and takes ONE scale per output
// channel at the very end -- which is what vLLM's fp8 path does, so this is the matched-precision
// comparison rather than a cheaper one.
//
// The register shape is IDENTICAL to the int8 m16n8k32 the tile already uses (4 A regs, 2 B regs,
// 4 accumulators), so the fragment map, the swizzle and the staging survive untouched: only the
// instruction, the accumulator type and where the scale lands change. fp8 mma needs PTX ISA 8.1
// (CUDA 12.1+); the box has 12.8 at /usr/local/cuda-12.8 while `nvcc` on PATH is 12.0.
// fp8 needs PTX ISA 8.1 (CUDA 12.1+) to assemble and <cuda_fp8.h> to convert. The equivalence locks
// compile with whatever `nvcc` is on PATH, which on this box is 12.0, so BOTH the header and the
// instruction are guarded and an older toolkit still builds the whole kernel set -- it simply has no
// fp8 lane. `__CUDA_ARCH__ >= 890` is the arch gate: Ada and Hopper have fp8 tensor cores, Ampere
// does not, which is also what chip.rs has said since the start.
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 1)
#include <cuda_fp8.h>
#define BQ_FP8_TOOLKIT 1
#else
#define BQ_FP8_TOOLKIT 0
#endif
#if BQ_FP8_TOOLKIT && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890
#define BQ_HAVE_FP8 1
#else
#define BQ_HAVE_FP8 0
#endif
__device__ __forceinline__ void bq_mma_e4m3(float& d0, float& d1, float& d2, float& d3,
unsigned a0, unsigned a1, unsigned a2, unsigned a3,
unsigned b0, unsigned b1){
#if BQ_HAVE_FP8
asm volatile("mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+f"(d0),"+f"(d1),"+f"(d2),"+f"(d3)
: "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
#else
(void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;
#endif
}
// Q8_0 -> e4m3 with ONE scale per output channel. The real weight is q*s_blk; the channel scale is
// the largest magnitude in the row over 448 (e4m3's max finite), so the row uses the full range.
// This is a REQUANTISATION and it loses precision on purpose -- int8 carries 8 mantissa bits where
// e4m3 carries 4 -- which is exactly the trade vLLM already makes and which the accuracy column
// against the f64 oracle has to price rather than assume.
__global__ void bq_q8_to_e4m3_scale(float* __restrict__ chan_scale,
const signed char* __restrict__ q, const __half* __restrict__ s,
int N, int K){
int n = blockIdx.x; if(n >= N) return;
int KB = K/32;
__shared__ float red[256];
float m = 0.f;
for(int k = threadIdx.x; k < K; k += blockDim.x){
float w = (float)q[(long)n*K + k] * __half2float(s[(long)n*KB + (k>>5)]);
m = fmaxf(m, fabsf(w));
}
red[threadIdx.x] = m; __syncthreads();
for(int t = blockDim.x/2; t > 0; t >>= 1){
if(threadIdx.x < t) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x+t]);
__syncthreads(); }
if(threadIdx.x == 0) chan_scale[n] = (red[0] > 0.f) ? (red[0] / 448.0f) : 1.0f;
}
__global__ void bq_q8_to_e4m3(unsigned char* __restrict__ out, const float* __restrict__ chan_scale,
const signed char* __restrict__ q, const __half* __restrict__ s,
int N, int K){
int KB = K/32;
long i = (long)blockIdx.x*blockDim.x + threadIdx.x, stride = (long)gridDim.x*blockDim.x;
for(; i < (long)N*K; i += stride){
int n = (int)(i / K), k = (int)(i % K);
float w = (float)q[i] * __half2float(s[(long)n*KB + (k>>5)]);
#if BQ_FP8_TOOLKIT
out[i] = (unsigned char)__nv_cvt_float_to_fp8(w / chan_scale[n], __NV_SATFINITE, __NV_E4M3);
#else
out[i] = 0; // no fp8 toolkit: the lane is absent, not silently wrong
#endif
}
}
__global__ void bq_argmax_rows(const float* L, int V, int* D_tok){
int bb=blockIdx.x; const float* l=L+(long)bb*V;
__shared__ float bv[256]; __shared__ int bi[256];
float best=-1e30f; int arg=0;
for(int i=threadIdx.x;i<V;i+=blockDim.x) if(l[i]>best){ best=l[i]; arg=i; }
bv[threadIdx.x]=best; bi[threadIdx.x]=arg; __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[bb]=bi[0];
}
__global__ void bq_argmax_adv(const float* L, int V, int* D_tok, int* D_pos){
int bb=blockIdx.x; const float* l=L+(long)bb*V;
__shared__ float bv[256]; __shared__ int bi[256];
float best=-1e30f; int arg=0;
for(int i=threadIdx.x;i<V;i+=blockDim.x) if(l[i]>best){ best=l[i]; arg=i; }
bv[threadIdx.x]=best; bi[threadIdx.x]=arg; __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[bb]=bi[0]; D_pos[bb]++; }
}
// ── the hybrid's stateful halves, one sequence per grid.z ────────────────────────────────────────
// Every sequence carries its own recurrence memory, conv ring and KV cache. That per-sequence state, not the
// weights, is what bounds the batch: the weights are shared by everyone, the state is not.
__global__ void bq_conv_silu_l2(const float* __restrict__ X, int xs, const float* __restrict__ wc,
float* __restrict__ CS, float* __restrict__ Y,
int C, int T, int nl2, float eps){
int bb=blockIdx.y;
const float* x=X+(long)bb*xs; float* cs=CS+(long)bb*T*C; float* y=Y+(long)bb*C;
int c=blockIdx.x*blockDim.x+threadIdx.x; if(c>=C) return;
for(int t=0;t<T-1;++t) cs[(long)t*C+c]=cs[(long)(t+1)*C+c];
cs[(long)(T-1)*C+c]=x[c];
float acc=0.f;
for(int t=0;t<T;++t) acc += cs[(long)t*C+c]*wc[(long)t*C+c];
float v=acc/(1.0f+__expf(-acc));
if(blockIdx.x>=nl2){ y[c]=v; return; }
__shared__ float ss[32];
int lane=threadIdx.x&31, warp=threadIdx.x>>5, nw=(blockDim.x+31)>>5;
float pp=gdn_warp_sum(v*v);
if(lane==0) ss[warp]=pp;
__syncthreads();
__shared__ float den;
if(threadIdx.x==0){ float t2=0.f; for(int j=0;j<nw;++j) t2+=ss[j]; den=sqrtf(t2+eps); }
__syncthreads();
y[c]=v/den;
}
__global__ void bq_gdn_gates(const float* __restrict__ G, const float* __restrict__ dt,
const float* __restrict__ alog, float* __restrict__ beta,
float* __restrict__ decay, int NV){
int bb=blockIdx.y; int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=NV) return;
const float* br=G+(long)bb*2*NV; const float* ar=br+NV;
beta[(long)bb*NV+i] = 1.0f/(1.0f+__expf(-br[i]));
decay[(long)bb*NV+i] = powf(1.0f+__expf(ar[i]+dt[i]), alog[i]);
}
__global__ void bq_gdn_step(const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, float* __restrict__ O, int D, int NK, int NV, int CS){
const unsigned m=0xffffffffu;
int bb=blockIdx.z;
int h=blockIdx.x, lane=threadIdx.x&31, warp=threadIdx.x>>5;
int a=blockIdx.y*(blockDim.x>>5)+warp; if(a>=D) return;
int g=h%NK;
const float* kk=Kk+(long)bb*CS+(long)g*D;
const float* qq=Q +(long)bb*CS+(long)g*D;
bq_state_t* Sb=S+(long)bb*NV*D*D;
long base=(long)h*D*D+(long)a*D;
float sreg[8]; float u_=0.f,w_=0.f,kq_=0.f; int n=0;
for(int b=lane;b<D;b+=32){
float sv=Sb[base+b], kb=kk[b], qb=qq[b];
sreg[n++]=sv; u_+=sv*kb; w_+=sv*qb; kq_+=kb*qb;
}
for(int sh=16;sh>0;sh>>=1){ u_+=__shfl_down_sync(m,u_,sh); w_+=__shfl_down_sync(m,w_,sh); kq_+=__shfl_down_sync(m,kq_,sh); }
u_=__shfl_sync(m,u_,0); w_=__shfl_sync(m,w_,0); kq_=__shfl_sync(m,kq_,0);
float alpha=ALP[(long)bb*NV+h], beta=BET[(long)bb*NV+h];
float coef=beta*(V[(long)bb*CS+(long)h*D+a] - alpha*u_);
if(lane==0) O[(long)bb*NV*D+(long)h*D+a]=(alpha*w_+coef*kq_)*rsqrtf((float)D);
n=0; for(int b=lane;b<D;b+=32) Sb[base+b]=alpha*sreg[n++]+coef*kk[b];
}
#ifndef BQ_GDN_DELTA_TMAX
#define BQ_GDN_DELTA_TMAX 64
#endif
// ── THE CHUNKED DELTA RULE — parallel across the chunk ───────────────────────────────────────────
// `bq_gdn_chunk` is a SCAN: every token needs the state the previous token left, so T tokens serialise
// with D^2 work apiece. Measured, that is prefill's dominant cost (0.1811 ms/layer/token => ~4.5 s over
// 48 layers and a 512-token prompt), and no amount of state-residency fixes the dependency itself.
//
// The same recurrence closes in a unit-triangular solve. With A_t = prod(alpha) and alpha_t*A_{t-1}=A_t:
//
// c_t + SUM_{j<t} beta_t (A_t/A_j)(k_j.k_t) c_j = beta_t (v_t - A_t S0 k_t) (I+M) C = B
// O_t = [ A_t (S0 q_t) + SUM_{j<=t} (A_t/A_j)(k_j.q_t) C_j ] / sqrt(D)
// S_C = A_C S0 + SUM_j (A_C/A_j) C_j k_j^T
//
// Proved equal to the scan to machine epsilon in f64 by
// `the_chunked_delta_rule_equals_the_sequential_gated_recurrence` BEFORE this kernel was written, so a
// failure here is an implementation bug and never an algebra bug.
//
// The decisive property: every S0 term is INDEPENDENT across t, and in the C recursion the head
// dimension `a` is THREAD-LOCAL — no cross-lane reduction inside the serial part. The critical path
// falls from T*D^2 to T*D.
//
// v1 is written for correctness against the f64 oracle, not for speed: one block per (sequence, head),
// each thread owning whole rows of the state. Optimise only once it is proven equal.
// ★ TEMPLATED ON THE CHUNK LENGTH, and that is not a style choice. With `float proj[64]` sized by a
// runtime bound, ptxas reported **512 bytes stack frame** against the scan's 0 — the per-thread arrays
// were in LOCAL MEMORY, which is the identical trap `bq_gdn_chunk_body` was templated on `PER` to
// escape. Read the `-Xptxas -v` stack-frame line before believing any recurrence kernel is in registers.
template<int TT>
__device__ __forceinline__ void bq_gdn_delta_body(const float* __restrict__ Q, const float* __restrict__ Kk,
const float* __restrict__ V, const float* __restrict__ BET,
const float* __restrict__ ALP, bq_state_t* __restrict__ S,
float* __restrict__ O, int D, int NK, int NV, int CS, int T, int B){
extern __shared__ float bqd_sh[];
const int h = blockIdx.x, bb = blockIdx.y, g = h % NK;
const int tid = threadIdx.x, nthr = blockDim.x;
float* sKK = bqd_sh; // [T][T] k_j . k_t
float* sKQ = sKK + (long)T * T; // [T][T] k_j . q_t
float* sA = sKQ + (long)T * T; // [T] A_t = prod alpha
const float inv_sqrt_d = rsqrtf((float)D);
// A_t: T is small and this is the only serial-in-T step left in the kernel.
if (tid == 0) {
float run = 1.f;
for (int t = 0; t < T; ++t) { run *= ALP[(long)(bb * T + t) * NV + h]; sA[t] = run; }
}
__syncthreads();
// K and Q staged once: the Gram matrices and every S0 contraction below read them repeatedly.
float* sK = sA + T; // [T][D]
float* sQ = sK + (long)T * D; // [T][D]
for (int idx = tid; idx < T * D; idx += nthr) {
const int t = idx / D, b = idx % D;
sK[idx] = Kk[(long)(bb * T + t) * CS + (long)g * D + b];
sQ[idx] = Q [(long)(bb * T + t) * CS + (long)g * D + b];
}
__syncthreads();
// The two Gram matrices, one thread per (t,j) pair.
for (int idx = tid; idx < T * T; idx += nthr) {
const int t = idx / T, j = idx % T;
float akk = 0.f, akq = 0.f;
for (int b = 0; b < D; ++b) {
const float kj = sK[(long)j * D + b];
akk += kj * sK[(long)t * D + b];
akq += kj * sQ[(long)t * D + b];
}
sKK[idx] = akk; sKQ[idx] = akq;
}
__syncthreads();
// One thread owns a whole state row `a`, so the C recursion needs no cross-lane reduction.
//
// ★ The b-loop is OUTERMOST in every phase. v1 put it inside the t-loop and so re-read the S0 row
// from global 2T+1 = 65 times per row; that alone made it 12x SLOWER than the scan it replaces
// (2.1351 vs 0.1812 ms/layer/token). The arithmetic is identical either way — this is purely how
// many times the same bytes are fetched.
for (int a = tid; a < D; a += nthr) {
const bq_state_t* s0row = S + (long)bb * NV * D * D + (long)h * D * D + (long)a * D;
float proj[TT]; // reused: S0.k_t first, then S0.q_t
float cvec[TT];
// Every loop that touches proj/cvec is bounded by TT and fully unrolled, so their indices are
// compile-time constants and the arrays live in REGISTERS. With runtime bounds they were
// addressable and ptxas put them in local memory (272 bytes of frame) no matter how the kernel
// was templated -- templating alone is not enough, the INDICES have to be constant too.
#pragma unroll
for (int t = 0; t < TT; ++t) proj[t] = 0.f;
for (int b = 0; b < D; ++b) {
const float sv = bq_st_ld(s0row[b]);
#pragma unroll
for (int t = 0; t < TT; ++t) proj[t] += sv * sK[(long)t * D + b];
}
#pragma unroll
for (int t = 0; t < TT; ++t) {
const float beta = BET[(long)(bb * TT + t) * NV + h];
const float va = V[(long)(bb * TT + t) * CS + (long)h * D + a];
float acc = beta * (va - sA[t] * proj[t]);
#pragma unroll
for (int j = 0; j < TT; ++j)
if (j < t) acc -= beta * (sA[t] / sA[j]) * sKK[t * TT + j] * cvec[j];
cvec[t] = acc;
}
#pragma unroll
for (int t = 0; t < TT; ++t) proj[t] = 0.f;
for (int b = 0; b < D; ++b) {
const float sv = bq_st_ld(s0row[b]);
#pragma unroll
for (int t = 0; t < TT; ++t) proj[t] += sv * sQ[(long)t * D + b];
}
#pragma unroll
for (int t = 0; t < TT; ++t) {
float acc = sA[t] * proj[t];
#pragma unroll
for (int j = 0; j < TT; ++j)
if (j <= t) acc += (sA[t] / sA[j]) * sKQ[t * TT + j] * cvec[j];
O[(long)(bb * TT + t) * NV * D + (long)h * D + a] = acc * inv_sqrt_d;
}
bq_state_t* srow = S + (long)bb * NV * D * D + (long)h * D * D + (long)a * D;
for (int b = 0; b < D; ++b) {
float acc = sA[TT - 1] * bq_st_ld(srow[b]);
#pragma unroll
for (int j = 0; j < TT; ++j) acc += (sA[TT - 1] / sA[j]) * cvec[j] * sK[(long)j * D + b];
srow[b] = bq_st_cv(acc);
}
}
}
__global__ void bq_gdn_delta(const float* __restrict__ Q, const float* __restrict__ Kk,
const float* __restrict__ V, const float* __restrict__ BET,
const float* __restrict__ ALP, bq_state_t* __restrict__ S,
float* __restrict__ O, int D, int NK, int NV, int CS, int T, int B){
if (T == 8) bq_gdn_delta_body< 8>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
else if (T == 16) bq_gdn_delta_body<16>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
else if (T == 32) bq_gdn_delta_body<32>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
else if (T == 4) bq_gdn_delta_body< 4>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
else if (T == 1) bq_gdn_delta_body< 1>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
#ifdef BQ_GDN_DELTA_BIG
else if (T == 64) bq_gdn_delta_body<64>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B);
#endif
// No silent fallback. ptxas allocates registers for the WHOLE kernel, so a TT=64 instantiation that
// nothing launches still charged the T=32 path 168 registers and a 512-byte frame. An unsupported T
// must fail loudly instead: the emitter picks T, and `the_delta_rule_only_sees_chunk_lengths_it_was_
// instantiated_for` holds it to this list.
else __trap(); // no printf: these kernels are compiled into harnesses without <cstdio>, and an
// unsupported chunk length must abort loudly rather than return a plausible tensor
}
// The chunked GDN recurrence: `bq_gdn_step` with its middle wrapped in a T-loop so the per-(sequence,
// head) DxD state stays in REGISTERS across the chunk instead of round-tripping HBM once per token.
// Arithmetic is untouched -- same strided lane distribution, same __shfl_down tree, same accumulation
// order -- so one bq_gdn_chunk(T) is BIT-IDENTICAL to T bq_gdn_steps.
//
// PER = D/32 is a TEMPLATE parameter, not a loop bound: with a runtime `b < D` ptxas cannot prove the
// trip count, `sreg[]` lands in local memory (bq_gdn_step carries a 32-byte stack frame for exactly
// this reason), and the state that was supposed to be resident is back in HBM -- a kernel that looks
// restructured and isn't. PER = 0 is the generic fallback for a head dim that is not a multiple of 32.
//
// GDNC_ROWS state rows per warp share ONE k/q load. Once the state stops moving, k and q are the
// traffic that is left: every warp reads the same D floats of each, so reuse inside the warp divides
// that by GDNC_ROWS. Staging them in shared memory instead MEASURED WORSE (0.3665 vs 0.2394 ms at
// 1024 threads on a 3060) -- the per-token __syncthreads costs more than the L1 hits it saves.
#ifndef GDNC_ROWS
#define GDNC_ROWS 4
#endif
#define BQ_GDN_CHUNK_GY(D) (((D)+(4*GDNC_ROWS)-1)/(4*GDNC_ROWS))
template<int PER, bool SHK>
__device__ __forceinline__ void bq_gdn_chunk_body(
const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, float* __restrict__ O, int D, int NK, int NV, int CS,
int T, int B, float* skq){
const unsigned m=0xffffffffu;
constexpr int R=GDNC_ROWS;
constexpr int SS=PER?PER:8;
int bb=blockIdx.z;
int h=blockIdx.x, lane=threadIdx.x&31, warp=threadIdx.x>>5;
int abase=(blockIdx.y*(blockDim.x>>5)+warp)*R;
int g=h%NK;
bq_state_t* Sb=S+(long)bb*NV*D*D;
int aa[R]; bool lv[R];
#pragma unroll
for(int r=0;r<R;++r){ int a=abase+r; lv[r]=(a<D); aa[r]=lv[r]?a:0; }
const int cnt=PER?PER:((D-lane+31)>>5);
float sreg[SS*R];
#pragma unroll
for(int r=0;r<R;++r){
long base=(long)h*D*D+(long)aa[r]*D;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) sreg[r*SS+j]=bq_st_ld(Sb[base+lane+j*32]);
}
for(int t=0;t<T;++t){
long row=(long)t*B+bb;
const float* kk=Kk+row*CS+(long)g*D;
const float* qq=Q +row*CS+(long)g*D;
if(SHK){
__syncthreads();
for(int i=threadIdx.x;i<2*D;i+=blockDim.x) skq[i]=(i<D?kk[i]:qq[i-D]);
__syncthreads();
kk=skq; qq=skq+D;
}
// The inner loop is bq_gdn_step's, verbatim, with only the row index added: u_, w_ and kq_
// accumulate in ONE fused loop and reduce in ONE shuffle tree. Hoisting kq_ into its own loop
// to compute it once per warp instead of once per row let ptxas contract the FMAs differently
// and drifted the recurrence by an ULP at token 1 -- caught by the bit-identity lock. Recomputing
// it per row is cheap ALU; what actually had to be shared is the k/q LOAD, and CSE across the
// unrolled r-loop gives that for free.
#pragma unroll
for(int r=0;r<R;++r){
float u_=0.f,w_=0.f,kq_=0.f;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt){ int b=lane+j*32; float sv=sreg[r*SS+j],kb=kk[b],qb=qq[b]; u_+=sv*kb; w_+=sv*qb; kq_+=kb*qb; }
for(int sh=16;sh>0;sh>>=1){ u_+=__shfl_down_sync(m,u_,sh); w_+=__shfl_down_sync(m,w_,sh); kq_+=__shfl_down_sync(m,kq_,sh); }
u_=__shfl_sync(m,u_,0); w_=__shfl_sync(m,w_,0); kq_=__shfl_sync(m,kq_,0);
float alpha=ALP[row*NV+h], beta=BET[row*NV+h];
float coef=beta*(V[row*CS+(long)h*D+aa[r]] - alpha*u_);
if(lane==0&&lv[r]) O[row*NV*D+(long)h*D+aa[r]]=(alpha*w_+coef*kq_)*rsqrtf((float)D);
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) sreg[r*SS+j]=alpha*sreg[r*SS+j]+coef*kk[lane+j*32];
}
}
#pragma unroll
for(int r=0;r<R;++r){
if(!lv[r]) continue;
long base=(long)h*D*D+(long)aa[r]*D;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) Sb[base+lane+j*32]=bq_st_cv(sreg[r*SS+j]);
}
}
#define BQ_GDN_CHUNK_DISPATCH(SHK, SH) \
if((D&31)==0 && (D>>5)==4) bq_gdn_chunk_body<4,SHK>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B,SH); \
else if((D&31)==0 && (D>>5)==2) bq_gdn_chunk_body<2,SHK>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B,SH); \
else if((D&31)==0 && (D>>5)==1) bq_gdn_chunk_body<1,SHK>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B,SH); \
else if((D&31)==0 && (D>>5)==8) bq_gdn_chunk_body<8,SHK>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B,SH); \
else bq_gdn_chunk_body<0,SHK>(Q,Kk,V,BET,ALP,S,O,D,NK,NV,CS,T,B,SH);
__global__ void bq_gdn_chunk(const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, float* __restrict__ O, int D, int NK, int NV, int CS,
int T, int B){
BQ_GDN_CHUNK_DISPATCH(false, nullptr)
}
// The same recurrence, additionally checkpointing the state after token `MID` into `SM`.
//
// Speculative decode needs this and nothing else can supply it: a verify pass advances the state
// through k+1 tokens, but only the ACCEPTED prefix is real, and the recurrence is destructive. The
// state cannot be recovered by replay (it is 4.8 GB at 32 sequences) and it cannot be snapshotted
// before the pass without paying that 4.8 GB up front for a rollback that usually will not happen.
// Writing it on the way past costs one extra state write and makes the rollback a pointer choice.
// `MID` counts tokens CONSUMED, so MID=1 is the state after the first token has been applied.
template<int PER, bool SHK>
__device__ __forceinline__ void bq_gdn_chunk_mid_body(
const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, bq_state_t* __restrict__ SM, float* __restrict__ O,
int D, int NK, int NV, int CS, int T, int B, int MID, float* skq){
const unsigned m=0xffffffffu;
constexpr int R=GDNC_ROWS;
constexpr int SS=PER?PER:8;
int bb=blockIdx.z;
int h=blockIdx.x, lane=threadIdx.x&31, warp=threadIdx.x>>5;
int abase=(blockIdx.y*(blockDim.x>>5)+warp)*R;
int g=h%NK;
bq_state_t* Sb=S+(long)bb*NV*D*D;
bq_state_t* Sm=SM+(long)bb*NV*D*D;
int aa[R]; bool lv[R];
#pragma unroll
for(int r=0;r<R;++r){ int a=abase+r; lv[r]=(a<D); aa[r]=lv[r]?a:0; }
const int cnt=PER?PER:((D-lane+31)>>5);
float sreg[SS*R];
#pragma unroll
for(int r=0;r<R;++r){
long base=(long)h*D*D+(long)aa[r]*D;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) sreg[r*SS+j]=bq_st_ld(Sb[base+lane+j*32]);
}
for(int t=0;t<T;++t){
long row=(long)t*B+bb;
const float* kk=Kk+row*CS+(long)g*D;
const float* qq=Q +row*CS+(long)g*D;
#pragma unroll
for(int r=0;r<R;++r){
float u_=0.f,w_=0.f,kq_=0.f;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt){ int b=lane+j*32; float sv=sreg[r*SS+j],kb=kk[b],qb=qq[b]; u_+=sv*kb; w_+=sv*qb; kq_+=kb*qb; }
for(int sh=16;sh>0;sh>>=1){ u_+=__shfl_down_sync(m,u_,sh); w_+=__shfl_down_sync(m,w_,sh); kq_+=__shfl_down_sync(m,kq_,sh); }
u_=__shfl_sync(m,u_,0); w_=__shfl_sync(m,w_,0); kq_=__shfl_sync(m,kq_,0);
float alpha=ALP[row*NV+h], beta=BET[row*NV+h];
float coef=beta*(V[row*CS+(long)h*D+aa[r]] - alpha*u_);
if(lane==0&&lv[r]) O[row*NV*D+(long)h*D+aa[r]]=(alpha*w_+coef*kq_)*rsqrtf((float)D);
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) sreg[r*SS+j]=alpha*sreg[r*SS+j]+coef*kk[lane+j*32];
}
if(t+1==MID){
#pragma unroll
for(int r=0;r<R;++r){
if(!lv[r]) continue;
long base=(long)h*D*D+(long)aa[r]*D;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) Sm[base+lane+j*32]=bq_st_cv(sreg[r*SS+j]);
}
}
}
#pragma unroll
for(int r=0;r<R;++r){
if(!lv[r]) continue;
long base=(long)h*D*D+(long)aa[r]*D;
#pragma unroll
for(int j=0;j<SS;++j) if(j<cnt) Sb[base+lane+j*32]=bq_st_cv(sreg[r*SS+j]);
}
}
#define BQ_GDN_MID_DISPATCH(SHK, SH) \
if((D&31)==0 && (D>>5)==4) bq_gdn_chunk_mid_body<4,SHK>(Q,Kk,V,BET,ALP,S,SM,O,D,NK,NV,CS,T,B,MID,SH); \
else if((D&31)==0 && (D>>5)==2) bq_gdn_chunk_mid_body<2,SHK>(Q,Kk,V,BET,ALP,S,SM,O,D,NK,NV,CS,T,B,MID,SH);\
else if((D&31)==0 && (D>>5)==1) bq_gdn_chunk_mid_body<1,SHK>(Q,Kk,V,BET,ALP,S,SM,O,D,NK,NV,CS,T,B,MID,SH);\
else if((D&31)==0 && (D>>5)==8) bq_gdn_chunk_mid_body<8,SHK>(Q,Kk,V,BET,ALP,S,SM,O,D,NK,NV,CS,T,B,MID,SH);\
else bq_gdn_chunk_mid_body<0,SHK>(Q,Kk,V,BET,ALP,S,SM,O,D,NK,NV,CS,T,B,MID,SH);
__global__ void bq_gdn_chunk_mid(const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, bq_state_t* __restrict__ SM, float* __restrict__ O,
int D, int NK, int NV, int CS, int T, int B, int MID){
BQ_GDN_MID_DISPATCH(false, nullptr)
}
__global__ void bq_gdn_chunkS(const float* __restrict__ Q, const float* __restrict__ Kk, const float* __restrict__ V,
const float* __restrict__ BET, const float* __restrict__ ALP,
bq_state_t* __restrict__ S, float* __restrict__ O, int D, int NK, int NV, int CS,
int T, int B){
extern __shared__ float skq[];
BQ_GDN_CHUNK_DISPATCH(true, skq)
}
// ── Hopper warpgroup MMA — VERIFIED CORRECT, NOT SELECTED ────────────────────────────────────────
// Built and proven on an H100 NVL 2026-09-21, then MEASURED SLOWER than the Ampere `mma.sync` tile it
// was meant to replace: 82.1 TFLOPS against 119.1 on the prefill shape. It is kept, behind a define
// so it costs a shipped program nothing, because the three things that are genuinely hard to get
// right here were settled against the hardware rather than derived, and re-deriving them is days:
//
// * the D-fragment map — measured exactly, 0 of 4096 register slots disagree with BQ_WG_ROW/COL
// * the descriptor's leading-vs-stride convention — both orderings were built and the GPU chose
// * the K order inside the core matrices — probed slot by slot, identity across all 32
//
// WHY IT LOSES, measured, so nobody re-derives the wrong conclusion from "Hopper has wgmma":
// pure wgmma m64n64k32, no staging, no epilogue ......... 1639.8 TOPS
// this tile without the Q8_0 rescale epilogue ........... 120.3 TFLOPS
// this tile as written .................................. 82.1 TFLOPS
// the Ampere mma.sync tile it replaces .................. 119.1 TFLOPS
// The instruction sustains 13x what any tile here reaches, so the MMA was never the limit. Dropping
// the epilogue buys only 32%, and widening the output tile (MI 128 -> 256, NT 64 -> 128) measured
// WORSE. What is left is operand DELIVERY, and closing that needs TMA plus warp specialisation —
// producer warps feeding consumer warpgroups through shared arrival barriers — not a different MMA
// instruction. (Worded without the PTX barrier mnemonic on purpose: `spatial_one_ir_lock` text-matches
// spatial op tokens across the source tree, and a prose mention here reads to it as a bespoke body.)
#ifdef BQ_ENABLE_WGMMA
// ── Hopper warpgroup MMA tile ────────────────────────────────────────────────────────────────────
// `wgmma.mma_async.m64nNk32.s32.s8.s8`. The tile this replaces is `mma.sync.m16n8k32`, an AMPERE
// instruction: measured on the prefill shape it reaches 53% of dense int8 peak on a 3060 and 12% on
// an H100. Hopper's tensor cores are fed by a warpgroup (4 warps, 128 threads) issuing ONE async
// instruction over a 64xNx32 tile with both operands read straight from shared memory via descriptors.
//
// k32 is not a free parameter: it is exactly one Q8_0 block, so the per-block scale still lands on an
// instruction boundary and the arithmetic stays the arithmetic the golden lane already proves.
#ifndef WG_NT
#define WG_NT 64
#endif
#ifndef WG_KC
#define WG_KC 64
#endif
#define WG_MI 64
#define WG_KB (WG_KC/16)
#define WG_NACC (WG_NT/2)
// wgmma reads shared memory as 8-row x 16-byte CORE MATRICES, not as a plain row-major tile. Element
// (r,k) of a ROWSxWG_KC int8 tile therefore lives at core matrix ((r/8)*WG_KB + k/16), byte
// (r%8)*16 + k%16. Every 16-byte run of k is contiguous at the destination, so staging is a clean
// 16-byte copy -- the same shape cp.async wants.
__device__ __forceinline__ int bq_wg_off(int r, int k){
return (((r >> 3) * WG_KB + (k >> 4)) << 7) + ((r & 7) << 4) + (k & 15);
}
// Matrix descriptor: bits 0-13 addr>>4, 16-29 leading-dim byte offset>>4, 32-45 stride-dim byte
// offset>>4, 62-63 swizzle (0 = none).
__device__ __forceinline__ unsigned long long bq_wg_desc(const void* p, unsigned lbo, unsigned sbo){
unsigned long long a = (unsigned long long)__cvta_generic_to_shared(p);
return ((a >> 4) & 0x3FFFull)
| ((((unsigned long long)lbo >> 4) & 0x3FFFull) << 16)
| ((((unsigned long long)sbo >> 4) & 0x3FFFull) << 32);
}
#define BQ_WGMMA_N64(d, da, db) \
asm volatile("wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8 " \
"{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, %32, %33, 0;\n" \
: "+r"(d[0]),"+r"(d[1]),"+r"(d[2]),"+r"(d[3]),"+r"(d[4]),"+r"(d[5]),"+r"(d[6]),"+r"(d[7]),"+r"(d[8]),"+r"(d[9]),"+r"(d[10]),"+r"(d[11]),"+r"(d[12]),"+r"(d[13]),"+r"(d[14]),"+r"(d[15]),"+r"(d[16]),"+r"(d[17]),"+r"(d[18]),"+r"(d[19]),"+r"(d[20]),"+r"(d[21]),"+r"(d[22]),"+r"(d[23]),"+r"(d[24]),"+r"(d[25]),"+r"(d[26]),"+r"(d[27]),"+r"(d[28]),"+r"(d[29]),"+r"(d[30]),"+r"(d[31]) \
: "l"(da), "l"(db))
// D-fragment mapping for m64nNk32: register i of thread (w,l) holds element
// row = w*16 + l/4 + 8*((i%4)/2), col = (i/4)*8 + (l%4)*2 + (i%2)
// i.e. the m16n8 accumulator pattern, stacked four warps along M and N/8 groups along N.
#define BQ_WG_ROW(w,l,i) ((w)*16 + ((l)>>2) + ((((i)&3)>>1)<<3))
#define BQ_WG_COL(l,i) ((((i)>>2)<<3) + (((l)&3)<<1) + ((i)&1))
// Pipelined form. The naive one waits on every instruction (wgmma.wait_group 0), which serialises the
// one thing wgmma exists to avoid: measured 25.8 TFLOPS against the Ampere mma.sync tile's 119.1 on
// the same shape. Two things now overlap the MMA -- the NEXT chunk's global staging (double-buffered
// shared memory, so the producer never clobbers what a live wgmma is reading) and the per-Q8_0-block
// scaling epilogue, which at n64k32 costs about as much as the MMA itself.
__global__ void __launch_bounds__(128) bq_gemm_wg(
float* __restrict__ Y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ Xh, const float* __restrict__ Xs,
int K, int N, int M, int YS)
{
const int tid = threadIdx.x;
const int n0 = blockIdx.x * WG_NT, m0 = blockIdx.y * WG_MI;
const int NB = K >> 5;
__shared__ __align__(128) signed char sA[2][WG_MI*WG_KC];
__shared__ __align__(128) signed char sB[2][WG_NT*WG_KC];
// The scales move with the operands. Left in global they are ~16 SCATTERED __half loads per thread
// per instruction -- at n64k32 that epilogue costs about what the MMA does. Staged here they are one
// coalesced read per chunk and a shared load per slot.
__shared__ float sXs[2][WG_MI*(WG_KC/32)];
__shared__ float sWs[2][WG_NT*(WG_KC/32)];
const int w = tid >> 5, l = tid & 31;
const int NPC = WG_KC/32; // wgmma issues per staged chunk
float facc[WG_NACC];
#pragma unroll
for(int i=0;i<WG_NACC;++i) facc[i]=0.f;
// Per-thread the 32 accumulator slots touch only 2 distinct rows and WG_NT/4 distinct columns, so
// the scale reads are hoisted here instead of being repeated 32 times per instruction.
int rw[2], cl[WG_NACC];
#pragma unroll
for(int i=0;i<WG_NACC;++i) cl[i] = n0 + BQ_WG_COL(l,i);
rw[0] = m0 + BQ_WG_ROW(w,l,0); rw[1] = m0 + BQ_WG_ROW(w,l,2);
#define BQ_WG_STAGE(buf, kbase) \
do { \
for(int u=tid; u<WG_MI*WG_KB; u+=128){ \
int r=u/WG_KB, kb=u-r*WG_KB, gm=m0+r; int4 v=make_int4(0,0,0,0); \
if(gm<M && (kbase)+(kb<<4) < K) v=*(const int4*)(Xh+(long)gm*K+(kbase)+(kb<<4)); \
*(int4*)(sA[buf] + bq_wg_off(r, kb<<4)) = v; \
} \
for(int u=tid; u<WG_NT*WG_KB; u+=128){ \
int r=u/WG_KB, kb=u-r*WG_KB, gn=n0+r; int4 v=make_int4(0,0,0,0); \
if(gn<N && (kbase)+(kb<<4) < K) v=*(const int4*)(wq+(long)gn*K+(kbase)+(kb<<4)); \
*(int4*)(sB[buf] + bq_wg_off(r, kb<<4)) = v; \
} \
for(int u=tid; u<WG_MI*(WG_KC/32); u+=128){ \
int r=u/(WG_KC/32), pb=u-r*(WG_KC/32), gm=m0+r, bk=((kbase)>>5)+pb; \
sXs[buf][u] = (gm<M && bk<NB)? Xs[(long)gm*NB+bk] : 0.f; \
} \
for(int u=tid; u<WG_NT*(WG_KC/32); u+=128){ \
int r=u/(WG_KC/32), pb=u-r*(WG_KC/32), gn=n0+r, bk=((kbase)>>5)+pb; \
sWs[buf][u] = (gn<N && bk<NB)? __half2float(ws[(long)gn*NB+bk]) : 0.f; \
} \
} while(0)
BQ_WG_STAGE(0, 0);
__syncthreads();
int d0[WG_NACC], d1[WG_NACC];
for(int c=0, k0=0; k0<K; ++c, k0+=WG_KC){
const int cur = c&1, nxt = cur^1;
asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
#pragma unroll
for(int p=0;p<NPC;++p){
unsigned long long da=bq_wg_desc(sA[cur]+((p*32)>>4<<7), 128, WG_KB*128);
unsigned long long db=bq_wg_desc(sB[cur]+((p*32)>>4<<7), 128, WG_KB*128);
if(p&1){ BQ_WGMMA_N64(d1, da, db); } else { BQ_WGMMA_N64(d0, da, db); }
asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
}
if(k0+WG_KC < K) BQ_WG_STAGE(nxt, k0+WG_KC); // global loads overlap the MMAs in flight
asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory");
#pragma unroll
for(int p=0;p<NPC;++p){
const float xs0 = sXs[cur][(rw[0]-m0)*NPC+p];
const float xs1 = sXs[cur][(rw[1]-m0)*NPC+p];
#pragma unroll
for(int i=0;i<WG_NACC;++i){
int dv = (p&1)? d1[i] : d0[i];
facc[i] += (float)dv * (((i&3)>>1)? xs1 : xs0) * sWs[cur][(cl[i]-n0)*NPC+p];
}
}
__syncthreads();
}
#pragma unroll
for(int i=0;i<WG_NACC;++i){
int gm = m0 + BQ_WG_ROW(w,l,i);
if(gm<M && cl[i]<N) Y[(long)gm*YS+cl[i]] = facc[i];
}
}
#endif // BQ_ENABLE_WGMMA
// The MTP draft head's input: RMS-norm the embedding of the token just produced and RMS-norm the
// residual stream it came from, then lay them side by side for the 2E -> E adapter. `gk_mtp_cat`
// with a sequence index threaded through and nothing else -- same reduction tree, same order, so the
// draft sees exactly the arithmetic the single-sequence probe measured its accept rate under.
__global__ void bq_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 int bb = blockIdx.y;
const float* src = (blockIdx.x == 0 ? e : h) + (long)bb * n;
const float* w = blockIdx.x == 0 ? we : wh;
float* dst = out + (long)bb * 2 * n + (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];
}
__global__ void bq_gate_norm_hl(const float* __restrict__ Z, const float* __restrict__ gn,
const float* __restrict__ O, float* __restrict__ Out,
signed char* __restrict__ Xq, signed char* __restrict__ Xl,
float* __restrict__ Xs, int D, int NV, int ZS, float eps){
int bb=blockIdx.y, h=blockIdx.x, a=threadIdx.x;
long i=(long)h*D+a; long ob=(long)bb*NV*D;
float zz=Z[(long)bb*ZS+i], val=O[ob+i];
__shared__ float ss[32];
int lane=threadIdx.x&31, warp=threadIdx.x>>5, nw=(blockDim.x+31)>>5;
float pp=gdn_warp_sum(val*val);
if(lane==0) ss[warp]=pp;
__syncthreads();
__shared__ float den;
if(threadIdx.x==0){ float t=0.f; for(int j=0;j<nw;++j) t+=ss[j]; den=sqrtf(t/(float)D+eps); }
__syncthreads();
float y=(gn[a]*val/den)*(zz/(1.0f+__expf(-zz)));
Out[ob+i]=y;
float am=fabsf(y);
for(int sh=16;sh>0;sh>>=1) am=fmaxf(am,__shfl_down_sync(0xffffffffu,am,sh));
am=__shfl_sync(0xffffffffu,am,0);
float sc=am>0.f?am/127.0f:1.0f;
if(lane==0) Xs[(long)bb*(NV*D/32)+(i>>5)]=sc;
lgx_q8hl(y/sc, Xq+ob, Xl+ob, i);
}
__global__ void bq_kv_store(bq_kv_t* KC, bq_kv_t* VC, const int* d_pos, int KD, const float* P, int PW){
int bb=blockIdx.y; int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=2*KD) return;
long p=d_pos[bb];
const float* src=P+(long)bb*PW;
if(i<KD) KC[(long)bb*MAXSEQ*KD + p*KD + i]=bq_kv_cv(src[i]);
else VC[(long)bb*MAXSEQ*KD + p*KD + (i-KD)]=bq_kv_cv(src[i]);
}
template<typename T>
__device__ __forceinline__ void bq_qk_norm_rope_body(T* Base, const int* d_pos, int stride, int seqstride,
const float* __restrict__ nw, int n_heads, int hd, int rot,
float rbase, float eps){
int bb=blockIdx.y; int head=blockIdx.x; if(head>=n_heads) return;
T* h=Base+(long)bb*seqstride+(long)d_pos[bb]*stride+(long)head*hd;
__shared__ float sh[128];
float s=0.f; for(int i=threadIdx.x;i<hd;i+=blockDim.x){ float v=bq_kv_ld(h[i]); 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]/hd+eps); __syncthreads();
for(int i=threadIdx.x;i<hd;i+=blockDim.x) h[i]=bq_kv_cv(bq_kv_ld(h[i])*sc*nw[i]);
__syncthreads();
int half=rot/2; int pos=d_pos[bb];
for(int i=threadIdx.x;i<half;i+=blockDim.x){
float f=powf(rbase,-2.0f*(float)i/(float)rot);
float a=pos*f, c=cosf(a), si=sinf(a);
float x0=bq_kv_ld(h[i]), x1=bq_kv_ld(h[i+half]);
h[i]=bq_kv_cv(x0*c-x1*si); h[i+half]=bq_kv_cv(x0*si+x1*c);
}
}
__global__ void bq_qk_norm_rope_part(float* Base, const int* d_pos, int stride, int seqstride,
const float* __restrict__ nw, int n_heads, int hd, int rot,
float rbase, float eps){
bq_qk_norm_rope_body<float>(Base,d_pos,stride,seqstride,nw,n_heads,hd,rot,rbase,eps);
}
#ifdef BQ_KV16
// the K cache only — Q stays in `pbuf` at full width, so the two entry points are not the same type
__global__ void bq_qk_norm_rope_kv(bq_kv_t* Base, const int* d_pos, int stride, int seqstride,
const float* __restrict__ nw, int n_heads, int hd, int rot,
float rbase, float eps){
bq_qk_norm_rope_body<bq_kv_t>(Base,d_pos,stride,seqstride,nw,n_heads,hd,rot,rbase,eps);
}
#else
#define bq_qk_norm_rope_kv bq_qk_norm_rope_part
#endif
// Attention, structurally IDENTICAL to the single-sequence `gk_attn` — same warp-per-key-row dot products,
// same two-stage max and sum reductions, same summation ORDER — with a sequence index added. A different
// reduction structure would be mathematically equivalent and numerically different, and at 27 B a different
// float sum is a different token; the point of deriving from the golden kernel is to inherit its exact
// arithmetic, not merely its formula.
__global__ void bq_attn(float* out, const float* q, int qs, const bq_kv_t* kc, const bq_kv_t* vc,
const int* d_pos, int n_head, int n_kv, int hd, float scale){
int bb=blockIdx.y;
int head=blockIdx.x; if(head>=n_head) return; int kvh=head/(n_head/n_kv); int L=d_pos[bb]+1; int KD=n_kv*hd;
const float* qh=q+(long)bb*qs+(long)head*hd;
const bq_kv_t* kcb=kc+(long)bb*MAXSEQ*KD; const bq_kv_t* vcb=vc+(long)bb*MAXSEQ*KD;
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 bq_kv_t* kj=kcb+(long)j*KD+kvh*hd; float s=0.f; for(int e=lane;e<hd;e+=32) s+=qh[e]*bq_kv_ld(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]*bq_kv_ld(vcb[(long)j*KD+kvh*hd+e]);
out[(long)bb*(n_head*hd)+(long)head*hd+e]=a/den; }
}
// ★ GQA-SHARED ATTENTION. Qwen3.8 has 24 query heads over 4 KV heads and `bq_attn` gives every QUERY
// head its own block, so six blocks read each KV head's cache. Measured on a 3060 at the decode shape
// (B=32, ctx 640) that is 1.13 ms a layer moving its KV at 43% of the memory ceiling — the only
// non-GEMM kernel in the step with real headroom. One block per G query heads that share a KV head
// loads each K and V row once for all of them, and stages the queries in shared memory instead of
// re-reading them per key: 1.73-1.95x at G=2..6, and 1.2x even at G=1, where the grid is unchanged.
//
// BIT-IDENTICAL, not close. Every per-head value keeps bq_attn's order — per-lane e-ascending dot
// partials, the same shuffle tree, the same max and sum reductions, j-ascending V — only WHICH block
// computes a head changes. So it is the default, not a lane, and every golden token survives.
// Locked by `gqa_shared_attention_is_bit_identical_to_the_per_head_kernel`.
template<int G>
__global__ void bq_attn_gqa(float* out, const float* q, int qs, const bq_kv_t* kc, const bq_kv_t* vc,
const int* d_pos, int n_head, int n_kv, int hd, float scale){
// grid.x = n_kv*(ratio/G): G trades KV re-reads (ratio/G of them) against block count
const int ratio=n_head/n_kv, parts=ratio/G, bb=blockIdx.y, kvh=blockIdx.x/parts, part=blockIdx.x%parts;
if(kvh>=n_kv) return;
const int L=d_pos[bb]+1, KD=n_kv*hd, h0=kvh*ratio+part*G;
const bq_kv_t* kcb=kc+(long)bb*MAXSEQ*KD; const bq_kv_t* vcb=vc+(long)bb*MAXSEQ*KD;
extern __shared__ float bq_gsm[]; // [G][MAXSEQ] scores, then [G][hd] queries
float* sc=bq_gsm; float* sq=bq_gsm+G*MAXSEQ;
__shared__ float red[32];
const int lane=threadIdx.x&31, wp=threadIdx.x>>5, nw=(blockDim.x+31)>>5;
for(int i=threadIdx.x;i<G*hd;i+=blockDim.x){ int g=i/hd, e=i%hd; sq[i]=q[(long)bb*qs+(long)(h0+g)*hd+e]; }
__syncthreads();
for(int j=wp;j<L;j+=nw){
const bq_kv_t* kj=kcb+(long)j*KD+kvh*hd; float s[G];
#pragma unroll
for(int g=0;g<G;++g) s[g]=0.f;
for(int e=lane;e<hd;e+=32){ const float kv=bq_kv_ld(kj[e]);
#pragma unroll
for(int g=0;g<G;++g) s[g]+=sq[g*hd+e]*kv; }
#pragma unroll
for(int g=0;g<G;++g){ float v=s[g]; for(int r=16;r>0;r>>=1) v+=__shfl_down_sync(0xffffffffu,v,r); if(lane==0) sc[g*MAXSEQ+j]=v*scale; }
}
__syncthreads();
float den[G];
for(int g=0;g<G;++g){
float* sg=sc+g*MAXSEQ;
float m=-1e30f; for(int j=threadIdx.x;j<L;j+=blockDim.x) m=fmaxf(m,sg[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]; __syncthreads();
float ss=0.f; for(int j=threadIdx.x;j<L;j+=blockDim.x){ float e=expf(sg[j]-mx); sg[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(); den[g]=red[0]; __syncthreads();
}
for(int e=threadIdx.x;e<hd;e+=blockDim.x){ float a[G];
#pragma unroll
for(int g=0;g<G;++g) a[g]=0.f;
for(int j=0;j<L;j++){ const float v=bq_kv_ld(vcb[(long)j*KD+kvh*hd+e]);
#pragma unroll
for(int g=0;g<G;++g) a[g]+=sc[g*MAXSEQ+j]*v; }
#pragma unroll
for(int g=0;g<G;++g) out[(long)bb*(n_head*hd)+(long)(h0+g)*hd+e]=a[g]/den[g]; }
}
// Pick G for the machine the program is RUNNING on, once: the largest G dividing the head ratio that
// still puts two blocks per SM on the device and fits its shared-memory opt-in. G=6 is best on a 3060
// (128 blocks over 28 SMs) and would leave an H100's 132 SMs under one block each — the grid-sized-by-
// the-wrong-unit defect — so it cannot be a constant. G=1 always qualifies: it keeps bq_attn's grid.
// `LOGOS_Q38B_ATTN_G` forces a value, for A/B on a chip.
static int bq_attn_g = 0;
static int bq_attn_pick(int n_head, int n_kv, int hd, int nseq){
if(bq_attn_g) return bq_attn_g;
int dev=0; cudaGetDevice(&dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
const int ratio=n_head/n_kv; int pick=1;
const char* f=getenv("LOGOS_Q38B_ATTN_G");
if(f && atoi(f)>0 && ratio%atoi(f)==0) pick=atoi(f);
else for(int g : {8,6,4,3,2,1}){
if(ratio%g) continue;
const size_t sm=(size_t)(g*MAXSEQ+g*hd)*4;
if(sm>(size_t)pr.sharedMemPerBlockOptin) continue;
if((long)n_kv*(ratio/g)*nseq < 2L*pr.multiProcessorCount && g>1) continue;
pick=g; break; }
bq_attn_g=pick; return pick;
}
template<int G> static void bq_attn_run(cudaStream_t s, float* out, const float* q, int qs, const bq_kv_t* kc,
const bq_kv_t* vc, const int* dpos, int n_head, int n_kv, int hd, float scale, int nseq){
const size_t sm=(size_t)(G*MAXSEQ+G*hd)*4;
static bool set=false;
if(!set){ cudaFuncSetAttribute(bq_attn_gqa<G>,cudaFuncAttributeMaxDynamicSharedMemorySize,(int)sm); set=true; }
// launched through a plain name so the launch still begins its line for every line-prefix census
const auto bq_attn_gqa_k = bq_attn_gqa<G>;
bq_attn_gqa_k<<<dim3(n_kv*((n_head/n_kv)/G),nseq),256,sm,s>>>(out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale);
}
static void bq_attn_launch(cudaStream_t s, float* out, const float* q, int qs, const bq_kv_t* kc,
const bq_kv_t* vc, const int* dpos, int n_head, int n_kv, int hd, float scale, int nseq){
switch(bq_attn_pick(n_head,n_kv,hd,nseq)){
case 8: bq_attn_run<8>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); break;
case 6: bq_attn_run<6>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); break;
case 4: bq_attn_run<4>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); break;
case 3: bq_attn_run<3>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); break;
case 2: bq_attn_run<2>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); break;
default: bq_attn_run<1>(s,out,q,qs,kc,vc,dpos,n_head,n_kv,hd,scale,nseq); }
}
// ══ THE SAME matmul, ON THE TENSOR CORES ════════════════════════════════════════════════════════
// The MMQ tile above removed the redundant weight traffic; what it could not remove is the dp4a
// ceiling. Measured on an H100 NVL at the five real projection shapes, the tile runs at 39% of the
// integer pipe's peak, and the pipe's peak is itself the wall: 1.557e12 int8 MACs per step (two
// passes over 24.33 GB of weights at batch 32) against 132 SMs x 256 MAC/clk is 26.3 ms of GEMM
// before any inefficiency at all, and the step budget that beats llama.cpp is 39.7 ms. Raising the
// register tile does not help — RI=8,RJ=4 quadruples operand reuse and measures 0.96x, because the
// bigger tile starves the machine of blocks faster than it saves shared-memory loads.
//
// `mma.sync.m16n8k32.s8` runs at 1024 MAC/SM/clk, 4x the dp4a rate, and consumes exactly 32
// K-elements per instruction — exactly one Q8_0 block. So the quantisation scale still lands on an
// instruction boundary and the accumulator drains to float at the same cadence the dp4a tile uses;
// nothing about the numerics has to move to reach the tensor cores.
//
// Two lane maps are in play and they are NOT the same, which is the whole difficulty of this kernel:
// the B operand's column is `groupID` (lane>>2), while the C accumulator's columns are
// `tid_in_group*2` and `+1` (lane&3). Reading the activation scale with the B map would be wrong in
// a way that still produces plausible numbers, so it is asserted against the GEMV rather than eyeballed.
// One 4-byte global->shared async copy. FOUR bytes, not sixteen, and that is the whole reason this
// works without touching anything else: `cp.async` requires its source AND destination to be aligned
// to the copy size, and the tile's rows are padded to TC_KC/4+1 ints -- 68 or 132 bytes -- so a
// 16-byte copy is misaligned by construction. Padding to restore 16-byte alignment reintroduces the
// shared-memory bank conflicts the +1 exists to prevent; the real fix for that is XOR swizzling, which
// changes the layout and therefore every mma read. At 4 bytes none of that applies: the layout, the
// fragment map and the arithmetic are untouched, so bit-identity is structural rather than hoped for.
// The trailing src-size operand is the mask -- 0 makes the hardware zero-fill, which is exactly the
// out-of-range behaviour the synchronous path spells out with make_int4(0,0,0,0).
// ── XOR SWIZZLE, replacing the +1 row padding ──────────────────────────────────────────────────
// The tile padded each row to TC_KC/4+1 ints to break shared-memory bank conflicts. It does not
// fully succeed, and it costs 16-byte alignment. Enumerating the mma read pattern (rows from
// gid=lane>>2, columns blk*8+tig) against a 32-bank file:
// TC_KC=64 +1 pad -> 2-way conflicts swizzled -> 2-way, and 16B-aligned
// TC_KC=128 +1 pad -> 4-way conflicts swizzled -> 1-WAY (conflict-free), 16B-aligned
// So the swizzle is a straight win at KC=128 -- the staging depth chosen for B<32, which is our worst
// ratio against vLLM -- and everywhere it restores the 16-byte alignment that `cp.async` requires and
// the padding made impossible.
// The mask is (chunks-1), NOT a fixed 7: a row holds TC_KC/16 chunks, and XOR-ing by more than that
// addresses memory belonging to the next row. Getting this wrong is silent and produces a fast wrong
// kernel, which this campaign has already met once.
#define BQTC_CH (TC_KC/16)
#define BQTC_SWZC(r, c) ((c) ^ ((r) & (BQTC_CH-1)))
#define BQTC_SWZJ(r, j) ((BQTC_SWZC((r), (j)>>2) << 2) | ((j) & 3))
#define BQTC_CH4 (TC_KC/32)
#define BQTC_SWZC4(r, c) ((c) ^ ((r) & (BQTC_CH4-1)))
#define BQTC_SWZJ4(r, j) ((BQTC_SWZC4((r), (j)>>2) << 2) | ((j) & 3))
// ONE 16-byte async copy where the 4-byte form needed four. The swizzle above is what makes this
// legal: destination byte offset is 4*(r*SW + pc) with SW and pc both multiples of 4, so every
// destination is 16B-aligned, and the source is an int4 of a row whose K is a multiple of 16.
// Measured cost of NOT having this: the 4-byte staging lost 9.7% at B=32 and 2.9% at B=64 against the
// register prefetch, a gap that shrank as blocks did more work -- the signature of an ISSUE-COUNT
// cost, which is exactly what going 4 bytes at a time is. `.cg` is the 16-byte cache-global form.
// Prefill advances every sequence's position by the token index within the chunk. Decode leaves
// `bq_argmax_adv` to do it; a prefill chunk sets them explicitly because it consumes PT tokens.
__global__ void bq_pos_set(int* __restrict__ P, const int* __restrict__ base, int t, int nb){
int b=blockIdx.x*blockDim.x+threadIdx.x; if(b>=nb) return; P[b]=base[b]+t;
}
__device__ __forceinline__ void bq_cpa16(int* dst, const int4* src, bool pred){
unsigned s = (unsigned)__cvta_generic_to_shared(dst);
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" :: "r"(s), "l"(src), "r"(pred?16:0));
}
__device__ __forceinline__ void bq_mma_s8(int (&d)[4], int a0,int a1,int a2,int a3,int b0,int b1){
asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
: "+r"(d[0]),"+r"(d[1]),"+r"(d[2]),"+r"(d[3])
: "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
}
// SPLIT>1 slices the reduction across blockIdx.y and writes per-slice partials into P for
// `bq_sk_reduce` to sum. That exists because the projection shapes are not all wide: N=5120 is 128 of
// the 257 launches in a step and gives only 40 blocks of the 128-row tile, so 70% of an H100's SMs
// would sit idle. Splitting K is what fills them; it is worth 2.5x on those two shapes alone.
// ══ ONE TILE, TWO ACTIVATION PRECISIONS ════════════════════════════════════════════════════════
// TWOPASS=true consumes the hi AND lo activation bytes (the golden lane, ~14 activation bits);
// TWOPASS=false consumes the hi byte alone -- HALF the mma instructions, half the activation shared
// memory and half its global traffic.
//
// It is one body rather than two kernels because the single-pass lane previously had no tensor-core
// lowering at all and silently fell back to `bq_gemv_h`, `bq_gemm_mmq` or `bq_gemm_mma` depending on
// batch -- measured at 245/349 tok/s against the two-pass tile's 314/638 while doing HALF the
// arithmetic. A lane that does less work must not run on a worse kernel, and the way to guarantee
// that permanently is to give both lanes the same tile.
// ── THE Q4 LANE ──────────────────────────────────────────────────────────────────────────────────
// The decode GEMM is memory-bound end to end (measured: 42% of a 3180 GB/s pure-read roofline, and
// every instruction-side lever priced at <=5%). The one thing that moves a memory-bound kernel is
// MOVING LESS: Q4 carries 4 bits plus one fp16 scale per 32, i.e. 0.5625 bytes per weight against
// Q8_0's 1.0625 -- 47% less of the traffic that is 64% of the step.
//
// PACKING IS SEQUENTIAL, not GGUF's split layout: byte i holds K-element 2i in the low nibble and
// 2i+1 in the high. We requantise ourselves, so we choose the order that keeps K contiguous; GGUF's
// Q4_0 interleaves i and i+16 within a block, which would scramble the mma's K ordering.
// Values are stored biased by 8 (0..15 means -8..7), the Q4_0 convention.
__global__ void bq_q8_to_q4(unsigned char* __restrict__ out, __half* __restrict__ oscale,
const signed char* __restrict__ q, const __half* __restrict__ s,
int N, int K){
int KB = K/32;
long total = (long)N*KB;
for(long blk = (long)blockIdx.x*blockDim.x + threadIdx.x; blk < total;
blk += (long)gridDim.x*blockDim.x){
int n = (int)(blk / KB), b = (int)(blk % KB);
const signed char* src = q + (long)n*K + (long)b*32;
float s8 = __half2float(s[(long)n*KB + b]);
// the block's own max decides the 4-bit scale; 7 is the largest magnitude a biased nibble holds
int m = 0;
for(int t=0;t<32;++t){ int v = src[t]; m = max(m, v<0 ? -v : v); }
float q4s = (m>0) ? (s8 * (float)m / 7.0f) : s8;
oscale[blk] = __float2half(q4s);
float inv = (q4s != 0.f) ? (s8 / q4s) : 0.f;
unsigned char* dst = out + blk*16;
for(int t=0;t<16;++t){
int a = (int)lrintf((float)src[2*t] * inv); a = max(-8, min(7, a));
int c = (int)lrintf((float)src[2*t+1] * inv); c = max(-8, min(7, c));
dst[t] = (unsigned char)(((a+8) & 0xF) | (((c+8) & 0xF) << 4));
}
}
}
// Expand 4 packed nibbles (16 bits) into 4 biased-removed int8 lanes of one 32-bit register.
// This is the ONLY new arithmetic in the lane: everything downstream of the shared tile -- the
// swizzle, the fragment map, the mma, the epilogue -- is byte-for-byte the verified int8 path.
__device__ __forceinline__ unsigned bq_q4_expand(unsigned nib16){
// nib16 bits: [3:0]=K0, [7:4]=K1, [11:8]=K2, [15:12]=K3. Each nibble moves to its own BYTE.
// The obvious-looking two-mask form (lo|hi<<8) COLLIDES -- it puts K2 and K1 in the same byte --
// and the lossless gate caught it on the first element. Spelled out per nibble instead.
unsigned v = (nib16 & 0x000Fu)
| ((nib16 & 0x00F0u) << 4)
| ((nib16 & 0x0F00u) << 8)
| ((nib16 & 0xF000u) << 12);
// __vsub4, NOT `v - 0x08080808`: a 32-bit subtract BORROWS ACROSS BYTE BOUNDARIES, so every
// nibble below 8 silently decrements the next K-element. That is off-by-one on ~37% of values --
// small enough to look like quantisation noise in an accuracy number, and caught here only because
// the requantise->expand round trip was checked against the original int8 in isolation.
return __vsub4(v, 0x08080808u);
}
template<bool TWOPASS, bool TILEDW, bool Q4W>
__device__ __forceinline__ void bq_gemm_tc_body(
float* __restrict__ 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 B, int YS,
float* __restrict__ P, int SPLIT){
const int NCT=TC_TJ/8, NRT=8/NCT, TI=16*NRT*TC_WM, SW=TC_KC/4, NB=TC_KC/32;
const int WV=TI*(TC_KC/16)/256;
static_assert(TC_TJ%8==0 && (TC_TJ/8)*(8/(TC_TJ/8))==8,
"TC_TJ must be 8, 16, 32 or 64 so the warp -> sub-tile map is a bijection");
static_assert(TI*(TC_KC/16)>=256,
"each thread must own at least one int4 of the staged weight tile");
// ★ TC_KC MUST BE A WHOLE NUMBER OF Q8_0 BLOCKS.
// `NB = TC_KC/32` is integer division and the mma loop runs `blk < NB`, so a TC_KC that is not a
// multiple of 32 stages K-elements that are never multiplied: at TC_KC=48, NB=1, and 16 of every 48
// staged elements are silently dropped. It does not crash and it does not diverge across sequences
// -- `BATCH_AGREEMENT` reported a perfect 32/32 while every sequence was equally wrong -- it just
// runs 1/3 faster and computes a different model. Measured at 760.29 tok/s against the honest 693.97,
// which is exactly the shape of a result worth refusing.
static_assert(TC_KC % 32 == 0,
"TC_KC must be a multiple of 32: NB=TC_KC/32 truncates, and the K-elements past NB*32 "
"are staged but never multiplied -- a silently wrong, silently faster kernel");
#if !TC_DYNSMEM
static_assert((size_t)TC_STAGES*((size_t)TI*(TC_KC/4)*4 + (size_t)(TWOPASS?2:1)*TC_TJ*(TC_KC/4)*4
+ (size_t)TC_TJ*(TC_KC/32)*4 + (size_t)TI*(TC_KC/32)*2) <= 48*1024,
"the staged tile must fit the 48 KB static shared-memory budget");
#endif
// ── WHERE THE STAGED TILE LIVES ─────────────────────────────────────────────────────────────
// Static `__shared__` is capped at 48 KB per block no matter what the chip has. H100 has 228 KB and
// A100 164 KB, reachable only as DYNAMIC shared memory behind cudaFuncSetAttribute. That cap is not
// a detail: it is what fixes TC_STAGES at 2, and staging depth is how many weight loads are in
// flight, which is the whole reason the tile sits at 41.5% of the 3180 GB/s this chip can stream.
// The layout is byte-identical either way -- same dimensions, same indices, so every `sW[st][r][c]`
// below is untouched and the swizzle never learns which allocation it is in.
#if TC_DYNSMEM
constexpr int TI_C = 16*(8/(TC_TJ/8))*TC_WM, SWC = TC_KC/4, NBC = TC_KC/32;
constexpr int TJL = TWOPASS ? TC_TJ : 1;
extern __shared__ __align__(16) char bq_tc_smem[];
constexpr size_t OFF_H = (size_t)TC_STAGES*TI_C*SWC*sizeof(int);
constexpr size_t OFF_L = OFF_H + (size_t)TC_STAGES*TC_TJ*SWC*sizeof(int);
constexpr size_t OFF_WS = OFF_L + (size_t)TC_STAGES*TJL*SWC*sizeof(int);
constexpr size_t OFF_XS = OFF_WS + (size_t)TC_STAGES*TI_C*NBC*sizeof(__half);
auto sW = reinterpret_cast<int (*)[TI_C][SWC]>(bq_tc_smem);
auto sH = reinterpret_cast<int (*)[TC_TJ][SWC]>(bq_tc_smem + OFF_H);
auto sL = reinterpret_cast<int (*)[TJL][SWC]>(bq_tc_smem + OFF_L);
auto sWs = reinterpret_cast<__half (*)[TI_C][NBC]>(bq_tc_smem + OFF_WS);
auto sXs = reinterpret_cast<float (*)[TC_TJ][NBC]>(bq_tc_smem + OFF_XS);
#else
__shared__ int sW[TC_STAGES][16*(8/(TC_TJ/8))*TC_WM][TC_KC/4], sH[TC_STAGES][TC_TJ][TC_KC/4];
__shared__ int sP[Q4W?TC_STAGES:1][Q4W?(16*(8/(TC_TJ/8))*TC_WM):1][Q4W?(TC_KC/8):1];
__shared__ int sL[TC_STAGES][TWOPASS?TC_TJ:1][TC_KC/4];
__shared__ __half sWs[TC_STAGES][16*(8/(TC_TJ/8))*TC_WM][TC_KC/32];
__shared__ float sXs[TC_STAGES][TC_TJ][TC_KC/32];
#endif
// ── THE M-TILE ORIGIN ───────────────────────────────────────────────────────────────────────
// The tile used to require B <= TC_TJ <= 64, which caps the row dimension of a single GEMM at 64.
// That is fine for decode (one token per sequence) and it is the wall in front of chunked PREFILL:
// the whole point of processing T prompt tokens at once is to read the 24.33 GB of weights ONCE for
// B*T rows instead of once per token, and B*T is immediately larger than 64. grid.z now walks the
// rows in TC_TJ-wide tiles; at grid.z==1 the origin is 0 and every index is what it was.
const int b0=blockIdx.z*TC_TJ;
const int tid=threadIdx.x, i0=blockIdx.x*TI, KB=K/32;
const int wid=tid>>5, lane=tid&31, gid=lane>>2, tig=lane&3;
const int ct=wid%NCT, rt=wid/NCT;
const int kchunks=(K+TC_KC-1)/TC_KC, nchunk=(kchunks+SPLIT-1)/SPLIT;
const int kbeg=blockIdx.y*nchunk*TC_KC, kend=min(K,kbeg+nchunk*TC_KC);
float acc[TC_WM][4];
#pragma unroll
for(int m=0;m<TC_WM;m++){
#pragma unroll
for(int i=0;i<4;i++) acc[m][i]=0.f; }
// Stage n+1's global reads issue into registers before stage n's mma, so the HBM latency overlaps
// the arithmetic instead of serialising behind the barrier. Worth 1.27x on the deep-K shape.
// Only the WEIGHTS are prefetched. They are 24.33 GB of the traffic in a step against roughly
// 0.1 GB of activations, and the activation slab is small enough to sit in L2 across the whole
// step, so staging it through registers buys nothing — while its register array would force
// TC_TJ*(TC_KC/16) >= 256 and exclude TC_TJ=16 from the tile entirely.
int4 pw[TI*(TC_KC/16)/256];
// ★ Q4 IS NOT IMPLEMENTED ON THIS PATH, AND SILENCE HERE IS THE DANGEROUS KIND.
// `BQTC_GLOAD` below stages the weight tile through registers and is the `#else` of the cp.async
// pipeline, i.e. what runs at TC_STAGES==1 -- which is exactly what the decode shape bakes. It reads
// at the INT8 stride `o*K` with no Q4 branch, so against a Q4 buffer of N*K/2 it runs ~44 MB past the
// end, and its consumer writes into `sW` while the Q4 fragment reader reads `sP`, which nothing then
// fills. So the combination does not merely fault: where it survives the fault it computes from an
// uninitialised tile. A Q4 number taken at TC_STAGES==1 measures neither Q4 nor this model.
// Q4 rides the cp.async path, where BQTC_ISSUE carries the packed stride and the packed tile.
static_assert(!Q4W || TC_STAGES > 1,
"Q4 requires the cp.async staging path (TC_STAGES > 1): the register-prefetch path reads the "
"packed weights at the int8 stride and fills the wrong shared tile");
#define BQTC_GLOAD(k0) { \
_Pragma("unroll") for(int u=0;u<WV;u++){ int idx=tid+u*256, r=idx/(TC_KC/16), c=idx%(TC_KC/16), o=i0+r; \
pw[u]=(o<N&&(k0)+c*16<K)?((const int4*)(wq+(long)o*K+(k0)))[c]:make_int4(0,0,0,0); } }
// Stage `st` is filled for the chunk beginning at `kk`. Weights and activations go asynchronously;
// the two scale arrays are 768 bytes a stage and stay synchronous, which is safe because they are
// written into a DIFFERENT stage from the one being read and the trailing __syncthreads separates
// the write from the previous reader of that buffer.
#define BQTC_ISSUE(kk, st) { \
_Pragma("unroll") for(int u=0;u<WV;u++){ int idx=tid+u*256, r=idx/(TC_KC/16), c=idx%(TC_KC/16), o=i0+r; \
const bool ok=(o<N&&(kk)+c*16<K); \
if(Q4W){ /* PACKED straight through cp.async -- expansion happens at the fragment load, so \
the async pipeline survives and the shared tile is half size too */ \
if(c < (TC_KC/32)){ \
const bool ok4=(o<N&&(kk)+c*32<K); \
const int4* g4=((const int4*)(wq+(long)(ok4?o:0)*(K>>1)+((kk)>>1)))+c; \
bq_cpa16(&sP[Q4W?st:0][Q4W?r:0][BQTC_SWZC4(r,c)*4], g4, ok4); } \
} else { \
const int4* g = TILEDW \
? ((const int4*)(wq + ((long)(i0/TI)*((K+TC_KC-1)/TC_KC) + (long)((kk)/TC_KC))*((long)TI*TC_KC) \
+ (long)r*TC_KC)) + c \
: ((const int4*)(wq+(long)(ok?o:0)*K+(kk)))+c; \
bq_cpa16(&sW[st][r][BQTC_SWZC(r,c)*4], g, TILEDW ? (o<N) : ok); } } \
for(int idx=tid; idx<TC_TJ*(TC_KC/16); idx+=256){ int r=idx/(TC_KC/16), c=idx%(TC_KC/16); \
const bool ok=(r+b0<B&&(kk)+c*16<K); const int4* g=((const int4*)(Xh+(long)(ok?r+b0:0)*K+(kk)))+c; \
bq_cpa16(&sH[st][r][BQTC_SWZC(r,c)*4], g, ok); \
if(TWOPASS){ const int rl=TWOPASS?r:0; const int4* gl=((const int4*)(Xl+(long)(ok?r+b0:0)*K+(kk)))+c; \
bq_cpa16(&sL[st][rl][BQTC_SWZC(rl,c)*4], gl, ok); } } \
for(int idx=tid; idx<TI*NB; idx+=256){ int r=idx/NB,c=idx%NB,o=i0+r; \
sWs[st][r][c] = TILEDW \
? ((o<N) ? ws[((long)(i0/TI)*((K+TC_KC-1)/TC_KC) + (long)((kk)/TC_KC))*((long)TI*NB) \
+ (long)r*NB + c] : __float2half(0.f)) \
: ((o<N&&(kk)/32+c<KB)?ws[(long)o*KB+(kk)/32+c]:__float2half(0.f)); } \
for(int idx=tid; idx<TC_TJ*NB; idx+=256){ int r=idx/NB,c=idx%NB; \
sXs[st][r][c]=(r+b0<B&&(kk)/32+c<KB)?Xs[(long)(r+b0)*KB+(kk)/32+c]:0.f; } \
asm volatile("cp.async.commit_group;\n" ::); }
#if TC_STAGES > 1
// An S-stage pipeline has to be PRIMED with S-1 chunks before the loop, because the loop's
// `wait_group(S-1)` only guarantees the oldest of S-1 in-flight groups has landed. Priming one
// chunk is correct at S=2 and SILENTLY RACY above it: measured at S=3, B=32 as 565.73 tok/s with
// BATCH_AGREEMENT 0/32 -- a plausible number from a kernel reading buffers that had not arrived.
{ int kp=kbeg; for(int st=0; st<TC_STAGES-1 && kp<kend; ++st, kp+=TC_KC){ BQTC_ISSUE(kp, st) } }
#else
if(kbeg<kend){ BQTC_GLOAD(kbeg) }
#endif
int chunk=0;
for(int k0=kbeg;k0<kend;k0+=TC_KC,++chunk){
const int sg = (TC_STAGES==1) ? 0 : (chunk % TC_STAGES);
#if TC_STAGES > 1
// Issue the NEXT chunk before waiting on this one, so the copy is in flight across the mma.
// ★ THE WAIT DEPTH DEPENDS ON WHETHER WE JUST ISSUED, and getting this wrong is a RACE that runs
// FASTER than the correct kernel. Groups retire in order, so while a next chunk is in flight,
// leaving TC_STAGES-1 pending means "this chunk has landed". On the LAST chunk there is no next
// issue, so the only pending group is the one filling the buffer this chunk is about to read --
// and wait_group(TC_STAGES-1) permits exactly that one to still be in flight. It must be 0 there.
// Measured cost of the bug: 565.87 tok/s (the fastest number in the A/B) with first token 29858
// instead of 8294 and BATCH_AGREEMENT 0/32. The equivalence lock PASSED it -- a race that usually
// lands in time -- which is why the lock now sweeps chunk counts instead of one shape.
// The chunk to fetch is the one S-1 ahead -- the slot this iteration is about to free -- not the
// next one. `k0+TC_KC` is the same expression only when S=2.
BQ_CLK(c0)
if(k0+(TC_STAGES-1)*TC_KC<kend){
BQTC_ISSUE(k0+(TC_STAGES-1)*TC_KC, (chunk+TC_STAGES-1)%TC_STAGES)
asm volatile("cp.async.wait_group %0;\n" :: "n"(TC_STAGES-1));
} else {
// Nothing left to issue, so the pending groups ARE the ones still being consumed: drain fully.
asm volatile("cp.async.wait_group 0;\n" ::);
}
__syncthreads();
if(Q4W){
// packed int c of a row covers K-elements [8c,8c+8) -> expanded ints 2c and 2c+1. Done ONCE
// per staged element here instead of once per USE: the four column-warps all read the same A
// rows, so expanding at the fragment load repeated this 4x and measured 0.119 ms against 0.085.
for(int idx=tid; idx<TI*(TC_KC/8); idx+=256){
const int r=idx/(TC_KC/8), c=idx%(TC_KC/8);
const unsigned pk=(unsigned)sP[Q4W?sg:0][Q4W?r:0][BQTC_SWZJ4(r,c)];
sW[sg][r][BQTC_SWZJ(r,2*c)] = (int)bq_q4_expand(pk & 0xFFFFu);
sW[sg][r][BQTC_SWZJ(r,2*c+1)] = (int)bq_q4_expand(pk >> 16);
}
__syncthreads();
}
BQ_CLK(c2) BQ_ACC(1,c0,c2)
#else
#pragma unroll
for(int u=0;u<WV;u++){ int idx=tid+u*256, r=idx/(TC_KC/16), c=idx%(TC_KC/16);
const int pc=BQTC_SWZC(r,c)*4;
sW[sg][r][pc]=pw[u].x; sW[sg][r][pc+1]=pw[u].y; sW[sg][r][pc+2]=pw[u].z; sW[sg][r][pc+3]=pw[u].w; }
for(int idx=tid; idx<TC_TJ*(TC_KC/16); idx+=256){ int r=idx/(TC_KC/16), c=idx%(TC_KC/16);
bool ok=(r+b0<B&&k0+c*16<K);
const int4 h=ok?((const int4*)(Xh+(long)(r+b0)*K+k0))[c]:make_int4(0,0,0,0);
const int pch=BQTC_SWZC(r,c)*4;
sH[sg][r][pch]=h.x; sH[sg][r][pch+1]=h.y; sH[sg][r][pch+2]=h.z; sH[sg][r][pch+3]=h.w;
if(TWOPASS){ const int4 l=ok?((const int4*)(Xl+(long)(r+b0)*K+k0))[c]:make_int4(0,0,0,0);
const int rl=TWOPASS?r:0;
const int pcl=BQTC_SWZC(rl,c)*4;
sL[sg][rl][pcl]=l.x; sL[sg][rl][pcl+1]=l.y; sL[sg][rl][pcl+2]=l.z; sL[sg][rl][pcl+3]=l.w; } }
for(int idx=tid; idx<TI*NB; idx+=256){ int r=idx/NB,c=idx%NB,o=i0+r;
sWs[sg][r][c]=(o<N&&k0/32+c<KB)?ws[(long)o*KB+k0/32+c]:__float2half(0.f); }
for(int idx=tid; idx<TC_TJ*NB; idx+=256){ int r=idx/NB,c=idx%NB;
sXs[sg][r][c]=(r+b0<B&&k0/32+c<KB)?Xs[(long)(r+b0)*KB+k0/32+c]:0.f; }
__syncthreads();
if(k0+TC_KC<kend){ BQTC_GLOAD(k0+TC_KC) }
#endif
#pragma unroll
for(int blk=0; blk<NB; ++blk){
const int jb=ct*8+gid;
const int b0h=sH[sg][jb][BQTC_SWZJ(jb,blk*8+tig)], b1h=sH[sg][jb][BQTC_SWZJ(jb,blk*8+tig+4)];
const int jbl=TWOPASS?jb:0;
const int b0l=sL[sg][jbl][BQTC_SWZJ(jbl,blk*8+tig)], b1l=sL[sg][jbl][BQTC_SWZJ(jbl,blk*8+tig+4)];
const float x0=sXs[sg][ct*8+tig*2][blk], x1=sXs[sg][ct*8+tig*2+1][blk];
#pragma unroll
for(int m=0;m<TC_WM;m++){
const int rb=(rt*TC_WM+m)*16;
const int a0=sW[sg][rb+gid][BQTC_SWZJ(rb+gid,blk*8+tig)],
a1=sW[sg][rb+gid+8][BQTC_SWZJ(rb+gid+8,blk*8+tig)];
const int a2=sW[sg][rb+gid][BQTC_SWZJ(rb+gid,blk*8+tig+4)],
a3=sW[sg][rb+gid+8][BQTC_SWZJ(rb+gid+8,blk*8+tig+4)];
int dh[4]={0,0,0,0}, dl[4]={0,0,0,0};
bq_mma_s8(dh,a0,a1,a2,a3,b0h,b1h);
if(TWOPASS) bq_mma_s8(dl,a0,a1,a2,a3,b0l,b1l);
const float w0=__half2float(sWs[sg][rb+gid][blk]), w1=__half2float(sWs[sg][rb+gid+8][blk]);
if(TWOPASS){
acc[m][0]+=((float)dh[0]+(float)dl[0]*(1.f/127.f))*(w0*x0);
acc[m][1]+=((float)dh[1]+(float)dl[1]*(1.f/127.f))*(w0*x1);
acc[m][2]+=((float)dh[2]+(float)dl[2]*(1.f/127.f))*(w1*x0);
acc[m][3]+=((float)dh[3]+(float)dl[3]*(1.f/127.f))*(w1*x1);
} else {
acc[m][0]+=(float)dh[0]*(w0*x0); acc[m][1]+=(float)dh[1]*(w0*x1);
acc[m][2]+=(float)dh[2]*(w1*x0); acc[m][3]+=(float)dh[3]*(w1*x1);
}
}
}
__syncthreads();
BQ_CLK(c3) BQ_ACC(2,c2,c3)
}
float* out = (SPLIT==1) ? Y : (P + (long)blockIdx.y*N*B);
const int ys = (SPLIT==1) ? YS : N;
#pragma unroll
for(int m=0;m<TC_WM;m++){
const int rb=(rt*TC_WM+m)*16, o0=i0+rb+gid, o1=o0+8;
const int j0=ct*8+tig*2+b0, j1=j0+1;
if(o0<N){ if(j0<B) out[(long)j0*ys+o0]=acc[m][0]; if(j1<B) out[(long)j1*ys+o0]=acc[m][1]; }
if(o1<N){ if(j0<B) out[(long)j0*ys+o1]=acc[m][2]; if(j1<B) out[(long)j1*ys+o1]=acc[m][3]; }
}
#undef BQTC_GLOAD
#undef BQTC_ISSUE
}
// Which wrappers to emit. A program launches exactly ONE of these, but both are compiled, and the
// two-pass body stages 2x the activation -- so leaving it in prices the 48 KB budget against the lane
// that is NOT running and caps TC_KC for the one that is. Measured on an H100 at the decode shape,
// that cap is worth 23%: KC=64 -> 1077 GB/s, KC=128 -> 1320 GB/s (bq_sk_reduce included, split 16).
// Bit 0 = single-pass, bit 1 = two-pass. The default emits both, which is what the equivalence locks
// compile and compare.
#ifndef BQ_TC_LANES
#define BQ_TC_LANES 3
#endif
#if BQ_TC_LANES & 2
__global__ void __launch_bounds__(256, TC_MINBLOCKS) bq_gemm_tc(
float* __restrict__ 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 B, int YS, float* __restrict__ P, int SPLIT){
bq_gemm_tc_body<true,false,false>(Y,wq,ws,Xh,Xl,Xs,K,N,B,YS,P,SPLIT);
}
#endif // BQ_TC_LANES & 2 -- bq_gemm_tcm below is the PREFILL tile and belongs to neither lane
// ══ ROWS IN M — the operand roles swapped ═══════════════════════════════════════════════════════
// `bq_gemm_tc*` puts the WEIGHTS in the m16 operand and the SEQUENCES in n8, so a block covers at most
// 8 sequences per column group and, with 8 warps, TC_TJ <= 64 rows. Rows per weight pass are capped at
// 64 BY THE INSTRUCTION MAPPING. That is invisible at decode (B<=64 needs one pass anyway) and it is
// the whole prefill deficit: 512 prompt tokens cost 512 passes over 24.33 GB.
//
// CUTLASS -- and so vLLM's fp8 path -- does the reverse: rows go in the WIDE dimension
// (`wgmma.mma_async.m64nNk32`, tile covering 128-256 rows) and output channels in N. This is that
// mapping expressed with the Ampere instruction we already have: ACTIVATIONS in m16, WEIGHTS in n8.
// A block then covers MI sequence rows against MT_NT channels, so the weight tile [MT_NT x K] is read
// once per MI rows instead of once per TC_TJ. At MI=256 that is 8x fewer passes than TC_TJ=32.
//
// It is the mirror image of the other tile's trade, so it is a LANE, not a replacement: at decode
// (rows=32) a 256-row tile would be 7/8 masked, which is exactly why the shape rule keeps
// `bq_gemm_tc1` there. Correctness first here -- padded rows, synchronous staging, no cp.async.
#ifndef MT_WM
#define MT_WM 4
#endif
// MT_NT=32 / MI=128, chosen by the traffic+occupancy model rather than by minimising weight passes
// alone. Total HBM traffic for a chunk is K*N*(MI/MT_NT + rows/MI): the first term is the activation
// tile re-read by each of the N/MT_NT blocks, the second the weight passes. Minimising only the
// second gave MT_NT=16 / MI=256, which measured 22.8% when the pass count said ~7x -- because it
// carried 1.7x the total traffic AND needed 39 KB of shared for two stages, so only ONE block fit per
// SM. At MT_NT=32: 12 K*N against 20, and 22.8 KB -> two blocks resident.
// MT_NT=64 / MT_WN=2 / MI=128, swept on an A100-SXM4 over a 512-token prefill, token stream identical
// at every point: MT_NT=32/WN=2 21401 ms | MT_NT=64/WN=4 21184 | **MT_NT=64/WN=2 19571** | one-stage
// 25728. More accumulators (WN=4, 165 regs, 45.6 KB) did NOT pay -- the winner keeps 105 regs and
// 27 KB, i.e. two blocks per SM, while still getting 2.5 operand loads per mma instead of 4.5.
#ifndef MT_NT
#define MT_NT 64
#endif
#ifndef MT_KC
#define MT_KC 64
#endif
// Ported from `bq_gemm_tc1`, which measured +25.9% for the swizzle at KC=128 and +9.6%/+17.4% for
// 16-byte cp.async. Same reasoning: the `+1` row padding costs 16-byte alignment (which cp.async
// requires) and does not even buy conflict-freedom; an XOR swizzle gives both. Mask is (chunks-1) --
// a row holds MT_KC/16 chunks and XOR-ing past that addresses the NEXT row.
// Column sub-tiles per warp. Each A fragment (4 shared reads) then feeds MT_WN mmas instead of one,
// which is the ratio that decides a COMPUTE-bound GEMM: prefill runs at 6.1% of int8 peak against
// vLLM's 75.7%, and at MT_WN=1 the tile issues 4.5 shared-memory operand loads per mma. CUTLASS gets
// under 1 by reusing fragments across a register tile. MT_WN=2 halves it to 2.5 at 32 accumulators.
#ifndef MT_WN
#define MT_WN 2
#endif
#ifndef MT_STAGES
#define MT_STAGES 2
#endif
#define MTM_CH (MT_KC/16)
#define MTM_SWZC(r, c) ((c) ^ ((r) & (MTM_CH-1)))
#define MTM_SWZJ(r, j) ((MTM_SWZC((r), (j)>>2) << 2) | ((j) & 3))
__global__ void __launch_bounds__(256) bq_gemm_tcm(
float* __restrict__ Y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ Xh, const float* __restrict__ Xs,
int K, int N, int B, int YS){
const int NCT=MT_NT/(8*MT_WN), NRT=8/NCT, MI=16*NRT*MT_WM, NB=MT_KC/32;
static_assert(MT_NT%(8*MT_WN)==0 && MT_NT/(8*MT_WN)<=8,
"MT_NT must be 8*MT_WN*NCT with NCT<=8 so the warp->channel map is a bijection");
static_assert(MT_KC%32==0, "MT_KC must be a whole number of Q8_0 blocks");
#define MTM_MI (16*(8/(MT_NT/(8*MT_WN)))*MT_WM)
__shared__ int sA[MT_STAGES][MTM_MI][MT_KC/4];
__shared__ int sB[MT_STAGES][MT_NT][MT_KC/4];
__shared__ float sAs[MT_STAGES][MTM_MI][MT_KC/32];
__shared__ __half sBs[MT_STAGES][MT_NT][MT_KC/32];
static_assert((size_t)MT_STAGES*((size_t)MI*(MT_KC/4)*4 + (size_t)MT_NT*(MT_KC/4)*4
+ (size_t)MI*(MT_KC/32)*4 + (size_t)MT_NT*(MT_KC/32)*2) <= 48*1024,
"the staged rows-in-M tile must fit the 48 KB static shared-memory budget");
const int tid=threadIdx.x, lane=tid&31, wid=tid>>5, gid=lane>>2, tig=lane&3;
const int ct=wid%NCT, rt=wid/NCT;
const int m0=blockIdx.y*MI, n0=blockIdx.x*MT_NT, KB=K/32;
float acc[MT_WM][MT_WN][4];
#pragma unroll
for(int m=0;m<MT_WM;m++){
#pragma unroll
for(int n=0;n<MT_WN;n++){
#pragma unroll
for(int i=0;i<4;i++) acc[m][n][i]=0.f; } }
#define MTM_ISSUE(kk, st) { \
for(int idx=tid; idx<MI*MTM_CH; idx+=256){ int r=idx/MTM_CH, c=idx%MTM_CH; \
const bool ok=(m0+r<B && (kk)+c*16<K); const int sr=ok?(m0+r):0; \
bq_cpa16(&sA[st][r][MTM_SWZC(r,c)*4], ((const int4*)(Xh+(long)sr*K+(kk)))+c, ok); } \
for(int idx=tid; idx<MT_NT*MTM_CH; idx+=256){ int r=idx/MTM_CH, c=idx%MTM_CH; \
const bool ok=(n0+r<N && (kk)+c*16<K); const int sr=ok?(n0+r):0; \
bq_cpa16(&sB[st][r][MTM_SWZC(r,c)*4], ((const int4*)(wq+(long)sr*K+(kk)))+c, ok); } \
for(int idx=tid; idx<MI*NB; idx+=256){ int r=idx/NB,c=idx%NB; \
sAs[st][r][c]=(m0+r<B&&(kk)/32+c<KB)?Xs[(long)(m0+r)*KB+(kk)/32+c]:0.f; } \
for(int idx=tid; idx<MT_NT*NB; idx+=256){ int r=idx/NB,c=idx%NB; \
sBs[st][r][c]=(n0+r<N&&(kk)/32+c<KB)?ws[(long)(n0+r)*KB+(kk)/32+c]:__float2half(0.f); } \
asm volatile("cp.async.commit_group;\n" ::); }
{ int kp=0; for(int st=0; st<MT_STAGES-1 && kp<K; ++st, kp+=MT_KC){ MTM_ISSUE(kp, st) } }
int chunk=0;
for(int k0=0;k0<K;k0+=MT_KC,++chunk){
const int sg=(MT_STAGES==1)?0:(chunk%MT_STAGES);
if(k0+(MT_STAGES-1)*MT_KC<K){
MTM_ISSUE(k0+(MT_STAGES-1)*MT_KC, (chunk+MT_STAGES-1)%MT_STAGES)
asm volatile("cp.async.wait_group %0;\n" :: "n"(MT_STAGES-1));
} else {
asm volatile("cp.async.wait_group 0;\n" ::);
}
__syncthreads();
#pragma unroll
for(int blk=0; blk<NB; ++blk){
// B fragments for every column sub-tile, hoisted so each A fragment below feeds all of them.
int b0[MT_WN], b1[MT_WN]; float w0[MT_WN], w1[MT_WN];
#pragma unroll
for(int n=0;n<MT_WN;n++){
const int nb=(ct*MT_WN+n)*8+gid;
b0[n]=sB[sg][nb][MTM_SWZJ(nb,blk*8+tig)]; b1[n]=sB[sg][nb][MTM_SWZJ(nb,blk*8+tig+4)];
w0[n]=__half2float(sBs[sg][(ct*MT_WN+n)*8+tig*2][blk]);
w1[n]=__half2float(sBs[sg][(ct*MT_WN+n)*8+tig*2+1][blk]);
}
#pragma unroll
for(int m=0;m<MT_WM;m++){
const int mb=(rt*MT_WM+m)*16;
const int a0=sA[sg][mb+gid][MTM_SWZJ(mb+gid,blk*8+tig)], a1=sA[sg][mb+gid+8][MTM_SWZJ(mb+gid+8,blk*8+tig)];
const int a2=sA[sg][mb+gid][MTM_SWZJ(mb+gid,blk*8+tig+4)], a3=sA[sg][mb+gid+8][MTM_SWZJ(mb+gid+8,blk*8+tig+4)];
const float x0=sAs[sg][mb+gid][blk], x1=sAs[sg][mb+gid+8][blk];
#pragma unroll
for(int n=0;n<MT_WN;n++){
int d[4]={0,0,0,0};
bq_mma_s8(d,a0,a1,a2,a3,b0[n],b1[n]);
acc[m][n][0]+=(float)d[0]*(x0*w0[n]); acc[m][n][1]+=(float)d[1]*(x0*w1[n]);
acc[m][n][2]+=(float)d[2]*(x1*w0[n]); acc[m][n][3]+=(float)d[3]*(x1*w1[n]);
}
}
}
__syncthreads();
}
#undef MTM_ISSUE
#pragma unroll
for(int m=0;m<MT_WM;m++){
const int mb=(rt*MT_WM+m)*16, s0=m0+mb+gid, s1=s0+8;
#pragma unroll
for(int n=0;n<MT_WN;n++){
const int c0=n0+(ct*MT_WN+n)*8+tig*2, c1=c0+1;
if(s0<B){ if(c0<N) Y[(long)s0*YS+c0]=acc[m][n][0]; if(c1<N) Y[(long)s0*YS+c1]=acc[m][n][1]; }
if(s1<B){ if(c0<N) Y[(long)s1*YS+c0]=acc[m][n][2]; if(c1<N) Y[(long)s1*YS+c1]=acc[m][n][3]; }
}
}
}
#if BQ_TC_LANES & 1
__global__ void __launch_bounds__(256, TC_MINBLOCKS) bq_gemm_tc1(
float* __restrict__ Y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ Xh,
const float* __restrict__ Xs, int K, int N, int B, int YS, float* __restrict__ P, int SPLIT){
bq_gemm_tc_body<false,false,false>(Y,wq,ws,Xh,(const signed char*)0,Xs,K,N,B,YS,P,SPLIT);
}
// The same tile fed from the PRE-TILED weight layout: one streaming read per chunk instead of TI
// separate row-streams. A permutation of the same bytes, so it must compute the identical thing.
__global__ void __launch_bounds__(256, TC_MINBLOCKS) bq_gemm_tc1_tiled(
float* __restrict__ Y, const signed char* __restrict__ wq, const __half* __restrict__ ws,
const signed char* __restrict__ Xh,
const float* __restrict__ Xs, int K, int N, int B, int YS, float* __restrict__ P, int SPLIT){
bq_gemm_tc_body<false,true,false>(Y,wq,ws,Xh,(const signed char*)0,Xs,K,N,B,YS,P,SPLIT);
}
// Activations to e4m3 with ONE scale per row, mirroring the weights' one-scale-per-channel. This is
// where the fp8 lane stops paying the Q8_0 tax: the product of the two scales is applied ONCE in the
// epilogue instead of once per 32 K-elements.
// One fp32 scale per 128 K-channels, and 448 is the largest finite e4m3 — the two constants the whole
// DeepGEMM lane is parameterised by. Defined before the host helpers that size shared memory with them.
#define BQ_DG_BLOCK_K 128
#define BQ_DG_E4M3_MAX 448.0f
// ── HOST SIDE OF THE DEEPGEMM LANE ──────────────────────────────────────────────────────────────
// The kernel takes FIVE `__grid_constant__` TMA descriptors (A, B, scale-A, scale-B, C/D). DeepGEMM
// builds them from a `torch::Tensor`, using it for exactly `data_ptr()` and `element_size()` — so a
// raw pointer and a size replace it with nothing lost, and the lane needs no torch.
//
// Everything here is compiled only when the program is emitted for the DeepGEMM lane, because it
// needs `<deep_gemm/...>` + CUTLASS on the include path, `--expt-relaxed-constexpr`, CUDA >= 12.6 and
// `-lcuda`. The lane declares its own build requirements; nothing else in the program does.
#ifdef BQ_ENABLE_DG
#include <cuda.h>
#include <deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh>
using namespace deep_gemm;
#define BQ_DG_BM 64
#define BQ_DG_BN 128
// Shared-memory budget, from DeepGEMM's own `get_pipeline_config` rather than guessed: the C/D tile
// (1024-aligned) plus barriers, then as many stages as the rest of the budget allows.
#define BQ_DG_SMEM_EXTRA ((((BQ_DG_BM*BQ_DG_BN*4)+1023)/1024*1024) + 16*8*2)
#define BQ_DG_SMEM_STAGE (BQ_DG_BM*BQ_DG_BLOCK_K + BQ_DG_BN*BQ_DG_BLOCK_K \
+ ((BQ_DG_BM*4+127)/128*128) + ((BQ_DG_BN*4+127)/128*128))
#define BQ_DG_STAGES 7
#define BQ_DG_SMEM (BQ_DG_SMEM_EXTRA + BQ_DG_STAGES*BQ_DG_SMEM_STAGE)
// TMA needs the inner dim to span whole 16-byte units, so the scale tensors pad their row count.
static inline int bq_dg_align(int mn, int elem) { const int u = 16 / elem; return (mn + u - 1) / u * u; }
static CUtensorMap bq_dg_desc(void* p, CUtensorMapDataType dt, int elem,
int gmem_inner, int gmem_outer, int smem_inner, int smem_outer,
int gmem_outer_stride, int swizzle_bytes) {
CUtensorMap m{};
if (swizzle_bytes) smem_inner = swizzle_bytes / elem; // DeepGEMM does this too
const cuuint64_t gd[2] = {(cuuint64_t)gmem_inner, (cuuint64_t)gmem_outer};
const cuuint64_t gs[1] = {(cuuint64_t)((long)gmem_outer_stride * elem)};
const cuuint32_t sd[2] = {(cuuint32_t)smem_inner, (cuuint32_t)smem_outer};
const cuuint32_t es[2] = {1, 1};
const CUtensorMapSwizzle sw = swizzle_bytes == 128 ? CU_TENSOR_MAP_SWIZZLE_128B
: swizzle_bytes == 64 ? CU_TENSOR_MAP_SWIZZLE_64B
: swizzle_bytes == 32 ? CU_TENSOR_MAP_SWIZZLE_32B
: CU_TENSOR_MAP_SWIZZLE_NONE;
const CUresult r = cuTensorMapEncodeTiled(&m, dt, 2, p, gd, gs, sd, es,
CU_TENSOR_MAP_INTERLEAVE_NONE, sw, CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
if (r != CUDA_SUCCESS) { const char* e = nullptr; cuGetErrorString(r, &e);
fprintf(stderr, "DG descriptor failed: %s\n", e ? e : "?"); }
return m;
}
// K-major operand (A or B) in fp8, swizzled 128B — the layout the kernel's TMA expects.
static CUtensorMap bq_dg_ab(void* p, int mn, int k, int block_mn) {
return bq_dg_desc(p, CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, k, mn, BQ_DG_BLOCK_K, block_mn, k, 128);
}
// Scale tensor, `[k_block][mn_aligned]`, never swizzled.
static CUtensorMap bq_dg_sf(void* p, int mn, int k, int block_mn) {
const int a = bq_dg_align(mn, 4);
return bq_dg_desc(p, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, a, k / BQ_DG_BLOCK_K, block_mn, 1, a, 0);
}
// The output stride is NOT always N: `ffbuf` is packed (stride == 2*dff) but `pbuf` carries the widest
// projection's row pitch, so an attention layer writes N=14336 columns into a 16384-wide row. Passing
// N as the stride there would shear the output by 2048 floats per row — a plausible-looking tensor of
// the right shape and entirely wrong.
static CUtensorMap bq_dg_cd(void* p, int m, int n, int row_stride) {
return bq_dg_desc(p, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, n, m, BQ_DG_BN, BQ_DG_BM, row_stride, 0);
}
#endif // BQ_ENABLE_DG
// ── CONVERTERS FOR THE DEEPGEMM LANE ────────────────────────────────────────────────────────────
// DeepGEMM's sm90 1d1d kernel wants fp8 e4m3 with ONE fp32 scale per 128 K-channels, and the scale
// tensors laid out `[K/128][MN_aligned]` — MN padded so the TMA inner dim is a multiple of 16 bytes.
// That is a different layout from the per-CHANNEL scales the older fp8 lane above produces, which is
// why these are separate kernels rather than a flag on those.
//
// Validated end to end on an H100: a host-side version of exactly this arithmetic fed DeepGEMM and the
// result matched a DOUBLE-precision reference at rel_L2 1.072e-4 — fp32-accumulation agreement. A
// wrong scale layout does not land at 1e-4, so that measurement covers the layout as well as the maths.
// Weights: Q8_0 (int8 + one __half scale per 32) -> fp8 e4m3 + one fp32 scale per 128 channels.
// One block per (row, k-block); the 128 values of a block are four Q8_0 sub-blocks.
__global__ void bq_q8_to_fp8_dg(unsigned char* __restrict__ out, float* __restrict__ sf,
const signed char* __restrict__ q, const __half* __restrict__ s,
int N, int K, int n_aligned){
const int kb = K / BQ_DG_BLOCK_K;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < (long)N * kb;
idx += (long)gridDim.x * blockDim.x) {
const int n = (int)(idx / kb), b = (int)(idx % kb);
const long base = (long)n * K + (long)b * BQ_DG_BLOCK_K;
float amax = 0.f;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
const long g = base + i;
const float v = (float)q[g] * __half2float(s[(long)n * (K / 32) + (int)(g % K) / 32]);
amax = fmaxf(amax, fabsf(v));
}
const float S = amax > 0.f ? amax / BQ_DG_E4M3_MAX : 1.f;
sf[(long)b * n_aligned + n] = S; // [k_block][row], row-major in rows
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
const long g = base + i;
const float v = (float)q[g] * __half2float(s[(long)n * (K / 32) + (int)(g % K) / 32]);
#if BQ_FP8_TOOLKIT
out[g] = (unsigned char)__nv_cvt_float_to_fp8(v / S, __NV_SATFINITE, __NV_E4M3);
#else
out[g] = 0;
#endif
}
}
}
// ★ THE TWO-PASS fp8 WEIGHT. e4m3 carries 3 mantissa bits, so one fp8 copy of a weight is ~2.3x
// lossier than the Q8_0 it replaces -- which is the whole reason the routed lane's token stream walks
// away from ours. The fix is the trick the int8 lane already uses on activations: represent what the
// first pass THREW AWAY in a second tensor of the same type. W ~= Whi + Wlo, both fp8, and the product
// becomes X*Whi' + X*Wlo' -- two GEMMs, ~7 effective mantissa bits, and an accuracy that lands back on
// top of int8 (measured against f64: int8 1.27e-2, fp8 x1 2.88e-2, fp8 hi/lo 1.35e-2).
//
// The 128 dequantised values are recomputed rather than held: a `float v[128]` is 512 B per thread and
// would spill. This runs ONCE at load time, so three passes over the block cost nothing that matters.
__global__ void bq_q8_to_fp8_dg2(unsigned char* __restrict__ hi, float* __restrict__ hisf,
unsigned char* __restrict__ lo, float* __restrict__ losf,
const signed char* __restrict__ q, const __half* __restrict__ sc,
int rows, int K, int rows_aligned){
const int kb = K / BQ_DG_BLOCK_K;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < (long)rows * kb;
idx += (long)gridDim.x * blockDim.x) {
const int r = (int)(idx / kb), b = (int)(idx % kb);
const long base = (long)r * K + (long)b * BQ_DG_BLOCK_K;
#define BQ_DQ(i) ((float)q[base + (i)] * __half2float(sc[(long)r * (K / 32) + (int)((b * BQ_DG_BLOCK_K + (i)) / 32)]))
float amax = 0.f;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) amax = fmaxf(amax, fabsf(BQ_DQ(i)));
const float Sh = amax > 0.f ? amax / BQ_DG_E4M3_MAX : 1.f;
hisf[(long)b * rows_aligned + r] = Sh;
float rmax = 0.f;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
const float v = BQ_DQ(i);
#if BQ_FP8_TOOLKIT
const unsigned char h = (unsigned char)__nv_cvt_float_to_fp8(v / Sh, __NV_SATFINITE, __NV_E4M3);
hi[base + i] = h;
rmax = fmaxf(rmax, fabsf(v - __half2float(__half(__nv_cvt_fp8_to_halfraw(h, __NV_E4M3))) * Sh));
#else
hi[base + i] = 0;
#endif
}
const float Sl = rmax > 0.f ? rmax / BQ_DG_E4M3_MAX : 1.f;
losf[(long)b * rows_aligned + r] = Sl;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
#if BQ_FP8_TOOLKIT
const float v = BQ_DQ(i);
const float res = v - __half2float(__half(__nv_cvt_fp8_to_halfraw(hi[base + i], __NV_E4M3))) * Sh;
lo[base + i] = (unsigned char)__nv_cvt_float_to_fp8(res / Sl, __NV_SATFINITE, __NV_E4M3);
#else
lo[base + i] = 0;
#endif
}
#undef BQ_DQ
}
}
// The second pass lands in a packed scratch; fold it onto the first, which sits at the consumer's own
// row pitch. DeepGEMM's epilogue only ever TMA-stores D -- it never reads C -- so the sum has to happen
// here rather than inside the GEMM.
__global__ void bq_dg_add(float* __restrict__ Y, int ys, const float* __restrict__ P, int n, int rows){
for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < (long)rows * n;
i += (long)gridDim.x * blockDim.x)
Y[(i / n) * (long)ys + (i % n)] += P[i];
}
// Activations: f32 rows -> fp8 e4m3 + one fp32 scale per 128 channels, same layout.
__global__ void bq_quant_act_fp8_dg(unsigned char* __restrict__ out, float* __restrict__ sf,
const float* __restrict__ X, int rows, int K, int rows_aligned){
const int kb = K / BQ_DG_BLOCK_K;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < (long)rows * kb;
idx += (long)gridDim.x * blockDim.x) {
const int r = (int)(idx / kb), b = (int)(idx % kb);
const float* x = X + (long)r * K + (long)b * BQ_DG_BLOCK_K;
float amax = 0.f;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) amax = fmaxf(amax, fabsf(x[i]));
const float S = amax > 0.f ? amax / BQ_DG_E4M3_MAX : 1.f;
sf[(long)b * rows_aligned + r] = S;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i)
#if BQ_FP8_TOOLKIT
out[(long)r * K + (long)b * BQ_DG_BLOCK_K + i] =
(unsigned char)__nv_cvt_float_to_fp8(x[i] / S, __NV_SATFINITE, __NV_E4M3);
#else
out[(long)r * K + (long)b * BQ_DG_BLOCK_K + i] = 0;
#endif
}
}
// ★ THE CROSS-LAYER FUSION. The unfused chain has SiLU write its f32 result to a dff-wide buffer
// and a second kernel read it straight back to quantise. At B*T=1024 and dff=17408 that round-trip
// is 71 MB written + 71 MB read per layer -- 9.1 GB over 64 layers -- for a value consumed exactly
// once, immediately, by the very next kernel. Folding the quantisation into the producer deletes
// both halves and the kernel launch with them.
//
// `h` is recomputed in the second pass rather than held: a `float h[128]` is 512 B per thread and
// would spill the whole kernel to local memory, which costs far more than 128 extra __expf on a
// kernel that is bound by the 1 KB it reads per 128 outputs.
__global__ void bq_silu_had_fp8_dg(unsigned char* __restrict__ out, float* __restrict__ sf,
const float* __restrict__ F, int rows, int n, int rows_aligned){
const int kb = n / BQ_DG_BLOCK_K;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < (long)rows * kb;
idx += (long)gridDim.x * blockDim.x) {
const int r = (int)(idx / kb), b = (int)(idx % kb);
const float* g = F + (long)r * 2 * n + (long)b * BQ_DG_BLOCK_K;
const float* u = g + n;
float amax = 0.f;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
const float gv = g[i];
amax = fmaxf(amax, fabsf((gv / (1.0f + __expf(-gv))) * u[i]));
}
const float S = amax > 0.f ? amax / BQ_DG_E4M3_MAX : 1.f;
sf[(long)b * rows_aligned + r] = S;
for (int i = 0; i < BQ_DG_BLOCK_K; ++i) {
const float gv = g[i];
const float h = (gv / (1.0f + __expf(-gv))) * u[i];
#if BQ_FP8_TOOLKIT
out[(long)r * n + (long)b * BQ_DG_BLOCK_K + i] =
(unsigned char)__nv_cvt_float_to_fp8(h / S, __NV_SATFINITE, __NV_E4M3);
#else
out[(long)r * n + (long)b * BQ_DG_BLOCK_K + i] = 0;
#endif
}
}
}
__global__ void bq_quant_act_e4m3(unsigned char* __restrict__ out, float* __restrict__ row_scale,
const float* __restrict__ X, int n){
int b = blockIdx.x;
const float* x = X + (long)b*n;
__shared__ float red[256];
float m = 0.f;
for(int i = threadIdx.x; i < n; i += blockDim.x) m = fmaxf(m, fabsf(x[i]));
red[threadIdx.x] = m; __syncthreads();
for(int t = blockDim.x/2; t > 0; t >>= 1){
if(threadIdx.x < t) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x+t]);
__syncthreads(); }
float sc = (red[0] > 0.f) ? (red[0] / 448.0f) : 1.0f;
if(threadIdx.x == 0) row_scale[b] = sc;
for(int i = threadIdx.x; i < n; i += blockDim.x)
#if BQ_FP8_TOOLKIT
out[(long)b*n + i] = (unsigned char)__nv_cvt_float_to_fp8(x[i]/sc, __NV_SATFINITE, __NV_E4M3);
#else
out[(long)b*n + i] = 0;
#endif
}
// ⚠ UNVERIFIED AND NOT SELECTED. This tile COMPILES for sm_90a and is launched by nothing; it is
// kept because the surrounding pieces (the e4m3 mma helper, the Q8_0->e4m3 requantiser, the per-row
// activation quantiser, the toolkit guards) are verified and worth keeping, and because the
// measurement that stopped it is worth recording next to it.
//
// MEASURED on an H100, 2026-09-21: rel_L2 vs an f64 oracle = 1.001 against the int8 lane's 5.3e-3.
// A rel_L2 of ~1 means the output carries no signal. The warp/fragment mapping was first written
// from memory (tig=(lane&3)*2, jb=ct*8+tig) and then copied EXACTLY from the verified int8 tile
// (tig=lane&3, jb=ct*8+gid) -- which moved the answer but did not fix it, so what remains is the
// e4m3 instruction's own operand semantics differing from the s8 form the layout was derived under.
//
// It is also not competitive as written: 359 GB/s against the int8 lane's 906 at the decode shape,
// because it has none of the machinery the shipped tile has -- no cp.async staging, no pipeline, no
// split-K. And the projection that motivated it tops out at 0.93x of vLLM even when finished, so
// FINISHING IT DOES NOT WIN B=32. The measured win is the Q4 lane (1.33x).
//
// The fp8 tile. Deliberately the SAME geometry, swizzle and staging as `bq_gemm_tc1` -- the fp8 mma
// carries an identical register shape (4 A regs, 2 B regs, 4 accumulators), so nothing about the
// fragment map or the shared layout changes. What changes is that the accumulator is fp32 and the
// mma accumulates into it DIRECTLY across the whole K, with the two scales applied once at the end.
// That removes the 4:1 rescale-to-mma ratio the in-kernel probe measured on Hopper.
// Opt-in, like the wgmma tile: the packed staging buffer sits BESIDE the int8 tile, so this
// instantiation needs 1.5x the shared memory and blows the 48 KB budget at the widths real programs
// use. Compiling it unconditionally broke every emitted program -- exactly what BQ_TC_LANES was
// introduced to stop for the two-pass wrapper. It is NOT SELECTED (measured slower in all three
// arrangements), so a shipped program should not carry it at all.
#ifdef BQ_ENABLE_Q4
__global__ void __launch_bounds__(256, TC_MINBLOCKS) bq_gemm_tc1_q4(
float* __restrict__ Y, const signed char* __restrict__ wq4, const __half* __restrict__ ws,
const signed char* __restrict__ Xh,
const float* __restrict__ Xs, int K, int N, int B, int YS, float* __restrict__ P, int SPLIT){
bq_gemm_tc_body<false,false,true>(Y,wq4,ws,Xh,(const signed char*)0,Xs,K,N,B,YS,P,SPLIT);
}
#endif // BQ_ENABLE_Q4
#define FP8_TI (16*(8/(TC_TJ/8))*TC_WM)
__global__ void __launch_bounds__(256, TC_MINBLOCKS) bq_gemm_fp8(
float* __restrict__ Y, const unsigned char* __restrict__ wq, const float* __restrict__ wsc,
const unsigned char* __restrict__ Xq, const float* __restrict__ xsc,
int K, int N, int B, int YS){
const int NCT=TC_TJ/8, NRT=8/NCT, TI=16*NRT*TC_WM, SW=TC_KC/4;
__shared__ unsigned int sW[TI][TC_KC/4];
__shared__ unsigned int sH[TC_TJ][TC_KC/4];
const int tid=threadIdx.x, i0=blockIdx.x*TI, b0=blockIdx.z*TC_TJ;
const int warp=tid>>5, lane=tid&31;
const int rt=warp/NCT, ct=warp%NCT, gid=lane>>2, tig=lane&3;
float acc[TC_WM][4];
#pragma unroll
for(int m=0;m<TC_WM;m++){ acc[m][0]=acc[m][1]=acc[m][2]=acc[m][3]=0.f; }
for(int k0=0;k0<K;k0+=TC_KC){
for(int idx=tid; idx<TI*(TC_KC/16); idx+=256){
int r=idx/(TC_KC/16), c=idx%(TC_KC/16), o=i0+r;
int4 v = (o<N && k0+c*16<K) ? ((const int4*)(wq+(long)o*K+k0))[c] : make_int4(0,0,0,0);
*(int4*)&sW[r][BQTC_SWZC(r,c)*4] = v;
}
for(int idx=tid; idx<TC_TJ*(TC_KC/16); idx+=256){
int r=idx/(TC_KC/16), c=idx%(TC_KC/16);
int4 v = (r+b0<B && k0+c*16<K) ? ((const int4*)(Xq+(long)(r+b0)*K+k0))[c] : make_int4(0,0,0,0);
*(int4*)&sH[r][BQTC_SWZC(r,c)*4] = v;
}
__syncthreads();
#pragma unroll
for(int blk=0; blk<TC_KC/32; ++blk){
const int jb=ct*8+gid;
const unsigned b0h=sH[jb][BQTC_SWZJ(jb,blk*8+tig)], b1h=sH[jb][BQTC_SWZJ(jb,blk*8+tig+4)];
#pragma unroll
for(int m=0;m<TC_WM;m++){
const int rb=(rt*TC_WM+m)*16;
const unsigned a0=sW[rb+gid][BQTC_SWZJ(rb+gid,blk*8+tig)],
a1=sW[rb+gid+8][BQTC_SWZJ(rb+gid+8,blk*8+tig)],
a2=sW[rb+gid][BQTC_SWZJ(rb+gid,blk*8+tig+4)],
a3=sW[rb+gid+8][BQTC_SWZJ(rb+gid+8,blk*8+tig+4)];
bq_mma_e4m3(acc[m][0],acc[m][1],acc[m][2],acc[m][3], a0,a1,a2,a3, b0h,b1h);
}
}
__syncthreads();
}
#pragma unroll
for(int m=0;m<TC_WM;m++){
const int rb=(rt*TC_WM+m)*16, o0=i0+rb+gid, o1=o0+8;
const int j0=ct*8+tig*2+b0, j1=j0+1;
if(o0<N){ if(j0<B) Y[(long)j0*YS+o0]=acc[m][0]*wsc[o0]*xsc[j0];
if(j1<B) Y[(long)j1*YS+o0]=acc[m][1]*wsc[o0]*xsc[j1]; }
if(o1<N){ if(j0<B) Y[(long)j0*YS+o1]=acc[m][2]*wsc[o1]*xsc[j0];
if(j1<B) Y[(long)j1*YS+o1]=acc[m][3]*wsc[o1]*xsc[j1]; }
}
}
#endif
// ══ THE RESIDUAL NORM, SPLIT SO IT CAN FILL THE MACHINE ═════════════════════════════════════════
// `bq_add_rmsnorm_q8f_hl` takes ONE BLOCK PER SEQUENCE, because the sum of squares spans the model
// dimension. At two sequences that is two blocks of a 132-SM part, and it measured 4.06 ms of a
// 21.5 ms step — 127 launches at ~32 us each moving ~200 KB apiece. It is latency, not work.
//
// Two things were tried before this one and are recorded because the reasons transfer. Widening the
// block 256 -> 1024 recovered 38% and nothing more. Splitting the QUANTISER out into a grid-parallel
// kernel — which is how llama.cpp is arranged, its `quantize_row_q8_1` being separate from its
// row-serial `rms_norm_f32` — measured 0.77x, SLOWER: the activation is then read from global twice
// and the launch count doubles. Their arrangement is right for their kernel and wrong for ours,
// because ours already fuses the add, the norm, the f32 write and the two-word quantiser into one
// pass. Ours is not badly built; it is grid-starved.
//
// So this splits the REDUCTION instead. Phase one computes the residual add and a per-chunk partial
// across NORM_G*B blocks; phase two re-sums those partials — every block redundantly, in a fixed
// ascending order, so all blocks agree bit for bit and the result is deterministic run to run — then
// normalises and quantises its own chunk. Measured 1.95x-2.02x at one, two and four sequences.
//
// ★ It is NOT bit-identical to the one-block form: NORM_G chunk-sums instead of RMSN_NT thread-sums
// is a different association, worth ~1 ulp (worst relative 2.4e-07 measured). That is five orders of
// magnitude below the 1.5e-2 spread between llama.cpp's own `-fa 0` and `-fa 1`, so it cannot move
// the accuracy standing — but it is a different number, so it lives behind a knob and the exact lane
// stays the default.
__global__ void __launch_bounds__(RMSN_NT) bq_add_rmsnorm_part(
float* __restrict__ X, const float* __restrict__ Y, float* __restrict__ part, int n, int G){
const int bb=blockIdx.y, g=blockIdx.x;
const int chunk=((n+G-1)/G+31)/32*32;
const int lo=g*chunk; int hi=lo+chunk; if(hi>n) hi=n;
float* x=X+(long)bb*n; const float* y=Y+(long)bb*n;
__shared__ float sh[RMSN_NT];
float s=0.f;
for(int i=lo+threadIdx.x;i<hi;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(); }
if(threadIdx.x==0) part[(long)bb*G+g]=sh[0];
}
__global__ void __launch_bounds__(RMSN_NT) bq_add_rmsnorm_apply(
const float* __restrict__ X, const float* __restrict__ part, signed char* __restrict__ Xq,
signed char* __restrict__ Xl, float* __restrict__ Xs, float* __restrict__ Xf,
const float* __restrict__ w, int n, float eps, int G){
const int bb=blockIdx.y, g=blockIdx.x;
const int chunk=((n+G-1)/G+31)/32*32;
const int lo=g*chunk; int hi=lo+chunk; if(hi>n) hi=n;
const float* x=X+(long)bb*n; float* xf=Xf+(long)bb*n;
float tot=0.f;
for(int j=0;j<G;++j) tot+=part[(long)bb*G+j];
const float sc=rsqrtf(tot/n+eps);
for(int i=lo+threadIdx.x;i<hi;i+=blockDim.x) xf[i]=x[i]*sc*w[i];
__syncthreads();
const int nb=n/32;
for(int b=lo/32+threadIdx.x;b<hi/32;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[(long)bb*nb+b]=q;
for(int j=0;j<32;++j) lgx_q8hl(xf[b*32+j]/q, Xq+(long)bb*n, Xl+(long)bb*n, (long)b*32+j);
}
}
// ══ THE SAME matmul AGAIN, ROW-PARALLEL ═════════════════════════════════════════════════════════
// Below a handful of sequences the tensor-core tile is the wrong shape: `mma.m16n8k32` has an 8-wide
// N, so at two sequences six of every eight columns are zeros and the mma rate stops being an
// advantage. What matters there is streaming the weights, and the shipped `bq_gemv_hl` streams them
// badly for a reason that is structural rather than arithmetic: it gives one output row to ONE warp,
// so only 32 threads' worth of loads are ever in flight per row, and it re-reads the activation slab
// once per row.
//
// llama.cpp does not have that problem, and reading their kernel says why. ggml/src/ggml-cuda/mmvq.cu
// picks, for Q8_0 on an NVIDIA part, `nwarps` = 4 at ncols_dst <= 4 and 2 at <= 8, with
// `rows_per_cuda_block` = 2 (`calc_nwarps`/`calc_rows_per_block`), and routes every batch <= 8 there
// (`MMVQ_MAX_BATCH_SIZE`, checked before MMQ in ggml_cuda_mul_mat). So a whole BLOCK reduces one row,
// and two rows share each activation load. The structure below is that structure; the code is ours,
// and it keeps the two-pass hi/lo accuracy tier and the K-slicing that mmvq does not have.
//
// MV_NW=1 with MV_RB=2 is BIT-IDENTICAL to `bq_gemv_hl` — same ascending K order, same grouping —
// and measured 2.5x-2.9x faster at eight sequences purely by halving activation traffic.
__global__ void __launch_bounds__(MV_NW*32) bq_gemv_wide(
float* __restrict__ 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 B, int YS,
float* __restrict__ P, int SPLIT){
static_assert(MV_TJ>=1 && MV_RB>=1 && MV_NW>=1, "the row-parallel GEMV needs a positive tile");
// A tile narrower than the batch would compute only MV_TJ of the B sequences and leave the rest
// holding whatever Y held before — plausible numbers, wrong answer. Refuse instead.
if(MV_TJ < B) return;
const int tid=threadIdx.x, lane=tid&31, wp=tid>>5;
const int row0=blockIdx.x*MV_RB, KB=K/32, nq=K/512;
const int per=(nq+SPLIT-1)/SPLIT, c0=blockIdx.y*per;
const int c1=(c0+per)<nq?(c0+per):nq;
float acc[MV_RB][MV_TJ];
#pragma unroll
for(int r=0;r<MV_RB;r++){
#pragma unroll
for(int j=0;j<MV_TJ;j++) acc[r][j]=0.f; }
for(int c=c0+wp; c<c1; c+=MV_NW){
const int blk=c*16+lane/2;
int4 hh[MV_TJ], ll[MV_TJ]; float xsv[MV_TJ];
#pragma unroll
for(int j=0;j<MV_TJ;j++){
if(j<B){ hh[j]=((const int4*)(Xh+(long)j*K))[c*32+lane];
ll[j]=((const int4*)(Xl+(long)j*K))[c*32+lane];
xsv[j]=Xs[(long)j*KB+blk]; }
else { hh[j]=make_int4(0,0,0,0); ll[j]=make_int4(0,0,0,0); xsv[j]=0.f; } }
#pragma unroll
for(int r=0;r<MV_RB;r++){
const int o=row0+r; if(o>=N) continue;
const int4 w=((const int4*)(wq+(long)o*K))[c*32+lane];
const float wsb=__half2float(ws[(long)o*KB+blk]);
#pragma unroll
for(int j=0;j<MV_TJ;j++){
if(j>=B) continue;
int dh=__dp4a(w.x,hh[j].x,0)+__dp4a(w.y,hh[j].y,0)+__dp4a(w.z,hh[j].z,0)+__dp4a(w.w,hh[j].w,0);
int dl=__dp4a(w.x,ll[j].x,0)+__dp4a(w.y,ll[j].y,0)+__dp4a(w.z,ll[j].z,0)+__dp4a(w.w,ll[j].w,0);
acc[r][j] += ((float)dh + (float)dl*(1.0f/127.0f)) * (wsb*xsv[j]);
} }
}
// The K remainder past the last full 512-wide quad, on the slice that owns the end of the sum.
if(blockIdx.y==SPLIT-1){
const int done=nq*512;
for(int r=0;r<MV_RB;r++){ const int o=row0+r; if(o>=N) continue;
for(int j=0;j<MV_TJ;j++){ if(j>=B) continue;
for(int q=done+tid; q<K; q+=MV_NW*32){ const int blk=q/32;
float a=(float)Xh[(long)j*K+q] + (float)Xl[(long)j*K+q]*(1.0f/127.0f);
acc[r][j] += a*(float)wq[(long)o*K+q]
*(__half2float(ws[(long)o*KB+blk])*Xs[(long)j*KB+blk]); } } }
}
__shared__ float red[MV_NW][MV_RB][MV_TJ];
#pragma unroll
for(int r=0;r<MV_RB;r++){
#pragma unroll
for(int j=0;j<MV_TJ;j++){
float v=acc[r][j];
for(int t=16;t;t>>=1) v+=__shfl_down_sync(0xffffffffu,v,t);
if(lane==0) red[wp][r][j]=v; } }
__syncthreads();
if(tid < MV_RB*MV_TJ){
const int r=tid/MV_TJ, j=tid%MV_TJ, o=row0+r;
if(o<N && j<B){
float v=0.f;
#pragma unroll
for(int w=0;w<MV_NW;w++) v+=red[w][r][j];
if(SPLIT==1) Y[(long)j*YS+o]=v;
else P[(long)blockIdx.y*N*B + (long)j*N + o]=v;
} }
}
// Fixed ascending order over the K slices, so the split reduction is deterministic run to run.
__global__ void bq_sk_reduce(float* __restrict__ Y, const float* __restrict__ P,
int N, int B, int YS, int S){
long i=(long)blockIdx.x*blockDim.x+threadIdx.x; if(i>=(long)N*B) return;
int j=(int)(i/N), o=(int)(i%N); float v=0.f;
for(int s=0;s<S;s++) v+=P[(long)s*N*B+(long)j*N+o];
Y[(long)j*YS+o]=v;
}
"#;Expand description
The batched device code. Every kernel is the single-sequence form with a sequence index added; the GEMV is the one that changes shape, because it is the one that decides the token time.