// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections.
// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's
// native path), in a Metal layout (no paging: unified memory):
//
// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur;
// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp);
// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last
// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion
// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell);
// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key).
// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b],
// the tail block n_bid scores `dead` and gets +1e9 when it has cells.
// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest
// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select).
// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l,
// acc = p . v -> partials
// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep
// planes) for attn_output.
//
// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table).
#define QH 24
#define QKV 2
#define QD 256
#define QROT 64
#define IH 4
#define ID 128
#define QBUDGET 2051
#define QCH 64
constant int QSA_NTOK [[function_constant(3)]];
// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE,
// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row
// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own -
// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE.
constant bool QSA_MROPE [[function_constant(24)]];
constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE;
// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16.
constant bool QSA_Q_F16_FC [[function_constant(25)]];
constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false;
// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose
// queries nobody reads). Undefined: 0.
constant int QSA_ROW0 [[function_constant(28)]];
constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0;
struct QsaArgs {
uint pos0; // the window's first cell (relative to the sequence)
uint pos_base; // absolute position of cell 0 (RoPE)
uint max_cells;
uint n_tok;
float eps;
};
// cache layout per layer
// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half
// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128]
static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) {
v = simd_sum(v);
if (lane == 0) part[sg] = v;
threadgroup_barrier(mem_flags::mem_threadgroup);
float t = 0.0f;
for (int i = 0; i < nsg; ++i) t += part[i];
return t;
}
// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32)
static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row,
threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps,
device const float2* cs, uint absp, bool rope, device const int* mtab,
bool reuse = false) {
const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8);
const float rs = rsqrt(ss / float(w) + eps);
const float y = tid < w ? (rs * x) * g[tid] : 0.0f;
float o = y;
if (rope) {
if (tid < QROT) row[tid] = y;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid < QROT) {
const int p = tid & 31;
const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp;
const float2 c = cs[pp * 32 + p];
const float a = row[p], b = row[p + 32];
o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x;
}
}
if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup);
(void)xhi_unused;
return o;
}
// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE]
// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own
// [n_tok][rows] buffers.
constant int QSA_KV_STRIDE [[function_constant(26)]];
constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0;
// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads;
// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 =
// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored.
kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]],
device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]],
device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]],
device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]],
device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]],
device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]],
device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]],
device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]],
device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]],
device const int* mtab [[buffer(20)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup float part[8];
threadgroup float row[QROT];
threadgroup float amax[8];
const int r = int(tg.x) + qsa_row0;
// the row strides of kcur / vcur, ikraw, iqraw
const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID;
const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID;
if (r == 32) {
// the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail;
// every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel.
const int t_self = int(tg.y);
if (t_self >= 3) {
const uint cell = a.pos0 + uint(t_self);
if ((cell & 3) == 3) {
float m = 0.0f;
if (tid < ID) {
m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) +
ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f;
}
const uint b = cell >> 2;
const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab);
if (tid < ID) pooled[b * ID + tid] = o;
}
return;
}
if (t_self != 0) return;
const int n_pre = min(int(a.n_tok), 3);
for (int t = 0; t < n_pre; ++t) {
const uint cell = a.pos0 + uint(t);
const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f;
if (tid < ID) tail[(cell & 3) * ID + tid] = x;
threadgroup_barrier(mem_flags::mem_device);
if (cell == 0) {
const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true);
if (tid < ID) dead[tid] = o;
}
if ((cell & 3) == 3) {
float m = 0.0f;
if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f;
const uint b = cell >> 2;
const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true);
if (tid < ID) pooled[b * ID + tid] = o;
}
threadgroup_barrier(mem_flags::mem_device);
}
const int t_tail0 = max(n_pre, int(a.n_tok) - 4);
for (int t = t_tail0; t < int(a.n_tok); ++t) {
const uint cell = a.pos0 + uint(t);
if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid];
}
return;
}
const int t = int(tg.y);
const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell;
if (r < QH) {
const float x = qfull[(t * QH + r) * 2 * QD + tid];
const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab);
if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o);
else qcur[(t * QH + r) * QD + tid] = o;
} else if (r < QH + 2 * QKV) {
const bool isv = r >= QH + QKV;
const int h = r - QH - (isv ? QKV : 0);
float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid];
if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab);
// int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads)
float m = simd_max(fabs(x));
if (lane == 0) amax[sg] = m;
threadgroup_barrier(mem_flags::mem_threadgroup);
const int g = tid >> 6;
m = max(amax[2 * g], amax[2 * g + 1]);
const half sh = half(m / 127.0f);
const float sf = float(sh);
int q = 0;
if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127);
const uint rowi = cell * QKV + uint(h);
(isv ? vq : kq)[rowi * QD + tid] = char(q);
if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh;
} else {
if (a.pos0 + a.n_tok <= uint(QBUDGET)) return;
const int h = r - QH - 2 * QKV;
const float x = tid < ID ? iqraw[t * s_iq + h * ID + tid] : 0.0f;
const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab);
if (tid < ID) qidx[(t * IH + h) * ID + tid] = o;
}
}
// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block.
kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]],
device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]],
constant QsaArgs& a [[buffer(4)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
const int t = int(tg.y);
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2;
const uint b = tg.x * 8 + sg;
if (b > n_bid) return;
device const float* key = b == n_bid ? dead : pooled + b * ID;
device const float* qp = qidx + (t * IH) * ID + lane;
float k[4], q[IH][4];
#pragma clang loop unroll(full)
for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i];
#pragma clang loop unroll(full)
for (int h = 0; h < IH; ++h) {
#pragma clang loop unroll(full)
for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i];
}
float s = 0.0f;
#pragma clang loop unroll(full)
for (int h = 0; h < IH; ++h) {
float d = 0.0f;
#pragma clang loop unroll(full)
for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d);
d = simd_sum(d);
s += max(d, 0.0f);
}
if (b == n_bid && (n_kv & 3) != 0) s += 1e9f;
if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s;
}
// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score =
// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the
// threshold key T and how many cells with key == T to take (the lowest cells first).
static inline uint okey(float f) {
const uint u = as_type<uint>(f);
return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
}
kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]],
constant QsaArgs& a [[buffer(2)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup atomic_uint hist[256];
threadgroup uint sel[4];
threadgroup uint2 pref[32];
threadgroup uint2 totals;
const int t = int(tg.x);
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1;
device int* out = ids + t * QBUDGET;
device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2);
if (n_kv <= QBUDGET) {
for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c);
return;
}
const bool cached = nb <= 4096;
uint keys[4];
if (cached)
for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]);
// radix select of the `need`-th largest cell key
uint prefix = 0, pmask = 0, need = QBUDGET;
for (int pass = 0; pass < 4; ++pass) {
const int shift = 24 - 8 * pass;
if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed);
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint b = tid; b < nb; b += 1024) {
const uint k = cached ? keys[b / 1024] : okey(sc[b]);
if ((k & pmask) == prefix) {
const uint w = b < n_bid ? 4u : (n_kv & 3u);
if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0) {
uint counts[8], sum = 0;
for (uint j = 0; j < 8; ++j) {
counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed);
sum += counts[j];
}
uint acc = simd_prefix_exclusive_sum(sum);
if (acc < need && acc + sum >= need) {
for (uint j = 0; j < 8; ++j) {
if (acc + counts[j] >= need) {
sel[0] = 255 - (lane * 8 + j);
sel[1] = need - acc;
break;
}
acc += counts[j];
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
prefix |= sel[0] << shift;
pmask |= 255u << shift;
need = sel[1];
}
// T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending.
// Pass over blocks in order of index with a prefix scan of each block's selected cell count.
const uint T = prefix;
uint base = 0, eq_taken = 0;
for (uint b0 = 0; b0 < nb; b0 += 1024) {
const uint b = b0 + tid;
uint w = 0, weq = 0;
if (b < nb) {
const uint k = cached ? keys[b / 1024] : okey(sc[b]);
const uint cells = b < n_bid ? 4u : (n_kv & 3u);
if (k > T) w = cells;
else if (k == T) weq = cells;
}
const uint eq_before = simd_prefix_exclusive_sum(weq);
const uint gt_before = simd_prefix_exclusive_sum(w);
if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0) {
const uint2 count = pref[lane];
const uint2 before = uint2(simd_prefix_exclusive_sum(count.x),
simd_prefix_exclusive_sum(count.y));
pref[lane] = before;
if (lane == 31) totals = before + count;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint2 group_before = pref[sg], count = totals;
simdgroup_barrier(mem_flags::mem_threadgroup);
const uint eq0 = eq_taken + group_before.x + eq_before;
const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u;
const uint mine = w + take_eq;
// Clipping the equal-key prefix preserves the lowest-cell tie break.
const uint remaining = need - min(need, eq_taken);
const uint off = base + group_before.y + gt_before +
min(remaining, group_before.x + eq_before);
for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c);
base += count.y + min(remaining, count.x);
eq_taken += count.x;
}
}
// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads.
kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]],
device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]],
device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]],
device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
constexpr int HPG = QH / QKV; // 12 query heads per kv head
threadgroup float qs[HPG][QD];
threadgroup float s[HPG][QCH];
threadgroup int cid[QCH];
const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z);
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
const int c0 = chunk * QCH;
const int nchunks = (QBUDGET + QCH - 1) / QCH;
device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2);
if (uint(c0) >= width) {
// an empty chunk: m = -inf, l = 0 (the merge skips it)
for (int h = 0; h < HPG; ++h)
if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; }
return;
}
const int nc = min(QCH, int(width) - c0);
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid];
if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0;
threadgroup_barrier(mem_flags::mem_threadgroup);
// scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each)
threadgroup const float* q_lane = &qs[0][lane * 8];
for (int c = sg; c < nc; c += 8) {
const uint rowi = uint(cid[c]) * QKV + uint(g);
device const char* kr = kq + rowi * QD + lane * 8;
const char4 k0 = *(device const char4*)(kr);
const char4 k1 = *(device const char4*)(kr + 4);
const float scl = float(ksc[rowi * 4 + (lane >> 3)]);
const float kv[8] = {
float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl,
float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl
};
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) {
threadgroup const float* qh = q_lane + h * QD;
float d = 0.0f;
#pragma clang loop unroll(full)
for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d);
d = simd_sum(d);
if (lane == 0) s[h][c] = d * 0.0625f;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8)
threadgroup float mh[HPG], lh[HPG];
// nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop)
const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc;
for (int h = sg; h < HPG; h += 8) {
const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY;
const float m = simd_max(max(s0, s1));
const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f;
if (in0) s[h][lane] = p0;
if (in1) s[h][lane + 32] = p1;
const float l = simd_sum(p0 + p1);
if (lane == 0) { mh[h] = m; lh[h] = l; }
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j
float acc[HPG];
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) acc[h] = 0.0f;
// four cells' V loads in flight at once; the fma order per head stays c, c + 1, ...
auto vcell = [&](int c) {
const uint rowi = uint(cid[c]) * QKV + uint(g);
return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]);
};
int c = 0;
for (; c + 4 <= nc; c += 4) {
const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3);
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) {
float x = fma(s[h][c], v0, acc[h]);
x = fma(s[h][c + 1], v1, x);
x = fma(s[h][c + 2], v2, x);
acc[h] = fma(s[h][c + 3], v3, x);
}
}
for (; c < nc; ++c) {
const float v = vcell(c);
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]);
}
#pragma clang loop unroll(full)
for (int h = 0; h < HPG; ++h) {
device float* o = pp + (g * HPG + h) * (QD + 2);
o[tid] = acc[h];
if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; }
}
}
constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f};
constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f};
static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a,
int t, int h, ushort tid) {
// prepared planes (xprep layout) of the n = QH * QD vector
const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid;
const int j = tid & 15;
device float* raw = xp + uint(t) * n + d;
raw[0] = y;
raw[plane] = y * qsa_q2s[j & 7];
raw[2 * plane] = y * qsa_p4s[j & 3];
raw[3 * plane] = (j & 1) ? y * 0.00390625f : y;
float s0 = 0.0f, s1 = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < 16; k += 2) {
s0 += simd_shuffle_down(y, ushort(k));
s1 += simd_shuffle_down(y, ushort(k + 1));
}
if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1;
}
// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token
kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]],
device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
const int h = int(tg.x), t = int(tg.y);
const int nchunks = (QBUDGET + QCH - 1) / QCH;
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
const int nc = int((width + QCH - 1) / QCH);
constexpr uint STRIDE = QH * (QD + 2);
device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2);
// the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64)
threadgroup float mk[64], lk[64], ek[64], L_sh;
if (tid < uint(nc)) {
const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD);
mk[tid] = ml.x;
lk[tid] = ml.y;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float M = -INFINITY;
for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]);
M = simd_max(M);
if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid == 0) {
float L = 0.0f;
for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L);
L_sh = L;
}
float o = 0.0f;
device const float* q = pp + tid;
for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o);
threadgroup_barrier(mem_flags::mem_threadgroup);
o = o / L_sh;
const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid];
const float y = o * (1.0f / (1.0f + exp(-gate)));
qsa_write_prepared(y, xp, a, t, h, tid);
}