pub const PRELUDE: &str = r#"
#include <cstddef>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
// An element as the storage type `T` holds it — `f32` as computed, binary16 and bfloat16 rounded to the nearest, ties
// to even — and as the `f32` every kernel computes in, widened exactly.
template<typename T> __device__ __forceinline__ T lk_to(float v){ return v; }
template<> __device__ __forceinline__ __half lk_to<__half>(float v){ return __float2half_rn(v); }
template<> __device__ __forceinline__ __nv_bfloat16 lk_to<__nv_bfloat16>(float v){ return __float2bfloat16_rn(v); }
__device__ __forceinline__ float lk_from(float v){ return v; }
__device__ __forceinline__ float lk_from(__half v){ return __half2float(v); }
__device__ __forceinline__ float lk_from(__nv_bfloat16 v){ return __bfloat162float(v); }
template<typename T> __device__ __forceinline__ float lk_kept(float v){ return lk_from(lk_to<T>(v)); }
// Two and four consecutive elements from `p` as `f32`: one load of their bytes, eight-byte aligned for four binary16
// values as for two `f32`.
__device__ __forceinline__ float2 lk_from2(const float* p){ return *(const float2*)p; }
__device__ __forceinline__ float2 lk_from2(const __half* p){ return __half22float2(*(const __half2*)p); }
__device__ __forceinline__ float4 lk_from4(const float* p){ return *(const float4*)p; }
__device__ __forceinline__ float4 lk_from4(const __half* p){
uint2 u = *(const uint2*)p;
float2 a = __half22float2(*(const __half2*)&u.x), b = __half22float2(*(const __half2*)&u.y);
return make_float4(a.x, a.y, b.x, b.y);
}
// One element of a result's array as it is written: `f32`, or binary16 rounded to the nearest.
struct lk_ref {
void* p; bool h;
__device__ __forceinline__ void operator=(float v) const { if(h) *(__half*)p = __float2half_rn(v); else *(float*)p = v; }
};
// A result's array — `f32`, or binary16 that every value written is rounded to — taken from a typed pointer, so a kernel
// writes either through one pointer type. It can be offset and written, never read.
struct lk_out {
void* p = nullptr; bool h = false;
lk_out() = default;
__host__ __device__ lk_out(std::nullptr_t) {}
__host__ __device__ lk_out(float* q) : p(q) {}
__host__ __device__ lk_out(__half* q) : p(q), h(true) {}
__host__ __device__ explicit operator bool() const { return p != nullptr; }
__device__ __forceinline__ void* at(long i) const { return h ? (void*)((__half*)p + i) : (void*)((float*)p + i); }
__device__ __forceinline__ lk_out operator+(long i) const { lk_out o; o.p = at(i); o.h = h; return o; }
__device__ __forceinline__ lk_ref operator[](long i) const { return lk_ref{at(i), h}; }
};
// A result written into the row a 1-based position names, or the buffer itself when no position is given.
#define LK_ROW(p, row, pitch) ((p) + ((row) ? (long)(*(row) - 1) * (long)(pitch) : 0L))
__device__ __forceinline__ float lk_warp_sum(float v){
for(int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
return __shfl_sync(0xffffffffu, v, 0);
}
__device__ __forceinline__ double lk_warp_sum_d(double v){
for(int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
return __shfl_sync(0xffffffffu, v, 0);
}
__device__ __forceinline__ float lk_warp_max(float v){
for(int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, o));
return __shfl_sync(0xffffffffu, v, 0);
}
__device__ __forceinline__ double lk_warp_sum(double v){ return lk_warp_sum_d(v); }
// The `V` values each lane holds summed over the warp's lanes, `V` a power of two to 32: lane `l` returns the total of
// value `l % V`. At a distance `W` below `V` each lane keeps one half of its values and sends the other, the halves
// summed apart; at and past `V` the whole set is summed. Starting at a distance `W` below 16 sums over each run of
// `2·W` lanes apart.
template<int V, int W = 16>
__device__ __forceinline__ float lk_lane_sums(float (&v)[V]){
if constexpr(W >= V){
#pragma unroll
for(int i = 0; i < V; ++i) v[i] += __shfl_xor_sync(0xffffffffu, v[i], W);
} else {
bool up = threadIdx.x & W;
#pragma unroll
for(int i = 0; i < W; ++i){
float send = up ? v[i] : v[i + W], keep = up ? v[i + W] : v[i];
v[i] = keep + __shfl_xor_sync(0xffffffffu, send, W);
}
}
if constexpr(W > 1) return lk_lane_sums<V, W / 2>(v);
else return v[0];
}
// The precise numerics (Q 0) sum in f64, the quantized in f32.
template<int Q> struct lk_acc { typedef float t; };
template<> struct lk_acc<0> { typedef double t; };
// The products that read their activation `f32`: every one under the precise numerics (Q 0), and under Q 3 — the
// latent's expansion halves under the quantized numerics, which sum in f32.
template<int Q> __host__ __device__ constexpr bool lk_f32_acts(){ return Q == 0 || Q == 3; }
// Every thread of the block gets the total over `vw` warps, in the accumulator `A` (`float` or `double`); `red` is 32 of
// them in shared memory. Warp `w` of the block's `nw` plays the warps `w`, `w + nw`, … below `vw`, value `j` of each
// lane its lane's of the `j`-th: each warp summed over its lanes, then the warps' sums in order — the tree a block of
// `vw` warps sums its values by.
template<int VW, typename A>
__device__ __forceinline__ A lk_warps_sum(const A (&v)[VW], int vw, A* red){
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = blockDim.x >> 5;
A s[VW];
#pragma unroll
for(int j = 0; j < VW; ++j) s[j] = lk_warp_sum(v[j]);
__syncthreads();
#pragma unroll
for(int j = 0; j < VW; ++j) if(lane == 0 && warp + j * nw < vw) red[warp + j * nw] = s[j];
__syncthreads();
return lk_warp_sum(lane < vw ? red[lane] : (A)0);
}
// Every thread of the block gets the block's total.
template<typename A>
__device__ __forceinline__ A lk_block_sum(A v, A* red){
const A one[1] = {v};
return lk_warps_sum<1>(one, (blockDim.x + 31) >> 5, red);
}
__device__ __forceinline__ float lk_block_max(float v, float* red){
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
v = lk_warp_max(v);
__syncthreads();
if(lane == 0) red[warp] = v;
__syncthreads();
return lk_warp_max(lane < nw ? red[lane] : -INFINITY);
}
// An activation: ACT 0 silu `x / (1 + e^-x)`, 1 sigmoid `1 / (1 + e^-x)`, 2 gelu in ggml's tanh form
// `x / (1 + e^-2z)`, `z = √(2/π)·(x + 0.044715·x³)`.
template<int ACT>
__device__ __forceinline__ float lk_act(float v){
if(ACT == 0) return v / (1.0f + expf(-v));
if(ACT == 1) return 1.0f / (1.0f + expf(-v));
float z = 0.7978845608028654f * (v + 0.044715f * v * v * v);
return v / (1.0f + expf(-2.0f * z));
}
// The two-word scale of values whose largest magnitude is `a`, and value `v`'s two words under scale `s`, by quotients:
// `v ≈ s·(h + l/128)` (the two-word activation quantizer's); by products with the scale's reciprocal `r` where a kernel
// takes many values under one scale; four words packed into one, the first lowest.
__device__ __forceinline__ float lk_two_scale(float a){ return (a > 1e-30f ? a : 1e-30f) / 127.0f; }
__device__ __forceinline__ void lk_two_words(float v, float s, int& h, int& l){
h = __float2int_rn(v / s); h = h > 127 ? 127 : (h < -128 ? -128 : h);
l = __float2int_rn(__fmaf_rn(-(float)h, s, v) / (s * 0.0078125f)); l = l > 127 ? 127 : (l < -128 ? -128 : l);
}
__device__ __forceinline__ void lk_two_words_r(float v, float s, float r, int& h, int& l){
h = __float2int_rn(v * r); h = h > 127 ? 127 : (h < -128 ? -128 : h);
l = __float2int_rn(__fmaf_rn(-(float)h, s, v) * r * 128.f); l = l > 127 ? 127 : (l < -128 ? -128 : l);
}
__device__ __forceinline__ int lk_int8x4(int a, int b, int c, int d){
return (int)((unsigned)(a & 255) | ((unsigned)(b & 255) << 8) | ((unsigned)(c & 255) << 16) | ((unsigned)(d & 255) << 24));
}
// Quantize the 32-value group this warp holds, one value per lane, as group `g`.
template<int Q>
__device__ __forceinline__ void lk_quant_lane(float v, signed char* __restrict__ qh, signed char* __restrict__ ql,
float* __restrict__ qs, long i, long g){
float a = lk_warp_max(fabsf(v));
if(Q == 1){
float s = a > 0.f ? a / 127.0f : 1.0f;
if((threadIdx.x & 31) == 0) qs[g] = s;
qh[i] = (signed char)__float2int_rn(v / s);
} else {
float s = lk_two_scale(a);
int h, l;
lk_two_words(v, s, h, l);
if((threadIdx.x & 31) == 0) qs[g] = s;
qh[i] = (signed char)h; ql[i] = (signed char)l;
}
}
// Quantize a pair of values `i`, `i + 1` of a 32-value group sixteen lanes of a warp hold, a pair a lane, as group
// `i >> 5` — the group's largest magnitude taken over the sixteen lanes on every lane (a lane `on` false holds none of
// it), and each value and the scale then what `lk_quant_lane` makes of the group.
template<int Q>
__device__ __forceinline__ void lk_quant_pair(float x, float y, bool on, signed char* __restrict__ qh, signed char* __restrict__ ql,
float* __restrict__ qs, long i){
float a = on ? fmaxf(fabsf(x), fabsf(y)) : 0.f;
#pragma unroll
for(int o = 8; o > 0; o >>= 1) a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, o));
if(!on) return;
if(Q == 1){
float s = a > 0.f ? a / 127.0f : 1.0f;
if((threadIdx.x & 15) == 0) qs[i >> 5] = s;
qh[i] = (signed char)__float2int_rn(x / s);
qh[i + 1] = (signed char)__float2int_rn(y / s);
} else {
float s = lk_two_scale(a);
int h, l;
lk_two_words(x, s, h, l);
qh[i] = (signed char)h; ql[i] = (signed char)l;
lk_two_words(y, s, h, l);
qh[i + 1] = (signed char)h; ql[i + 1] = (signed char)l;
if((threadIdx.x & 15) == 0) qs[i >> 5] = s;
}
}
// Where an attention quantizes its output for the projection that reads it: its words, low words and scales.
struct lk_qout { signed char* q; signed char* l; float* s; const float* g; int ga; };
// Values `i`, `i + 1` of an attention's output: gated by `lk_act` (`ga`) of the gate's own values `g[i]`, `g[i + 1]` where
// `g` is given, as a product with the activation computes it; written `f32` where `out` is given; quantized under `QO`
// into `qo` (`lk_quant_pair`) — every lane of the warp calls it, `on` the lanes that hold a pair.
template<int QO>
__device__ __forceinline__ void lk_attn_store(float* __restrict__ out, lk_qout qo, long i, bool on, float x, float y){
if(on && qo.g){
const float gx = qo.g[i], gy = qo.g[i + 1];
x *= qo.ga == 0 ? lk_act<0>(gx) : qo.ga == 1 ? lk_act<1>(gx) : lk_act<2>(gx);
y *= qo.ga == 0 ? lk_act<0>(gy) : qo.ga == 1 ? lk_act<1>(gy) : lk_act<2>(gy);
}
if(on && out) *(float2*)(out + i) = make_float2(x, y);
if(QO > 0) lk_quant_pair<QO>(x, y, on, qo.q, qo.l, qo.s, i);
}
// One step of a Q8_0 dot's partial sum, as its two factors: the partial `p` becomes `fmaf(a, b, p)` for the int8 dot
// `d` (and, in two words, the low words' `e`) of a block of weight scale `w` and activation scale `s`, each operation
// spelled out, so every kernel that sums as the GEMV sums takes the same steps (`lk_q8_acc` takes one).
template<int Q>
__device__ __forceinline__ float2 lk_q8_term(float w, float s, float d, float e){
if(Q == 1) return make_float2(__fmul_rn(w, s), d);
return make_float2(w, __fmaf_rn(s, d, __fmul_rn(__fmul_rn(s, 0.0078125f), e)));
}
template<int Q>
__device__ __forceinline__ float lk_q8_acc(float acc, float w, float s, float d, float e){
const float2 t = lk_q8_term<Q>(w, s, d, e);
return __fmaf_rn(t.x, t.y, acc);
}
// The four partials of a quantized Q8_0 dot summed in the one order every kernel that sums as the GEMV sums takes.
__device__ __forceinline__ float lk_q8_sum4(float p0, float p1, float p2, float p3){ return (p0 + p2) + (p1 + p3); }
// The int8 dot of one 32-value block two lanes hold sixteen values each of (`w` and `x` their words), exact, in both.
__device__ __forceinline__ int lk_q8_block_dot(int4 w, int4 x){
int d = __dp4a(w.x, x.x, 0) + __dp4a(w.y, x.y, 0) + __dp4a(w.z, x.z, 0) + __dp4a(w.w, x.w, 0);
return d + __shfl_xor_sync(0xffffffffu, d, 1);
}
// The scratch, in `float2`, a warp of a quantized Q8_0 GEMV takes for `U` dots at once over more than one activation
// row: a round's steps of up to eight dots at a time, a dot's sixteen 20 apart, so a half-warp's chain reads fall on
// distinct banks.
template<int U> __host__ __device__ constexpr int lk_q8_scratch(){ return (U < 8 ? U : 8) * 20; }
// Round `c` of rows `o … o + W − 1` dotted with one activation row (`lk_q8_rows_dot`): its first `n` blocks' steps
// into each row's partials in lanes `2q`, chained in order from pairs `q`, `q + 4`, `q + 8`, `q + 12` by `shfl_down`;
// a whole round's sixteen (`WHOLE`) need no test of the blocks past the row's end.
template<int Q, int W, bool WHOLE>
__device__ __forceinline__ void lk_q8_round(float (&acc)[W], const signed char* __restrict__ wq, const __half* __restrict__ ws, long o, int K,
const signed char* __restrict__ xq, const signed char* __restrict__ xl, const float* __restrict__ xs,
int c, int n){
const int lane = threadIdx.x & 31, j = lane >> 1, nb = K >> 5, blk = c * 16 + j;
const bool in = WHOLE || j < n;
const int4 zero = make_int4(0, 0, 0, 0);
const int4 h = in ? ((const int4*)xq)[c * 32 + lane] : zero;
const int4 l = Q == 2 && in ? ((const int4*)xl)[c * 32 + lane] : zero;
const float s = in ? xs[blk] : 0.f;
#pragma unroll
for(int r = 0; r < W; ++r){
const int4 w = in ? ((const int4*)(wq + (o + r) * K))[c * 32 + lane] : zero;
const float wsc = in ? __half2float(ws[(o + r) * nb + blk]) : 0.f;
const int d = lk_q8_block_dot(w, h), e = Q == 2 ? lk_q8_block_dot(w, l) : 0;
const float2 t = lk_q8_term<Q>(wsc, s, (float)d, (float)e);
if(in) acc[r] = __fmaf_rn(t.x, t.y, acc[r]);
#pragma unroll
for(int k = 1; k < 4; ++k){
const float a = __shfl_down_sync(0xffffffffu, t.x, 8 * k), b = __shfl_down_sync(0xffffffffu, t.y, 8 * k);
if(WHOLE || j + 4 * k < n) acc[r] = __fmaf_rn(a, b, acc[r]);
}
}
}
// Rows `o … o + W − 1` of a Q8_0 matrix (int8 values `wq`, a half scale per 32 in `ws`) dotted with the `M` activation
// rows of `K` values — one warp of a block of at most `WARPS` — row `r`'s dots at `out + r·M`, in every lane. Under the
// precise numerics (Q 0) the activations are `f32` and the sums `f64`, four values a lane at a time. Quantized (int8
// words, two under Q 2), each 32-value block's dot is exact in int32 and its step (`lk_q8_term`) goes into partial
// `b mod 4` of block `b`, a partial's blocks in order; the four partials sum as `lk_q8_sum4` sums them. The row goes
// 512 values (16 blocks) a round, sixteen bytes a lane: lanes `2j`, `2j + 1` block `j` of each (`lk_q8_block_dot`),
// so a round's steps of partial `q` are pairs `q`, `q + 4`, `q + 8` and `q + 12`'s. With one activation row each row's
// partials chain in lanes `2q`, reading the steps by `shfl_down`; with more, the round's steps go through the warp's
// scratch (`lk_q8_scratch`), up to eight dots' a time, and lane `4u + q` chains the time's dot `u`'s partial `q`. The
// blocks past the row's end are none of its partials'.
template<int Q, int M, int W, int WARPS = 8>
__device__ __forceinline__ void lk_q8_rows_dot(const signed char* __restrict__ wq, const __half* __restrict__ ws, long o, int K,
const float* __restrict__ x, const signed char* __restrict__ xq,
const signed char* __restrict__ xl, const float* __restrict__ xs, float* out){
const int lane = threadIdx.x & 31, nb = K >> 5;
if constexpr(Q == 0){
double acc[W][M];
#pragma unroll
for(int r = 0; r < W; ++r)
#pragma unroll
for(int t = 0; t < M; ++t) acc[r][t] = 0.0;
for(int g = lane; g < (K >> 2); g += 32){
char4 w[W];
double sc[W];
#pragma unroll
for(int r = 0; r < W; ++r){ w[r] = ((const char4*)(wq + (o + r) * K))[g]; sc[r] = (double)__half2float(ws[(o + r) * nb + (g >> 3)]); }
#pragma unroll
for(int t = 0; t < M; ++t){
const float* xt = x + (long)t * K + 4 * g;
float x0 = xt[0], x1 = xt[1], x2 = xt[2], x3 = xt[3];
#pragma unroll
for(int r = 0; r < W; ++r) acc[r][t] += sc[r] * ((double)w[r].x * x0 + (double)w[r].y * x1 + (double)w[r].z * x2 + (double)w[r].w * x3);
}
}
#pragma unroll
for(int r = 0; r < W; ++r)
#pragma unroll
for(int t = 0; t < M; ++t) out[r * M + t] = (float)lk_warp_sum_d(acc[r][t]);
} else if constexpr(M == 1){
float acc[W];
#pragma unroll
for(int r = 0; r < W; ++r) acc[r] = 0.f;
const int full = nb >> 4, rest = nb & 15;
for(int c = 0; c < full; ++c) lk_q8_round<Q, W, true>(acc, wq, ws, o, K, xq, xl, xs, c, 16);
if(rest) lk_q8_round<Q, W, false>(acc, wq, ws, o, K, xq, xl, xs, full, rest);
#pragma unroll
for(int r = 0; r < W; ++r){
float v = acc[r];
v += __shfl_down_sync(0xffffffffu, v, 4);
v += __shfl_down_sync(0xffffffffu, v, 2);
out[r] = __shfl_sync(0xffffffffu, v, 0);
}
} else {
constexpr int U = W * M, G = U < 8 ? U : 8, R = (U + G - 1) / G;
__shared__ float2 lk_q8s[WARPS][lk_q8_scratch<U>()];
const int warp = threadIdx.x >> 5, j = lane >> 1, u0 = lane >> 2, q = lane & 3;
if(warp >= WARPS) __trap();
float2* sm = lk_q8s[warp];
float acc[R];
#pragma unroll
for(int i = 0; i < R; ++i) acc[i] = 0.f;
const int4 zero = make_int4(0, 0, 0, 0);
for(int c = 0; c * 16 < nb; ++c){
const int blk = c * 16 + j, n = min(16, nb - c * 16);
const bool in = j < n;
int4 w[W];
float wsc[W];
#pragma unroll
for(int r = 0; r < W; ++r){
w[r] = in ? ((const int4*)(wq + (o + r) * K))[c * 32 + lane] : zero;
wsc[r] = in ? __half2float(ws[(o + r) * nb + blk]) : 0.f;
}
float2 st[U];
#pragma unroll
for(int t = 0; t < M; ++t){
const int4 h = in ? ((const int4*)(xq + (long)t * K))[c * 32 + lane] : zero;
const int4 l = Q == 2 && in ? ((const int4*)(xl + (long)t * K))[c * 32 + lane] : zero;
const float s = in ? xs[(long)t * nb + blk] : 0.f;
#pragma unroll
for(int r = 0; r < W; ++r){
const int d = lk_q8_block_dot(w[r], h), e = Q == 2 ? lk_q8_block_dot(w[r], l) : 0;
st[r * M + t] = lk_q8_term<Q>(wsc[r], s, (float)d, (float)e);
}
}
// Up to eight dots' steps a time: pair `j` writes its block's (the even lane an even dot's, the odd an odd's),
// then lane `4u + q` chains dot `u`'s partial `q` over blocks `q`, `q + 4`, `q + 8`, `q + 12`.
#pragma unroll
for(int i = 0; i < R; ++i){
#pragma unroll
for(int g = 0; g < G; ++g) if(i * G + g < U && (lane & 1) == (g & 1)) sm[g * 20 + j] = st[i * G + g];
__syncwarp();
if(u0 < G && i * G + u0 < U){
#pragma unroll
for(int k = 0; k < 4; ++k) if(q + 4 * k < n){ const float2 v = sm[u0 * 20 + q + 4 * k]; acc[i] = __fmaf_rn(v.x, v.y, acc[i]); }
}
__syncwarp();
}
}
#pragma unroll
for(int i = 0; i < R; ++i){
float v = acc[i];
v += __shfl_down_sync(0xffffffffu, v, 2);
v += __shfl_down_sync(0xffffffffu, v, 1);
#pragma unroll
for(int g = 0; g < G; ++g) if(i * G + g < U) out[i * G + g] = __shfl_sync(0xffffffffu, v, 4 * g);
}
}
}
"#;Expand description
The helpers every kernel shares: the storage types, row addressing, warp and block reductions, the activation quantizer.