//! Full LFM2 forward in WGSL — correctness-first (naive kernels, f32, one dispatch per op), verified
//! against the CPU [`crate::reference::RefModel`] oracle. Optimisation (Q4_0 dequant-GEMV, fused
//! norm+matmul, flash-decode attention, on-GPU argmax, one-sync/token) layers on top once correct.
//!
//! Kernels here beyond the GEMV in [`crate::Gemv`]: rmsnorm (one workgroup per vector), the gated
//! short-conv (one thread per channel, ring buffer), RoPE, single-query attention (one workgroup per
//! query head), SwiGLU silu·mul, and the residual add.
use crate::weights::Op;
use crate::{GpuCtx, Weights};
use anyhow::Result;
use wgpu::util::DeviceExt;
/// Wall-clock for the engine's (diagnostic-only) timing instrumentation. `wclock::Instant`
/// is unimplemented on bare `wasm32-unknown-unknown` and PANICS at `now()`; in the browser the
/// timings simply degrade to zero — the numbers feed `OSFKB_STAGE_TRACE`, never correctness.
mod wclock {
#[cfg(not(target_arch = "wasm32"))]
pub use std::time::Instant;
#[cfg(target_arch = "wasm32")]
#[derive(Clone, Copy)]
pub struct Instant;
#[cfg(target_arch = "wasm32")]
impl Instant {
pub fn now() -> Self {
Instant
}
pub fn elapsed(&self) -> std::time::Duration {
std::time::Duration::ZERO
}
}
}
/// Per-sequence KV capacity (prompt + generated tokens).
///
/// 4096 was too small for a real serving workload: the claim-extractor's system prompt alone is
/// ~2.5k tokens, so its SDK's default `max_tokens=20_000` could not be honoured — the cap covers
/// prompt AND generation.
const MAX_T: usize = 16384;
/// Paged-KV geometry. `PAGE_ROW` is DERIVED, not restated: it is the block-table row stride, so a
/// sequence can address `PAGE_ROW * PAGE_BLK` tokens and that product must be exactly `MAX_T`.
pub const PAGE_BLK: usize = 16; // tokens per KV block
pub const PAGE_ROW: usize = MAX_T / PAGE_BLK; // block-table row stride
/// Rows in the image-embedding arena on a multimodal model — the ceiling on image tokens ALIVE at
/// once across every batched sequence, not per request.
///
/// 4096 rows is ~16 pages of a dense document at NuExtract-3's geometry (a 620×400 invoice is 228),
/// and costs `rows · hidden · 4` = 42 MB at hidden 2560. Text models allocate nothing.
const DEFAULT_IMG_TOKENS: usize = 4096;
/// The block table gives each sequence ONE row of `PAGE_ROW` entries, each mapping a `PAGE_BLK`-token
/// block — so `PAGE_ROW * PAGE_BLK` is exactly the addressable per-sequence capacity. If this ever
/// drifts from `MAX_T`, the scheduler admits sequences the shaders cannot address, and the overflow
/// reads whatever sits past the row: another sequence's KV. Fail the build instead.
const _: () = assert!(
PAGE_ROW * PAGE_BLK == MAX_T,
"block-table geometry drifted: PAGE_ROW * PAGE_BLK must equal MAX_T"
);
/// Substitute the paged-KV geometry into a shader source.
///
/// The three paged-KV shaders declare `const PROW: u32 = {{PROW}}u;` — which is NOT valid WGSL. That
/// is the point: an unpatched source fails to compile at pipeline creation, loudly, instead of
/// silently addressing the block table with the wrong stride and reading another sequence's KV. The
/// stride used to be a `256u` literal repeated in three shaders next to a Rust constant, with a
/// comment saying they MUST match — which is a hazard, not a guarantee.
fn page_geom(src: &str) -> String {
debug_assert!(
src.contains("{{PROW}}"),
"paged-KV shader has no {{{{PROW}}}} placeholder — did someone re-hardcode the stride?"
);
src.replace("{{PROW}}", &PAGE_ROW.to_string())
}
/// Server sequence slots (block-table rows 1..=MAX_SLOTS; also conv ring-state slots).
pub const MAX_SLOTS: usize = 192;
/// Max chained decode steps per submit (the on-GPU token-feedback loop that removes the per-step
/// CPU sync — the batch analogue of the solo loop's submit chaining).
pub const CHAIN_MAX: usize = 8;
const CACHE_MAX: usize = 128; // prefix-cache depth: max prompt position whose conv-state we snapshot
// ── WGSL ────────────────────────────────────────────────────────────────────────────────────────
/// RMSNorm: one workgroup per vector. `kp = (vec_len, n_vecs)`; vector v = workgroup_id.x.
/// `out[v,i] = x[v,i] / sqrt(mean_i(x²)+eps) · w[i]`.
const RMSNORM: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (vec_len, n_vecs, eps, _)
var<workgroup> partial: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let n = u32(kp.x);
let eps = kp.z;
let vbase = wid.x * n;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { let v = x[vbase + i]; s = s + v * v; }
partial[lid.x] = s;
workgroupBarrier();
for (var stride = 128u; stride > 0u; stride = stride >> 1u) {
if (lid.x < stride) { partial[lid.x] = partial[lid.x] + partial[lid.x + stride]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(partial[0] / f32(n) + eps);
for (var i = lid.x; i < n; i = i + 256u) { y[vbase + i] = x[vbase + i] * inv * w[i]; }
}
"#;
/// Mean-subtract: `y[i] = x[i] - mean(x)`. The one primitive LayerNorm needs beyond RMSNorm —
/// NON-parametric LayerNorm(x) = RMSNorm(x - mean(x)), so an arch that norms with LayerNorm (OLMo1)
/// centers into a scratch, then feeds it to the SAME RMSNorm-fused projection kernels (with ones
/// weights): `centered * inv_rms(centered) * 1 = (x-mean)/sqrt(var+eps)` exactly. The residual stream
/// (raw `x`) is untouched — only the projection input is centered. `kp=(vec_len, _, _, _)`.
const CENTER: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> kp: vec4<u32>; // (vec_len, _, _, _)
var<workgroup> partial: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let n = kp.x;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { s = s + x[i]; }
partial[lid.x] = s;
workgroupBarrier();
for (var stride = 128u; stride > 0u; stride = stride >> 1u) {
if (lid.x < stride) { partial[lid.x] = partial[lid.x] + partial[lid.x + stride]; }
workgroupBarrier();
}
let mean = partial[0] / f32(n);
for (var i = lid.x; i < n; i = i + 256u) { y[i] = x[i] - mean; }
}
"#;
/// Full AFFINE LayerNorm in one pass: `y[i] = (x[i] − mean)/sqrt(var + eps)·w[i] + b[i]`, where the
/// norm buffer `nw` packs `[weight(n) | bias(n)]` back-to-back (so no new struct field is needed —
/// affine arches just store a 2n-wide norm buffer). Used by affine-LayerNorm arches (StableLM) whose
/// `+bias` cannot be folded into the RMSNorm-fused kernels; the result is fed to PLAIN gemvs (the
/// projection sees the fully-normed input), so this path is backend-agnostic. `kp=(n, _, eps, _)`.
const LAYERNORM_AFFINE: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> nw: array<f32>; // [weight(n) | bias(n)]
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (n, _, eps, _)
var<workgroup> psum: array<f32, 256>;
var<workgroup> psq: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let n = u32(kp.x);
let eps = kp.z;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { s = s + x[i]; }
psum[lid.x] = s;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lid.x < st) { psum[lid.x] = psum[lid.x] + psum[lid.x + st]; }
workgroupBarrier();
}
let mean = psum[0] / f32(n);
workgroupBarrier();
var v = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { let d = x[i] - mean; v = v + d * d; }
psq[lid.x] = v;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lid.x < st) { psq[lid.x] = psq[lid.x] + psq[lid.x + st]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(psq[0] / f32(n) + eps);
for (var i = lid.x; i < n; i = i + 256u) { y[i] = (x[i] - mean) * inv * nw[i] + nw[n + i]; }
}
"#;
/// Add a (possibly offset) bias slice into a vector in place: `y[i] += b[off + i]`. Phi-2 packs its
/// per-layer biases (fc1, and the parallel-combined o_proj+fc2) into the layer's `qkv_bias` buffer
/// and its lm_head bias into `embedding_norm`, so a single kernel reading at an offset covers them
/// all without any per-layer struct field. `kp = (n, off, _, _)`.
const BIAS_ADD: &str = r#"
@group(0) @binding(0) var<storage, read> b: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> kp: vec4<u32>; // (n, offset, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x;
if (i >= kp.x) { return; }
y[i] = y[i] + b[kp.y + i];
}
"#;
/// OLMo2's FULL-VECTOR qk-norm: RMSNorm applied to the WHOLE q projection `[qd]` and the whole k
/// projection `[kd]` (not the engine's per-head qk-norm). Runs in place on the fused `qkv` buffer —
/// workgroup 0 normalizes q `[0..qd]` with `qn`, workgroup 1 normalizes k `[qd..qd+kd]` with `kn`; v
/// is left untouched. After this the qkrc kernel does rope with NO per-head norm (`arch_has_qk_norm`
/// = false). `mp=(qd, kd, _, _)`, `ep.x=eps`.
const QK_FULLNORM: &str = r#"
@group(0) @binding(0) var<storage, read_write> qkv: array<f32>; // [qd + 2*kd] (q|k|v)
@group(0) @binding(1) var<storage, read> nw: array<f32>; // [qd + kd] = [q_norm | k_norm]
@group(0) @binding(2) var<uniform> mp: vec4<u32>; // (qd, kd, _, _)
@group(0) @binding(3) var<uniform> ep: vec4<f32>; // (eps, _, _, _)
var<workgroup> partial: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let qd = mp.x; let kd = mp.y;
// wid.x=0 → q at [0..qd] with weights nw[0..qd]; wid.x=1 → k at [qd..qd+kd] with nw[qd..qd+kd].
// The qkv slice and the weight slice share the SAME base, so `qk-norm applies inv AND the full
// weight here` — qkrc then gets ones (no per-head norm) and does rope only.
var n = qd; var base = 0u;
if (wid.x == 1u) { n = kd; base = qd; }
var s = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { let x = qkv[base + i]; s = s + x * x; }
partial[lid.x] = s;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lid.x < st) { partial[lid.x] = partial[lid.x] + partial[lid.x + st]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(partial[0] / f32(n) + ep.x);
for (var i = lid.x; i < n; i = i + 256u) {
qkv[base + i] = qkv[base + i] * inv * nw[base + i];
}
}
"#;
/// Mamba depthwise causal conv1d, ONE decode step (one thread per channel): `y[c] =
/// SiLU(bias[c] + Σ_k cw[c,k]·window[c,k])` where `window = [conv_state(d_conv-1) | x[c]]`, then the
/// per-channel `conv_state` shifts (drops oldest, appends `x[c]`). `mp=(d_inner, d_conv, _, _)`.
const MAMBA_CONV: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>; // x_ssm [d_inner]
@group(0) @binding(1) var<storage, read> cw: array<f32>; // conv_w [d_inner·d_conv]
@group(0) @binding(2) var<storage, read> cb: array<f32>; // conv_bias [d_inner]
@group(0) @binding(3) var<storage, read_write> cs: array<f32>; // conv_state [d_inner·(d_conv-1)]
@group(0) @binding(4) var<storage, read_write> y: array<f32>; // out [d_inner]
@group(0) @binding(5) var<uniform> mp: vec4<u32>; // (d_inner, d_conv, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let c = gid.x; let di = mp.x; let dc = mp.y;
if (c >= di) { return; }
let dcm1 = dc - 1u;
var acc = cb[c];
for (var k = 0u; k < dcm1; k = k + 1u) { acc = acc + cw[c * dc + k] * cs[c * dcm1 + k]; }
acc = acc + cw[c * dc + dcm1] * x[c];
for (var k = 0u; k + 1u < dcm1; k = k + 1u) { cs[c * dcm1 + k] = cs[c * dcm1 + k + 1u]; }
cs[c * dcm1 + dcm1 - 1u] = x[c];
y[c] = acc / (1.0 + exp(-acc));
}
"#;
/// Mamba selective-scan, ONE decode step (one thread per channel c). `dt = softplus(dt_raw[c] +
/// dt_bias[c])`; for each state s: `h = exp(dt·A[c,s])·h + dt·B[s]·x[c]` (A = −exp(A_log)), and
/// `y[c] = Σ_s C[s]·h + D[c]·x[c]`, gated by `SiLU(z[c])`. `dbl` is the `x_proj` output
/// `[dt_rank | B(d_state) | C(d_state)]`; `par = [A_log(d_inner·d_state) | D(d_inner) | dt_bias]`.
/// `mp=(d_inner, d_state, dt_rank, _)`.
const MAMBA_SSM: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>; // x_conv [d_inner]
@group(0) @binding(1) var<storage, read> dtr: array<f32>; // dt_proj out [d_inner]
@group(0) @binding(2) var<storage, read> dbl: array<f32>; // [dt_rank | B | C]
@group(0) @binding(3) var<storage, read> z: array<f32>; // in_proj out; gate at [z_off + c]
@group(0) @binding(4) var<storage, read> par: array<f32>; // [A_log | D | dt_bias]
@group(0) @binding(5) var<storage, read_write> st: array<f32>; // ssm_state [d_inner·d_state]
@group(0) @binding(6) var<storage, read_write> y: array<f32>; // out [d_inner]
@group(0) @binding(7) var<uniform> mp: vec4<u32>; // (d_inner, d_state, dt_rank, z_off)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let c = gid.x; let di = mp.x; let ds = mp.y; let dtrk = mp.z;
if (c >= di) { return; }
let dr = dtr[c] + par[di * ds + di + c]; // dt_raw + dt_bias
let dt = max(dr, 0.0) + log(1.0 + exp(-abs(dr))); // softplus (stable)
let xc = x[c];
var acc = 0.0;
for (var s = 0u; s < ds; s = s + 1u) {
let a = -exp(par[c * ds + s]);
let dA = exp(a * dt);
let h = dA * st[c * ds + s] + dt * dbl[dtrk + s] * xc;
st[c * ds + s] = h;
acc = acc + dbl[dtrk + ds + s] * h;
}
acc = acc + par[di * ds + c] * xc; // + D[c]·x[c]
let zc = z[mp.w + c];
y[c] = acc * (zc / (1.0 + exp(-zc)));
}
"#;
/// RWKV token-shift + per-channel interpolation, ONE decode step (one thread per channel). Produces
/// the mixed inputs `xk/xv/xr[i] = x[i]·mix + x_prev[i]·(1−mix)` for the three projections (mix
/// packed `[time_mix_k | time_mix_v | time_mix_r]`), then updates the token-shift state `x_prev = x`.
/// The channel-mix has no value stream — its mix packs `[time_mix_k | 0 | time_mix_r]` and `outv`
/// (garbage) is ignored. `mp=(hidden, _, _, _)`.
const RWKV_TOKENSHIFT: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>; // norm(cur) [hidden]
@group(0) @binding(1) var<storage, read_write> xp: array<f32>; // token-shift state [hidden]
@group(0) @binding(2) var<storage, read> mix: array<f32>; // [mk|mv|mr] (3·hidden)
@group(0) @binding(3) var<storage, read_write> outk: array<f32>; // [hidden]
@group(0) @binding(4) var<storage, read_write> outv: array<f32>; // [hidden]
@group(0) @binding(5) var<storage, read_write> outr: array<f32>; // [hidden]
@group(0) @binding(6) var<uniform> mp: vec4<u32>; // (hidden, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = mp.x; let i = gid.x;
if (i >= n) { return; }
let xi = x[i]; let prev = xp[i];
let mk = mix[i]; let mv = mix[n + i]; let mr = mix[2u * n + i];
outk[i] = xi * mk + prev * (1.0 - mk);
outv[i] = xi * mv + prev * (1.0 - mv);
outr[i] = xi * mr + prev * (1.0 - mr);
xp[i] = xi; // this token becomes the previous token for the next step
}
"#;
/// RWKV-4 WKV linear-attention recurrence, ONE decode step (one thread per channel). Numerically
/// stable running exponential-decay weighted sum (the HF `rwkv_linear_attention_cpu` step): state is
/// `[num | den | max]` (max initialized to −1e38). `time_decay` (td) is pre-folded to `−exp(decay)`,
/// `time_first` (tf) is the bonus. `mp=(hidden, _, _, _)`.
const RWKV_WKV: &str = r#"
@group(0) @binding(0) var<storage, read> k: array<f32>; // key [hidden]
@group(0) @binding(1) var<storage, read> v: array<f32>; // value [hidden]
@group(0) @binding(2) var<storage, read> td: array<f32>; // -exp(time_decay) [hidden]
@group(0) @binding(3) var<storage, read> tf: array<f32>; // time_first [hidden]
@group(0) @binding(4) var<storage, read_write> st: array<f32>; // [num | den | max] (3·hidden)
@group(0) @binding(5) var<storage, read_write> wkv: array<f32>; // out [hidden]
@group(0) @binding(6) var<uniform> mp: vec4<u32>; // (hidden, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = mp.x; let i = gid.x;
if (i >= n) { return; }
let aa = st[i]; let bb = st[n + i]; let pp = st[2u * n + i];
let ck = k[i]; let cv = v[i];
// wkv at t: combine the running state with the current (key+bonus)-weighted value.
let ww = tf[i] + ck;
let q = max(pp, ww);
let e1 = exp(pp - q);
let e2 = exp(ww - q);
wkv[i] = (e1 * aa + e2 * cv) / (e1 * bb + e2);
// advance the state by one decay step (key without the bonus).
let ww2 = pp + td[i];
let q2 = max(ww2, ck);
let e1b = exp(ww2 - q2);
let e2b = exp(ck - q2);
st[i] = e1b * aa + e2b * cv;
st[n + i] = e1b * bb + e2b;
st[2u * n + i] = q2;
}
"#;
/// RWKV receptance gate: `y[i] = base + sigmoid(a[i])·b[i]`, where `base = y[i]` when `mp.y==1`
/// (accumulate into the residual, channel-mix) else `0` (overwrite, time-mix `r·wkv`). One thread
/// per channel. `mp=(hidden, add, _, _)`.
const RWKV_SIGMUL: &str = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>; // receptance (pre-sigmoid) [hidden]
@group(0) @binding(1) var<storage, read> b: array<f32>; // wkv or kv [hidden]
@group(0) @binding(2) var<storage, read_write> y: array<f32>; // out [hidden]
@group(0) @binding(3) var<uniform> mp: vec4<u32>; // (hidden, add, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = mp.x; let i = gid.x;
if (i >= n) { return; }
let g = 1.0 / (1.0 + exp(-a[i]));
var base = 0.0;
if (mp.y == 1u) { base = y[i]; }
y[i] = base + g * b[i];
}
"#;
/// Elementwise vector copy `y[i] = x[i]` (distinct buffers). Used to write RWKV's block-0 `ln0`
/// result back into the residual stream: LAYERNORM_AFFINE cannot alias src=dst within one dispatch
/// (wgpu forbids read + read_write on the same buffer), so ln0 lands in a scratch and is copied in.
/// `mp=(n, _, _, _)`.
const COPY_VEC: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> mp: vec4<u32>; // (n, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x;
if (i >= mp.x) { return; }
y[i] = x[i];
}
"#;
/// On-GPU embed gather: `cur[i] = embed[tokens[pos]*h + i]` — reads the current token from a GPU
/// buffer (no CPU token offset), so the decode loop never returns to the CPU between tokens.
const GATHER: &str = r#"
@group(0) @binding(0) var<storage, read> embed: array<f32>;
@group(0) @binding(1) var<storage, read> tokens: array<u32>;
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (pos, prompt_len, vocab, hidden)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
cur[i] = embed[tokens[posu.x] * h + i];
}
"#;
/// GLU activation·multiply for the tiled-GEMM MLP (`OSFKB_WIDE_MLP_GEMM`): when the wide
/// prefill plan runs gate and up as two SEPARATE tiled Q4 GEMMs (instead of the per-column
/// fused `mlp_gate_q4_k`), this folds them — `y = act(gate)·up` over `[k·im]` — feeding the
/// down projection. `dims.y` selects SiLU (0) vs gelu_pytorch_tanh (1), the same flag the
/// fused kernel reads from `epsm.y`.
const MLP_ACT_MUL: &str = r#"
@group(0) @binding(0) var<storage, read> g: array<f32>; // gate_raw [k·im]
@group(0) @binding(1) var<storage, read> u: array<f32>; // up_raw [k·im]
@group(0) @binding(2) var<storage, read_write> y: array<f32>; // out [k·im]
@group(0) @binding(3) var<uniform> dims: vec4<u32>; // (total, gelu, up_offset, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x; if (i >= dims.x) { return; }
// dims.z: read the up stream at an offset — the Gemma-4 PLE injection binds the whole
// per-layer-inputs buffer and selects its layer's slice here. 0 (every MLP call) is exact.
let gate = g[i]; let upv = u[i + dims.z];
var act: f32;
if (dims.y == 2u) {
// EXACT (erf) GELU — Abramowitz-Stegun 7.1.26 erf (max err ~1.5e-7). Falcon's `gelu`.
let s = sign(gate);
let ax = abs(gate) * 0.7071067811865476;
let t = 1.0 / (1.0 + 0.3275911 * ax);
let e = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-ax * ax);
act = 0.5 * gate * (1.0 + s * e);
} else if (dims.y == 3u) {
// ReLU. Nemotron's squared-ReLU MLP is relu(g)² = relu(g)·g, reached by binding both the
// gate and up streams to the same up-projection: act(g)·u = relu(u)·u = relu(u)².
act = max(gate, 0.0);
} else if (dims.y != 0u) {
// gelu_pytorch_tanh / gelu_new (the tanh approximation) — Phi-2, Gemma.
let g3 = gate * gate * gate;
let t = clamp(0.7978845608028654 * (gate + 0.044715 * g3), -20.0, 20.0);
act = 0.5 * gate * (1.0 + tanh(t));
} else {
act = gate / (1.0 + exp(-gate));
}
y[i] = act * upv;
}
"#;
/// f16 embed gather (`OSFKB_Q1_WEIGHTS`, Bonsai-class): the token-gather table is stored f16
/// because a 248k×5120 f32 table is 5 GB — over Metal's 4 GB per-buffer cap — while f16 is
/// 2.5 GB. Same copy, widened to f32 into the residual stream.
const GATHER_F16: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> embed_lo: array<f16>;
@group(0) @binding(1) var<storage, read> tokens: array<u32>;
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (pos, prompt_len, vocab, hidden)
@group(0) @binding(4) var<storage, read> embed_hi: array<f16>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
// One f16 table, TWO bindings: wgpu's Vulkan backend caps a storage BINDING at 2 GiB-4
// (Metal allows 4 GiB), and 248k x 5120 f16 is 2.54 GB. The host binds the same buffer as
// two row-aligned sub-ranges split at vocab/2; rows < split come from lo, the rest from hi.
let tok = tokens[posu.x];
let vs = posu.z / 2u;
if (tok < vs) {
cur[i] = f32(embed_lo[tok * h + i]);
} else {
cur[i] = f32(embed_hi[(tok - vs) * h + i]);
}
}
"#;
/// Column-indexed embed gather (MTP serving draft): `cur[col·h + i] = embed[toks[col]·h + i]`.
const GATHER_COLS: &str = r#"
@group(0) @binding(0) var<storage, read> embed: array<f32>;
@group(0) @binding(1) var<storage, read> toks: array<u32>;
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (_, _, vocab, hidden)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
let col = gid.y;
cur[col * h + i] = embed[toks[col] * h + i];
}
"#;
/// On-GPU argmax: `tokens[pos+1] = argmax(logits)` — but only once we're past the prompt
/// (`pos+1 >= prompt_len`), so prefill keeps the supplied prompt. One workgroup over the vocab.
const ARGMAX: &str = r#"
@group(0) @binding(0) var<storage, read> logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> tokens: array<u32>;
@group(0) @binding(2) var<uniform> posu: vec4<u32>; // (pos, prompt_len, vocab, hidden)
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let v = posu.z;
var bv = -1e30;
var bi = 0u;
for (var i = lid.x; i < v; i = i + 256u) { if (logits[i] > bv) { bv = logits[i]; bi = i; } }
vmax[lid.x] = bv; imax[lid.x] = bi;
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u && posu.x + 1u >= posu.y) { tokens[posu.x + 1u] = imax[0]; }
}
"#;
/// Parallel argmax stage 1: 256 workgroups each scan a vocab slice → (max,idx) into `pv`/`pi[256]`.
/// (The old single-workgroup argmax used 1 of ~40 GPU cores; this fills the GPU.)
const ARGMAX_BLOCK: &str = r#"
@group(0) @binding(0) var<storage, read> logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> pv: array<f32>;
@group(0) @binding(2) var<storage, read_write> pi: array<u32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (pos, prompt_len, vocab, hidden)
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let v = posu.z;
let per = (v + 255u) / 256u;
let start = wid.x * per;
var bv = -1e30; var bi = 0u;
for (var k = lid.x; k < per; k = k + 256u) {
let idx = start + k;
if (idx < v && logits[idx] > bv) { bv = logits[idx]; bi = idx; }
}
vmax[lid.x] = bv; imax[lid.x] = bi;
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u) { pv[wid.x] = vmax[0]; pi[wid.x] = imax[0]; }
}
"#;
/// Parallel argmax stage 2: 1 workgroup merges the 256 block winners → `tokens[pos+1]` (past prompt).
const ARGMAX_MERGE: &str = r#"
@group(0) @binding(0) var<storage, read> pv: array<f32>;
@group(0) @binding(1) var<storage, read> pi: array<u32>;
@group(0) @binding(2) var<storage, read_write> tokens: array<u32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (pos, prompt_len, vocab, hidden)
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
vmax[lid.x] = pv[lid.x]; imax[lid.x] = pi[lid.x];
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u && posu.x + 1u >= posu.y) { tokens[posu.x + 1u] = imax[0]; }
}
"#;
/// Column-slot variant of [`ARGMAX_MERGE`]: writes `chosen[colu.x]` (batched serving — one merge
/// per batch column). Comparisons only, so a fresh compilation cannot drift FP behavior.
const ARGMAX_MERGE_COL: &str = r#"
@group(0) @binding(0) var<storage, read> pv: array<f32>;
@group(0) @binding(1) var<storage, read> pi: array<u32>;
@group(0) @binding(2) var<storage, read_write> chosen: array<u32>; // [steps, stride]
@group(0) @binding(3) var<storage, read> cnt: array<u32>;
@group(0) @binding(4) var<uniform> colu: vec4<u32>; // (col, step_stride, _, _)
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
vmax[lid.x] = pv[lid.x]; imax[lid.x] = pi[lid.x];
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u) { chosen[cnt[0] * colu.y + colu.x] = imax[0]; }
}
"#;
/// Step-counter bump (chained decode): one dispatch at the end of each chained step; dispatch
/// ordering within an encoder makes it a sequencing point between steps.
const CNT_BUMP: &str = r#"
@group(0) @binding(0) var<storage, read_write> cnt: array<u32>;
@compute @workgroup_size(1)
fn main() { cnt[0] = cnt[0] + 1u; }
"#;
// Expert-dedup dispatch-order sort: TRIED AND REVERTED (2026-07-05). A 1-thread insertion sort
// of the 40 (col,slot) pairs per MoE layer cost ~75 µs/dispatch (dependent global reads) —
// +3.2 ms/round for NO measurable dedup gain: the concurrent per-pair weight sweeps already
// share L2 without ordering (same finding as the banded head — the small-width verify is
// latency/occupancy-bound, not weight-traffic-bound). See the memory topic for the numbers.
/// Gated-DeltaNet depthwise conv, SLOTTED + k-column (the CONV_MIX_K discipline): one thread per
/// conv channel walks the k columns IN ORDER — same-slot columns (prefill of one sequence) update
/// the ring sequentially, distinct slots are independent. `mixed` is the fused-projection output
/// `[k, in_rows]`; only the first `conv_dim` rows are convolved. `dims=(conv_dim, K, k, stride)`,
/// `dims2=(in_rows, _, _, _)`; slot = cmeta btrow (0 solo, 1..=MAX_SLOTS server, trash parked).
const DN_CONV_K: &str = r#"
@group(0) @binding(0) var<storage, read> mixed: array<f32>; // [k, in_rows]
@group(0) @binding(1) var<storage, read> cw: array<f32>; // [conv_dim, K] oldest-first
@group(0) @binding(2) var<storage, read_write> ring: array<f32>; // [slots, conv_dim, K]
@group(0) @binding(3) var<storage, read_write> conved: array<f32>; // [k, conv_dim]
@group(0) @binding(4) var<storage, read> cmeta: array<u32>; // [steps, stride, 2]
@group(0) @binding(5) var<storage, read> cnt: array<u32>;
@group(0) @binding(6) var<uniform> dims: vec4<u32>; // (conv_dim, K, k, stride)
@group(0) @binding(7) var<uniform> dims2: vec4<u32>; // (in_rows, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let cd = dims.x; let kk = dims.y; let k = dims.z;
let c = gid.x;
if (c >= cd) { return; }
for (var p = 0u; p < k; p = p + 1u) {
let slot = cmeta[(cnt[0] * dims.w + p) * 2u + 1u];
let base = (slot * cd + c) * kk;
var acc = 0.0;
for (var t = 0u; t + 1u < kk; t = t + 1u) {
let nxt = ring[base + t + 1u];
ring[base + t] = nxt;
acc = acc + nxt * cw[c * kk + t];
}
let xn = mixed[p * dims2.x + c];
ring[base + kk - 1u] = xn;
acc = acc + xn * cw[c * kk + kk - 1u];
conved[p * cd + c] = acc / (1.0 + exp(-acc));
}
}
"#;
/// Gated-DeltaNet recurrent step + fused gated RMSNorm, SLOTTED + k-column: one workgroup per
/// v-head; columns processed IN ORDER (recurrence). Mirrors `deltanet::DeltaNetRef::core`
/// op-for-op (the primitive's gates pinned the math). Requires dk ≤ 256, dv ≤ 128.
/// `dims=(nv, nk, dk, dv)`, `dims2=(in_rows, conv_dim, k, stride)`, `epsu=(rms_eps,_,_,_)`.
/// `gpar = [A_log(nv) | dt_bias(nv) | norm_w(dv)]`; z/b/a read from `mixed` at row offsets.
const DN_STEP_K: &str = r#"
@group(0) @binding(0) var<storage, read> conved: array<f32>; // [k, conv_dim] = q|k|v
@group(0) @binding(1) var<storage, read> mixed: array<f32>; // [k, in_rows] (z|b|a tail)
@group(0) @binding(2) var<storage, read> gpar: array<f32>;
@group(0) @binding(3) var<storage, read_write> st: array<f32>; // [slots, nv, dk, dv]
@group(0) @binding(4) var<storage, read_write> core: array<f32>; // [k, nv·dv]
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (nv, nk, dk, dv)
@group(0) @binding(8) var<uniform> dims2: vec4<u32>; // (in_rows, conv_dim, k, stride)
@group(0) @binding(9) var<uniform> epsu: vec4<f32>;
var<workgroup> qsh: array<f32, 256>;
var<workgroup> ksh: array<f32, 256>;
var<workgroup> osh: array<f32, 128>;
var<workgroup> scal: array<f32, 3>;
var<workgroup> red2: array<f32, 256>;
fn softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
if (x < -20.0) { return exp(x); }
return log(1.0 + exp(x));
}
@compute @workgroup_size(128)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let nv = dims.x; let nk = dims.y; let dk = dims.z; let dv = dims.w;
let in_rows = dims2.x; let cd = dims2.y; let k = dims2.z;
let kh = h / (nv / nk);
let j = lid.x;
for (var p = 0u; p < k; p = p + 1u) {
let slot = cmeta[(cnt[0] * dims2.w + p) * 2u + 1u];
let cbase = p * cd;
for (var i = j; i < dk; i = i + 128u) {
qsh[i] = conved[cbase + kh * dk + i];
ksh[i] = conved[cbase + nk * dk + kh * dk + i];
}
workgroupBarrier();
// Parallel q/k RMS norms (tree-reduced): the serial form ran THREE dk-loops on
// thread 0 while 127 lanes idled — every step, every head; the fresh decode profile
// put dn.step at 26.5% and this section was most of its critical path.
var psq = 0.0;
var psk = 0.0;
for (var i = j; i < dk; i = i + 128u) {
psq = psq + qsh[i] * qsh[i];
psk = psk + ksh[i] * ksh[i];
}
red2[j] = psq;
red2[128u + j] = psk;
workgroupBarrier();
for (var sred = 64u; sred > 0u; sred = sred >> 1u) {
if (j < sred) {
red2[j] = red2[j] + red2[j + sred];
red2[128u + j] = red2[128u + j] + red2[128u + j + sred];
}
workgroupBarrier();
}
let invq = 1.0 / sqrt(red2[0] + 1e-6);
let invk = 1.0 / sqrt(red2[128] + 1e-6);
let scale = 1.0 / sqrt(f32(dk));
for (var i = j; i < dk; i = i + 128u) {
qsh[i] = qsh[i] * invq * scale;
ksh[i] = ksh[i] * invk;
}
if (j == 0u) {
let b = mixed[p * in_rows + cd + nv * dv + h];
let a = mixed[p * in_rows + cd + nv * dv + nv + h];
let g = -exp(gpar[h]) * softplus(a + gpar[nv + h]);
scal[0] = exp(g);
scal[1] = 1.0 / (1.0 + exp(-b));
}
workgroupBarrier();
if (j < dv) {
let sbase = ((slot * nv + h) * dk) * dv + j;
let decay = scal[0];
var kv = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
kv = kv + ksh[i] * st[sbase + i * dv];
}
kv = kv * decay;
let vv = conved[cbase + 2u * nk * dk + h * dv + j];
let delta = (vv - kv) * scal[1];
var o = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
let sv = st[sbase + i * dv] * decay + ksh[i] * delta;
st[sbase + i * dv] = sv;
o = o + qsh[i] * sv;
}
osh[j] = o;
}
workgroupBarrier();
// Parallel o RMS norm (tree-reduced) — the same fix the q/k norms got: the serial
// form ran a dv-length loop on thread 0 per step per head.
var pso = 0.0;
if (j < dv) { pso = osh[j] * osh[j]; }
red2[j] = pso;
workgroupBarrier();
for (var sred = 64u; sred > 0u; sred = sred >> 1u) {
if (j < sred) { red2[j] = red2[j] + red2[j + sred]; }
workgroupBarrier();
}
if (j == 0u) {
scal[2] = 1.0 / sqrt(red2[0] / f32(dv) + epsu.x);
}
workgroupBarrier();
if (j < dv) {
let zz = mixed[p * in_rows + cd + h * dv + j];
core[p * nv * dv + h * dv + j] =
osh[j] * scal[2] * gpar[2u * nv + j] * (zz / (1.0 + exp(-zz)));
}
workgroupBarrier();
}
}
"#;
/// PARALLEL-COLUMN variant of [`DN_STEP_K`] for DECODE batches (grid `(nv, k)`): every
/// column is a DISTINCT sequence slot (asserted host-side), so the recurrence has no
/// cross-column order — 8× the occupancy of the sequential kernel at k=16. Prefill
/// chunks (same-slot columns) MUST keep the sequential kernel.
/// PARALLEL-COLUMN DeltaNet step for BATCH plans (grid `(nv, k)`): slot-owner discipline —
/// the first column of each slot processes all of that slot\u{2019}s columns sequentially, so
/// mixed decode+prefill batches stay correct while pure-decode batches get full
/// occupancy. Uniformity: `wid.y` is workgroup-uniform, so the early return and the
/// in-loop `continue` are uniform per workgroup (barriers inside remain legal).
/// [`DN_STEP_PK`] with the two dk-serial walks parallelized over `ic` sub-lanes (WG =
/// 128*IC). At decode (k=1) the PK grid is nv workgroups (~48) of 4 warps = ~2.4 warps/SM
/// on the V100S — 4% occupancy, and each thread chained ~2*dk dependent state loads; the
/// band=1 profile put dn_step at 4.3ms/token vs ~0.5ms of state I/O. Here lane (ic, j)
/// owns the i-slice {ic, ic+IC, ...}: kv and o become IC-way tree reductions (fixed order,
/// deterministic; NOT bitwise vs PK — both plan paths swap together under OSFKB_DN_PK2).
/// State writes stay disjoint per (i, j); slot-owner column walk unchanged.
fn dn_step_pk2_src(ic: u32) -> String {
let wg = 128 * ic;
format!(
r#"
@group(0) @binding(0) var<storage, read> conved: array<f32>;
@group(0) @binding(1) var<storage, read> mixed: array<f32>;
@group(0) @binding(2) var<storage, read> gpar: array<f32>;
@group(0) @binding(3) var<storage, read_write> st: array<f32>;
@group(0) @binding(4) var<storage, read_write> core: array<f32>;
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> dims: vec4<u32>;
@group(0) @binding(8) var<uniform> dims2: vec4<u32>;
@group(0) @binding(9) var<uniform> epsu: vec4<f32>;
var<workgroup> qsh: array<f32, 256>;
var<workgroup> ksh: array<f32, 256>;
var<workgroup> osh: array<f32, 128>;
var<workgroup> scal: array<f32, 3>;
var<workgroup> red2: array<f32, {wg2}>;
var<workgroup> kvp: array<f32, {wg}>;
fn softplus(x: f32) -> f32 {{
if (x > 20.0) {{ return x; }}
if (x < -20.0) {{ return exp(x); }}
return log(1.0 + exp(x));
}}
@compute @workgroup_size({wg})
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let h = wid.x;
let nv = dims.x; let nk = dims.y; let dk = dims.z; let dv = dims.w;
let in_rows = dims2.x; let cd = dims2.y; let k = dims2.z;
let kh = h / (nv / nk);
let t = lid.x;
let j = t % 128u;
let icl = t / 128u;
let p0 = wid.y;
let my_slot = cmeta[(cnt[0] * dims2.w + p0) * 2u + 1u];
for (var q = 0u; q < p0; q = q + 1u) {{
if (cmeta[(cnt[0] * dims2.w + q) * 2u + 1u] == my_slot) {{ return; }}
}}
for (var p = p0; p < k; p = p + 1u) {{
let slot = cmeta[(cnt[0] * dims2.w + p) * 2u + 1u];
if (slot != my_slot) {{ continue; }}
let cbase = p * cd;
for (var i = t; i < dk; i = i + {wg}u) {{
qsh[i] = conved[cbase + kh * dk + i];
ksh[i] = conved[cbase + nk * dk + kh * dk + i];
}}
workgroupBarrier();
var psq = 0.0;
var psk = 0.0;
for (var i = t; i < dk; i = i + {wg}u) {{
psq = psq + qsh[i] * qsh[i];
psk = psk + ksh[i] * ksh[i];
}}
red2[t] = psq;
red2[{wg}u + t] = psk;
workgroupBarrier();
for (var sred = {wgh}u; sred > 0u; sred = sred >> 1u) {{
if (t < sred) {{
red2[t] = red2[t] + red2[t + sred];
red2[{wg}u + t] = red2[{wg}u + t] + red2[{wg}u + t + sred];
}}
workgroupBarrier();
}}
let invq = 1.0 / sqrt(red2[0] + 1e-6);
let invk = 1.0 / sqrt(red2[{wg}] + 1e-6);
let scale = 1.0 / sqrt(f32(dk));
for (var i = t; i < dk; i = i + {wg}u) {{
qsh[i] = qsh[i] * invq * scale;
ksh[i] = ksh[i] * invk;
}}
if (t == 0u) {{
let b = mixed[p * in_rows + cd + nv * dv + h];
let a = mixed[p * in_rows + cd + nv * dv + nv + h];
let g = -exp(gpar[h]) * softplus(a + gpar[nv + h]);
scal[0] = exp(g);
scal[1] = 1.0 / (1.0 + exp(-b));
}}
workgroupBarrier();
// phase A: kv[j] = sum_i ksh[i]*st[i][j], i split over {ic} sub-lanes
var kvpart = 0.0;
if (j < dv) {{
let sbase = ((slot * nv + h) * dk) * dv + j;
for (var i = icl; i < dk; i = i + {ic}u) {{
kvpart = kvpart + ksh[i] * st[sbase + i * dv];
}}
}}
kvp[t] = kvpart;
workgroupBarrier();
// phase B: delta from the reduced kv; each ic-slice updates its own st rows and
// accumulates a partial o; o reduced over ic.
if (j < dv) {{
let sbase = ((slot * nv + h) * dk) * dv + j;
let decay = scal[0];
var kv = 0.0;
for (var c = 0u; c < {ic}u; c = c + 1u) {{
kv = kv + kvp[c * 128u + j];
}}
kv = kv * decay;
let vv = conved[cbase + 2u * nk * dk + h * dv + j];
let delta = (vv - kv) * scal[1];
var o = 0.0;
for (var i = icl; i < dk; i = i + {ic}u) {{
let sv = st[sbase + i * dv] * decay + ksh[i] * delta;
st[sbase + i * dv] = sv;
o = o + qsh[i] * sv;
}}
kvp[t] = o;
}} else {{
kvp[t] = 0.0;
}}
workgroupBarrier();
if (icl == 0u && j < dv) {{
var o = 0.0;
for (var c = 0u; c < {ic}u; c = c + 1u) {{
o = o + kvp[c * 128u + j];
}}
osh[j] = o;
}}
workgroupBarrier();
var pso = 0.0;
if (icl == 0u && j < dv) {{ pso = osh[j] * osh[j]; }}
red2[t] = pso;
workgroupBarrier();
for (var sred = {wgh}u; sred > 0u; sred = sred >> 1u) {{
if (t < sred) {{ red2[t] = red2[t] + red2[t + sred]; }}
workgroupBarrier();
}}
if (t == 0u) {{
scal[2] = 1.0 / sqrt(red2[0] / f32(dv) + epsu.x);
}}
workgroupBarrier();
if (icl == 0u && j < dv) {{
let zz = mixed[p * in_rows + cd + h * dv + j];
core[p * nv * dv + h * dv + j] =
osh[j] * scal[2] * gpar[2u * nv + j] * (zz / (1.0 + exp(-zz)));
}}
workgroupBarrier();
}}
}}
"#,
wg = wg,
wg2 = 2 * wg,
wgh = wg / 2,
ic = ic,
)
}
const DN_STEP_PK: &str = r#"
@group(0) @binding(0) var<storage, read> conved: array<f32>; // [k, conv_dim] = q|k|v
@group(0) @binding(1) var<storage, read> mixed: array<f32>; // [k, in_rows] (z|b|a tail)
@group(0) @binding(2) var<storage, read> gpar: array<f32>;
@group(0) @binding(3) var<storage, read_write> st: array<f32>; // [slots, nv, dk, dv]
@group(0) @binding(4) var<storage, read_write> core: array<f32>; // [k, nv·dv]
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (nv, nk, dk, dv)
@group(0) @binding(8) var<uniform> dims2: vec4<u32>; // (in_rows, conv_dim, k, stride)
@group(0) @binding(9) var<uniform> epsu: vec4<f32>;
var<workgroup> qsh: array<f32, 256>;
var<workgroup> ksh: array<f32, 256>;
var<workgroup> osh: array<f32, 128>;
var<workgroup> scal: array<f32, 3>;
var<workgroup> red2: array<f32, 256>;
fn softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
if (x < -20.0) { return exp(x); }
return log(1.0 + exp(x));
}
@compute @workgroup_size(128)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let nv = dims.x; let nk = dims.y; let dk = dims.z; let dv = dims.w;
let in_rows = dims2.x; let cd = dims2.y; let k = dims2.z;
let kh = h / (nv / nk);
let j = lid.x;
// One workgroup per (head, column). The FIRST column of each slot owns the slot: it walks
// its own column and every later same-slot column IN ORDER (prefill chunks stay sequential);
// other instances exit. Decode batches (distinct slots) get full (nv × k) parallelism.
let p0 = wid.y;
let my_slot = cmeta[(cnt[0] * dims2.w + p0) * 2u + 1u];
for (var q = 0u; q < p0; q = q + 1u) {
if (cmeta[(cnt[0] * dims2.w + q) * 2u + 1u] == my_slot) { return; }
}
for (var p = p0; p < k; p = p + 1u) {
let slot = cmeta[(cnt[0] * dims2.w + p) * 2u + 1u];
if (slot != my_slot) { continue; }
let cbase = p * cd;
for (var i = j; i < dk; i = i + 128u) {
qsh[i] = conved[cbase + kh * dk + i];
ksh[i] = conved[cbase + nk * dk + kh * dk + i];
}
workgroupBarrier();
// Parallel q/k RMS norms (tree-reduced): the serial form ran THREE dk-loops on
// thread 0 while 127 lanes idled — every step, every head; the fresh decode profile
// put dn.step at 26.5% and this section was most of its critical path.
var psq = 0.0;
var psk = 0.0;
for (var i = j; i < dk; i = i + 128u) {
psq = psq + qsh[i] * qsh[i];
psk = psk + ksh[i] * ksh[i];
}
red2[j] = psq;
red2[128u + j] = psk;
workgroupBarrier();
for (var sred = 64u; sred > 0u; sred = sred >> 1u) {
if (j < sred) {
red2[j] = red2[j] + red2[j + sred];
red2[128u + j] = red2[128u + j] + red2[128u + j + sred];
}
workgroupBarrier();
}
let invq = 1.0 / sqrt(red2[0] + 1e-6);
let invk = 1.0 / sqrt(red2[128] + 1e-6);
let scale = 1.0 / sqrt(f32(dk));
for (var i = j; i < dk; i = i + 128u) {
qsh[i] = qsh[i] * invq * scale;
ksh[i] = ksh[i] * invk;
}
if (j == 0u) {
let b = mixed[p * in_rows + cd + nv * dv + h];
let a = mixed[p * in_rows + cd + nv * dv + nv + h];
let g = -exp(gpar[h]) * softplus(a + gpar[nv + h]);
scal[0] = exp(g);
scal[1] = 1.0 / (1.0 + exp(-b));
}
workgroupBarrier();
if (j < dv) {
let sbase = ((slot * nv + h) * dk) * dv + j;
let decay = scal[0];
var kv = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
kv = kv + ksh[i] * st[sbase + i * dv];
}
kv = kv * decay;
let vv = conved[cbase + 2u * nk * dk + h * dv + j];
let delta = (vv - kv) * scal[1];
var o = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
let sv = st[sbase + i * dv] * decay + ksh[i] * delta;
st[sbase + i * dv] = sv;
o = o + qsh[i] * sv;
}
osh[j] = o;
}
workgroupBarrier();
// Parallel o RMS norm (tree-reduced) — the same fix the q/k norms got: the serial
// form ran a dv-length loop on thread 0 per step per head.
var pso = 0.0;
if (j < dv) { pso = osh[j] * osh[j]; }
red2[j] = pso;
workgroupBarrier();
for (var sred = 64u; sred > 0u; sred = sred >> 1u) {
if (j < sred) { red2[j] = red2[j] + red2[j + sred]; }
workgroupBarrier();
}
if (j == 0u) {
scal[2] = 1.0 / sqrt(red2[0] / f32(dv) + epsu.x);
}
workgroupBarrier();
if (j < dv) {
let zz = mixed[p * in_rows + cd + h * dv + j];
core[p * nv * dv + h * dv + j] =
osh[j] * scal[2] * gpar[2u * nv + j] * (zz / (1.0 + exp(-zz)));
}
workgroupBarrier();
}
}
"#;
/// CHUNK-PARALLEL DeltaNet recurrence for WIDE (prefill) plans — the production port of
/// `deltanet.rs`'s gated CHUNK form (see `DeltaNetRef::core_chunk` for the derivation:
/// `S_r = Γ_r S₀ + Σ_{i≤r} (Γ_r/Γ_i) k_i u_iᵀ`, entering-state term folded into the
/// pseudo-value solve). Same bindings and slot-owner discipline as [`DN_STEP_PK`]; the owner
/// workgroup walks its slot's columns in chunks of ≤16, touching the recurrent state THREE
/// times per chunk (2 reads + 1 write) instead of 3 per column — the arithmetic-intensity
/// flip that unblocks prefill (sequential DN was the measured 250 prompt-tok/s wall).
/// Shared budget ≈ 27 KiB (k/q/u staging + C×C pair matrices) — fits Metal's 32 KiB floor
/// and every desktop adapter, vendor-neutral WGSL only. Output values differ from the
/// sequential kernel in last ulps (different summation grouping — the point of the form);
/// gates are tolerance + argmax anchors plus real-model text equality, not bitwise.
/// CHUNK-PARALLEL DeltaNet recurrence for WIDE (prefill) plans — codegen'd (see
/// `DeltaNetRef::core_chunk` for the algebra; same bindings + slot-owner discipline as
/// [`DN_STEP_PK`]). v2 rework, both informed by measurement: (1) ALL per-position private
/// accumulators are STATIC registers (dynamically-indexed private arrays land in DRAM-backed
/// Naga scratch — the 4.2× GEMM-unroll lesson); (2) barriers are BATCHED per chunk (~12) —
/// the v1 per-position tree reductions ran ~150 barriers/chunk and dominated the kernel
/// (1.57 ms/layer at 64 cols). The pseudo-value solve is lane-local (u_r[j] reads only its
/// own lane's prior writes), so the whole triangular chain runs barrier-free. Shared ≈ 30 KiB
/// (Metal 32 KiB floor ✓). Requires dk ≤ 128, dv ≤ 128 (the plan builder asserts).
/// CHUNK-PARALLEL DeltaNet recurrence, dv-SPLIT v3 (see the in-kernel comment): grid
/// `(nv·2, k)`, two WGs per head, RAW o + per-half Σo² partials out; [`DN_CHUNK_ONORM`]
/// finishes the gated RMSNorm + z-gate. Same slot-owner discipline and chunk algebra as v2
/// (batched barriers, static accumulators — the measured lessons).
fn dn_chunk_pk_src() -> String {
r##"@group(0) @binding(0) var<storage, read> conved: array<f32>; // [k, conv_dim] = q|k|v
@group(0) @binding(1) var<storage, read> mixed: array<f32>; // [k, in_rows] (z|b|a tail)
@group(0) @binding(2) var<storage, read> gpar: array<f32>;
@group(0) @binding(3) var<storage, read_write> st: array<f32>; // [slots, nv, dk, dv]
@group(0) @binding(4) var<storage, read_write> core: array<f32>; // [k, nv·dv] RAW o (pre-norm)
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (nv, nk, dk, dv)
@group(0) @binding(8) var<uniform> dims2: vec4<u32>; // (in_rows, conv_dim, k, stride)
@group(0) @binding(9) var<storage, read_write> opart: array<f32>; // [k, nv, 2] Σo² per half
const CMAX = 16u;
var<workgroup> ksh: array<f32, 2048>; // [C, dk] normalized k
var<workgroup> qsh: array<f32, 2048>; // [C, dk] normalized+scaled q
var<workgroup> ush: array<f32, 2048>; // [C, dv] pseudo-values (this half's lanes)
var<workgroup> amat: array<f32, 256>; // [C, C] decay-ratio'd KKᵀ (strict lower)
var<workgroup> qkm: array<f32, 256>; // [C, C] decay-ratio'd QKᵀ (lower incl. diag)
var<workgroup> oval: array<f32, 1024>; // [8, 128] o² staging / P1c partials
var<workgroup> invs: array<f32, 40>; // [invq 16 | invk 16 | Σo² 8]
var<workgroup> lgs: array<f32, 16>;
var<workgroup> bets: array<f32, 16>;
fn softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
if (x < -20.0) { return exp(x); }
return log(1.0 + exp(x));
}
@compute @workgroup_size(128)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
// dv-SPLIT: two WGs per head (grid (nv·2, k)); each owns half the dv lanes — 2× the
// occupancy of the 1-WG/head form (32 WGs on 80 SMs was the measured dn wall). The
// gated RMSNorm needs Σo² over ALL dv, so o is written RAW and each half publishes its
// partial to `opart`; DN_CHUNK_ONORM applies the norm + z-gate afterwards (the state
// update never depends on the norm). dk-space phases (P1c/P2) run redundantly in both
// halves (~10% duplicated work — cheaper than a cross-WG handshake).
let h = wid.x / 2u;
let half = wid.x % 2u;
let nv = dims.x; let nk = dims.y; let dk = dims.z; let dv = dims.w;
let half_len = (dv + 1u) / 2u;
let in_rows = dims2.x; let cd = dims2.y; let k = dims2.z;
let kh = h / (nv / nk);
let j = lid.x;
let jg = half * half_len + j;
let lane = j < half_len && jg < dv;
let p0 = wid.y;
let my_slot = cmeta[(cnt[0] * dims2.w + p0) * 2u + 1u];
for (var q = 0u; q < p0; q = q + 1u) {
if (cmeta[(cnt[0] * dims2.w + q) * 2u + 1u] == my_slot) { return; }
}
var scan = p0;
loop {
var pcols: array<u32, 16>;
var cc = 0u;
{
var p = scan;
for (; p < k && cc < CMAX; p = p + 1u) {
if (cmeta[(cnt[0] * dims2.w + p) * 2u + 1u] == my_slot) {
pcols[cc] = p;
cc = cc + 1u;
}
}
}
if (cc == 0u) { return; }
for (var r = 0u; r < cc; r = r + 1u) {
let cbase = pcols[r] * cd;
for (var i = j; i < dk; i = i + 128u) {
qsh[r * dk + i] = conved[cbase + kh * dk + i];
ksh[r * dk + i] = conved[cbase + nk * dk + kh * dk + i];
}
}
if (j == 0u) {
var lg_run = 0.0;
for (var r = 0u; r < cc; r = r + 1u) {
let p = pcols[r];
let b = mixed[p * in_rows + cd + nv * dv + h];
let a = mixed[p * in_rows + cd + nv * dv + nv + h];
let g = -exp(gpar[h]) * softplus(a + gpar[nv + h]);
lg_run = lg_run + g;
lgs[r] = lg_run;
bets[r] = 1.0 / (1.0 + exp(-b));
}
}
workgroupBarrier();
{
let r = j / 8u;
let sub = j % 8u;
var sq = 0.0;
var sk = 0.0;
if (r < cc) {
for (var i = sub; i < dk; i = i + 8u) {
let vq = qsh[r * dk + i];
let vk = ksh[r * dk + i];
sq = sq + vq * vq;
sk = sk + vk * vk;
}
}
oval[j] = sq;
oval[128u + j] = sk;
}
workgroupBarrier();
if (j < 2u * CMAX) {
let r = j % CMAX;
let base = (j / CMAX) * 128u + r * 8u;
var t = 0.0;
for (var e = 0u; e < 8u; e = e + 1u) { t = t + oval[base + e]; }
invs[j] = 1.0 / sqrt(t + 1e-6);
}
workgroupBarrier();
{
let scale = 1.0 / sqrt(f32(dk));
for (var r = 0u; r < cc; r = r + 1u) {
let iq = invs[r] * scale;
let ik = invs[CMAX + r];
for (var i = j; i < dk; i = i + 128u) {
qsh[r * dk + i] = qsh[r * dk + i] * iq;
ksh[r * dk + i] = ksh[r * dk + i] * ik;
}
}
}
workgroupBarrier();
for (var pp = j; pp < cc * cc; pp = pp + 128u) {
let r = pp / cc;
let i = pp % cc;
if (i <= r) {
var kk_d = 0.0;
var qk_d = 0.0;
for (var d = 0u; d < dk; d = d + 1u) {
kk_d = kk_d + ksh[i * dk + d] * ksh[r * dk + d];
qk_d = qk_d + ksh[i * dk + d] * qsh[r * dk + d];
}
let ratio = exp(lgs[r] - lgs[i]);
if (i < r) { amat[r * CMAX + i] = ratio * kk_d; }
qkm[r * CMAX + i] = ratio * qk_d;
}
}
workgroupBarrier();
let sbase = ((my_slot * nv + h) * dk) * dv + jg;
var acck0 = 0.0;
var accq0 = 0.0;
var acck1 = 0.0;
var accq1 = 0.0;
var acck2 = 0.0;
var accq2 = 0.0;
var acck3 = 0.0;
var accq3 = 0.0;
var acck4 = 0.0;
var accq4 = 0.0;
var acck5 = 0.0;
var accq5 = 0.0;
var acck6 = 0.0;
var accq6 = 0.0;
var acck7 = 0.0;
var accq7 = 0.0;
var acck8 = 0.0;
var accq8 = 0.0;
var acck9 = 0.0;
var accq9 = 0.0;
var acck10 = 0.0;
var accq10 = 0.0;
var acck11 = 0.0;
var accq11 = 0.0;
var acck12 = 0.0;
var accq12 = 0.0;
var acck13 = 0.0;
var accq13 = 0.0;
var acck14 = 0.0;
var accq14 = 0.0;
var acck15 = 0.0;
var accq15 = 0.0;
if (lane) {
for (var i = 0u; i < dk; i = i + 1u) {
let s = st[sbase + i * dv];
acck0 = acck0 + ksh[0u * dk + i] * s;
accq0 = accq0 + qsh[0u * dk + i] * s;
acck1 = acck1 + ksh[1u * dk + i] * s;
accq1 = accq1 + qsh[1u * dk + i] * s;
acck2 = acck2 + ksh[2u * dk + i] * s;
accq2 = accq2 + qsh[2u * dk + i] * s;
acck3 = acck3 + ksh[3u * dk + i] * s;
accq3 = accq3 + qsh[3u * dk + i] * s;
acck4 = acck4 + ksh[4u * dk + i] * s;
accq4 = accq4 + qsh[4u * dk + i] * s;
acck5 = acck5 + ksh[5u * dk + i] * s;
accq5 = accq5 + qsh[5u * dk + i] * s;
acck6 = acck6 + ksh[6u * dk + i] * s;
accq6 = accq6 + qsh[6u * dk + i] * s;
acck7 = acck7 + ksh[7u * dk + i] * s;
accq7 = accq7 + qsh[7u * dk + i] * s;
acck8 = acck8 + ksh[8u * dk + i] * s;
accq8 = accq8 + qsh[8u * dk + i] * s;
acck9 = acck9 + ksh[9u * dk + i] * s;
accq9 = accq9 + qsh[9u * dk + i] * s;
acck10 = acck10 + ksh[10u * dk + i] * s;
accq10 = accq10 + qsh[10u * dk + i] * s;
acck11 = acck11 + ksh[11u * dk + i] * s;
accq11 = accq11 + qsh[11u * dk + i] * s;
acck12 = acck12 + ksh[12u * dk + i] * s;
accq12 = accq12 + qsh[12u * dk + i] * s;
acck13 = acck13 + ksh[13u * dk + i] * s;
accq13 = accq13 + qsh[13u * dk + i] * s;
acck14 = acck14 + ksh[14u * dk + i] * s;
accq14 = accq14 + qsh[14u * dk + i] * s;
acck15 = acck15 + ksh[15u * dk + i] * s;
accq15 = accq15 + qsh[15u * dk + i] * s;
}
}
if (lane) {
if (0u < cc) {
var u = conved[pcols[0u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[0u]) * acck0;
ush[0u * dv + j] = u * bets[0u];
}
if (1u < cc) {
var u = conved[pcols[1u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[1u]) * acck1;
u = u - amat[16u] * ush[0u * dv + j];
ush[1u * dv + j] = u * bets[1u];
}
if (2u < cc) {
var u = conved[pcols[2u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[2u]) * acck2;
u = u - amat[32u] * ush[0u * dv + j];
u = u - amat[33u] * ush[1u * dv + j];
ush[2u * dv + j] = u * bets[2u];
}
if (3u < cc) {
var u = conved[pcols[3u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[3u]) * acck3;
u = u - amat[48u] * ush[0u * dv + j];
u = u - amat[49u] * ush[1u * dv + j];
u = u - amat[50u] * ush[2u * dv + j];
ush[3u * dv + j] = u * bets[3u];
}
if (4u < cc) {
var u = conved[pcols[4u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[4u]) * acck4;
u = u - amat[64u] * ush[0u * dv + j];
u = u - amat[65u] * ush[1u * dv + j];
u = u - amat[66u] * ush[2u * dv + j];
u = u - amat[67u] * ush[3u * dv + j];
ush[4u * dv + j] = u * bets[4u];
}
if (5u < cc) {
var u = conved[pcols[5u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[5u]) * acck5;
u = u - amat[80u] * ush[0u * dv + j];
u = u - amat[81u] * ush[1u * dv + j];
u = u - amat[82u] * ush[2u * dv + j];
u = u - amat[83u] * ush[3u * dv + j];
u = u - amat[84u] * ush[4u * dv + j];
ush[5u * dv + j] = u * bets[5u];
}
if (6u < cc) {
var u = conved[pcols[6u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[6u]) * acck6;
u = u - amat[96u] * ush[0u * dv + j];
u = u - amat[97u] * ush[1u * dv + j];
u = u - amat[98u] * ush[2u * dv + j];
u = u - amat[99u] * ush[3u * dv + j];
u = u - amat[100u] * ush[4u * dv + j];
u = u - amat[101u] * ush[5u * dv + j];
ush[6u * dv + j] = u * bets[6u];
}
if (7u < cc) {
var u = conved[pcols[7u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[7u]) * acck7;
u = u - amat[112u] * ush[0u * dv + j];
u = u - amat[113u] * ush[1u * dv + j];
u = u - amat[114u] * ush[2u * dv + j];
u = u - amat[115u] * ush[3u * dv + j];
u = u - amat[116u] * ush[4u * dv + j];
u = u - amat[117u] * ush[5u * dv + j];
u = u - amat[118u] * ush[6u * dv + j];
ush[7u * dv + j] = u * bets[7u];
}
if (8u < cc) {
var u = conved[pcols[8u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[8u]) * acck8;
u = u - amat[128u] * ush[0u * dv + j];
u = u - amat[129u] * ush[1u * dv + j];
u = u - amat[130u] * ush[2u * dv + j];
u = u - amat[131u] * ush[3u * dv + j];
u = u - amat[132u] * ush[4u * dv + j];
u = u - amat[133u] * ush[5u * dv + j];
u = u - amat[134u] * ush[6u * dv + j];
u = u - amat[135u] * ush[7u * dv + j];
ush[8u * dv + j] = u * bets[8u];
}
if (9u < cc) {
var u = conved[pcols[9u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[9u]) * acck9;
u = u - amat[144u] * ush[0u * dv + j];
u = u - amat[145u] * ush[1u * dv + j];
u = u - amat[146u] * ush[2u * dv + j];
u = u - amat[147u] * ush[3u * dv + j];
u = u - amat[148u] * ush[4u * dv + j];
u = u - amat[149u] * ush[5u * dv + j];
u = u - amat[150u] * ush[6u * dv + j];
u = u - amat[151u] * ush[7u * dv + j];
u = u - amat[152u] * ush[8u * dv + j];
ush[9u * dv + j] = u * bets[9u];
}
if (10u < cc) {
var u = conved[pcols[10u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[10u]) * acck10;
u = u - amat[160u] * ush[0u * dv + j];
u = u - amat[161u] * ush[1u * dv + j];
u = u - amat[162u] * ush[2u * dv + j];
u = u - amat[163u] * ush[3u * dv + j];
u = u - amat[164u] * ush[4u * dv + j];
u = u - amat[165u] * ush[5u * dv + j];
u = u - amat[166u] * ush[6u * dv + j];
u = u - amat[167u] * ush[7u * dv + j];
u = u - amat[168u] * ush[8u * dv + j];
u = u - amat[169u] * ush[9u * dv + j];
ush[10u * dv + j] = u * bets[10u];
}
if (11u < cc) {
var u = conved[pcols[11u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[11u]) * acck11;
u = u - amat[176u] * ush[0u * dv + j];
u = u - amat[177u] * ush[1u * dv + j];
u = u - amat[178u] * ush[2u * dv + j];
u = u - amat[179u] * ush[3u * dv + j];
u = u - amat[180u] * ush[4u * dv + j];
u = u - amat[181u] * ush[5u * dv + j];
u = u - amat[182u] * ush[6u * dv + j];
u = u - amat[183u] * ush[7u * dv + j];
u = u - amat[184u] * ush[8u * dv + j];
u = u - amat[185u] * ush[9u * dv + j];
u = u - amat[186u] * ush[10u * dv + j];
ush[11u * dv + j] = u * bets[11u];
}
if (12u < cc) {
var u = conved[pcols[12u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[12u]) * acck12;
u = u - amat[192u] * ush[0u * dv + j];
u = u - amat[193u] * ush[1u * dv + j];
u = u - amat[194u] * ush[2u * dv + j];
u = u - amat[195u] * ush[3u * dv + j];
u = u - amat[196u] * ush[4u * dv + j];
u = u - amat[197u] * ush[5u * dv + j];
u = u - amat[198u] * ush[6u * dv + j];
u = u - amat[199u] * ush[7u * dv + j];
u = u - amat[200u] * ush[8u * dv + j];
u = u - amat[201u] * ush[9u * dv + j];
u = u - amat[202u] * ush[10u * dv + j];
u = u - amat[203u] * ush[11u * dv + j];
ush[12u * dv + j] = u * bets[12u];
}
if (13u < cc) {
var u = conved[pcols[13u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[13u]) * acck13;
u = u - amat[208u] * ush[0u * dv + j];
u = u - amat[209u] * ush[1u * dv + j];
u = u - amat[210u] * ush[2u * dv + j];
u = u - amat[211u] * ush[3u * dv + j];
u = u - amat[212u] * ush[4u * dv + j];
u = u - amat[213u] * ush[5u * dv + j];
u = u - amat[214u] * ush[6u * dv + j];
u = u - amat[215u] * ush[7u * dv + j];
u = u - amat[216u] * ush[8u * dv + j];
u = u - amat[217u] * ush[9u * dv + j];
u = u - amat[218u] * ush[10u * dv + j];
u = u - amat[219u] * ush[11u * dv + j];
u = u - amat[220u] * ush[12u * dv + j];
ush[13u * dv + j] = u * bets[13u];
}
if (14u < cc) {
var u = conved[pcols[14u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[14u]) * acck14;
u = u - amat[224u] * ush[0u * dv + j];
u = u - amat[225u] * ush[1u * dv + j];
u = u - amat[226u] * ush[2u * dv + j];
u = u - amat[227u] * ush[3u * dv + j];
u = u - amat[228u] * ush[4u * dv + j];
u = u - amat[229u] * ush[5u * dv + j];
u = u - amat[230u] * ush[6u * dv + j];
u = u - amat[231u] * ush[7u * dv + j];
u = u - amat[232u] * ush[8u * dv + j];
u = u - amat[233u] * ush[9u * dv + j];
u = u - amat[234u] * ush[10u * dv + j];
u = u - amat[235u] * ush[11u * dv + j];
u = u - amat[236u] * ush[12u * dv + j];
u = u - amat[237u] * ush[13u * dv + j];
ush[14u * dv + j] = u * bets[14u];
}
if (15u < cc) {
var u = conved[pcols[15u] * cd + 2u * nk * dk + h * dv + jg] - exp(lgs[15u]) * acck15;
u = u - amat[240u] * ush[0u * dv + j];
u = u - amat[241u] * ush[1u * dv + j];
u = u - amat[242u] * ush[2u * dv + j];
u = u - amat[243u] * ush[3u * dv + j];
u = u - amat[244u] * ush[4u * dv + j];
u = u - amat[245u] * ush[5u * dv + j];
u = u - amat[246u] * ush[6u * dv + j];
u = u - amat[247u] * ush[7u * dv + j];
u = u - amat[248u] * ush[8u * dv + j];
u = u - amat[249u] * ush[9u * dv + j];
u = u - amat[250u] * ush[10u * dv + j];
u = u - amat[251u] * ush[11u * dv + j];
u = u - amat[252u] * ush[12u * dv + j];
u = u - amat[253u] * ush[13u * dv + j];
u = u - amat[254u] * ush[14u * dv + j];
ush[15u * dv + j] = u * bets[15u];
}
}
workgroupBarrier();
// P5 phase 0: raw o out + per-half Σo² partials (norm applied by ONORM).
{
if (0u < cc && lane) {
var o = exp(lgs[0u]) * accq0;
o = o + qkm[0u] * ush[0u * dv + j];
core[pcols[0u] * nv * dv + h * dv + jg] = o;
oval[0u * 128u + j] = o * o;
} else if (j < 128u) { oval[0u * 128u + j] = 0.0; }
if (1u < cc && lane) {
var o = exp(lgs[1u]) * accq1;
o = o + qkm[16u] * ush[0u * dv + j];
o = o + qkm[17u] * ush[1u * dv + j];
core[pcols[1u] * nv * dv + h * dv + jg] = o;
oval[1u * 128u + j] = o * o;
} else if (j < 128u) { oval[1u * 128u + j] = 0.0; }
if (2u < cc && lane) {
var o = exp(lgs[2u]) * accq2;
o = o + qkm[32u] * ush[0u * dv + j];
o = o + qkm[33u] * ush[1u * dv + j];
o = o + qkm[34u] * ush[2u * dv + j];
core[pcols[2u] * nv * dv + h * dv + jg] = o;
oval[2u * 128u + j] = o * o;
} else if (j < 128u) { oval[2u * 128u + j] = 0.0; }
if (3u < cc && lane) {
var o = exp(lgs[3u]) * accq3;
o = o + qkm[48u] * ush[0u * dv + j];
o = o + qkm[49u] * ush[1u * dv + j];
o = o + qkm[50u] * ush[2u * dv + j];
o = o + qkm[51u] * ush[3u * dv + j];
core[pcols[3u] * nv * dv + h * dv + jg] = o;
oval[3u * 128u + j] = o * o;
} else if (j < 128u) { oval[3u * 128u + j] = 0.0; }
if (4u < cc && lane) {
var o = exp(lgs[4u]) * accq4;
o = o + qkm[64u] * ush[0u * dv + j];
o = o + qkm[65u] * ush[1u * dv + j];
o = o + qkm[66u] * ush[2u * dv + j];
o = o + qkm[67u] * ush[3u * dv + j];
o = o + qkm[68u] * ush[4u * dv + j];
core[pcols[4u] * nv * dv + h * dv + jg] = o;
oval[4u * 128u + j] = o * o;
} else if (j < 128u) { oval[4u * 128u + j] = 0.0; }
if (5u < cc && lane) {
var o = exp(lgs[5u]) * accq5;
o = o + qkm[80u] * ush[0u * dv + j];
o = o + qkm[81u] * ush[1u * dv + j];
o = o + qkm[82u] * ush[2u * dv + j];
o = o + qkm[83u] * ush[3u * dv + j];
o = o + qkm[84u] * ush[4u * dv + j];
o = o + qkm[85u] * ush[5u * dv + j];
core[pcols[5u] * nv * dv + h * dv + jg] = o;
oval[5u * 128u + j] = o * o;
} else if (j < 128u) { oval[5u * 128u + j] = 0.0; }
if (6u < cc && lane) {
var o = exp(lgs[6u]) * accq6;
o = o + qkm[96u] * ush[0u * dv + j];
o = o + qkm[97u] * ush[1u * dv + j];
o = o + qkm[98u] * ush[2u * dv + j];
o = o + qkm[99u] * ush[3u * dv + j];
o = o + qkm[100u] * ush[4u * dv + j];
o = o + qkm[101u] * ush[5u * dv + j];
o = o + qkm[102u] * ush[6u * dv + j];
core[pcols[6u] * nv * dv + h * dv + jg] = o;
oval[6u * 128u + j] = o * o;
} else if (j < 128u) { oval[6u * 128u + j] = 0.0; }
if (7u < cc && lane) {
var o = exp(lgs[7u]) * accq7;
o = o + qkm[112u] * ush[0u * dv + j];
o = o + qkm[113u] * ush[1u * dv + j];
o = o + qkm[114u] * ush[2u * dv + j];
o = o + qkm[115u] * ush[3u * dv + j];
o = o + qkm[116u] * ush[4u * dv + j];
o = o + qkm[117u] * ush[5u * dv + j];
o = o + qkm[118u] * ush[6u * dv + j];
o = o + qkm[119u] * ush[7u * dv + j];
core[pcols[7u] * nv * dv + h * dv + jg] = o;
oval[7u * 128u + j] = o * o;
} else if (j < 128u) { oval[7u * 128u + j] = 0.0; }
workgroupBarrier();
if (j < 8u) {
if ((0u * 8u + j) < cc) {
var ms = 0.0;
for (var e = 0u; e < half_len; e = e + 1u) { ms = ms + oval[j * 128u + e]; }
opart[(pcols[0u * 8u + j] * nv + h) * 2u + half] = ms;
}
}
workgroupBarrier();
}
// P5 phase 1: raw o out + per-half Σo² partials (norm applied by ONORM).
{
if (8u < cc && lane) {
var o = exp(lgs[8u]) * accq8;
o = o + qkm[128u] * ush[0u * dv + j];
o = o + qkm[129u] * ush[1u * dv + j];
o = o + qkm[130u] * ush[2u * dv + j];
o = o + qkm[131u] * ush[3u * dv + j];
o = o + qkm[132u] * ush[4u * dv + j];
o = o + qkm[133u] * ush[5u * dv + j];
o = o + qkm[134u] * ush[6u * dv + j];
o = o + qkm[135u] * ush[7u * dv + j];
o = o + qkm[136u] * ush[8u * dv + j];
core[pcols[8u] * nv * dv + h * dv + jg] = o;
oval[0u * 128u + j] = o * o;
} else if (j < 128u) { oval[0u * 128u + j] = 0.0; }
if (9u < cc && lane) {
var o = exp(lgs[9u]) * accq9;
o = o + qkm[144u] * ush[0u * dv + j];
o = o + qkm[145u] * ush[1u * dv + j];
o = o + qkm[146u] * ush[2u * dv + j];
o = o + qkm[147u] * ush[3u * dv + j];
o = o + qkm[148u] * ush[4u * dv + j];
o = o + qkm[149u] * ush[5u * dv + j];
o = o + qkm[150u] * ush[6u * dv + j];
o = o + qkm[151u] * ush[7u * dv + j];
o = o + qkm[152u] * ush[8u * dv + j];
o = o + qkm[153u] * ush[9u * dv + j];
core[pcols[9u] * nv * dv + h * dv + jg] = o;
oval[1u * 128u + j] = o * o;
} else if (j < 128u) { oval[1u * 128u + j] = 0.0; }
if (10u < cc && lane) {
var o = exp(lgs[10u]) * accq10;
o = o + qkm[160u] * ush[0u * dv + j];
o = o + qkm[161u] * ush[1u * dv + j];
o = o + qkm[162u] * ush[2u * dv + j];
o = o + qkm[163u] * ush[3u * dv + j];
o = o + qkm[164u] * ush[4u * dv + j];
o = o + qkm[165u] * ush[5u * dv + j];
o = o + qkm[166u] * ush[6u * dv + j];
o = o + qkm[167u] * ush[7u * dv + j];
o = o + qkm[168u] * ush[8u * dv + j];
o = o + qkm[169u] * ush[9u * dv + j];
o = o + qkm[170u] * ush[10u * dv + j];
core[pcols[10u] * nv * dv + h * dv + jg] = o;
oval[2u * 128u + j] = o * o;
} else if (j < 128u) { oval[2u * 128u + j] = 0.0; }
if (11u < cc && lane) {
var o = exp(lgs[11u]) * accq11;
o = o + qkm[176u] * ush[0u * dv + j];
o = o + qkm[177u] * ush[1u * dv + j];
o = o + qkm[178u] * ush[2u * dv + j];
o = o + qkm[179u] * ush[3u * dv + j];
o = o + qkm[180u] * ush[4u * dv + j];
o = o + qkm[181u] * ush[5u * dv + j];
o = o + qkm[182u] * ush[6u * dv + j];
o = o + qkm[183u] * ush[7u * dv + j];
o = o + qkm[184u] * ush[8u * dv + j];
o = o + qkm[185u] * ush[9u * dv + j];
o = o + qkm[186u] * ush[10u * dv + j];
o = o + qkm[187u] * ush[11u * dv + j];
core[pcols[11u] * nv * dv + h * dv + jg] = o;
oval[3u * 128u + j] = o * o;
} else if (j < 128u) { oval[3u * 128u + j] = 0.0; }
if (12u < cc && lane) {
var o = exp(lgs[12u]) * accq12;
o = o + qkm[192u] * ush[0u * dv + j];
o = o + qkm[193u] * ush[1u * dv + j];
o = o + qkm[194u] * ush[2u * dv + j];
o = o + qkm[195u] * ush[3u * dv + j];
o = o + qkm[196u] * ush[4u * dv + j];
o = o + qkm[197u] * ush[5u * dv + j];
o = o + qkm[198u] * ush[6u * dv + j];
o = o + qkm[199u] * ush[7u * dv + j];
o = o + qkm[200u] * ush[8u * dv + j];
o = o + qkm[201u] * ush[9u * dv + j];
o = o + qkm[202u] * ush[10u * dv + j];
o = o + qkm[203u] * ush[11u * dv + j];
o = o + qkm[204u] * ush[12u * dv + j];
core[pcols[12u] * nv * dv + h * dv + jg] = o;
oval[4u * 128u + j] = o * o;
} else if (j < 128u) { oval[4u * 128u + j] = 0.0; }
if (13u < cc && lane) {
var o = exp(lgs[13u]) * accq13;
o = o + qkm[208u] * ush[0u * dv + j];
o = o + qkm[209u] * ush[1u * dv + j];
o = o + qkm[210u] * ush[2u * dv + j];
o = o + qkm[211u] * ush[3u * dv + j];
o = o + qkm[212u] * ush[4u * dv + j];
o = o + qkm[213u] * ush[5u * dv + j];
o = o + qkm[214u] * ush[6u * dv + j];
o = o + qkm[215u] * ush[7u * dv + j];
o = o + qkm[216u] * ush[8u * dv + j];
o = o + qkm[217u] * ush[9u * dv + j];
o = o + qkm[218u] * ush[10u * dv + j];
o = o + qkm[219u] * ush[11u * dv + j];
o = o + qkm[220u] * ush[12u * dv + j];
o = o + qkm[221u] * ush[13u * dv + j];
core[pcols[13u] * nv * dv + h * dv + jg] = o;
oval[5u * 128u + j] = o * o;
} else if (j < 128u) { oval[5u * 128u + j] = 0.0; }
if (14u < cc && lane) {
var o = exp(lgs[14u]) * accq14;
o = o + qkm[224u] * ush[0u * dv + j];
o = o + qkm[225u] * ush[1u * dv + j];
o = o + qkm[226u] * ush[2u * dv + j];
o = o + qkm[227u] * ush[3u * dv + j];
o = o + qkm[228u] * ush[4u * dv + j];
o = o + qkm[229u] * ush[5u * dv + j];
o = o + qkm[230u] * ush[6u * dv + j];
o = o + qkm[231u] * ush[7u * dv + j];
o = o + qkm[232u] * ush[8u * dv + j];
o = o + qkm[233u] * ush[9u * dv + j];
o = o + qkm[234u] * ush[10u * dv + j];
o = o + qkm[235u] * ush[11u * dv + j];
o = o + qkm[236u] * ush[12u * dv + j];
o = o + qkm[237u] * ush[13u * dv + j];
o = o + qkm[238u] * ush[14u * dv + j];
core[pcols[14u] * nv * dv + h * dv + jg] = o;
oval[6u * 128u + j] = o * o;
} else if (j < 128u) { oval[6u * 128u + j] = 0.0; }
if (15u < cc && lane) {
var o = exp(lgs[15u]) * accq15;
o = o + qkm[240u] * ush[0u * dv + j];
o = o + qkm[241u] * ush[1u * dv + j];
o = o + qkm[242u] * ush[2u * dv + j];
o = o + qkm[243u] * ush[3u * dv + j];
o = o + qkm[244u] * ush[4u * dv + j];
o = o + qkm[245u] * ush[5u * dv + j];
o = o + qkm[246u] * ush[6u * dv + j];
o = o + qkm[247u] * ush[7u * dv + j];
o = o + qkm[248u] * ush[8u * dv + j];
o = o + qkm[249u] * ush[9u * dv + j];
o = o + qkm[250u] * ush[10u * dv + j];
o = o + qkm[251u] * ush[11u * dv + j];
o = o + qkm[252u] * ush[12u * dv + j];
o = o + qkm[253u] * ush[13u * dv + j];
o = o + qkm[254u] * ush[14u * dv + j];
o = o + qkm[255u] * ush[15u * dv + j];
core[pcols[15u] * nv * dv + h * dv + jg] = o;
oval[7u * 128u + j] = o * o;
} else if (j < 128u) { oval[7u * 128u + j] = 0.0; }
workgroupBarrier();
if (j < 8u) {
if ((1u * 8u + j) < cc) {
var ms = 0.0;
for (var e = 0u; e < half_len; e = e + 1u) { ms = ms + oval[j * 128u + e]; }
opart[(pcols[1u * 8u + j] * nv + h) * 2u + half] = ms;
}
}
workgroupBarrier();
}
if (lane) {
let lgc = lgs[cc - 1u];
let gc = exp(lgc);
for (var i = 0u; i < dk; i = i + 1u) {
var s = st[sbase + i * dv] * gc;
if (0u < cc) { s = s + exp(lgc - lgs[0u]) * ksh[0u * dk + i] * ush[0u * dv + j]; }
if (1u < cc) { s = s + exp(lgc - lgs[1u]) * ksh[1u * dk + i] * ush[1u * dv + j]; }
if (2u < cc) { s = s + exp(lgc - lgs[2u]) * ksh[2u * dk + i] * ush[2u * dv + j]; }
if (3u < cc) { s = s + exp(lgc - lgs[3u]) * ksh[3u * dk + i] * ush[3u * dv + j]; }
if (4u < cc) { s = s + exp(lgc - lgs[4u]) * ksh[4u * dk + i] * ush[4u * dv + j]; }
if (5u < cc) { s = s + exp(lgc - lgs[5u]) * ksh[5u * dk + i] * ush[5u * dv + j]; }
if (6u < cc) { s = s + exp(lgc - lgs[6u]) * ksh[6u * dk + i] * ush[6u * dv + j]; }
if (7u < cc) { s = s + exp(lgc - lgs[7u]) * ksh[7u * dk + i] * ush[7u * dv + j]; }
if (8u < cc) { s = s + exp(lgc - lgs[8u]) * ksh[8u * dk + i] * ush[8u * dv + j]; }
if (9u < cc) { s = s + exp(lgc - lgs[9u]) * ksh[9u * dk + i] * ush[9u * dv + j]; }
if (10u < cc) { s = s + exp(lgc - lgs[10u]) * ksh[10u * dk + i] * ush[10u * dv + j]; }
if (11u < cc) { s = s + exp(lgc - lgs[11u]) * ksh[11u * dk + i] * ush[11u * dv + j]; }
if (12u < cc) { s = s + exp(lgc - lgs[12u]) * ksh[12u * dk + i] * ush[12u * dv + j]; }
if (13u < cc) { s = s + exp(lgc - lgs[13u]) * ksh[13u * dk + i] * ush[13u * dv + j]; }
if (14u < cc) { s = s + exp(lgc - lgs[14u]) * ksh[14u * dk + i] * ush[14u * dv + j]; }
if (15u < cc) { s = s + exp(lgc - lgs[15u]) * ksh[15u * dk + i] * ush[15u * dv + j]; }
st[sbase + i * dv] = s;
}
}
workgroupBarrier();
scan = pcols[cc - 1u] + 1u;
if (scan >= k) { return; }
}
}
"##
.to_string()
}
/// Deferred gated RMSNorm for the dv-split chunk kernel: `core` holds RAW o and `opart` the
/// two halves' Σo²; this applies `o·(1/√(Σo²/dv+ε))·norm_w·silu(z)` in place. Grid `(nv, k)`.
const DN_CHUNK_ONORM: &str = r#"
@group(0) @binding(0) var<storage, read_write> core: array<f32>; // [k, nv·dv] RAW → normed
@group(0) @binding(1) var<storage, read> mixed: array<f32>; // [k, in_rows] (z tail)
@group(0) @binding(2) var<storage, read> gpar: array<f32>; // [.. | norm_w(dv)]
@group(0) @binding(3) var<storage, read> opart: array<f32>; // [k, nv, 2]
@group(0) @binding(4) var<uniform> dims: vec4<u32>; // (nv, nk, dk, dv)
@group(0) @binding(5) var<uniform> dims2: vec4<u32>; // (in_rows, conv_dim, k, _)
@group(0) @binding(6) var<uniform> epsu: vec4<f32>;
@compute @workgroup_size(128)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let p = wid.y;
let nv = dims.x; let dv = dims.w;
let in_rows = dims2.x; let cd = dims2.y;
let j = lid.x;
if (j >= dv) { return; }
let ms = (opart[(p * nv + h) * 2u] + opart[(p * nv + h) * 2u + 1u]) / f32(dv);
let inv = 1.0 / sqrt(ms + epsu.x);
let o = core[p * nv * dv + h * dv + j];
let zz = mixed[p * in_rows + cd + h * dv + j];
core[p * nv * dv + h * dv + j] = o * inv * gpar[2u * nv + j] * (zz / (1.0 + exp(-zz)));
}
"#;
/// Elementwise `y = y · sigmoid(g)` — the Qwen3.5 attention output gate. Grid (n/64, k).
/// State-RESIDENT DeltaNet step for the small-batch verify band (k ≤ 8, same-slot chains):
/// the recurrent state S (dk×dv f32 = 64 KB/head) is split across TWO workgroups by dv-half
/// (32 KB shared each), loaded ONCE, updated in shared across all same-slot columns, stored
/// ONCE — vs the global-memory variant's 2-read+1-write per column (30 MB/layer at 5 cols).
/// The output is written RAW; [`DN_ONORM`] applies the g-RMSNorm + z-gate in a follow-up
/// dispatch (the norm needs all dv lanes, which now live in two WGs). All reduction orders
/// (q/k norms, per-i state walks) are kept EXACTLY as [`DN_STEP_PK`] — bitwise by construction.
const DN_STEP_RES: &str = r#"
@group(0) @binding(0) var<storage, read> conved: array<f32>; // [k, conv_dim] = q|k|v
@group(0) @binding(1) var<storage, read> mixed: array<f32>; // [k, in_rows] (z|b|a tail)
@group(0) @binding(2) var<storage, read> gpar: array<f32>;
@group(0) @binding(3) var<storage, read_write> st: array<f32>; // [slots, nv, dk, dv]
@group(0) @binding(4) var<storage, read_write> core: array<f32>; // [k, nv·dv] RAW o
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (nv, nk, dk, dv)
@group(0) @binding(8) var<uniform> dims2: vec4<u32>; // (in_rows, conv_dim, k, stride)
@group(0) @binding(9) var<uniform> epsu: vec4<f32>;
var<workgroup> S: array<f32, 8192>; // [dk, 64] resident half-state
var<workgroup> qsh: array<f32, 256>;
var<workgroup> ksh: array<f32, 256>;
var<workgroup> scal: array<f32, 3>;
var<workgroup> red2: array<f32, 256>;
fn softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
if (x < -20.0) { return exp(x); }
return log(1.0 + exp(x));
}
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x / 2u;
let half = wid.x % 2u;
let nv = dims.x; let nk = dims.y; let dk = dims.z; let dv = dims.w;
let in_rows = dims2.x; let cd = dims2.y; let k = dims2.z;
let kh = h / (nv / nk);
let jl = lid.x;
let j = half * 64u + jl;
let p0 = wid.y;
let my_slot = cmeta[(cnt[0] * dims2.w + p0) * 2u + 1u];
for (var q = 0u; q < p0; q = q + 1u) {
if (cmeta[(cnt[0] * dims2.w + q) * 2u + 1u] == my_slot) { return; }
}
let sbase = ((my_slot * nv + h) * dk) * dv + j;
if (j < dv) {
for (var i = 0u; i < dk; i = i + 1u) { S[i * 64u + jl] = st[sbase + i * dv]; }
}
for (var p = p0; p < k; p = p + 1u) {
let slot = cmeta[(cnt[0] * dims2.w + p) * 2u + 1u];
if (slot != my_slot) { continue; }
let cbase = p * cd;
for (var i = jl; i < dk; i = i + 64u) {
qsh[i] = conved[cbase + kh * dk + i];
ksh[i] = conved[cbase + nk * dk + kh * dk + i];
}
workgroupBarrier();
if (jl == 0u) {
var sq = 0.0;
var sk = 0.0;
for (var i = 0u; i < dk; i = i + 1u) { sq = sq + qsh[i] * qsh[i]; }
for (var i = 0u; i < dk; i = i + 1u) { sk = sk + ksh[i] * ksh[i]; }
let invq = 1.0 / sqrt(sq + 1e-6);
let invk = 1.0 / sqrt(sk + 1e-6);
let scale = 1.0 / sqrt(f32(dk));
for (var i = 0u; i < dk; i = i + 1u) {
qsh[i] = qsh[i] * invq * scale;
ksh[i] = ksh[i] * invk;
}
let b = mixed[p * in_rows + cd + nv * dv + h];
let a = mixed[p * in_rows + cd + nv * dv + nv + h];
let g = -exp(gpar[h]) * softplus(a + gpar[nv + h]);
scal[0] = exp(g);
scal[1] = 1.0 / (1.0 + exp(-b));
}
workgroupBarrier();
if (j < dv) {
let decay = scal[0];
var kv = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
kv = kv + ksh[i] * (S[i * 64u + jl] * decay);
}
let vv = conved[cbase + 2u * nk * dk + h * dv + j];
let delta = (vv - kv) * scal[1];
var o = 0.0;
for (var i = 0u; i < dk; i = i + 1u) {
let sv = S[i * 64u + jl] * decay + ksh[i] * delta;
S[i * 64u + jl] = sv;
o = o + qsh[i] * sv;
}
core[p * nv * dv + h * dv + j] = o;
}
workgroupBarrier();
}
if (j < dv) {
for (var i = 0u; i < dk; i = i + 1u) { st[sbase + i * dv] = S[i * 64u + jl]; }
}
}
"#;
/// The g-RMSNorm + z-gate epilogue split out of [`DN_STEP_RES`] (needs all dv lanes of a head
/// in one group). Reads RAW o from `core` and rewrites it in place — the ms scan keeps the
/// sequential 0..dv order of [`DN_STEP_PK`], so results are bitwise-identical. Grid (nv, k).
const DN_ONORM: &str = r#"
@group(0) @binding(0) var<storage, read> mixed: array<f32>; // [k, in_rows]
@group(0) @binding(1) var<storage, read> gpar: array<f32>;
@group(0) @binding(2) var<storage, read_write> core: array<f32>; // [k, nv·dv] RAW → normed
@group(0) @binding(3) var<uniform> dims: vec4<u32>; // (nv, in_rows, cd, dv)
@group(0) @binding(4) var<uniform> epsu: vec4<f32>;
var<workgroup> osh: array<f32, 128>;
var<workgroup> scal: array<f32, 1>;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let p = wid.y;
let nv = dims.x; let dv = dims.w;
let in_rows = dims.y; let cd = dims.z;
for (var j = lid.x; j < dv; j = j + 64u) {
osh[j] = core[p * nv * dv + h * dv + j];
}
workgroupBarrier();
if (lid.x == 0u) {
var ms = 0.0;
for (var jj = 0u; jj < dv; jj = jj + 1u) { ms = ms + osh[jj] * osh[jj]; }
scal[0] = 1.0 / sqrt(ms / f32(dv) + epsu.x);
}
workgroupBarrier();
for (var j = lid.x; j < dv; j = j + 64u) {
let zz = mixed[p * in_rows + cd + h * dv + j];
core[p * nv * dv + h * dv + j] =
osh[j] * scal[0] * gpar[2u * nv + j] * (zz / (1.0 + exp(-zz)));
}
}
"#;
const GATE_MUL_K: &str = r#"
@group(0) @binding(0) var<storage, read> g: array<f32>; // [k, n]
@group(0) @binding(1) var<storage, read_write> y: array<f32>; // [k, n]
@group(0) @binding(2) var<uniform> dims: vec4<u32>; // (n, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = dims.x; let i = gid.x;
if (i >= n) { return; }
let o = gid.y * n + i;
y[o] = y[o] * (1.0 / (1.0 + exp(-g[o])));
}
"#;
/// Shared-expert gate: per column, `s = sigmoid(w · xn)` then `y *= s` — one workgroup per
/// column (reduce 256-wide, then scale). `dims=(hidden, _, _, _)`.
const SGATE_SCALE_K: &str = r#"
@group(0) @binding(0) var<storage, read> w: array<f32>; // [hidden]
@group(0) @binding(1) var<storage, read> xn: array<f32>; // [k, hidden] (ffn-normed)
@group(0) @binding(2) var<storage, read_write> y: array<f32>; // [k, hidden] (shared-expert out)
@group(0) @binding(3) var<uniform> dims: vec4<u32>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = dims.x; let col = wid.x;
var sacc = 0.0;
for (var i = lid.x; i < h; i = i + 256u) { sacc = sacc + w[i] * xn[col * h + i]; }
red[lid.x] = sacc;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lid.x < st) { red[lid.x] = red[lid.x] + red[lid.x + st]; }
workgroupBarrier();
}
let sg = 1.0 / (1.0 + exp(-red[0]));
for (var i = lid.x; i < h; i = i + 256u) { y[col * h + i] = y[col * h + i] * sg; }
}
"#;
/// Elementwise residual add `y += x` (the MoE block joins the stream AFTER gating). Grid (n/64, k).
const ADD_K: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> dims: vec4<u32>; // (n, _, _, _)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = dims.x; let i = gid.x;
if (i >= n) { return; }
let o = gid.y * n + i;
y[o] = y[o] + x[o];
}
"#;
/// Qwen3.5 FUSED MoE router: absorbs the standalone ffn-RMSNorm and the shared-expert-gate
/// dispatches (3 tiny dispatches → 1 — at small batch widths each extra dispatch costs
/// hundreds of µs of pure drain latency on V100/Vulkan, the measured single-stream wall).
/// Bitwise contract: the Σx² reduction + `x·inv·w` expression replicate the RMSNORM kernel
/// exactly, and the gate dot replicates SGATE_SCALE_K's strided+tree shape exactly, so
/// `normed`, the router selection, and `g` carry the same bits as the three-dispatch chain.
/// Router split, stage 1 — NORM + shared-gate (per-column WG; the RMSNORM/SGATE reduction
/// shapes are replicated exactly, so `normed` and `g` keep their bits).
const ROUTER_NORM_SGATE: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>; // [k, h] RAW residual
@group(0) @binding(1) var<storage, read> wn: array<f32>; // [h] ffn_norm
@group(0) @binding(2) var<storage, read> sgw: array<f32>; // [h] shared-expert gate
@group(0) @binding(3) var<storage, read_write> normed: array<f32>; // [k, h] out
@group(0) @binding(4) var<storage, read_write> g: array<f32>; // [k] out
@group(0) @binding(5) var<uniform> mp: vec4<u32>; // (_, h, _, _)
@group(0) @binding(6) var<uniform> epsf: vec4<f32>; // (eps, _, _, _)
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = mp.y;
let col = wid.x;
let lx = lid.x;
var sq = 0.0;
for (var i = lx; i < h; i = i + 256u) { let v = x[col * h + i]; sq = sq + v * v; }
red[lx] = sq;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(red[0] / f32(h) + epsf.x);
workgroupBarrier();
for (var i = lx; i < h; i = i + 256u) { normed[col * h + i] = x[col * h + i] * inv * wn[i]; }
storageBarrier();
workgroupBarrier();
var sacc = 0.0;
for (var i = lx; i < h; i = i + 256u) { sacc = sacc + sgw[i] * normed[col * h + i]; }
red[lx] = sacc;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
if (lx == 0u) { g[col] = 1.0 / (1.0 + exp(-red[0])); }
}
"#;
/// Router split, stage 2 — COALESCED logits: 16 lanes per expert, 16 experts per workgroup
/// (grid (⌈E/16⌉, k)). The old per-thread serial h-loop read `wr` with a 2048-float stride
/// between threads — ~500 µs/layer of uncoalesced traffic, the single largest kernel of the
/// small-batch verify. Lane-tree accumulation (ulp-level shift; selection ties are decided on
/// well-separated probabilities in practice and every plan uses the same kernel).
const ROUTER_LOGITS: &str = r#"
@group(0) @binding(0) var<storage, read> wr: array<vec4<f32>>; // [E, h/4] router
@group(0) @binding(1) var<storage, read> normed: array<vec4<f32>>; // [k, h/4]
@group(0) @binding(2) var<storage, read_write> logits: array<f32>; // [k, E]
@group(0) @binding(3) var<uniform> mp: vec4<u32>; // (E, h, _, _)
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let e_n = mp.x; let h4 = mp.y / 4u;
let col = wid.y;
let e = wid.x * 16u + lid.x / 16u;
let lane = lid.x % 16u;
var acc = 0.0;
if (e < e_n) {
for (var j = lane; j < h4; j = j + 16u) {
acc = acc + dot(wr[e * h4 + j], normed[col * h4 + j]);
}
}
red[lid.x] = acc;
workgroupBarrier();
if (lane < 8u) { red[lid.x] = red[lid.x] + red[lid.x + 8u]; }
workgroupBarrier();
if (lane < 4u) { red[lid.x] = red[lid.x] + red[lid.x + 4u]; }
workgroupBarrier();
if (lane < 2u) { red[lid.x] = red[lid.x] + red[lid.x + 2u]; }
workgroupBarrier();
if (lane == 0u && e < e_n) {
logits[col * e_n + e] = red[lid.x] + red[lid.x + 1u];
}
}
"#;
/// Router split, stage 3 — softmax + top-k pick from the logits buffer (the original router's
/// softmax/selection body verbatim; one WG per column).
/// Subgroup-32 router pick (NVIDIA): same softmax reduction TREE as [`ROUTER_PICK`] (bitwise-
/// identical denominator), top-K via parallel subgroup argmax with min-index tie-break — the
/// exact winner set of the sequential strict-`>` scan. Kills the 256-thread WG's serial
/// 2048-iteration top-K loop + 16 wide barriers (measured 110 µs → target ~4 µs per layer).
const ROUTER_PICK_SG32: &str = r#"
@group(0) @binding(0) var<storage, read> logits: array<f32>; // [k, E]
@group(0) @binding(1) var<storage, read> pes: array<f32>; // [E]
@group(0) @binding(2) var<storage, read_write> sel: array<u32>; // [k, K]
@group(0) @binding(3) var<storage, read_write> wsel: array<f32>; // [k, K]
@group(0) @binding(4) var<uniform> mp: vec4<u32>; // (E, _, K, _)
var<workgroup> probs: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(32)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(subgroup_invocation_id) sid: u32) {
let e_n = mp.x; let kk = mp.z;
let col = wid.x;
for (var i = sid; i < 256u; i = i + 32u) {
var logit = -1e30;
if (i < e_n) { logit = logits[col * e_n + i]; }
probs[i] = logit;
red[i] = logit;
}
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
for (var i = sid; i < st; i = i + 32u) { red[i] = max(red[i], red[i + st]); }
workgroupBarrier();
}
let mx = red[0];
workgroupBarrier();
for (var i = sid; i < 256u; i = i + 32u) {
var e_val = 0.0;
if (i < e_n) { e_val = exp(probs[i] - mx); }
// pads: -1.0 in probs (never selectable) but 0.0 in red (denominator padding).
if (i < e_n) { probs[i] = e_val; } else { probs[i] = -1.0; }
red[i] = e_val;
}
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
for (var i = sid; i < st; i = i + 32u) { red[i] = red[i] + red[i + st]; }
workgroupBarrier();
}
let den = red[0];
workgroupBarrier();
var wsum = 0.0;
for (var s = 0u; s < kk; s = s + 1u) {
// Lane-local argmax over its strided slots (ascending scan, strict > keeps the
// smallest index on ties — matching the sequential reference).
var lbv = -1.0; var lbi = 0xffffffffu;
for (var i = sid; i < 256u; i = i + 32u) {
let v = probs[i];
if (v > lbv) { lbv = v; lbi = i; }
}
let gbv = subgroupMax(lbv);
var cand = 0xffffffffu;
if (lbv == gbv) { cand = lbi; }
let gbi = subgroupMin(cand);
if (sid == 0u) {
sel[col * kk + s] = gbi;
wsel[col * kk + s] = gbv / den;
}
wsum = wsum + gbv / den;
if (gbi % 32u == sid) { probs[gbi] = -1.0; }
workgroupBarrier();
}
if (sid == 0u) {
for (var s = 0u; s < kk; s = s + 1u) {
let eid = sel[col * kk + s];
wsel[col * kk + s] = (wsel[col * kk + s] / wsum) * pes[eid];
}
}
}
"#;
const ROUTER_PICK: &str = r#"
@group(0) @binding(0) var<storage, read> logits: array<f32>; // [k, E]
@group(0) @binding(1) var<storage, read> pes: array<f32>; // [E]
@group(0) @binding(2) var<storage, read_write> sel: array<u32>; // [k, K]
@group(0) @binding(3) var<storage, read_write> wsel: array<f32>; // [k, K]
@group(0) @binding(4) var<uniform> mp: vec4<u32>; // (E, _, K, _)
var<workgroup> probs: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let e_n = mp.x; let kk = mp.z;
let col = wid.x;
let lx = lid.x;
var logit = -1e30;
if (lx < e_n) { logit = logits[col * e_n + lx]; }
probs[lx] = logit;
workgroupBarrier();
red[lx] = probs[lx];
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = max(red[lx], red[lx + st]); }
workgroupBarrier();
}
let mx = red[0];
workgroupBarrier();
var e_val = 0.0;
if (lx < e_n) { e_val = exp(probs[lx] - mx); }
probs[lx] = e_val;
red[lx] = e_val;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let den = red[0];
workgroupBarrier();
if (lx == 0u) {
var wsum = 0.0;
for (var s = 0u; s < kk; s = s + 1u) {
var bi = 0u; var bv = -1.0;
for (var e = 0u; e < e_n; e = e + 1u) {
if (probs[e] > bv) { bv = probs[e]; bi = e; }
}
sel[col * kk + s] = bi;
wsel[col * kk + s] = bv / den;
wsum = wsum + bv / den;
probs[bi] = -1.0; // consumed
}
for (var s = 0u; s < kk; s = s + 1u) {
let eid = sel[col * kk + s];
wsel[col * kk + s] = (wsel[col * kk + s] / wsum) * pes[eid];
}
}
}
"#;
const MOE_ROUTER_FUSED: &str = r#"
@group(0) @binding(0) var<storage, read> wr: array<f32>; // [E, h] router
@group(0) @binding(1) var<storage, read> pes: array<f32>; // [E] per-expert scale
@group(0) @binding(2) var<storage, read> x: array<f32>; // [k, h] RAW residual
@group(0) @binding(3) var<storage, read> wn: array<f32>; // [h] ffn_norm
@group(0) @binding(4) var<storage, read> sgw: array<f32>; // [h] shared-expert gate
@group(0) @binding(5) var<storage, read_write> sel: array<u32>; // [k, K]
@group(0) @binding(6) var<storage, read_write> wsel: array<f32>; // [k, K]
@group(0) @binding(7) var<storage, read_write> normed: array<f32>; // [k, h] out
@group(0) @binding(8) var<storage, read_write> g: array<f32>; // [k] sigmoid gate out
@group(0) @binding(9) var<uniform> mp: vec4<u32>; // (E, h, K, _)
@group(0) @binding(10) var<uniform> epsf: vec4<f32>; // (eps, _, _, _)
var<workgroup> probs: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let e_n = mp.x; let h = mp.y; let kk = mp.z;
let col = wid.x;
let lx = lid.x;
// (0) RMSNorm (exact RMSNORM shape): strided Σx² + 128→1 tree.
var sq = 0.0;
for (var i = lx; i < h; i = i + 256u) { let v = x[col * h + i]; sq = sq + v * v; }
red[lx] = sq;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(red[0] / f32(h) + epsf.x);
workgroupBarrier();
for (var i = lx; i < h; i = i + 256u) { normed[col * h + i] = x[col * h + i] * inv * wn[i]; }
storageBarrier();
workgroupBarrier();
// (1) logits: thread e computes W_r[e] · normed[col]
var logit = -1e30;
if (lx < e_n) {
var acc = 0.0;
for (var j = 0u; j < h; j = j + 1u) { acc = acc + wr[lx * h + j] * normed[col * h + j]; }
logit = acc;
}
probs[lx] = logit;
workgroupBarrier();
// (2) softmax over all E (f32): max-reduce, exp, sum-reduce
red[lx] = probs[lx];
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = max(red[lx], red[lx + st]); }
workgroupBarrier();
}
let mx = red[0];
workgroupBarrier();
var e_val = 0.0;
if (lx < e_n) { e_val = exp(probs[lx] - mx); }
probs[lx] = e_val;
red[lx] = e_val;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let den = red[0];
workgroupBarrier();
// (3) shared-expert gate (exact SGATE_SCALE_K shape): strided dot + tree → sigmoid.
var sacc = 0.0;
for (var i = lx; i < h; i = i + 256u) { sacc = sacc + sgw[i] * normed[col * h + i]; }
red[lx] = sacc;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
if (lx == 0u) { g[col] = 1.0 / (1.0 + exp(-red[0])); }
// (4) thread 0: k passes of stable argmax over the probabilities, renorm, scale.
if (lx == 0u) {
var wsum = 0.0;
for (var s = 0u; s < kk; s = s + 1u) {
var bi = 0u; var bv = -1.0;
for (var e = 0u; e < e_n; e = e + 1u) {
if (probs[e] > bv) { bv = probs[e]; bi = e; }
}
sel[col * kk + s] = bi;
wsel[col * kk + s] = bv / den;
wsum = wsum + bv / den;
probs[bi] = -1.0; // consumed
}
for (var s = 0u; s < kk; s = s + 1u) {
let eid = sel[col * kk + s];
wsel[col * kk + s] = (wsel[col * kk + s] / wsum) * pes[eid];
}
}
}
"#;
/// Grammar-masked parallel argmax stage 1: like [`ARGMAX_BLOCK`] but a 1-bit-per-token `mask` gates
/// which tokens are eligible. Keeps the logits on the GPU (no 256 KB/token readback) — the CPU only
/// uploads the ~8 KB bitmask of grammar-allowed tokens and reads back the single chosen id.
const ARGMAX_BLOCK_MASKED: &str = r#"
@group(0) @binding(0) var<storage, read> logits: array<f32>;
@group(0) @binding(1) var<storage, read_write> pv: array<f32>;
@group(0) @binding(2) var<storage, read_write> pi: array<u32>;
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (_, _, vocab, _)
@group(0) @binding(4) var<storage, read> mask: array<u32>; // 1 bit/token: grammar-allowed
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let v = posu.z;
let per = (v + 255u) / 256u;
let start = wid.x * per;
var bv = -1e30; var bi = 0u;
for (var k = lid.x; k < per; k = k + 256u) {
let idx = start + k;
if (idx < v) {
let allowed = (mask[idx >> 5u] >> (idx & 31u)) & 1u;
if (allowed == 1u && logits[idx] > bv) { bv = logits[idx]; bi = idx; }
}
}
vmax[lid.x] = bv; imax[lid.x] = bi;
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u) { pv[wid.x] = vmax[0]; pi[wid.x] = imax[0]; }
}
"#;
/// Compact the grammar's one-u32-per-token mask into a dense allowed-id list: `ids[0..count]` (order
/// arbitrary — the sparse argmax tie-breaks on the token id, not the slot). Feeds the grammar-sparse
/// lm_head, which then reads ONLY the allowed rows of the 65 536-row tied head instead of all of them.
const MASK_COMPACT: &str = r#"
@group(0) @binding(0) var<storage, read> mask: array<u32>; // 1 u32/token: !=0 ⇒ allowed
@group(0) @binding(1) var<storage, read_write> ids: array<u32>; // compacted allowed ids
@group(0) @binding(2) var<storage, read_write> smeta: array<atomic<u32>>; // [0] = running count
@group(0) @binding(3) var<uniform> posu: vec4<u32>; // (_, _, vocab, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x; if (i >= posu.z) { return; }
if (mask[i] != 0u) { let k = atomicAdd(&smeta[0], 1u); ids[k] = i; }
}
"#;
/// Turn the compacted count into the sparse lm_head's INDIRECT workgroup counts, freeze it for the
/// argmax (`smeta[4]`), and re-zero the running counter for the next token — all on-GPU, so the
/// decode loop never reads the count back. `NR` must match the sparse head's rows-per-workgroup.
const SPARSE_ARGS: &str = r#"
@group(0) @binding(0) var<storage, read_write> smeta: array<atomic<u32>>; // [0]=count [1..3]=indirect [4]=frozen
const NR: u32 = 8u;
@compute @workgroup_size(1)
fn main() {
let c = atomicLoad(&smeta[0]);
atomicStore(&smeta[4], c);
atomicStore(&smeta[1], (c + NR - 1u) / NR);
atomicStore(&smeta[2], 1u);
atomicStore(&smeta[3], 1u);
atomicStore(&smeta[0], 0u);
}
"#;
/// Argmax over the sparse (svals, ids) pairs — one workgroup, `count` entries. Tie-break is the
/// LOWEST TOKEN ID (matching the CPU oracle's argmax), which also makes the result independent of
/// the compaction order. An empty count leaves `chosen = 0` — the established dead-end sentinel
/// (token 0 is then rejected by the FSM and the CPU re-walk truncates the query).
const SPARSE_ARGMAX: &str = r#"
@group(0) @binding(0) var<storage, read> svals: array<f32>;
@group(0) @binding(1) var<storage, read> ids: array<u32>;
@group(0) @binding(2) var<storage, read> smeta: array<u32>; // [4] = frozen count
@group(0) @binding(3) var<storage, read_write> chosen: array<u32>;
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let cnt = smeta[4];
var bv = -1e30; var bi = 0u;
for (var i = lid.x; i < cnt; i = i + 256u) {
let v = svals[i]; let id = ids[i];
if (v > bv || (v == bv && id < bi)) { bv = v; bi = id; }
}
vmax[lid.x] = bv; imax[lid.x] = bi;
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) {
let ov = vmax[lid.x + s]; let oi = imax[lid.x + s];
if (ov > vmax[lid.x] || (ov == vmax[lid.x] && oi < imax[lid.x])) { vmax[lid.x] = ov; imax[lid.x] = oi; }
}
workgroupBarrier();
}
if (lid.x == 0u) { chosen[0] = imax[0]; }
}
"#;
/// Masked argmax stage 2: 1 workgroup merges the 256 block winners → `chosen[0]` (the picked token id).
const ARGMAX_MERGE_CHOSEN: &str = r#"
@group(0) @binding(0) var<storage, read> pv: array<f32>;
@group(0) @binding(1) var<storage, read> pi: array<u32>;
@group(0) @binding(2) var<storage, read_write> chosen: array<u32>;
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
vmax[lid.x] = pv[lid.x]; imax[lid.x] = pi[lid.x];
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) { if (vmax[lid.x + s] > vmax[lid.x]) { vmax[lid.x] = vmax[lid.x + s]; imax[lid.x] = imax[lid.x + s]; } }
workgroupBarrier();
}
if (lid.x == 0u) { chosen[0] = imax[0]; }
}
"#;
/// TurboQuant stage-1 KV quantization (one workgroup per (kv-head, K|V)): reads the CURRENT
/// position's f32 K/V from the cache (qkrc just wrote it), rotates by the fixed orthogonal `rot`,
/// and writes 4-bit nearest-centroid indices + the per-vector scale. `kp=(nh,nkv,hd,pos+1)`;
/// wid.x = kv_head, wid.y = 0 for K / 1 for V. Scales layout: `[pos, nkv, 2]`.
const TQ_QUANT: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> kc: array<f16>;
@group(0) @binding(1) var<storage, read> vc: array<f16>;
@group(0) @binding(2) var<storage, read> rot: array<f32>; // [hd, hd] row-major
@group(0) @binding(3) var<storage, read> cents: array<f32>; // [16]
@group(0) @binding(4) var<storage, read> sk: array<f32>; // [hd, hd] QJL sketch
@group(0) @binding(5) var<storage, read_write> kvq: array<u32>; // [T, nkv, 2, hd/8]
@group(0) @binding(6) var<storage, read_write> kvs: array<f32>; // [T, nkv, 3] (kscale, vscale, gamma)
@group(0) @binding(7) var<storage, read_write> ksig: array<u32>; // [T, nkv, hd/32]
@group(0) @binding(8) var<uniform> kp: vec4<u32>; // (nh, nkv, hd, pos+1)
var<workgroup> vsh: array<f32, 256>; // the input vector (original domain)
var<workgroup> red: array<f32, 256>; // reductions / index staging
var<workgroup> dsh: array<f32, 256>; // rotated-domain dequantized vector
var<workgroup> rsh: array<f32, 256>; // key residual (original domain)
var<workgroup> csh: array<f32, 16>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nkv = kp.y; let hd = kp.z; let pos = kp.w - 1u;
let head = wid.x; let is_v = wid.y;
let j = lid.x;
if (j < 16u) { csh[j] = cents[j]; }
var xj = 0.0;
if (j < hd) {
let src = pos * nkv * hd + head * hd + j;
if (is_v == 0u) { xj = f32(kc[src]); } else { xj = f32(vc[src]); }
}
vsh[j] = xj;
red[j] = xj * xj;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (j < st) { red[j] = red[j] + red[j + st]; }
workgroupBarrier();
}
let norm = sqrt(red[0]);
let scale = norm / sqrt(f32(hd));
var inv = 0.0;
if (scale > 0.0) { inv = 1.0 / scale; }
if (j == 0u) { kvs[(pos * nkv + head) * 3u + is_v] = scale; }
workgroupBarrier();
// Rotate + nearest centroid; stash index (red) and dequantized rotated coord (dsh).
var best = 0u;
if (j < hd) {
var y = 0.0;
for (var t = 0u; t < hd; t = t + 1u) { y = y + rot[j * hd + t] * vsh[t]; }
let z = y * inv;
var bd = 1e30;
for (var c = 0u; c < 16u; c = c + 1u) {
let dd = abs(z - csh[c]);
if (dd < bd) { bd = dd; best = c; }
}
dsh[j] = csh[best] * scale;
}
red[j] = f32(best);
workgroupBarrier();
if (j < hd && (j % 8u) == 0u) {
var word = 0u;
for (var b = 0u; b < 8u; b = b + 1u) {
word = word | ((u32(red[j + b]) & 0xFu) << (b * 4u));
}
kvq[((pos * nkv + head) * 2u + is_v) * (hd / 8u) + j / 8u] = word;
}
if (is_v == 1u) { return; }
// ── Stage-2 for KEYS: residual r = x − Πᵀ·d̂ (original domain), γ = ‖r‖, QJL sign bits. ──
workgroupBarrier();
var rt = 0.0;
if (j < hd) {
var back = 0.0;
for (var u = 0u; u < hd; u = u + 1u) { back = back + rot[u * hd + j] * dsh[u]; }
rt = vsh[j] - back;
}
rsh[j] = rt;
red[j] = rt * rt;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (j < st) { red[j] = red[j] + red[j + st]; }
workgroupBarrier();
}
if (j == 0u) { kvs[(pos * nkv + head) * 3u + 2u] = sqrt(red[0]); }
workgroupBarrier();
var bit = 0.0;
if (j < hd) {
var sv = 0.0;
for (var t = 0u; t < hd; t = t + 1u) { sv = sv + sk[j * hd + t] * rsh[t]; }
if (sv >= 0.0) { bit = 1.0; }
}
red[j] = bit;
workgroupBarrier();
if (j < hd && (j % 32u) == 0u) {
var word = 0u;
for (var b = 0u; b < 32u; b = b + 1u) {
if (red[j + b] != 0.0) { word = word | (1u << b); }
}
ksig[(pos * nkv + head) * (hd / 32u) + j / 32u] = word;
}
}
"#;
/// TurboQuant attention (stage-1 K and V): rotate the QUERY once per head (orthogonality keeps
/// q·k exact in the rotated domain), score against dequantized 4-bit K coordinates, accumulate the
/// output in the ROTATED domain from 4-bit V, then un-rotate once. Bindings mirror [`ATTN_K`]
/// plus (rot, cents, kvq, kvs); same sliding-window/scale flags in `lw`.
const ATTN_TQ: &str = r#"
@group(0) @binding(0) var<storage, read> q: array<f32>; // [k, nh*hd]
@group(0) @binding(1) var<storage, read> kvq: array<u32>; // [T, nkv, 2, hd/8]
@group(0) @binding(2) var<storage, read> kvs: array<f32>; // [T, nkv, 3]
@group(0) @binding(3) var<storage, read> rot: array<f32>; // [hd, hd]
@group(0) @binding(4) var<storage, read> cents: array<f32>; // [16]
@group(0) @binding(5) var<storage, read> sk: array<f32>; // [hd, hd] QJL sketch
@group(0) @binding(6) var<storage, read> ksig: array<u32>; // [T, nkv, hd/32]
@group(0) @binding(7) var<storage, read_write> out: array<f32>; // [k, nh*hd]
@group(0) @binding(8) var<uniform> kp: vec4<u32>; // (nh, nkv, hd, pos0+1)
@group(0) @binding(9) var<uniform> lw: vec4<u32>; // (window, scale_one, _, _)
var<workgroup> sc: array<f32, 512>;
var<workgroup> red: array<f32, 256>;
var<workgroup> qsh: array<f32, 256>; // original-domain query
var<workgroup> rq: array<f32, 256>; // Π·q (rotated query — exact vs quantized-K MSE part)
var<workgroup> sq: array<f32, 256>; // S·q (QJL-domain query — residual correction)
var<workgroup> acc: array<f32, 256>; // rotated-domain output accumulator
var<workgroup> csh: array<f32, 16>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let col = wid.y;
let t = kp.w + col;
var s0 = 0u;
if (lw.x != 0u && t > lw.x) { s0 = t - lw.x; }
let span = t - s0;
let qh = wid.x;
let kv = qh / (nh / nkv);
let qbase = col * nh * hd + qh * hd;
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) { scaling = 1.0; }
// √(π/2)/d — the unbiased QJL coefficient. lw.z: 0 = off (MSE-only debug), 1 = +, 2 = −.
var coef = sqrt(1.5707963267948966) / f32(hd);
if (lw.z == 0u) { coef = 0.0; }
if (lw.z == 2u) { coef = -coef; }
let lx = lid.x;
if (lx < 16u) { csh[lx] = cents[lx]; }
var qj = 0.0;
if (lx < hd) { qj = q[qbase + lx]; }
qsh[lx] = qj;
workgroupBarrier();
var a1 = 0.0;
var a2 = 0.0;
if (lx < hd) {
for (var u = 0u; u < hd; u = u + 1u) {
a1 = a1 + rot[lx * hd + u] * qsh[u];
a2 = a2 + sk[lx * hd + u] * qsh[u];
}
}
rq[lx] = a1;
sq[lx] = a2;
acc[lx] = 0.0;
workgroupBarrier();
var lmax = -1e30;
for (var s = lx; s < span; s = s + 256u) {
let pos = s0 + s;
let base = ((pos * nkv + kv) * 2u) * (hd / 8u);
let ks = kvs[(pos * nkv + kv) * 3u];
let gamma = kvs[(pos * nkv + kv) * 3u + 2u];
var dot = 0.0;
for (var w = 0u; w < hd / 8u; w = w + 1u) {
let word = kvq[base + w];
for (var b = 0u; b < 8u; b = b + 1u) {
dot = dot + rq[w * 8u + b] * csh[(word >> (b * 4u)) & 0xFu];
}
}
dot = dot * ks;
// Unbiased residual correction: + coef·γ·⟨S·q, sign_s⟩.
var corr = 0.0;
for (var w = 0u; w < hd / 32u; w = w + 1u) {
let bits = ksig[(pos * nkv + kv) * (hd / 32u) + w];
for (var b = 0u; b < 32u; b = b + 1u) {
var sgn = -1.0;
if (((bits >> b) & 1u) == 1u) { sgn = 1.0; }
corr = corr + sq[w * 32u + b] * sgn;
}
}
let sv = (dot + coef * gamma * corr) * scaling;
sc[s] = sv;
lmax = max(lmax, sv);
}
red[lx] = lmax;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) { if (lx < st) { red[lx] = max(red[lx], red[lx + st]); } workgroupBarrier(); }
let mx = red[0];
workgroupBarrier();
var lsum = 0.0;
for (var s = lx; s < span; s = s + 256u) { let e = exp(sc[s] - mx); sc[s] = e; lsum = lsum + e; }
red[lx] = lsum;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) { if (lx < st) { red[lx] = red[lx] + red[lx + st]; } workgroupBarrier(); }
let den = red[0];
workgroupBarrier();
if (lx < hd) {
var a = 0.0;
let w = lx / 8u;
let b = lx % 8u;
for (var s = 0u; s < span; s = s + 1u) {
let pos = s0 + s;
let base = ((pos * nkv + kv) * 2u + 1u) * (hd / 8u);
let vs = kvs[(pos * nkv + kv) * 3u + 1u];
let word = kvq[base + w];
a = a + sc[s] * vs * csh[(word >> (b * 4u)) & 0xFu];
}
acc[lx] = a / den;
}
workgroupBarrier();
if (lx < hd) {
var o = 0.0;
for (var u = 0u; u < hd; u = u + 1u) { o = o + rot[u * hd + lx] * acc[u]; }
out[qbase + lx] = o;
}
}
"#;
/// Batched [`GATHER`]: column `gid.y` reads its FED token (`fed[col]`, uploaded per batch) —
/// `cur[col·h + i] = embed[fed[col]·h + i]`. `posu = (_, _, _, hidden)`.
/// [`GATHER_K`] for Q1/GGUF models: the f16 gather table bound as TWO row-aligned halves
/// (a single storage binding caps at 2 GiB-4 on wgpu-Vulkan; the 248k x 5120 f16 table is
/// 2.5 GB). Split at vocab/2 (posu.z), picked per column by token id; widened on read.
const GATHER_K_F16: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> embed_lo: array<f16>;
@group(0) @binding(1) var<storage, read> fed: array<u32>; // [steps, stride]
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<storage, read> cnt: array<u32>; // chained-step counter
@group(0) @binding(4) var<uniform> posu: vec4<u32>; // (step_stride, _, vocab, hidden)
@group(0) @binding(5) var<storage, read> embed_hi: array<f16>;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
let col = gid.y;
let tok = fed[cnt[0] * posu.x + col];
let vs = posu.z / 2u;
if (tok < vs) {
cur[col * h + i] = f32(embed_lo[tok * h + i]);
} else {
cur[col * h + i] = f32(embed_hi[(tok - vs) * h + i]);
}
}
"#;
const GATHER_K: &str = r#"
@group(0) @binding(0) var<storage, read> embed: array<f32>;
@group(0) @binding(1) var<storage, read> fed: array<u32>; // [steps, stride]
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<storage, read> cnt: array<u32>; // chained-step counter
@group(0) @binding(4) var<uniform> posu: vec4<u32>; // (step_stride, _, _, hidden)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
let col = gid.y;
cur[col * h + i] = embed[fed[cnt[0] * posu.x + col] * h + i];
}
"#;
/// [`GATHER_K`] for a multimodal prompt: a column's hidden row comes EITHER from the token
/// embedding table (text) OR from the vision tower's output (an image placeholder).
///
/// The substitution has to happen here — before layer 0 — because an image placeholder token id
/// carries no information: every image in every prompt expands to the same repeated id, so gathering
/// its embedding would feed the model the same vector for a photograph of a cat and a scan of an
/// invoice.
///
/// It is driven by DATA (`isrc`, one entry per column) rather than by which copies the host encodes,
/// which is the whole point: the command buffer stays independent of the request and so remains
/// reusable across steps (see `BatchPlan::cmdbuf`). Encoding per-image buffer copies instead would
/// have made every multimodal step re-record its command buffer.
const GATHER_K_MM: &str = r#"
@group(0) @binding(0) var<storage, read> embed: array<f32>;
@group(0) @binding(1) var<storage, read> fed: array<u32>; // [steps, stride]
@group(0) @binding(2) var<storage, read_write> cur: array<f32>;
@group(0) @binding(3) var<storage, read> cnt: array<u32>; // chained-step counter
@group(0) @binding(4) var<storage, read> img: array<f32>; // [rows, hidden] vision output
@group(0) @binding(5) var<storage, read> isrc: array<i32>; // [steps, stride]; -1 = text
// NOTE: indexed by the CHAINED-ROUND index exactly like `fed`, so it must be defined for EVERY
// round, not just round 0. Chained rounds feed the model its own generated tokens — always text —
// so rounds >= 1 are all -1. Leaving them uninitialized reads 0, which means "image row 0": a text
// request would silently be fed whatever image happens to sit at the base of the arena.
@group(0) @binding(6) var<uniform> posu: vec4<u32>; // (step_stride, _, _, hidden)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let h = posu.w; let i = gid.x; if (i >= h) { return; }
let col = gid.y;
let idx = cnt[0] * posu.x + col;
let row = isrc[idx];
if (row < 0) {
cur[col * h + i] = embed[fed[idx] * h + i];
} else {
cur[col * h + i] = img[u32(row) * h + i];
}
}
"#;
/// DeepSeek MLA assembly (one workgroup per head): from the raw q `[nh·qk_head_dim]`, the kv_b
/// expansion `[nh·(qk_nope+v_head_dim)]`, and the shared `k_pe [qk_rope]`, apply the DECOUPLED
/// GPT-J interleaved-pair RoPE (rotate pairs `(2i,2i+1)` by angle `pos·θ^(−2i/qk_rope)`, computed
/// in-kernel) to the pe slices, broadcast k_pe to every head, and write: `qn` (full roped q), the
/// KV cache k `[k_nope | roped k_pe]` and v `[v | zero-pad to qk_head_dim]`. Cache addressing
/// mirrors `qkrc` (nkv = nh). θ is baked in at build.
fn mla_assemble_src(theta: f32) -> String {
format!(
r#"enable f16;
@group(0) @binding(0) var<storage, read> qr: array<f32>; // [nh·qk_head_dim] raw q
@group(0) @binding(1) var<storage, read> kvb: array<f32>; // [nh·(qk_nope+v_head_dim)]
@group(0) @binding(2) var<storage, read> kpe: array<f32>; // [qk_rope] shared k_pe
@group(0) @binding(3) var<storage, read_write> qn: array<f32>; // [nh·qk_head_dim] out
@group(0) @binding(4) var<storage, read_write> kc: array<f16>; // KV cache
@group(0) @binding(5) var<storage, read_write> vc: array<f16>;
@group(0) @binding(6) var<storage, read> btab: array<u32>;
@group(0) @binding(7) var<storage, read> cmeta: array<u32>;
@group(0) @binding(8) var<storage, read> cnt: array<u32>;
@group(0) @binding(9) var<uniform> kp: vec4<u32>; // (nh, qk_nope, qk_rope, v_head_dim)
@group(0) @binding(10) var<uniform> sp: vec4<u32>; // (step_stride, kpe_off, _, _)
const THETA: f32 = {theta:.1};
const PBLK: u32 = 16u; const PLOG: u32 = 4u; const PMSK: u32 = 15u; const PROW: u32 = {{{{PROW}}}}u;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let nh = kp.x; let qkn = kp.y; let qkr = kp.z; let vhd = kp.w;
let qkh = qkn + qkr;
let head = wid.x; let j = lid.x;
let mcol = cnt[0] * sp.x + wid.y;
let lpos = cmeta[mcol * 2u];
let pos = btab[cmeta[mcol * 2u + 1u] * PROW + (lpos >> PLOG)] * PBLK + (lpos & PMSK);
let kd = nh * qkh; // cache stride per position (nkv = nh)
let qbase = head * qkh;
let kvbase = head * (qkn + vhd);
let cbase = pos * kd + head * qkh;
// q nope (copy) + k nope (→cache) + v (→cache, zero-padded past v_head_dim)
if (j < qkn) {{ qn[qbase + j] = qr[qbase + j]; kc[cbase + j] = f16(kvb[kvbase + j]); }}
if (j < qkh) {{ var vv = 0.0; if (j < vhd) {{ vv = kvb[kvbase + qkn + j]; }} vc[cbase + j] = f16(vv); }}
// pe rope on the pairs (2i, 2i+1): q per head, k_pe shared → broadcast to this head
if (j < qkr / 2u) {{
let ang = f32(lpos) * pow(THETA, -2.0 * f32(j) / f32(qkr));
let c = cos(ang); let s = sin(ang);
let qa = qr[qbase + qkn + 2u * j]; let qb = qr[qbase + qkn + 2u * j + 1u];
qn[qbase + qkn + 2u * j] = qa * c - qb * s;
qn[qbase + qkn + 2u * j + 1u] = qa * s + qb * c;
let ka = kpe[sp.y + 2u * j]; let kb = kpe[sp.y + 2u * j + 1u];
kc[cbase + qkn + 2u * j] = f16(ka * c - kb * s);
kc[cbase + qkn + 2u * j + 1u] = f16(ka * s + kb * c);
}}
}}
"#
)
}
/// Qwen2 q/k/v projection bias: `qkv[i] += bias[i % stride]`, broadcasting the `[stride]` bias
/// across the `k` columns of the fused qkv output. A standalone in-place pass inserted between the
/// qkv GEMV and qkrc, ONLY on the bias-carrying arch — every other arch's plan is unchanged.
const QKV_BIAS_ADD: &str = r#"
@group(0) @binding(0) var<storage, read_write> qkv: array<f32>; // [k, stride]
@group(0) @binding(1) var<storage, read> bias: array<f32>; // [stride]
@group(0) @binding(2) var<uniform> d: vec4<u32>; // (stride, total, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x;
if (i >= d.y) { return; }
qkv[i] = qkv[i] + bias[i % d.x];
}
"#;
/// Does this decoder architecture apply per-head qk-RMSNorm before RoPE? Qwen3/Qwen35/Gemma do;
/// GLM-OCR and Llama/Mistral do NOT (they select the no-norm kernel variant, which still reads the
/// ones q/k-norm buffers but skips the normalization).
pub(crate) fn arch_has_qk_norm(arch: crate::weights::Arch) -> bool {
use crate::weights::Arch;
!matches!(
arch,
Arch::GlmOcr
| Arch::Moshi
| Arch::Llama
| Arch::Qwen2
| Arch::Qwen2Moe
| Arch::Phi3
| Arch::Granite
| Arch::Mixtral
| Arch::Olmo
| Arch::StableLm
| Arch::Falcon
| Arch::Cohere
| Arch::Nemotron
| Arch::Phi2
| Arch::Olmo2
)
}
/// Batched [`qk_norm_rope_cache_src`]: grid ((nh+2nkv), k); column `wid.y` = position `kp.w-1+col`,
/// cos/sin at `col·hd` (per-position tables uploaded per batch), qkv/qn at column offsets, KV-cache
/// writes at the column's absolute position. Body identical per column ⇒ bitwise-equal to M=1.
fn qk_norm_rope_cache_k_src(eps: f32, qknorm: bool) -> String {
// GLM-OCR's decoder has NO qk-norm: the no-norm variant still READS qnw/knw (ones buffers), so
// the bind-group layout is identical (Naga prunes unused bindings, which would renumber them),
// but drops the 1/rms scaling — with ones weights the line is an exact identity.
let normline = if qknorm {
"if (is_q) { nv = xj * inv * qknw[j]; } else { nv = xj * inv * qknw[hd + j]; }"
} else {
"if (is_q) { nv = xj * qknw[j]; } else { nv = xj * qknw[hd + j]; }"
};
format!(
r#"enable f16;
@group(0) @binding(0) var<storage, read> qkv: array<f32>; // [k, qd+2kd]
@group(0) @binding(1) var<storage, read> qknw: array<f32>; // [q_norm(hd) | k_norm(hd)]
@group(0) @binding(2) var<storage, read> cs: array<f32>; // [k, hd]
@group(0) @binding(3) var<storage, read> sn: array<f32>; // [k, hd]
@group(0) @binding(4) var<storage, read_write> qn: array<f32>; // [k, qd]
@group(0) @binding(5) var<storage, read_write> kcache: array<f16>; // paged pool
@group(0) @binding(6) var<storage, read_write> vcache: array<f16>;
@group(0) @binding(7) var<storage, read> btab: array<u32>; // block tables [rows, PROW]
@group(0) @binding(8) var<storage, read> cmeta: array<u32>; // [steps, stride, 2]: (pos, btab row)
@group(0) @binding(9) var<storage, read> cnt: array<u32>; // chained-step counter (0 unchained)
@group(0) @binding(10) var<uniform> kp: vec4<u32>; // (nh, nkv, hd, _)
@group(0) @binding(11) var<uniform> af: vec4<u32>; // (v_norm, k_eq_v, step_stride, _)
const EPS: f32 = {eps};
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {prow}u;
var<workgroup> nrm: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid3: vec3<u32>) {{
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let col = wid.y;
let mcol = cnt[0] * af.z + wid.y;
let lpos = cmeta[mcol * 2u];
let pos = btab[cmeta[mcol * 2u + 1u] * PROW + (lpos >> PLOG)] * PBLK + (lpos & PMSK);
let qd = nh * hd; let kd = nkv * hd;
// af.y: 0 = plain q|k|v, 1 = K=V (no v rows in the concatenated buffer), 2 = Q-ONLY
// (Gemma-4 edge KV-shared consumers — dispatched with gx = nh, so the k/v branches below
// never run and the donor caches bound here are never written).
var qkv_stride = qd + 2u * kd;
if (af.y == 1u) {{ qkv_stride = qd + kd; }}
if (af.y == 2u) {{ qkv_stride = qd; }}
let qkv_off = col * qkv_stride;
let cs_off = mcol * hd;
let idx = wid.x; let j = lid3.x;
if (idx >= nh + 2u * nkv) {{ return; }}
if (idx >= nh + nkv) {{
// v head. Source: the V rows, or (K=V mode) the K PROJECTION rows (pre-norm, un-roped).
let head = idx - nh - nkv;
var src = qkv_off + qd + kd + head * hd + j;
if (af.y == 1u) {{ src = qkv_off + qd + head * hd + j; }}
var xj = 0.0; if (j < hd) {{ xj = qkv[src]; }}
if (af.x != 0u) {{
// Gemma-4: WEIGHTLESS RMSNorm on V (with_scale=false — pure x/rms(x)).
red[j] = xj * xj;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{ if (j < st) {{ red[j] = red[j] + red[j + st]; }} workgroupBarrier(); }}
let inv = 1.0 / sqrt(red[0] / f32(hd) + EPS);
xj = xj * inv;
}}
if (j < hd) {{ vcache[pos * kd + head * hd + j] = f16(xj); }}
return;
}}
var in_base = 0u; var is_q = true;
if (idx < nh) {{ in_base = idx * hd; }}
else {{ is_q = false; in_base = qd + (idx - nh) * hd; }}
var xj = 0.0; if (j < hd) {{ xj = qkv[qkv_off + in_base + j]; }}
red[j] = xj * xj;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{ if (j < st) {{ red[j] = red[j] + red[j + st]; }} workgroupBarrier(); }}
let inv = 1.0 / sqrt(red[0] / f32(hd) + EPS);
var nv = 0.0;
if (j < hd) {{ {normline} }}
nrm[j] = nv;
workgroupBarrier();
if (j < hd) {{
// Partial RoPE (af.w = rotary dim, 0 = full): rotation pairs live WITHIN the first rd
// dims — (j, j + rd/2) — and dims ≥ rd pass through unrotated (Qwen3.5 factor 0.25).
var rd = hd;
if (af.w != 0u) {{ rd = af.w; }}
var out = nv;
if (j < rd) {{
let half = rd / 2u;
var rot: f32;
if (j < half) {{ rot = -nrm[j + half]; }} else {{ rot = nrm[j - half]; }}
out = nv * cs[cs_off + j] + rot * sn[cs_off + j];
}}
if (is_q) {{ qn[col * qd + in_base + j] = out; }}
else {{ let head = idx - nh; kcache[pos * kd + head * hd + j] = f16(out); }}
}}
}}
"#,
prow = PAGE_ROW,
normline = normline,
)
}
/// ROPE-FROM-POS variant of [`qk_norm_rope_cache_k_src`]: cos/sin computed IN-KERNEL from the
/// column's cmeta position (angle = pos·θ^(−2i/rd)) instead of CPU-staged tables — the enabler
/// for GPU-resident MTP rounds, where positions are data-dependent (pos0 += 1+j) and no host
/// can stage tables ahead. Bindings renumbered (cs/sn dropped — Naga omits unused bindings).
/// Numerics: GPU pow/cos/sin vs host libm differ in ulps — batch==solo stays exact when BOTH
/// run this variant (same module); vs the table path it's tolerance-class. Qwen-family only
/// (single θ; Gemma's dual-theta local rope keeps the table kernel). `OSFKB_ROPE_GPU=1`.
fn qk_norm_rope_gpu_k_src(eps: f32, theta: f32, qknorm: bool) -> String {
let base = qk_norm_rope_cache_k_src(eps, qknorm);
for frag in [
"@group(0) @binding(2) var<storage, read> cs: array<f32>; // [k, hd]",
"@group(0) @binding(3) var<storage, read> sn: array<f32>; // [k, hd]",
" let cs_off = mcol * hd;",
" out = nv * cs[cs_off + j] + rot * sn[cs_off + j];",
] {
assert!(base.contains(frag), "qkrc source drifted: {frag}");
}
let renumbered = base
.replace(
"@group(0) @binding(2) var<storage, read> cs: array<f32>; // [k, hd]\n",
"",
)
.replace(
"@group(0) @binding(3) var<storage, read> sn: array<f32>; // [k, hd]\n",
"",
)
.replace("@group(0) @binding(4) var<storage, read_write> qn:", "@group(0) @binding(2) var<storage, read_write> qn:")
.replace("@group(0) @binding(5) var<storage, read_write> kcache:", "@group(0) @binding(3) var<storage, read_write> kcache:")
.replace("@group(0) @binding(6) var<storage, read_write> vcache:", "@group(0) @binding(4) var<storage, read_write> vcache:")
.replace("@group(0) @binding(7) var<storage, read> btab:", "@group(0) @binding(5) var<storage, read> btab:")
.replace("@group(0) @binding(8) var<storage, read> cmeta:", "@group(0) @binding(6) var<storage, read> cmeta:")
.replace("@group(0) @binding(9) var<storage, read> cnt:", "@group(0) @binding(7) var<storage, read> cnt:")
.replace("@group(0) @binding(10) var<uniform> kp:", "@group(0) @binding(8) var<uniform> kp:")
.replace("@group(0) @binding(11) var<uniform> af:", "@group(0) @binding(9) var<uniform> af:")
.replace(" let cs_off = mcol * hd;\n", "")
.replace(
" out = nv * cs[cs_off + j] + rot * sn[cs_off + j];",
&format!(
" var fi = j; if (j >= half) {{ fi = j - half; }}\n let ang = f32(lpos) * pow({theta:.1}, -2.0 * f32(fi) / f32(rd));\n out = nv * cos(ang) + rot * sin(ang);"
),
);
assert!(!renumbered.contains("cs_off"), "cs_off survived");
renumbered
}
/// Wide (256 < head_dim ≤ 512) twin of [`qk_norm_rope_cache_k_src`] — Gemma-4 edge
/// full-attention layers (hd 512). Thread `j` owns dims `j` and `j+256` of its head; the norm
/// reduction folds both. RoPE is FULL-width by construction (the proportional tables carry
/// zero-frequency identity pairs, so `rd == hd` and `af.w` is not consulted): dim `j < half`
/// pairs with `j+half`, and the second slot `j2 = j+256 ≥ half` always pairs downward.
/// Bindings, consts and `af` semantics identical to the base kernel.
fn qk_norm_rope_cache_k_wide_src(eps: f32, qknorm: bool) -> String {
let normline = if qknorm {
"if (is_q) { nv0 = x0 * inv * qknw[j]; nv1 = x1 * inv * qknw[j2]; } else { nv0 = x0 * inv * qknw[hd + j]; nv1 = x1 * inv * qknw[hd + j2]; }"
} else {
"if (is_q) { nv0 = x0 * qknw[j]; nv1 = x1 * qknw[j2]; } else { nv0 = x0 * qknw[hd + j]; nv1 = x1 * qknw[hd + j2]; }"
};
format!(
r#"enable f16;
@group(0) @binding(0) var<storage, read> qkv: array<f32>; // [k, qd+2kd]
@group(0) @binding(1) var<storage, read> qknw: array<f32>; // [q_norm(hd) | k_norm(hd)]
@group(0) @binding(2) var<storage, read> cs: array<f32>; // [k, hd]
@group(0) @binding(3) var<storage, read> sn: array<f32>; // [k, hd]
@group(0) @binding(4) var<storage, read_write> qn: array<f32>; // [k, qd]
@group(0) @binding(5) var<storage, read_write> kcache: array<f16>; // paged pool
@group(0) @binding(6) var<storage, read_write> vcache: array<f16>;
@group(0) @binding(7) var<storage, read> btab: array<u32>; // block tables [rows, PROW]
@group(0) @binding(8) var<storage, read> cmeta: array<u32>; // [steps, stride, 2]: (pos, btab row)
@group(0) @binding(9) var<storage, read> cnt: array<u32>; // chained-step counter (0 unchained)
@group(0) @binding(10) var<uniform> kp: vec4<u32>; // (nh, nkv, hd, _)
@group(0) @binding(11) var<uniform> af: vec4<u32>; // (v_norm, mode, step_stride, _)
const EPS: f32 = {eps};
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {prow}u;
var<workgroup> nrm: array<f32, 512>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid3: vec3<u32>) {{
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let col = wid.y;
let mcol = cnt[0] * af.z + wid.y;
let lpos = cmeta[mcol * 2u];
let pos = btab[cmeta[mcol * 2u + 1u] * PROW + (lpos >> PLOG)] * PBLK + (lpos & PMSK);
let qd = nh * hd; let kd = nkv * hd;
var qkv_stride = qd + 2u * kd;
if (af.y == 1u) {{ qkv_stride = qd + kd; }}
if (af.y == 2u) {{ qkv_stride = qd; }}
let qkv_off = col * qkv_stride;
let cs_off = mcol * hd;
let idx = wid.x; let j = lid3.x; let j2 = j + 256u;
if (idx >= nh + 2u * nkv) {{ return; }}
if (idx >= nh + nkv) {{
// v head. Source: the V rows, or (K=V mode) the K PROJECTION rows (pre-norm, un-roped).
let head = idx - nh - nkv;
var base = qkv_off + qd + kd + head * hd;
if (af.y == 1u) {{ base = qkv_off + qd + head * hd; }}
var x0 = 0.0; var x1 = 0.0;
if (j < hd) {{ x0 = qkv[base + j]; }}
if (j2 < hd) {{ x1 = qkv[base + j2]; }}
if (af.x != 0u) {{
// Gemma-4: WEIGHTLESS RMSNorm on V (with_scale=false — pure x/rms(x)).
red[j] = x0 * x0 + x1 * x1;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{ if (j < st) {{ red[j] = red[j] + red[j + st]; }} workgroupBarrier(); }}
let inv = 1.0 / sqrt(red[0] / f32(hd) + EPS);
x0 = x0 * inv;
x1 = x1 * inv;
}}
if (j < hd) {{ vcache[pos * kd + head * hd + j] = f16(x0); }}
if (j2 < hd) {{ vcache[pos * kd + head * hd + j2] = f16(x1); }}
return;
}}
var in_base = 0u; var is_q = true;
if (idx < nh) {{ in_base = idx * hd; }}
else {{ is_q = false; in_base = qd + (idx - nh) * hd; }}
var x0 = 0.0; var x1 = 0.0;
if (j < hd) {{ x0 = qkv[qkv_off + in_base + j]; }}
if (j2 < hd) {{ x1 = qkv[qkv_off + in_base + j2]; }}
red[j] = x0 * x0 + x1 * x1;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{ if (j < st) {{ red[j] = red[j] + red[j + st]; }} workgroupBarrier(); }}
let inv = 1.0 / sqrt(red[0] / f32(hd) + EPS);
var nv0 = 0.0; var nv1 = 0.0;
if (j < hd) {{ {normline} }}
nrm[j] = nv0;
nrm[j2] = nv1;
workgroupBarrier();
let half = hd / 2u;
if (j < hd) {{
var rot: f32;
if (j < half) {{ rot = -nrm[j + half]; }} else {{ rot = nrm[j - half]; }}
let o0 = nv0 * cs[cs_off + j] + rot * sn[cs_off + j];
if (is_q) {{ qn[col * qd + in_base + j] = o0; }}
else {{ let head = idx - nh; kcache[pos * kd + head * hd + j] = f16(o0); }}
}}
if (j2 < hd) {{
let o1 = nv1 * cs[cs_off + j2] + nrm[j2 - half] * sn[cs_off + j2];
if (is_q) {{ qn[col * qd + in_base + j2] = o1; }}
else {{ let head = idx - nh; kcache[pos * kd + head * hd + j2] = f16(o1); }}
}}
}}
"#,
prow = PAGE_ROW,
normline = normline,
)
}
/// Batched [`ATTN`]: grid (n_heads, k); column `wid.y` attends causally over its OWN context length
/// `T = kp.w + col` (position pos0+col sees keys 0..pos0+col). The k columns' K/V were all written by
/// the preceding batched-qkrc dispatch, so column p reads columns < p's keys — sequential semantics
/// preserved. Body identical per column ⇒ bitwise-equal to M=1.
///
/// TILED prefill attention (wide plans, full-attention layers only): grid `(nh, ⌈k/8⌉)` — a
/// WG owns 8 consecutive µbatch COLUMNS of one head and streams the KV span in 32-key chunks
/// staged ONCE in shared f16 for all 8 columns (the per-column [`ATTN_K`] re-read the span
/// per column: 192× the KV traffic at prefill widths — 30.8% of the stage at pos 1024).
/// Online softmax per column (running m/l, rescaled accumulator — flash discipline); thread
/// layout 8 groups × 32 lanes, each lane owns 4 output dims (static registers). MIXED-slot
/// tiles (different btab rows — µbatch boundaries between requests) stay CORRECT: the shared
/// stage uses column 0's block table, and any column whose btab row differs reads the pool
/// DIRECTLY for both K and V (data-source divergence only — every barrier is WG-uniform).
/// Sliding-window layers keep [`ATTN_K`] (per-column spans diverge under a window).
const ATTN_TILED: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> q: array<f32>; // [k, n_heads*head_dim]
@group(0) @binding(1) var<storage, read> kc: array<f16>; // paged pool [blocks*PBLK, nkv, hd]
@group(0) @binding(2) var<storage, read> vc: array<f16>;
@group(0) @binding(3) var<storage, read_write> out: array<f32>; // [k, n_heads*head_dim]
@group(0) @binding(4) var<storage, read> btab: array<u32>; // block tables [rows, PROW]
@group(0) @binding(5) var<storage, read> cmeta: array<u32>; // [steps, stride, 2]: (pos, btab row)
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> kp: vec4<u32>; // (n_heads, n_kv, head_dim, kcols)
@group(0) @binding(8) var<uniform> lw: vec4<u32>; // (window==0, scale_one, step_stride, _)
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {{PROW}}u;
const TC: u32 = 8u; // columns per workgroup
const CB: u32 = 32u; // keys per staged chunk
var<workgroup> ksh: array<f16, 4096>; // [CB, hd≤128] staged keys (column-0's sequence)
var<workgroup> vsh: array<f16, 4096>; // [CB, hd≤128] staged values
var<workgroup> esh: array<f32, 256>; // [TC, CB] per-column exp weights
var<workgroup> red: array<f32, 256>; // per-group reductions ([TC, 32])
var<workgroup> mrow: array<u32, 8>; // per-column btab row
var<workgroup> trow: array<u32, 8>; // per-column span t
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nh = kp.x; let nkv = kp.y; let hd = kp.z; let kcols = kp.w;
let qh = wid.x;
let kv = qh / (nh / nkv);
let lane = lid.x % 32u;
let cg = lid.x / 32u;
let col = wid.y * TC + cg;
let cvalid = col < kcols;
if (lid.x < TC) {
let cX = wid.y * TC + lid.x;
if (cX < kcols) {
let mcol = cnt[0] * lw.z + cX;
trow[lid.x] = cmeta[mcol * 2u] + 1u;
mrow[lid.x] = cmeta[mcol * 2u + 1u];
} else {
trow[lid.x] = 0u;
mrow[lid.x] = 0xFFFFFFFFu;
}
}
workgroupBarrier();
// Shared stage follows column 0's sequence; the tile's loop bound is the max span.
let brow0 = mrow[0];
var tmax = 0u;
for (var c = 0u; c < TC; c = c + 1u) { tmax = max(tmax, trow[c]); }
let t_c = trow[cg];
let shared_ok = cvalid && mrow[cg] == brow0;
let btbase_c = mrow[cg] * PROW;
let qbase = col * nh * hd + qh * hd;
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) { scaling = 1.0; }
var m = -1e30;
var l = 0.0;
// Lane owns 4 output dims: lane, lane+32, lane+64, lane+96 (hd ≤ 128).
var acc0 = 0.0;
var acc1 = 0.0;
var acc2 = 0.0;
var acc3 = 0.0;
for (var c0 = 0u; c0 < tmax; c0 = c0 + CB) {
// Stage this chunk of column-0's K/V (f16, coalesced: 16 f16 per thread).
for (var jj = 0u; jj < 16u; jj = jj + 1u) {
let idx = lid.x * 16u + jj;
let key = idx / 128u;
let d = idx % 128u;
let sp = c0 + key;
var kvv = vec2<f16>(f16(0.0), f16(0.0));
if (sp < tmax && d < hd) {
let pp = btab[brow0 * PROW + (sp >> PLOG)] * PBLK + (sp & PMSK);
kvv.x = kc[(pp * nkv + kv) * hd + d];
kvv.y = vc[(pp * nkv + kv) * hd + d];
}
ksh[idx] = kvv.x;
vsh[idx] = kvv.y;
}
workgroupBarrier();
// Per-column scores for the chunk: lane computes key = lane (0..31).
let key = lane;
let sp = c0 + key;
var sv = -1e30;
if (cvalid && sp < t_c && key < CB) {
var dot = 0.0;
if (shared_ok) {
for (var j = 0u; j < hd; j = j + 1u) {
dot = dot + q[qbase + j] * f32(ksh[key * 128u + j]);
}
} else {
let pp = btab[btbase_c + (sp >> PLOG)] * PBLK + (sp & PMSK);
for (var j = 0u; j < hd; j = j + 1u) {
dot = dot + q[qbase + j] * f32(kc[(pp * nkv + kv) * hd + j]);
}
}
sv = dot * scaling;
}
red[lid.x] = sv;
workgroupBarrier();
// Per-group max over 32 lanes (tree in the group's red slice; uniform barriers).
for (var st = 16u; st > 0u; st = st >> 1u) {
if (lane < st) { red[lid.x] = max(red[lid.x], red[lid.x + st]); }
workgroupBarrier();
}
let new_m = max(m, red[cg * 32u]);
workgroupBarrier();
var e = 0.0;
if (sv > -1e29) { e = exp(sv - new_m); }
esh[cg * 32u + lane] = e;
red[lid.x] = e;
workgroupBarrier();
for (var st = 16u; st > 0u; st = st >> 1u) {
if (lane < st) { red[lid.x] = red[lid.x] + red[lid.x + st]; }
workgroupBarrier();
}
let esum = red[cg * 32u];
let corr = exp(m - new_m);
// Accumulate V for the 4 owned dims (shared V, or pool V on a foreign-slot column).
let lim = min(CB, tmax - c0);
acc0 = acc0 * corr;
acc1 = acc1 * corr;
acc2 = acc2 * corr;
acc3 = acc3 * corr;
for (var i = 0u; i < lim; i = i + 1u) {
let ev = esh[cg * 32u + i];
if (ev != 0.0) {
if (shared_ok) {
acc0 = acc0 + ev * f32(vsh[i * 128u + lane]);
acc1 = acc1 + ev * f32(vsh[i * 128u + lane + 32u]);
acc2 = acc2 + ev * f32(vsh[i * 128u + lane + 64u]);
acc3 = acc3 + ev * f32(vsh[i * 128u + lane + 96u]);
} else {
let sp2 = c0 + i;
let pp = btab[btbase_c + (sp2 >> PLOG)] * PBLK + (sp2 & PMSK);
let vb = (pp * nkv + kv) * hd;
acc0 = acc0 + ev * f32(vc[vb + lane]);
acc1 = acc1 + ev * f32(vc[vb + lane + 32u]);
acc2 = acc2 + ev * f32(vc[vb + lane + 64u]);
acc3 = acc3 + ev * f32(vc[vb + lane + 96u]);
}
}
}
l = l * corr + esum;
m = new_m;
workgroupBarrier();
}
if (cvalid) {
if (lane < hd) { out[qbase + lane] = acc0 / l; }
if (lane + 32u < hd) { out[qbase + lane + 32u] = acc1 / l; }
if (lane + 64u < hd) { out[qbase + lane + 64u] = acc2 / l; }
if (lane + 96u < hd) { out[qbase + lane + 96u] = acc3 / l; }
}
}
"#;
/// [`ATTN_TILED`] codegen'd for an arbitrary `head_dim` — the shared K/V stage is `CB × hd`
/// f16 capped at 4096, so `CB = 4096/hd` (32 at hd=128, **16 at hd=256**) and each lane owns
/// `DPL = hd/32` output dims. The const [`ATTN_TILED`] hard-codes hd=128 (4 accumulators,
/// stride 128) and is disabled above it — which dropped every `head_dim=256` model (Qwen3.5 /
/// Bonsai) back to the per-column [`ATTN_K`] at prefill, the measured long-context wall. This
/// generates the right accumulator count + strides so the 8-column K/V sharing applies at 256
/// too. `hd` must be a multiple of 32. Body/semantics are otherwise identical to the const
/// (online softmax, shared-or-pool divergence, `{{PROW}}` for `page_geom`).
fn attn_tiled_src(hd: usize) -> String {
assert!(
hd.is_multiple_of(32) && hd <= 256,
"attn_tiled hd must be 32..256/32"
);
// CB is capped at 32: the score phase assigns ONE key per lane (`key = lane`, 32 lanes)
// and `esh` is [TC, 32]. The uncapped 4096/hd (64 at hd=64) staged 64-key chunks whose
// keys 32..63 were never scored and whose V-loop read the NEXT column's exp weights —
// scheduler prefill emitted wrong tokens for any llama-family prompt past 32 tokens
// (caught by tests/sched_long_prompt.rs; the hd=128/256 geometries were always aligned).
let cb = (4096 / hd).min(32);
let dpl = hd / 32;
let accs = (0..dpl)
.map(|d| format!(" var acc{d} = 0.0;"))
.collect::<Vec<_>>()
.join("\n");
let rescale = (0..dpl)
.map(|d| format!(" acc{d} = acc{d} * corr;"))
.collect::<Vec<_>>()
.join("\n");
let v_shared = (0..dpl)
.map(|d| {
format!(
" acc{d} = acc{d} + ev * f32(vsh[i * {hd}u + lane + {}u]);",
32 * d
)
})
.collect::<Vec<_>>()
.join("\n");
let v_pool = (0..dpl)
.map(|d| {
format!(
" acc{d} = acc{d} + ev * f32(vc[vb + lane + {}u]);",
32 * d
)
})
.collect::<Vec<_>>()
.join("\n");
let writes = (0..dpl)
.map(|d| {
format!(
" if (lane + {off}u < hd) {{ out[qbase + lane + {off}u] = acc{d} / l; }}",
off = 32 * d
)
})
.collect::<Vec<_>>()
.join("\n");
format!(
r#"enable f16;
@group(0) @binding(0) var<storage, read> q: array<f32>;
@group(0) @binding(1) var<storage, read> kc: array<f16>;
@group(0) @binding(2) var<storage, read> vc: array<f16>;
@group(0) @binding(3) var<storage, read_write> out: array<f32>;
@group(0) @binding(4) var<storage, read> btab: array<u32>;
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> kp: vec4<u32>;
@group(0) @binding(8) var<uniform> lw: vec4<u32>;
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {{{{PROW}}}}u;
const TC: u32 = 8u;
const CB: u32 = {cb}u;
const SHD: u32 = {hd}u;
var<workgroup> ksh: array<f16, 4096>;
var<workgroup> vsh: array<f16, 4096>;
var<workgroup> esh: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
var<workgroup> mrow: array<u32, 8>;
var<workgroup> trow: array<u32, 8>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let nh = kp.x; let nkv = kp.y; let hd = kp.z; let kcols = kp.w;
let qh = wid.x;
let kv = qh / (nh / nkv);
let lane = lid.x % 32u;
let cg = lid.x / 32u;
let col = wid.y * TC + cg;
let cvalid = col < kcols;
if (lid.x < TC) {{
let cX = wid.y * TC + lid.x;
if (cX < kcols) {{
let mcol = cnt[0] * lw.z + cX;
trow[lid.x] = cmeta[mcol * 2u] + 1u;
mrow[lid.x] = cmeta[mcol * 2u + 1u];
}} else {{
trow[lid.x] = 0u;
mrow[lid.x] = 0xFFFFFFFFu;
}}
}}
workgroupBarrier();
let brow0 = mrow[0];
var tmax = 0u;
for (var c = 0u; c < TC; c = c + 1u) {{ tmax = max(tmax, trow[c]); }}
let t_c = trow[cg];
let shared_ok = cvalid && mrow[cg] == brow0;
let btbase_c = mrow[cg] * PROW;
let qbase = col * nh * hd + qh * hd;
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) {{ scaling = 1.0; }}
var m = -1e30;
var l = 0.0;
{accs}
for (var c0 = 0u; c0 < tmax; c0 = c0 + CB) {{
for (var jj = 0u; jj < 16u; jj = jj + 1u) {{
let idx = lid.x * 16u + jj;
let key = idx / SHD;
let d = idx % SHD;
let sp = c0 + key;
var kvv = vec2<f16>(f16(0.0), f16(0.0));
if (sp < tmax && d < hd && key < CB) {{
let pp = btab[brow0 * PROW + (sp >> PLOG)] * PBLK + (sp & PMSK);
kvv.x = kc[(pp * nkv + kv) * hd + d];
kvv.y = vc[(pp * nkv + kv) * hd + d];
}}
ksh[idx] = kvv.x;
vsh[idx] = kvv.y;
}}
workgroupBarrier();
let key = lane;
let sp = c0 + key;
var sv = -1e30;
if (cvalid && sp < t_c && key < CB) {{
var dot = 0.0;
if (shared_ok) {{
for (var j = 0u; j < hd; j = j + 1u) {{
dot = dot + q[qbase + j] * f32(ksh[key * SHD + j]);
}}
}} else {{
let pp = btab[btbase_c + (sp >> PLOG)] * PBLK + (sp & PMSK);
for (var j = 0u; j < hd; j = j + 1u) {{
dot = dot + q[qbase + j] * f32(kc[(pp * nkv + kv) * hd + j]);
}}
}}
sv = dot * scaling;
}}
red[lid.x] = sv;
workgroupBarrier();
for (var st = 16u; st > 0u; st = st >> 1u) {{
if (lane < st) {{ red[lid.x] = max(red[lid.x], red[lid.x + st]); }}
workgroupBarrier();
}}
let new_m = max(m, red[cg * 32u]);
workgroupBarrier();
var e = 0.0;
if (sv > -1e29) {{ e = exp(sv - new_m); }}
esh[cg * 32u + lane] = e;
red[lid.x] = e;
workgroupBarrier();
for (var st = 16u; st > 0u; st = st >> 1u) {{
if (lane < st) {{ red[lid.x] = red[lid.x] + red[lid.x + st]; }}
workgroupBarrier();
}}
let esum = red[cg * 32u];
let corr = exp(m - new_m);
let lim = min(CB, tmax - c0);
{rescale}
for (var i = 0u; i < lim; i = i + 1u) {{
let ev = esh[cg * 32u + i];
if (ev != 0.0) {{
if (shared_ok) {{
{v_shared}
}} else {{
let sp2 = c0 + i;
let pp = btab[btbase_c + (sp2 >> PLOG)] * PBLK + (sp2 & PMSK);
let vb = (pp * nkv + kv) * hd;
{v_pool}
}}
}}
}}
l = l * corr + esum;
m = new_m;
workgroupBarrier();
}}
if (cvalid) {{
{writes}
}}
}}
"#
)
}
/// [`ATTN_K`] with vec4 loads — the QK dot ran hd scalar mul-adds per key from hd scalar f16
/// loads, and the V accumulate one scalar per (key, dim); every decoder head dim is %4 (64/128),
/// so both become hd/4 vec4 ops. Same online-softmax structure, same chunking, same paged
/// lookups; selected at engine build when `hd % 4 == 0` (`OSFKB_ATTN_V4=0` reverts). Used by
/// BOTH the M=1 and batched plans, so batch==solo bitwise invariance is preserved by
/// construction — the FP-order change (vec4 dot regrouping) lands on every path equally.
const ATTN_K_V4: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> q: array<vec4<f32>>; // [k, n_heads*head_dim]
@group(0) @binding(1) var<storage, read> kc: array<vec4<f16>>; // paged pool
@group(0) @binding(2) var<storage, read> vc: array<vec4<f16>>;
@group(0) @binding(3) var<storage, read_write> out: array<vec4<f32>>;
@group(0) @binding(4) var<storage, read> btab: array<u32>;
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> kp: vec4<u32>;
@group(0) @binding(8) var<uniform> lw: vec4<u32>;
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {{PROW}}u;
var<workgroup> sc: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let hd4 = hd / 4u;
let col = wid.y;
let mcol = cnt[0] * lw.z + wid.y;
let t = cmeta[mcol * 2u] + 1u;
let btbase = cmeta[mcol * 2u + 1u] * PROW;
var s0 = 0u;
if (lw.x != 0u && t > lw.x) { s0 = t - lw.x; }
let span = t - s0;
let qh = wid.x;
let kv = qh / (nh / nkv);
let qbase4 = (col * nh * hd + qh * hd) / 4u;
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) { scaling = 1.0; }
let lx = lid.x;
var m = -1e30;
var l = 0.0;
var acc = vec4<f32>(0.0); // thread lx owns dims [4lx, 4lx+4) (lx < hd/4)
for (var c0 = 0u; c0 < span; c0 = c0 + 256u) {
let sidx = c0 + lx;
var sv = -1e30;
if (sidx < span) {
let sp = s0 + sidx;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
let kbase4 = (pp * nkv + kv) * hd4;
var dot4 = 0.0;
for (var j = 0u; j < hd4; j = j + 1u) {
dot4 = dot4 + dot(q[qbase4 + j], vec4<f32>(kc[kbase4 + j]));
}
sv = dot4 * scaling;
}
red[lx] = sv;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = max(red[lx], red[lx + st]); }
workgroupBarrier();
}
let new_m = max(m, red[0]);
workgroupBarrier();
var e = 0.0;
if (sidx < span) { e = exp(sv - new_m); }
sc[lx] = e;
red[lx] = e;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let corr = exp(m - new_m);
if (lx < hd4) {
acc = acc * corr;
let lim = min(256u, span - c0);
for (var i = 0u; i < lim; i = i + 1u) {
let sp = s0 + c0 + i;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
acc = acc + sc[i] * vec4<f32>(vc[(pp * nkv + kv) * hd4 + lx]);
}
}
l = l * corr + red[0];
m = new_m;
workgroupBarrier();
}
if (lx < hd4) { out[qbase4 + lx] = acc / l; }
}
"#;
/// [`ATTN_K_V4`] with the V-ACCUMULATION PARALLELIZED (OSFKB_ATTN_V4PAR=1, hd<=64 only).
/// The V4 kernel's V phase activates only `hd4` of 256 threads and each walks up to 256 keys
/// SERIALLY — measured ~70x off the bandwidth roofline at deep positions (the 6.7ms@2400
/// residual). Here the 256 threads split 16 dim-slots x 16 key-lanes: each lane accumulates a
/// strided 1/16th of the chunk's keys into a private vec4, then a fixed-order shared tree
/// folds the lanes (4KB extra workgroup storage; 6KB total, inside the 16KB budget). Lane
/// order + tree order are deterministic ⇒ reproducible under any load, but NOT bitwise vs the
/// serial walk — opt-in, e2e-quality-gated, exactly like the wide plans. The split-KV surgery
/// anchors (col/mcol, span, epilogue) are untouched, so `attn_split_patch` composes.
fn attn_k_v4_par_src() -> String {
let a_shared = "var<workgroup> sc: array<f32, 256>;";
let a_vloop = r#" let corr = exp(m - new_m);
if (lx < hd4) {
acc = acc * corr;
let lim = min(256u, span - c0);
for (var i = 0u; i < lim; i = i + 1u) {
let sp = s0 + c0 + i;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
acc = acc + sc[i] * vec4<f32>(vc[(pp * nkv + kv) * hd4 + lx]);
}
}
l = l * corr + red[0];"#;
for a in [a_shared, a_vloop] {
assert!(ATTN_K_V4.contains(a), "ATTN_K_V4 drifted: {a:?}");
}
ATTN_K_V4
.replace(
a_shared,
"var<workgroup> sc: array<f32, 256>;\nvar<workgroup> vpart: array<vec4<f32>, 256>;",
)
.replace(
a_vloop,
r#" let corr = exp(m - new_m);
let d = lx & 15u;
let kl = lx >> 4u;
var vp = vec4<f32>(0.0);
{
let lim = min(256u, span - c0);
for (var i = kl; i < lim; i = i + 16u) {
let sp = s0 + c0 + i;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
if (d < hd4) { vp = vp + sc[i] * vec4<f32>(vc[(pp * nkv + kv) * hd4 + d]); }
}
}
vpart[lx] = vp;
workgroupBarrier();
for (var st = 8u; st > 0u; st = st >> 1u) {
if (kl < st) { vpart[lx] = vpart[lx] + vpart[lx + st * 16u]; }
workgroupBarrier();
}
if (lx < hd4) { acc = acc * corr + vpart[lx]; }
l = l * corr + red[0];"#,
)
}
/// [`ATTN_K_V4`] with RB QUERY-COLUMN register blocking — the flash-prefill lever. One
/// workgroup owns `RB` ADJACENT µbatch columns of a head; each thread loads a key's K/V vec4s
/// ONCE and dots them against ALL `RB` columns (accumulates V into all), cutting the K/V DRAM
/// re-reads the per-column [`ATTN_K_V4`] pays by `RB×`. `RB` online-softmax states per thread.
/// Full-attention only (window==0 ⇒ every column starts at key 0, so the shared key stream is
/// valid); `lw.z` carries the column count for tail masking. The per-column key order + vec4
/// dot grouping are IDENTICAL to [`ATTN_K_V4`], so each column's output is bitwise-equal to the
/// plain-vec4 arm. `RB=2` measured the sweet spot on M4 (2.5× on attn); `RB=4` halves the K/V
/// reads again but drops the workgroup count (occupancy) — chosen by measurement per box.
fn attn_k_v4_rb_src(rb: usize) -> String {
// rb<=8 fits rb x 1KB f32 `sc` arrays + 1KB `red` in the 16KB workgroup-storage budget and
// stays BITWISE-equal to the plain-vec4 arm. rb>8 switches `sc` to f16 (rb x 0.5KB — 16
// fits at 9KB): the f16 rounding touches ONLY the V-accumulation weights (softmax probs in
// [0,1], rel ~5e-4), so it is NOT bitwise vs plain — same acceptance class as the wide
// plans, held by the e2e quality gate. Never silently: only OSFKB_PREFILL_RB>8 gets here.
assert!((2..=16).contains(&rb), "RB out of range");
let sc_ty = if rb > 8 { "f16" } else { "f32" };
let decls = (0..rb)
.map(|c| format!(
" let col{c} = cbase + {c}u;\n let c{c}v = col{c} < kcols;\n var t{c} = 0u;\n if (c{c}v) {{ t{c} = cmeta[(cnt[0] * kcols + col{c}) * 2u] + 1u; }}\n let q{c}base4 = (col{c} * nh * hd + qh * hd) / 4u;\n var m{c} = -1e30; var l{c} = 0.0; var acc{c} = vec4<f32>(0.0);"
))
.collect::<Vec<_>>().join("\n");
let span = (0..rb)
.map(|c| format!("t{c}"))
.collect::<Vec<_>>()
.join(", ");
let span_expr = (1..rb).fold("t0".to_string(), |a, c| format!("max({a}, t{c})"));
let _ = span;
let dots_decl = (0..rb)
.map(|c| format!(" var d{c} = 0.0;"))
.collect::<Vec<_>>()
.join("\n");
let dots = (0..rb)
.map(|c| format!(" d{c} = d{c} + dot(q[q{c}base4 + j], kv4);"))
.collect::<Vec<_>>()
.join("\n");
let sv_decl = (0..rb)
.map(|c| format!(" var sv{c} = -1e30;"))
.collect::<Vec<_>>()
.join("\n");
let sv_set = (0..rb)
.map(|c| {
if c == 0 {
" if (sidx < t0) { sv0 = d0 * scaling; }".to_string()
} else {
format!(" if (c{c}v && sidx < t{c}) {{ sv{c} = d{c} * scaling; }}")
}
})
.collect::<Vec<_>>()
.join("\n");
// Per-column softmax + running-state update, generated.
let softmax = (0..rb)
.map(|c| {
format!(
r#" red[lx] = sv{c};
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{
if (lx < st) {{ red[lx] = max(red[lx], red[lx + st]); }}
workgroupBarrier();
}}
let nm{c} = max(m{c}, red[0]);
workgroupBarrier();
var e{c} = 0.0; if (sv{c} > -1e29) {{ e{c} = exp(sv{c} - nm{c}); }}
sc{c}[lx] = {sc_ty}(e{c}); red[lx] = e{c};
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {{
if (lx < st) {{ red[lx] = red[lx] + red[lx + st]; }}
workgroupBarrier();
}}
let esum{c} = red[0];
workgroupBarrier();"#
)
})
.collect::<Vec<_>>()
.join("\n");
let corr = (0..rb)
.map(|c| format!(" let corr{c} = exp(m{c} - nm{c});"))
.collect::<Vec<_>>()
.join("\n");
let rescale = (0..rb)
.map(|c| format!(" acc{c} = acc{c} * corr{c};"))
.collect::<Vec<_>>()
.join("\n");
let vacc = (0..rb)
.map(|c| format!(" acc{c} = acc{c} + f32(sc{c}[i]) * v4;"))
.collect::<Vec<_>>()
.join("\n");
let update = (0..rb)
.map(|c| format!(" l{c} = l{c} * corr{c} + esum{c}; m{c} = nm{c};"))
.collect::<Vec<_>>()
.join("\n");
let writes = (0..rb)
.map(|c| {
if c == 0 {
" out[q0base4 + lx] = acc0 / l0;".to_string()
} else {
format!(" if (c{c}v) {{ out[q{c}base4 + lx] = acc{c} / l{c}; }}")
}
})
.collect::<Vec<_>>()
.join("\n");
let sc_decls = (0..rb)
.map(|c| format!("var<workgroup> sc{c}: array<{sc_ty}, 256>;"))
.collect::<Vec<_>>()
.join("\n");
format!(
r#"enable f16;
@group(0) @binding(0) var<storage, read> q: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read> kc: array<vec4<f16>>;
@group(0) @binding(2) var<storage, read> vc: array<vec4<f16>>;
@group(0) @binding(3) var<storage, read_write> out: array<vec4<f32>>;
@group(0) @binding(4) var<storage, read> btab: array<u32>;
@group(0) @binding(5) var<storage, read> cmeta: array<u32>;
@group(0) @binding(6) var<storage, read> cnt: array<u32>;
@group(0) @binding(7) var<uniform> kp: vec4<u32>;
@group(0) @binding(8) var<uniform> lw: vec4<u32>;
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {{{{PROW}}}}u;
const RB: u32 = {rb}u;
{sc_decls}
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let hd4 = hd / 4u;
let kcols = lw.z;
let cbase = wid.y * RB;
let btbase = cmeta[(cnt[0] * kcols + cbase) * 2u + 1u] * PROW;
let qh = wid.x;
let kv = qh / (nh / nkv);
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) {{ scaling = 1.0; }}
let lx = lid.x;
{decls}
let span = {span_expr};
for (var c0 = 0u; c0 < span; c0 = c0 + 256u) {{
let sidx = c0 + lx;
{sv_decl}
if (sidx < span) {{
let pp = btab[btbase + (sidx >> PLOG)] * PBLK + (sidx & PMSK);
let kbase4 = (pp * nkv + kv) * hd4;
{dots_decl}
for (var j = 0u; j < hd4; j = j + 1u) {{
let kv4 = vec4<f32>(kc[kbase4 + j]);
{dots}
}}
{sv_set}
}}
{softmax}
{corr}
if (lx < hd4) {{
{rescale}
let lim = min(256u, span - c0);
for (var i = 0u; i < lim; i = i + 1u) {{
let sp = c0 + i;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
let v4 = vec4<f32>(vc[(pp * nkv + kv) * hd4 + lx]);
{vacc}
}}
}}
{update}
workgroupBarrier();
}}
if (lx < hd4) {{
{writes}
}}
}}
"#
)
}
const ATTN_K: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> q: array<f32>; // [k, n_heads*head_dim]
@group(0) @binding(1) var<storage, read> kc: array<f16>; // paged pool [blocks*PBLK, nkv, hd]
@group(0) @binding(2) var<storage, read> vc: array<f16>;
@group(0) @binding(3) var<storage, read_write> out: array<f32>; // [k, n_heads*head_dim]
@group(0) @binding(4) var<storage, read> btab: array<u32>; // block tables [rows, PROW]
@group(0) @binding(5) var<storage, read> cmeta: array<u32>; // [steps, stride, 2]: (pos, btab row)
@group(0) @binding(6) var<storage, read> cnt: array<u32>; // chained-step counter (0 unchained)
@group(0) @binding(7) var<uniform> kp: vec4<u32>; // (n_heads, n_kv, head_dim, _)
@group(0) @binding(8) var<uniform> lw: vec4<u32>; // (window, scale_one, step_stride, _)
const PBLK: u32 = 16u;
const PLOG: u32 = 4u;
const PMSK: u32 = 15u;
const PROW: u32 = {{PROW}}u;
var<workgroup> sc: array<f32, 256>; // one CHUNK of scores/exps — span is UNBOUNDED
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nh = kp.x; let nkv = kp.y; let hd = kp.z;
let col = wid.y;
let mcol = cnt[0] * lw.z + wid.y;
let t = cmeta[mcol * 2u] + 1u;
let btbase = cmeta[mcol * 2u + 1u] * PROW;
var s0 = 0u;
if (lw.x != 0u && t > lw.x) { s0 = t - lw.x; }
let span = t - s0;
let qh = wid.x;
let kv = qh / (nh / nkv);
let qbase = col * nh * hd + qh * hd;
var scaling = 1.0 / sqrt(f32(hd));
if (lw.y != 0u) { scaling = 1.0; }
let lx = lid.x;
// Online softmax over 256-key chunks (flash-decode discipline): running max m, running sum
// l, output accumulator rescaled when the max advances — span limited only by the KV pool.
var m = -1e30;
var l = 0.0;
var acc = 0.0; // thread lx owns output dim lx (lx < hd)
for (var c0 = 0u; c0 < span; c0 = c0 + 256u) {
let sidx = c0 + lx;
var sv = -1e30;
if (sidx < span) {
let sp = s0 + sidx;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
let kbase = (pp * nkv + kv) * hd;
var dot = 0.0;
for (var j = 0u; j < hd; j = j + 1u) { dot = dot + q[qbase + j] * f32(kc[kbase + j]); }
sv = dot * scaling;
}
red[lx] = sv;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = max(red[lx], red[lx + st]); }
workgroupBarrier();
}
let new_m = max(m, red[0]);
workgroupBarrier();
var e = 0.0;
if (sidx < span) { e = exp(sv - new_m); }
sc[lx] = e;
red[lx] = e;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let corr = exp(m - new_m);
if (lx < hd) {
acc = acc * corr;
let lim = min(256u, span - c0);
for (var i = 0u; i < lim; i = i + 1u) {
let sp = s0 + c0 + i;
let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);
acc = acc + sc[i] * f32(vc[(pp * nkv + kv) * hd + lx]);
}
}
l = l * corr + red[0];
m = new_m;
workgroupBarrier();
}
if (lx < hd) { out[qbase + lx] = acc / l; }
}
"#;
/// SPLIT-KV flash decode, the LSE merge half: [`attn_split_patch`]'s partials — per
/// (column, head, split): `hd` floats of m-scaled accumulator, then `m` at `[hd]` and `l` at
/// `[hd+1]` (stride `hd+4`, vec4-aligned so the v4 partial can write its lanes directly) —
/// are folded in FIXED split order under the global max, then normalized once. Grid
/// `(n_heads, cols)`; `mp = (n_heads, head_dim, z, _)`. Empty splits carry `m = -1e30, l = 0`
/// and vanish through `exp(m - gm) = 0`; the split holding the true max contributes `l ≥ 1`,
/// so the denominator never underflows.
const ATTN_K_MERGE: &str = r#"
@group(0) @binding(0) var<storage, read> part: array<f32>;
@group(0) @binding(1) var<storage, read_write> out: array<f32>;
@group(0) @binding(2) var<uniform> mp: vec4<u32>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let nh = mp.x; let hd = mp.y; let z = mp.z;
let qh = wid.x;
let col = wid.y;
let lx = lid.x;
let stride = hd + 4u;
let base = (col * nh + qh) * z;
var gm = -1e30;
for (var sp = 0u; sp < z; sp = sp + 1u) { gm = max(gm, part[(base + sp) * stride + hd]); }
var gl = 0.0;
var o = 0.0;
for (var sp = 0u; sp < z; sp = sp + 1u) {
let ms = part[(base + sp) * stride + hd];
let ls = part[(base + sp) * stride + hd + 1u];
let wgt = exp(ms - gm);
gl = gl + ls * wgt;
if (lx < hd) { o = o + part[(base + sp) * stride + lx] * wgt; }
}
if (lx < hd) { out[col * nh * hd + qh * hd + lx] = o / gl; }
}
"#;
/// SPLIT-KV surgery over [`ATTN_K`]/[`ATTN_K_V4`] (fragments asserted so kernel drift is
/// caught at pipeline build): `grid.y` becomes `cols × z` (no Step/replay changes — the split
/// count rides the spare `lw.w` meta word), each split windows the span to its contiguous
/// slice and keeps the EXACT chunked online-softmax body, and the epilogue stores the RAW
/// running (acc, m, l) instead of normalizing — [`ATTN_K_MERGE`] folds the splits. This is
/// the FlashDecoding shape (the bonsai-webgpu space runs the same pair at z=8): decode
/// attention had `n_heads × cols` workgroups and NO KV-length parallelism, so a single
/// long-context stream left the GPU ~95% idle exactly where the plan spends its time.
fn attn_split_patch(base: &str, v4: bool) -> String {
let frags = [
" let col = wid.y;\n let mcol = cnt[0] * lw.z + wid.y;",
" let span = t - s0;",
if v4 {
" if (lx < hd4) { out[qbase4 + lx] = acc / l; }\n}"
} else {
" if (lx < hd) { out[qbase + lx] = acc / l; }\n}"
},
];
for f in frags {
assert!(base.contains(f), "attn kernel drifted: {f}");
}
let mut src = base
.replace(
" let col = wid.y;\n let mcol = cnt[0] * lw.z + wid.y;",
" let zs = max(lw.w, 1u);\n let col = wid.y / zs;\n let spid = wid.y % zs;\n let mcol = cnt[0] * lw.z + col;",
)
.replace(
" let span = t - s0;",
" let span_all = t - s0;\n let per = (span_all + zs - 1u) / zs;\n let w0 = min(spid * per, span_all);\n s0 = s0 + w0;\n let span = min(per, span_all - w0);",
);
if v4 {
src = src.replace(
" if (lx < hd4) { out[qbase4 + lx] = acc / l; }\n}",
" let pb4 = ((col * nh + qh) * zs + spid) * (hd4 + 1u);\n if (lx < hd4) { out[pb4 + lx] = acc; }\n if (lx == 0u) { out[pb4 + hd4] = vec4<f32>(m, l, 0.0, 0.0); }\n}",
);
} else {
src = src.replace(
" if (lx < hd) { out[qbase + lx] = acc / l; }\n}",
" let pb = ((col * nh + qh) * zs + spid) * (hd + 4u);\n if (lx < hd) { out[pb + lx] = acc; }\n if (lx == 0u) { out[pb + hd] = m; out[pb + hd + 1u] = l; }\n}",
);
}
src
}
/// Split count for the split-KV decode attention pair. `OSFKB_ATTN_SPLITKV=0` (or `1`)
/// reverts to the single-workgroup-per-(head,col) kernels; a number pins the split; unset
/// defaults to 8 (the measured knee class: n_heads × 8 ≈ 4 workgroups/core on this box, and
/// the empty-split cost at short spans is two tiny dispatches).
fn attn_splitkv_z() -> u32 {
match std::env::var("OSFKB_ATTN_SPLITKV").ok().as_deref() {
Some("0") | Some("1") => 1,
Some(v) => v.parse::<u32>().map_or(8, |z| z.clamp(2, 32)),
None => 8,
}
}
/// Sandwich-norm residual add (Gemma-3): `y[v·n + i] += rmsnorm(x[v·n + i]) · w[i]` — the block
/// output is normed BEFORE the residual add (`x + post_norm(block(x))`). One workgroup per vector
/// (column); grid.x = n_vecs (1 for M=1, k for the batched plan).
/// `kp.w` = Gemma-4 `layer_scalar`, multiplied onto the WHOLE stream after the add (0 = 1.0 —
/// every pre-existing caller passes 0 in that slot, and `·1.0` is bitwise-exact).
const RMSNORM_ADD: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (vec_len, n_vecs, eps, layer_scalar)
var<workgroup> partial: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let n = u32(kp.x);
let eps = kp.z;
var ls = kp.w;
if (ls == 0.0) { ls = 1.0; }
let vbase = wid.x * n;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 256u) { let v = x[vbase + i]; s = s + v * v; }
partial[lid.x] = s;
workgroupBarrier();
for (var stride = 128u; stride > 0u; stride = stride >> 1u) {
if (lid.x < stride) { partial[lid.x] = partial[lid.x] + partial[lid.x + stride]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(partial[0] / f32(n) + eps);
for (var i = lid.x; i < n; i = i + 256u) { y[vbase + i] = (y[vbase + i] + x[vbase + i] * inv * w[i]) * ls; }
}
"#;
/// [`RMSNORM`] with the Σx² reduction done via `subgroupAdd` + one tiny lane-0 pass over the
/// per-subgroup partials — 2 barriers instead of the 9-level shared tree. At decode (k=1)
/// each rmsnorm dispatch is ONE workgroup, so the barrier ladder was most of its measured
/// 21µs (129 dispatches/token ≈ 2.7ms). Reduction order is the device's fixed subgroup tree
/// (same determinism baseline as the production gemv kernels). Selected when the device has
/// subgroups; the tree kernel remains the fallback.
/// [`RMSNORM_SG`] that ALSO emits the packed-f16 twin the rp4/f16x GEMVs read (binding 4)
/// — producer-side packing deletes the separate pack_f16 dispatch for every consumer of
/// this output. Final pass maps PAIRS to threads so each lane packs one u32; the f32
/// values written are bitwise-identical to [`RMSNORM_SG`]'s.
#[cfg(test)]
mod rmsnorm_pk_tests {
use super::*;
/// PK twin contract: y bitwise-equal to RMSNORM_SG, y16 = pack2x16float(y pairs),
/// across MULTIPLE columns (the production k-plan shape).
#[test]
fn rmsnorm_pk_matches_sg_multicol() {
let Ok(ctx) = GpuCtx::new() else {
eprintln!("SKIP");
return;
};
if !ctx.subgroups {
eprintln!("SKIP: no subgroups");
return;
}
let (n, k) = (5120usize, 8usize);
let mut rng = 0xfeed_beef_dead_cafeu64;
let mut nx = || {
rng ^= rng << 13;
rng ^= rng >> 7;
rng ^= rng << 17;
((rng >> 33) as f32 / (1u64 << 31) as f32) - 0.5
};
let x: Vec<f32> = (0..n * k).map(|_| nx()).collect();
let w: Vec<f32> = (0..n).map(|_| 1.0 + nx() * 0.1).collect();
let xb = ctx.storage(&x);
let wb = ctx.storage(&w);
let y_sg = ctx.empty(n * k);
let y_pk = ctx.empty(n * k);
let y16 = ctx.empty(n * k / 2);
let meta = uni(
&ctx,
bytemuck::cast_slice(&[n as f32, k as f32, 1e-6f32, 0.0]),
);
let sg = pipeline(&ctx, "sg", RMSNORM_SG);
let pk = pipeline(&ctx, "pk", RMSNORM_SG_PK);
let bg_sg = make_bg(&ctx, &sg, &[&xb, &wb, &y_sg], &meta);
let bg_pk = make_bg(&ctx, &pk, &[&xb, &wb, &y_pk, &y16], &meta);
let mut enc = ctx.device.create_command_encoder(&Default::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&sg);
p.set_bind_group(0, &bg_sg, &[]);
p.dispatch_workgroups(k as u32, 1, 1);
p.set_pipeline(&pk);
p.set_bind_group(0, &bg_pk, &[]);
p.dispatch_workgroups(k as u32, 1, 1);
}
ctx.queue.submit([enc.finish()]);
let a = ctx.read(&y_sg, n * k).unwrap();
let b = ctx.read(&y_pk, n * k).unwrap();
let p16 = ctx.read_u32(&y16, n * k / 2).unwrap();
for i in 0..n * k {
assert!(a[i] == b[i], "y mismatch at {i}: {} vs {}", a[i], b[i]);
}
// WGSL pack2x16float rounding is implementation-defined (NVIDIA Vulkan differs
// from the half crate's RTE by 1 ULP on some values; Metal matches). The engine
// contract is GPU-self-consistent packing — allow 1 ULP vs the CPU reference.
for j in 0..n * k / 2 {
let lo = half::f16::from_f32(b[2 * j]).to_bits() as i32;
let hi = half::f16::from_f32(b[2 * j + 1]).to_bits() as i32;
let glo = (p16[j] & 0xffff) as i32;
let ghi = (p16[j] >> 16) as i32;
assert!(
(glo - lo).abs() <= 1 && (ghi - hi).abs() <= 1,
"y16 mismatch at {j}: {:#x} vs {lo:#x}/{hi:#x}",
p16[j]
);
}
eprintln!("PK == SG bitwise + pack OK ({} elems, {k} cols)", n * k);
}
}
const RMSNORM_SG_PK: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<storage, read_write> y16: array<u32>;
@group(0) @binding(4) var<uniform> kp: vec4<f32>; // (vec_len, n_vecs, eps, _)
var<workgroup> partial: array<f32, 32>;
@compute @workgroup_size(1024)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
@builtin(subgroup_invocation_id) siv: u32, @builtin(subgroup_size) ssz: u32) {
let n = u32(kp.x);
let eps = kp.z;
let vbase = wid.x * n;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 1024u) { let v = x[vbase + i]; s = s + v * v; }
s = subgroupAdd(s);
let sgid = lid.x / ssz;
if (siv == 0u) { partial[sgid] = s; }
workgroupBarrier();
if (lid.x == 0u) {
let nsg = 1024u / ssz;
var t = 0.0;
for (var c = 0u; c < nsg; c = c + 1u) { t = t + partial[c]; }
partial[0] = t;
}
workgroupBarrier();
let inv = 1.0 / sqrt(partial[0] / f32(n) + eps);
for (var i = lid.x * 2u; i < n; i = i + 2048u) {
let a = x[vbase + i] * inv * w[i];
let b = x[vbase + i + 1u] * inv * w[i + 1u];
y[vbase + i] = a;
y[vbase + i + 1u] = b;
y16[(vbase + i) / 2u] = pack2x16float(vec2<f32>(a, b));
}
}
"#;
const RMSNORM_SG: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (vec_len, n_vecs, eps, _)
var<workgroup> partial: array<f32, 32>;
@compute @workgroup_size(1024)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
@builtin(subgroup_invocation_id) siv: u32, @builtin(subgroup_size) ssz: u32) {
let n = u32(kp.x);
let eps = kp.z;
let vbase = wid.x * n;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 1024u) { let v = x[vbase + i]; s = s + v * v; }
s = subgroupAdd(s);
let sgid = lid.x / ssz;
if (siv == 0u) { partial[sgid] = s; }
workgroupBarrier();
if (lid.x == 0u) {
let nsg = 1024u / ssz;
var t = 0.0;
for (var c = 0u; c < nsg; c = c + 1u) { t = t + partial[c]; }
partial[0] = t;
}
workgroupBarrier();
let inv = 1.0 / sqrt(partial[0] / f32(n) + eps);
for (var i = lid.x; i < n; i = i + 1024u) { y[vbase + i] = x[vbase + i] * inv * w[i]; }
}
"#;
/// [`RMSNORM_ADD`]'s subgroup twin — same reduction swap as [`RMSNORM_SG`].
const RMSNORM_ADD_SG: &str = r#"
@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> y: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (vec_len, n_vecs, eps, layer_scalar)
var<workgroup> partial: array<f32, 32>;
@compute @workgroup_size(1024)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
@builtin(subgroup_invocation_id) siv: u32, @builtin(subgroup_size) ssz: u32) {
let n = u32(kp.x);
let eps = kp.z;
var ls = kp.w;
if (ls == 0.0) { ls = 1.0; }
let vbase = wid.x * n;
var s = 0.0;
for (var i = lid.x; i < n; i = i + 1024u) { let v = x[vbase + i]; s = s + v * v; }
s = subgroupAdd(s);
let sgid = lid.x / ssz;
if (siv == 0u) { partial[sgid] = s; }
workgroupBarrier();
if (lid.x == 0u) {
let nsg = 1024u / ssz;
var t = 0.0;
for (var c = 0u; c < nsg; c = c + 1u) { t = t + partial[c]; }
partial[0] = t;
}
workgroupBarrier();
let inv = 1.0 / sqrt(partial[0] / f32(n) + eps);
for (var i = lid.x; i < n; i = i + 1024u) { y[vbase + i] = (y[vbase + i] + x[vbase + i] * inv * w[i]) * ls; }
}
"#;
/// Gemma-4 edge PLE token gather: dequantize ONE row of the Q4_0 per-layer-embedding table
/// (already pre-scaled by sqrt(ple_dim) at load) into `out[cols]`. The token id is read from a
/// 1-element buffer so the M=1 record path can stage it queue-ordered per step. Layout mirrors
/// [`crate::weights::quantize_q4_0`]: block b of row r = scale `[r·nblk+b]` (f16) + 4 u32 words,
/// byte `y` of word `w` packing elements `w·4+y` (low nibble) and `w·4+y+16` (high nibble),
/// value `(q − 8)·d`.
const PLE_GATHER: &str = r#"enable f16;
@group(0) @binding(0) var<storage, read> scales: array<f16>;
@group(0) @binding(1) var<storage, read> quants: array<u32>;
@group(0) @binding(2) var<storage, read> fed: array<u32>; // [steps, stride] token ids
@group(0) @binding(3) var<storage, read_write> out: array<f32>; // [k, cols]
@group(0) @binding(4) var<storage, read> cnt: array<u32>; // chained-step counter
@group(0) @binding(5) var<uniform> d: vec4<u32>; // (cols, stride, _, _)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
let i = gid.x;
let cols = d.x;
if (i >= cols) { return; }
// Column = wid.y; the token id comes from the SAME per-round `fed` slice the embed gather
// reads, so chained rounds stay correct with no host round-trip. The M=1 plan binds its
// 1-element token buffer with stride 1 and the zero counter.
let col = wid.y;
let r = fed[cnt[0] * d.y + col];
let nblk = cols / 32u;
let blk = r * nblk + i / 32u;
let e = i % 32u;
var j = e;
var shift = 0u;
if (e >= 16u) { j = e - 16u; shift = 4u; }
let word = quants[blk * 4u + j / 4u];
let q = (word >> ((j % 4u) * 8u + shift)) & 0xFu;
out[col * cols + i] = (f32(q) - 8.0) * f32(scales[blk]);
}
"#;
/// Gemma-4 edge PLE prep: merge the CONTEXT stream into the token stream, in place. One
/// workgroup per layer-group `g` (`ple_dim ≤ 256` — checked at load):
/// `ple[g·pd + j] = (rmsnorm(ctx[g·pd + j] · hscale)·w[j] + ple[g·pd + j]) · 2^-0.5`,
/// where `ctx` is the per_layer_model_projection GEMV output, `hscale = hidden^-0.5` (HF scales
/// BEFORE the norm — the eps makes that ordering observable), `w` = per_layer_projection_norm,
/// and `ple` arrives holding the gathered token row (already ×sqrt(ple_dim) via the table).
const PLE_PREP: &str = r#"
@group(0) @binding(0) var<storage, read> ctxv: array<f32>;
@group(0) @binding(1) var<storage, read> w: array<f32>;
@group(0) @binding(2) var<storage, read_write> ple: array<f32>;
@group(0) @binding(3) var<uniform> kp: vec4<f32>; // (ple_dim, hscale, eps, _)
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let pd = u32(kp.x);
// kp.w = per-COLUMN row stride (n_layers·pd) for the batched shape; 0 in the M=1 plan
// (gy = 1, so wid.y = 0 and the term vanishes either way).
let base = wid.y * u32(kp.w) + wid.x * pd;
let j = lid.x;
var xs = 0.0;
if (j < pd) { xs = ctxv[base + j] * kp.y; }
red[j] = xs * xs;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (j < st) { red[j] = red[j] + red[j + st]; }
workgroupBarrier();
}
let inv = 1.0 / sqrt(red[0] / f32(pd) + kp.z);
if (j < pd) {
ple[base + j] = (xs * inv * w[j] + ple[base + j]) * 0.70710678118654752;
}
}
"#;
/// Gemma-4 edge PLE activation over k columns: `y[col·pd + j] = gelu_tanh(g[col·pd + j]) ·
/// ple[col·nlpd + off + j]` — the batched twin of the M=1 plan's MLP_ACT_MUL+offset reuse
/// (a scalar offset cannot express the per-column slice). gelu_pytorch_tanh is the only
/// activation any Gemma-4 ships.
const PLE_ACT_K: &str = r#"
@group(0) @binding(0) var<storage, read> g: array<f32>; // [k, pd]
@group(0) @binding(1) var<storage, read> ple: array<f32>; // [k, nlpd]
@group(0) @binding(2) var<storage, read_write> y: array<f32>; // [k, pd]
@group(0) @binding(3) var<uniform> d: vec4<u32>; // (total, off, pd, nlpd)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x; if (i >= d.x) { return; }
let col = i / d.z;
let j = i % d.z;
let gate = g[i];
let g3 = gate * gate * gate;
let t = clamp(0.7978845608028654 * (gate + 0.044715 * g3), -20.0, 20.0);
let act = 0.5 * gate * (1.0 + tanh(t));
y[i] = act * ple[col * d.w + d.y + j];
}
"#;
/// Wide (head_dim ≤ 512) twin of [`ATTN_K`], derived by string surgery with asserted fragments
/// (kernel drift is caught, like [`attn_splitkv_z`]'s split patch): thread `lx` owns output dims
/// `lx` AND `lx+256`. Only Gemma-4 edge full-attention layers (hd 512) ever dispatch it — every
/// other model keeps the untouched base kernel.
fn attn_k_wide_src() -> String {
let base = ATTN_K;
let frags = [
" var acc = 0.0; // thread lx owns output dim lx (lx < hd)\n",
" if (lx < hd) {\n acc = acc * corr;\n let lim = min(256u, span - c0);\n for (var i = 0u; i < lim; i = i + 1u) {\n let sp = s0 + c0 + i;\n let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);\n acc = acc + sc[i] * f32(vc[(pp * nkv + kv) * hd + lx]);\n }\n }\n",
" if (lx < hd) { out[qbase + lx] = acc / l; }\n",
];
for f in frags {
assert!(base.contains(f), "ATTN_K source drifted: {f}");
}
base.replace(frags[0], " var acc = 0.0; // thread lx owns output dims lx and lx+256\n var acc2 = 0.0;\n")
.replace(
frags[1],
" {\n acc = acc * corr;\n acc2 = acc2 * corr;\n let lim = min(256u, span - c0);\n let lx2 = lx + 256u;\n for (var i = 0u; i < lim; i = i + 1u) {\n let sp = s0 + c0 + i;\n let pp = btab[btbase + (sp >> PLOG)] * PBLK + (sp & PMSK);\n let vb = (pp * nkv + kv) * hd;\n if (lx < hd) { acc = acc + sc[i] * f32(vc[vb + lx]); }\n if (lx2 < hd) { acc2 = acc2 + sc[i] * f32(vc[vb + lx2]); }\n }\n }\n",
)
.replace(
frags[2],
" if (lx < hd) { out[qbase + lx] = acc / l; }\n if (lx + 256u < hd) { out[qbase + lx + 256u] = acc2 / l; }\n",
)
}
/// Column-slot variant of [`SPARSE_ARGMAX`]: writes `chosen[colu.x]` so the k verify columns share
/// one chosen buffer (binding offsets can't be 4-byte aligned; a per-column uniform can).
const SPARSE_ARGMAX_K: &str = r#"
@group(0) @binding(0) var<storage, read> svals: array<f32>;
@group(0) @binding(1) var<storage, read> ids: array<u32>;
@group(0) @binding(2) var<storage, read> smeta: array<u32>; // [4] = frozen count
@group(0) @binding(3) var<storage, read_write> chosen: array<u32>;
@group(0) @binding(4) var<uniform> colu: vec4<u32>; // (col, _, _, _)
var<workgroup> vmax: array<f32, 256>;
var<workgroup> imax: array<u32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let cnt = smeta[4];
var bv = -1e30; var bi = 0u;
for (var i = lid.x; i < cnt; i = i + 256u) {
let v = svals[i]; let id = ids[i];
if (v > bv || (v == bv && id < bi)) { bv = v; bi = id; }
}
vmax[lid.x] = bv; imax[lid.x] = bi;
workgroupBarrier();
for (var s = 128u; s > 0u; s = s >> 1u) {
if (lid.x < s) {
let ov = vmax[lid.x + s]; let oi = imax[lid.x + s];
if (ov > vmax[lid.x] || (ov == vmax[lid.x] && oi < imax[lid.x])) { vmax[lid.x] = ov; imax[lid.x] = oi; }
}
workgroupBarrier();
}
if (lid.x == 0u) { chosen[colu.x] = imax[0]; }
}
"#;
/// MoE router (Gemma-4 semantics, one workgroup per column): logits = W_r · x, softmax over ALL
/// experts in f32, top-k of the probabilities, renormalize the k weights to sum 1, then multiply by
/// the learned per-expert scale. Writes `sel[col·K+s]` (expert ids) + `wsel[col·K+s]` (weights).
/// Ties break to the LOWEST expert id (deterministic; matches a stable CPU argmax). E ≤ 256.
const MOE_ROUTER: &str = r#"
@group(0) @binding(0) var<storage, read> wr: array<f32>; // [E, h] router
@group(0) @binding(1) var<storage, read> pes: array<f32>; // [E] per-expert scale
@group(0) @binding(2) var<storage, read> x: array<f32>; // [n_cols, h] (ffn-normed input)
@group(0) @binding(3) var<storage, read_write> sel: array<u32>; // [n_cols, K]
@group(0) @binding(4) var<storage, read_write> wsel: array<f32>; // [n_cols, K]
@group(0) @binding(5) var<uniform> mp: vec4<u32>; // (E, h, K, _)
var<workgroup> probs: array<f32, 256>;
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let e_n = mp.x; let h = mp.y; let kk = mp.z;
let col = wid.x;
let lx = lid.x;
// (1) logits: thread e computes W_r[e] · x[col]
var logit = -1e30;
if (lx < e_n) {
var acc = 0.0;
for (var j = 0u; j < h; j = j + 1u) { acc = acc + wr[lx * h + j] * x[col * h + j]; }
logit = acc;
}
probs[lx] = logit;
workgroupBarrier();
// (2) softmax over all E (f32): max-reduce, exp, sum-reduce
red[lx] = probs[lx];
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = max(red[lx], red[lx + st]); }
workgroupBarrier();
}
let mx = red[0];
workgroupBarrier();
var e_val = 0.0;
if (lx < e_n) { e_val = exp(probs[lx] - mx); }
probs[lx] = e_val;
red[lx] = e_val;
workgroupBarrier();
for (var st = 128u; st > 0u; st = st >> 1u) {
if (lx < st) { red[lx] = red[lx] + red[lx + st]; }
workgroupBarrier();
}
let den = red[0];
workgroupBarrier();
// (3) thread 0: k passes of stable argmax over the probabilities, renorm, scale.
if (lx == 0u) {
var wsum = 0.0;
for (var s = 0u; s < kk; s = s + 1u) {
var bi = 0u; var bv = -1.0;
for (var e = 0u; e < e_n; e = e + 1u) {
if (probs[e] > bv) { bv = probs[e]; bi = e; }
}
sel[col * kk + s] = bi;
wsel[col * kk + s] = bv / den;
wsum = wsum + bv / den;
probs[bi] = -1.0; // consumed
}
for (var s = 0u; s < kk; s = s + 1u) {
let eid = sel[col * kk + s];
wsel[col * kk + s] = (wsel[col * kk + s] / wsum) * pes[eid];
}
}
}
"#;
/// DeepSeek-V3 aux-loss-free router (`DeepseekV3TopkRouter`). Unlike [`MOE_ROUTER`]'s softmax, this
/// scores each expert by `sigmoid(W_r·x)`, adds a learned `bias` for SELECTION only, then does
/// GROUP-LIMITED top-k: partition the `E` experts into `n_group` groups, score each group by the
/// SUM OF ITS TOP-2 corrected scores, keep the best `topk_group` groups, `masked_fill(0.0)` the
/// rest, and take the global top-k over the survivors. The emitted weights are the ORIGINAL
/// (un-biased) sigmoid scores of the chosen experts, renormalized and pre-multiplied by `scaling`
/// (`routed_scaling_factor`) — so the downstream `moe_gate`/`moe_down` need no per-expert scale.
/// One workgroup per token; the parallel part is the `E` logit dot-products (E ≤ 256, one/thread),
/// the group/select/weight logic is a short serial tail on thread 0. `mp=(E,h,K,n_group)`,
/// `mp2=(topk_group, bitcast(scaling), _, _)`.
const MOE_ROUTER_V3: &str = r#"
@group(0) @binding(0) var<storage, read> wr: array<f32>; // [E, h] router
@group(0) @binding(1) var<storage, read> bias: array<f32>; // [E] e_score_correction_bias
@group(0) @binding(2) var<storage, read> x: array<f32>; // [n_cols, h] ffn-normed
@group(0) @binding(3) var<storage, read_write> sel: array<u32>; // [n_cols, K]
@group(0) @binding(4) var<storage, read_write> wsel: array<f32>; // [n_cols, K]
@group(0) @binding(5) var<uniform> mp: vec4<u32>; // (E, h, K, n_group)
@group(0) @binding(6) var<uniform> mp2: vec4<u32>; // (topk_group, bitcast(scaling), _, _)
var<workgroup> score: array<f32, 256>; // unbiased sigmoid score
var<workgroup> sfc: array<f32, 256>; // score + bias (selection); consumed during top-k
var<workgroup> gscore: array<f32, 32>; // per-group top-2 sum (n_group ≤ 32)
var<workgroup> gmask: array<u32, 32>; // 1 = group survives
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let e_n = mp.x; let h = mp.y; let kk = mp.z; let ng = mp.w;
let tg = mp2.x;
let scaling = bitcast<f32>(mp2.y);
let col = wid.x;
let lx = lid.x;
// (1) per-expert logit → sigmoid score, plus the bias-corrected selection score.
if (lx < e_n) {
var acc = 0.0;
for (var j = 0u; j < h; j = j + 1u) { acc = acc + wr[lx * h + j] * x[col * h + j]; }
let s = 1.0 / (1.0 + exp(-acc));
score[lx] = s;
sfc[lx] = s + bias[lx];
}
workgroupBarrier();
if (lx == 0u) {
let epg = e_n / ng;
// (2) group score = sum of the top-2 corrected scores in each group.
for (var g = 0u; g < ng; g = g + 1u) {
var t1 = -1e30; var t2 = -1e30;
for (var i = 0u; i < epg; i = i + 1u) {
let val = sfc[g * epg + i];
if (val > t1) { t2 = t1; t1 = val; }
else if (val > t2) { t2 = val; }
}
gscore[g] = t1 + t2;
gmask[g] = 0u;
}
// (3) keep the top `tg` groups (stable argmax, ties → lower group id).
for (var s = 0u; s < tg; s = s + 1u) {
var bg = 0u; var bv = -1e30;
for (var g = 0u; g < ng; g = g + 1u) {
if (gmask[g] == 0u && gscore[g] > bv) { bv = gscore[g]; bg = g; }
}
gmask[bg] = 1u;
}
// (4) masked_fill(0.0) the experts in dropped groups (transformers' exact behavior).
for (var e = 0u; e < e_n; e = e + 1u) {
if (gmask[e / epg] == 0u) { sfc[e] = 0.0; }
}
// (5) global top-k over the masked scores; weight = un-biased sigmoid, renorm, × scaling.
var wsum = 0.0;
for (var s = 0u; s < kk; s = s + 1u) {
var bi = 0u; var bv = -1e30;
for (var e = 0u; e < e_n; e = e + 1u) {
if (sfc[e] > bv) { bv = sfc[e]; bi = e; }
}
sel[col * kk + s] = bi;
wsel[col * kk + s] = score[bi];
wsum = wsum + score[bi];
sfc[bi] = -1e30;
}
for (var s = 0u; s < kk; s = s + 1u) {
wsel[col * kk + s] = (wsel[col * kk + s] / wsum) * scaling;
}
}
}
"#;
/// Slot-indirected expert GLU gate: expert id = `sel[col·K + slot]` (chosen at dispatch-record time
/// only as a SLOT; the id itself is read on-GPU), weights indexed at `eid·mi` row offset into the
/// experts-concatenated gate/up matrices. Input is the ALREADY-normed `x` (the router's input) —
/// no fused norm here, unlike the dense MLP kernel. `dims=(mi, h, slot, K)`; epsm.y = gelu flag.
fn moe_gate_q4_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> s1: array<f16>;
@group(0) @binding(1) var<storage, read> q1: array<u32>;
@group(0) @binding(2) var<storage, read> s3: array<f16>;
@group(0) @binding(3) var<storage, read> q3: array<u32>;
@group(0) @binding(4) var<storage, read> x: array<vec4<f32>>; // [n_cols, h/4] normed
@group(0) @binding(5) var<storage, read> sel: array<u32>;
@group(0) @binding(6) var<storage, read_write> y: array<f32>; // [n_cols, mi]
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (mi, h, slot, K)
@group(0) @binding(8) var<uniform> epsm: vec4<f32>;
const WG: u32 = 32u;
var<workgroup> pg: array<f32, 32>;
var<workgroup> pu: array<f32, 32>;
@compute @workgroup_size(32)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid3: vec3<u32>) {{
let lid = lid3.x;
let mi = dims.x; let col = wid.y / dims.w;
let row_local = wid.x;
if (row_local >= mi) {{ return; }}
// FUSED slots: gy = cols × K, decomposed here (one dispatch replaces top_k serial ones);
// dims.z is unused. Output is [col, K, mi] so slots never collide.
let slot = wid.y % dims.w;
let eid = sel[col * dims.w + slot];
let row = eid * mi + row_local;
let nblk = dims.y / 32u;
let xoff = col * (dims.y / 4u);
var g = 0.0; var u = 0.0;
for (var b = lid; b < nblk; b = b + WG) {{
let d1 = f32(s1[row * nblk + b]); let d3 = f32(s3[row * nblk + b]);
let qb = (row * nblk + b) * 4u;
let xb = b * 8u;
for (var wi = 0u; wi < 4u; wi = wi + 1u) {{
let xlo = x[xoff + xb + wi]; let xhi = x[xoff + xb + 4u + wi];
g = g + d1 * (dot(q4_lo(q1[qb + wi]), xlo) + dot(q4_hi(q1[qb + wi]), xhi));
u = u + d3 * (dot(q4_lo(q3[qb + wi]), xlo) + dot(q4_hi(q3[qb + wi]), xhi));
}}
}}
pg[lid] = g; pu[lid] = u;
workgroupBarrier();
var stride = WG / 2u;
loop {{
if (lid < stride) {{ pg[lid] = pg[lid] + pg[lid + stride]; pu[lid] = pu[lid] + pu[lid + stride]; }}
workgroupBarrier();
if (stride == 1u) {{ break; }}
stride = stride / 2u;
}}
if (lid == 0u) {{
let gate = pg[0]; let upv = pu[0];
var act: f32;
if (epsm.y != 0.0) {{
let g3 = gate * gate * gate;
let targ = clamp(0.7978845608028654 * (gate + 0.044715 * g3), -20.0, 20.0);
act = 0.5 * gate * (1.0 + tanh(targ));
}} else {{
act = gate / (1.0 + exp(-gate));
}}
y[(col * dims.w + slot) * mi + row_local] = act * upv;
}}
}}
"#
)
}
/// Slot-indirected expert down projection with ROUTER-WEIGHTED accumulate:
/// `y[col,row] += wsel[col·K+slot] · (W_down[eid]·g[col])`, `eid = sel[col·K+slot]`, weights at
/// `eid·h` row offset. `dims=(h, mi, slot, K)`.
fn moe_down_q4_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> scales: array<f16>;
@group(0) @binding(1) var<storage, read> quants: array<u32>;
@group(0) @binding(2) var<storage, read> g: array<vec4<f32>>; // [n_cols, mi/4]
@group(0) @binding(3) var<storage, read> sel: array<u32>;
@group(0) @binding(4) var<storage, read> wsel: array<f32>;
@group(0) @binding(5) var<storage, read_write> y: array<f32>; // [n_cols, h]
@group(0) @binding(6) var<uniform> dims: vec4<u32>; // (h, mi, slot, K)
const WG: u32 = 64u;
var<workgroup> partial: array<f32, 64>;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let h = dims.x; let col = wid.y;
let row_local = wid.x;
if (row_local >= h) {{ return; }}
let nblk = dims.y / 32u;
// FUSED slots: loop the top-k experts INSIDE the kernel (dims.z unused). The per-slot
// accumulation into y keeps the exact order of the old one-dispatch-per-slot chain, so the
// result is bitwise-identical while the dispatch count drops top_k-fold.
var acc = 0.0;
for (var slot = 0u; slot < dims.w; slot = slot + 1u) {{
let eid = sel[col * dims.w + slot];
let row = eid * h + row_local;
let goff = (col * dims.w + slot) * (dims.y / 4u);
var a = 0.0;
for (var b = lid.x; b < nblk; b = b + WG) {{
let d = f32(scales[row * nblk + b]);
let qb = (row * nblk + b) * 4u;
let xb = b * 8u;
for (var wi = 0u; wi < 4u; wi = wi + 1u) {{
a = a + d * (dot(q4_lo(quants[qb + wi]), g[goff + xb + wi]) + dot(q4_hi(quants[qb + wi]), g[goff + xb + 4u + wi]));
}}
}}
partial[lid.x] = a;
workgroupBarrier();
for (var st = 32u; st > 0u; st = st >> 1u) {{
if (lid.x < st) {{ partial[lid.x] = partial[lid.x] + partial[lid.x + st]; }}
workgroupBarrier();
}}
if (lid.x == 0u) {{ acc = acc + wsel[col * dims.w + slot] * partial[0]; }}
workgroupBarrier();
}}
if (lid.x == 0u) {{
y[col * h + row_local] = y[col * h + row_local] + acc;
}}
}}
"#
)
}
/// Per-column-snapshot variant of [`DN_CONV_K`]: binding 8 captures the slot's conv ring AFTER
/// each column (`csnap[p]` = ring state once column p is consumed) — the speculative-verify
/// rollback restores `csnap[j]` on partial acceptance. Derived textually; the compute path is
/// byte-identical (the snapshot is write-only extra traffic, ~200 KB/column/layer).
fn dn_conv_k_snap_src() -> String {
let src = DN_CONV_K.to_string();
for frag in [
"@group(0) @binding(7) var<uniform> dims2: vec4<u32>; // (in_rows, _, _, _)",
" ring[base + kk - 1u] = xn;",
] {
assert!(src.contains(frag), "DN_CONV_K drifted: {frag}");
}
src.replace(
"@group(0) @binding(7) var<uniform> dims2: vec4<u32>; // (in_rows, _, _, _)",
"@group(0) @binding(7) var<uniform> dims2: vec4<u32>; // (in_rows, _, _, _)\n@group(0) @binding(8) var<storage, read_write> csnap: array<f32>; // [k, conv_dim, K]",
)
.replace(
" ring[base + kk - 1u] = xn;",
" ring[base + kk - 1u] = xn;\n for (var t = 0u; t < kk; t = t + 1u) { csnap[(p * cd + c) * kk + t] = ring[base + t]; }",
)
}
/// Per-column-snapshot variant of [`DN_STEP_PK`]: binding 10 captures the slot's recurrent
/// state slab AFTER each column (`snap[p]` = state once column p is applied). Same math.
/// Snapshot variant of [`DN_STEP_RES`]: after each column's update, the resident half-state
/// is also written to `snap[p]` (same layout/semantics as [`dn_step_pk_snap_src`]: state AFTER
/// column p, the rollback target for accepting through p).
fn dn_step_res_snap_src() -> String {
let src = DN_STEP_RES.to_string();
for frag in [
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;",
" core[p * nv * dv + h * dv + j] = o;",
] {
assert!(src.contains(frag), "DN_STEP_RES drifted: {frag}");
}
src.replace(
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;",
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;\n@group(0) @binding(10) var<storage, read_write> snap: array<f32>; // [k, nv, dk, dv]",
)
.replace(
" core[p * nv * dv + h * dv + j] = o;",
" core[p * nv * dv + h * dv + j] = o;\n for (var i = 0u; i < dk; i = i + 1u) { snap[((p * nv + h) * dk + i) * dv + j] = S[i * 64u + jl]; }",
)
}
fn dn_step_pk_snap_src() -> String {
let src = DN_STEP_PK.to_string();
for frag in [
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;",
" osh[j] = o;",
] {
assert!(src.contains(frag), "DN_STEP_PK drifted: {frag}");
}
src.replace(
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;",
"@group(0) @binding(9) var<uniform> epsu: vec4<f32>;\n@group(0) @binding(10) var<storage, read_write> snap: array<f32>; // [k, nv, dk, dv]",
)
.replace(
" osh[j] = o;",
" osh[j] = o;\n for (var i = 0u; i < dk; i = i + 1u) { snap[((p * nv + h) * dk + i) * dv + j] = st[sbase + i * dv]; }",
)
}
/// Qwen3.5 FUSED expert down-projection: absorbs the residual-ADD dispatch — the epilogue
/// finishes the block in place: `cur += (shared_out · g) + routed_acc`, the exact FP sequence
/// of the old sgate-scale → `y += acc` → `cur += y` chain (given `g` from the fused router),
/// so results stay bitwise while two dispatches per MoE layer disappear.
/// nbar-shaped expert GATE+UP (256-thread row-strips, x from L2, barriers only in the final
/// 16-lane trees): the MoE expert weights are the small-batch verify's dominant traffic
/// (k·top_k re-reads/layer) and the WG=32 kernel ran at ~44 GB/s — the row-strip shape measured
/// 2-3× on V100. Reduction is a 16-lane tree (vs 32) — batch==solo contracts hold because every
/// plan uses the same kernel on both sides; the CPU-reference gate is tolerance-based.
fn moe_gate_q4_nbar_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> s1: array<f16>;
@group(0) @binding(1) var<storage, read> q1: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> s3: array<f16>;
@group(0) @binding(3) var<storage, read> q3: array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> x: array<vec4<f32>>; // [n_cols, h/4] normed
@group(0) @binding(5) var<storage, read> sel: array<u32>;
@group(0) @binding(6) var<storage, read_write> y: array<f32>; // [n_cols, K, mi]
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (mi, h, _, K)
@group(0) @binding(8) var<uniform> epsm: vec4<f32>;
var<workgroup> rg: array<f32, 256>;
var<workgroup> ru: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid3: vec3<u32>) {{
let lid = lid3.x;
let mi = dims.x; let h = dims.y;
let col = wid.y / dims.w;
let slot = wid.y % dims.w;
let row_local = wid.x * 16u + lid / 16u;
let lane = lid % 16u;
let eid = sel[col * dims.w + slot];
let row = eid * mi + row_local;
let nblk = h / 32u;
let xoff = col * (h / 4u);
var g = 0.0;
var u = 0.0;
for (var b = lane; b < nblk; b = b + 16u) {{
var d1 = 0.0;
var d3 = 0.0;
var w1 = vec4<u32>();
var w3 = vec4<u32>();
if (row_local < mi) {{
d1 = f32(s1[row * nblk + b]);
d3 = f32(s3[row * nblk + b]);
w1 = q1[row * nblk + b];
w3 = q3[row * nblk + b];
}}
let xb = xoff + b * 8u;
let a0 = x[xb]; let a1 = x[xb + 1u];
let a2 = x[xb + 2u]; let a3 = x[xb + 3u];
let a4 = x[xb + 4u]; let a5 = x[xb + 5u];
let a6 = x[xb + 6u]; let a7 = x[xb + 7u];
var sg1 = dot(q4_lo(w1.x), a0) + dot(q4_hi(w1.x), a4);
sg1 = sg1 + dot(q4_lo(w1.y), a1) + dot(q4_hi(w1.y), a5);
sg1 = sg1 + dot(q4_lo(w1.z), a2) + dot(q4_hi(w1.z), a6);
sg1 = sg1 + dot(q4_lo(w1.w), a3) + dot(q4_hi(w1.w), a7);
g = g + d1 * sg1;
var su = dot(q4_lo(w3.x), a0) + dot(q4_hi(w3.x), a4);
su = su + dot(q4_lo(w3.y), a1) + dot(q4_hi(w3.y), a5);
su = su + dot(q4_lo(w3.z), a2) + dot(q4_hi(w3.z), a6);
su = su + dot(q4_lo(w3.w), a3) + dot(q4_hi(w3.w), a7);
u = u + d3 * su;
}}
rg[lid] = g;
ru[lid] = u;
workgroupBarrier();
if (lane < 8u) {{ rg[lid] = rg[lid] + rg[lid + 8u]; ru[lid] = ru[lid] + ru[lid + 8u]; }}
workgroupBarrier();
if (lane < 4u) {{ rg[lid] = rg[lid] + rg[lid + 4u]; ru[lid] = ru[lid] + ru[lid + 4u]; }}
workgroupBarrier();
if (lane < 2u) {{ rg[lid] = rg[lid] + rg[lid + 2u]; ru[lid] = ru[lid] + ru[lid + 2u]; }}
workgroupBarrier();
if (lane == 0u && row_local < mi) {{
let gate = rg[lid] + rg[lid + 1u];
let upv = ru[lid] + ru[lid + 1u];
var act: f32;
if (epsm.y != 0.0) {{
let g3 = gate * gate * gate;
let targ = clamp(0.7978845608028654 * (gate + 0.044715 * g3), -20.0, 20.0);
act = 0.5 * gate * (1.0 + tanh(targ));
}} else {{
act = gate / (1.0 + exp(-gate));
}}
y[(col * dims.w + slot) * mi + row_local] = act * upv;
}}
}}
"#
)
}
/// nbar-shaped fused expert DOWN (row-strips; slots accumulated per lane then ONE tree; the
/// epilogue joins the residual exactly like [`moe_down_q4_fused_src`]).
fn moe_down_q4_nbar_fused_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> scales: array<f16>;
@group(0) @binding(1) var<storage, read> quants: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> g: array<vec4<f32>>; // [n_cols, K, mi/4]
@group(0) @binding(3) var<storage, read> sel: array<u32>;
@group(0) @binding(4) var<storage, read> wsel: array<f32>;
@group(0) @binding(5) var<storage, read> y: array<f32>; // [n_cols, h] shared-out
@group(0) @binding(6) var<uniform> dims: vec4<u32>; // (h, mi, _, K)
@group(0) @binding(7) var<storage, read> sg: array<f32>; // [n_cols] shared gate
@group(0) @binding(8) var<storage, read_write> cur: array<f32>; // [n_cols, h] residual
var<workgroup> red: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid3: vec3<u32>) {{
let lid = lid3.x;
let h = dims.x; let mi = dims.y; let kx = dims.w;
let col = wid.y;
let row_local = wid.x * 16u + lid / 16u;
let lane = lid % 16u;
let nblk = mi / 32u;
var acc = 0.0;
for (var slot = 0u; slot < kx; slot = slot + 1u) {{
let eid = sel[col * kx + slot];
let row = eid * h + row_local;
let goff = (col * kx + slot) * (mi / 4u);
var a = 0.0;
for (var b = lane; b < nblk; b = b + 16u) {{
var d = 0.0;
var q = vec4<u32>();
if (row_local < h) {{
d = f32(scales[row * nblk + b]);
q = quants[row * nblk + b];
}}
let gb = goff + b * 8u;
var sv = dot(q4_lo(q.x), g[gb]) + dot(q4_hi(q.x), g[gb + 4u]);
sv = sv + dot(q4_lo(q.y), g[gb + 1u]) + dot(q4_hi(q.y), g[gb + 5u]);
sv = sv + dot(q4_lo(q.z), g[gb + 2u]) + dot(q4_hi(q.z), g[gb + 6u]);
sv = sv + dot(q4_lo(q.w), g[gb + 3u]) + dot(q4_hi(q.w), g[gb + 7u]);
a = a + d * sv;
}}
acc = acc + wsel[col * kx + slot] * a;
}}
red[lid] = acc;
workgroupBarrier();
if (lane < 8u) {{ red[lid] = red[lid] + red[lid + 8u]; }}
workgroupBarrier();
if (lane < 4u) {{ red[lid] = red[lid] + red[lid + 4u]; }}
workgroupBarrier();
if (lane < 2u) {{ red[lid] = red[lid] + red[lid + 2u]; }}
workgroupBarrier();
if (lane == 0u && row_local < h) {{
let o = col * h + row_local;
cur[o] = cur[o] + (y[o] * sg[col] + red[lid] + red[lid + 1u]);
}}
}}
"#
)
}
/// llama.cpp-shaped expert GATE+UP (WG=32=subgroup, 4 rows sharing activation loads,
/// subgroupAdd, zero barriers) — the expert projections are the largest per-layer kernel time
/// of the small-width verify (row-strip shape ≈ 370 µs/layer at 5 cols; this shape measured
/// 2-4× on every other GEMV family). Grid: (⌈mi/4⌉, n_cols·K).
/// llama.cpp-shaped expert GATE+UP GEMV — the kernel that actually runs the MoE on Vulkan/V100
/// (see `lcpp_family`), and the one that dominates the speculative verify: at a k-column verify it
/// moves `k × top_k × (w1+w3)` bytes with no cross-column reuse, because each column routes to its
/// own experts.
///
/// `pub` so `kernel-lab` can micro-benchmark the REAL production source rather than a copy — a lab
/// baseline that has drifted from what ships turns every "win" into a strawman comparison.
pub fn moe_gate_q4_lcpp_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> s1: array<f16>;
@group(0) @binding(1) var<storage, read> q1: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> s3: array<f16>;
@group(0) @binding(3) var<storage, read> q3: array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> x: array<vec4<f32>>; // [n_cols, h/4] normed
@group(0) @binding(5) var<storage, read> sel: array<u32>;
@group(0) @binding(6) var<storage, read_write> y: array<f32>; // [n_cols, K, mi]
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (mi, h, _, K)
@group(0) @binding(8) var<uniform> epsm: vec4<f32>;
@compute @workgroup_size(32)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(subgroup_invocation_id) sid: u32) {{
let mi = dims.x; let h = dims.y;
let col = wid.y / dims.w;
let slot = wid.y % dims.w;
let eid = sel[col * dims.w + slot];
let row0 = wid.x * 4u;
let nblk = h / 32u;
let xoff = col * (h / 4u);
var ag = vec4<f32>(0.0);
var au = vec4<f32>(0.0);
for (var b = sid; b < nblk; b = b + 32u) {{
let xb = xoff + b * 8u;
let v0 = x[xb]; let v1 = x[xb + 1u];
let v2 = x[xb + 2u]; let v3 = x[xb + 3u];
let v4 = x[xb + 4u]; let v5 = x[xb + 5u];
let v6 = x[xb + 6u]; let v7 = x[xb + 7u];
for (var r = 0u; r < 4u; r = r + 1u) {{
let row = eid * mi + min(row0 + r, mi - 1u);
{{
let d = f32(s1[row * nblk + b]);
let q = q1[row * nblk + b];
var sv = dot(q4_lo(q.x), v0) + dot(q4_hi(q.x), v4);
sv = sv + dot(q4_lo(q.y), v1) + dot(q4_hi(q.y), v5);
sv = sv + dot(q4_lo(q.z), v2) + dot(q4_hi(q.z), v6);
sv = sv + dot(q4_lo(q.w), v3) + dot(q4_hi(q.w), v7);
ag[r] = ag[r] + d * sv;
}}
{{
let d = f32(s3[row * nblk + b]);
let q = q3[row * nblk + b];
var sv = dot(q4_lo(q.x), v0) + dot(q4_hi(q.x), v4);
sv = sv + dot(q4_lo(q.y), v1) + dot(q4_hi(q.y), v5);
sv = sv + dot(q4_lo(q.z), v2) + dot(q4_hi(q.z), v6);
sv = sv + dot(q4_lo(q.w), v3) + dot(q4_hi(q.w), v7);
au[r] = au[r] + d * sv;
}}
}}
}}
let tg = subgroupAdd(ag);
let tu = subgroupAdd(au);
if (sid == 0u) {{
for (var r = 0u; r < 4u; r = r + 1u) {{
if (row0 + r < mi) {{
let gate = tg[r];
let upv = tu[r];
var act: f32;
if (epsm.y != 0.0) {{
let g3 = gate * gate * gate;
let targ = clamp(0.7978845608028654 * (gate + 0.044715 * g3), -20.0, 20.0);
act = 0.5 * gate * (1.0 + tanh(targ));
}} else {{
act = gate / (1.0 + exp(-gate));
}}
y[(col * dims.w + slot) * mi + row0 + r] = act * upv;
}}
}}
}}
}}
"#
)
}
/// llama.cpp-shaped fused expert DOWN (same shape; slots accumulated in-lane, residual epilogue
/// identical to the row-strip fused variant).
fn moe_down_q4_lcpp_fused_src() -> String {
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> scales: array<f16>;
@group(0) @binding(1) var<storage, read> quants: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> g: array<vec4<f32>>; // [n_cols, K, mi/4]
@group(0) @binding(3) var<storage, read> sel: array<u32>;
@group(0) @binding(4) var<storage, read> wsel: array<f32>;
@group(0) @binding(5) var<storage, read> y: array<f32>; // [n_cols, h] shared-out
@group(0) @binding(6) var<uniform> dims: vec4<u32>; // (h, mi, _, K)
@group(0) @binding(7) var<storage, read> sg: array<f32>; // [n_cols] shared gate
@group(0) @binding(8) var<storage, read_write> cur: array<f32>; // [n_cols, h] residual
@compute @workgroup_size(32)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(subgroup_invocation_id) sid: u32) {{
let h = dims.x; let mi = dims.y; let kx = dims.w;
let col = wid.y;
let row0 = wid.x * 4u;
let nblk = mi / 32u;
var acc = vec4<f32>(0.0);
for (var slot = 0u; slot < kx; slot = slot + 1u) {{
let eid = sel[col * kx + slot];
let w = wsel[col * kx + slot];
let goff = (col * kx + slot) * (mi / 4u);
for (var b = sid; b < nblk; b = b + 32u) {{
let gb = goff + b * 8u;
let v0 = g[gb]; let v1 = g[gb + 1u];
let v2 = g[gb + 2u]; let v3 = g[gb + 3u];
let v4 = g[gb + 4u]; let v5 = g[gb + 5u];
let v6 = g[gb + 6u]; let v7 = g[gb + 7u];
for (var r = 0u; r < 4u; r = r + 1u) {{
let row = eid * h + min(row0 + r, h - 1u);
let d = f32(scales[row * nblk + b]);
let q = quants[row * nblk + b];
var sv = dot(q4_lo(q.x), v0) + dot(q4_hi(q.x), v4);
sv = sv + dot(q4_lo(q.y), v1) + dot(q4_hi(q.y), v5);
sv = sv + dot(q4_lo(q.z), v2) + dot(q4_hi(q.z), v6);
sv = sv + dot(q4_lo(q.w), v3) + dot(q4_hi(q.w), v7);
acc[r] = acc[r] + w * d * sv;
}}
}}
}}
let tot = subgroupAdd(acc);
if (sid == 0u) {{
for (var r = 0u; r < 4u; r = r + 1u) {{
if (row0 + r < h) {{
let o = col * h + row0 + r;
cur[o] = cur[o] + (y[o] * sg[col] + tot[r]);
}}
}}
}}
}}
"#
)
}
fn moe_down_q4_fused_src() -> String {
let src = moe_down_q4_src();
// NOTE: `src` is the FORMATTED output (braces already single) — match post-format text.
for frag in [
"@group(0) @binding(5) var<storage, read_write> y: array<f32>; // [n_cols, h]",
" if (lid.x == 0u) {\n y[col * h + row_local] = y[col * h + row_local] + acc;\n }",
] {
assert!(src.contains(frag), "moe_down source drifted: {frag}");
}
src.replace(
"@group(0) @binding(5) var<storage, read_write> y: array<f32>; // [n_cols, h]",
"@group(0) @binding(5) var<storage, read> y: array<f32>; // [n_cols, h] shared-out (UNSCALED)\n@group(0) @binding(7) var<storage, read> sg: array<f32>; // [n_cols] shared gate\n@group(0) @binding(8) var<storage, read_write> cur: array<f32>; // [n_cols, h] residual",
)
.replace(
" if (lid.x == 0u) {\n y[col * h + row_local] = y[col * h + row_local] + acc;\n }",
" if (lid.x == 0u) {\n let o = col * h + row_local;\n cur[o] = cur[o] + (y[o] * sg[col] + acc);\n }",
)
}
/// GROUPED-MoE tile builder (one 256-thread WG): histogram the round's `n_pairs = ncols×K`
/// routed (col, slot) pairs by expert, prefix-sum, scatter into an expert-sorted order (shared
/// memory — n_pairs ≤ 1536 fits), then emit TILES of up to 8 same-expert pairs
/// (`tile_expert[t]`, `tile_pairs[t·8+i]`, 0xFFFFFFFF pads) plus the two indirect-dispatch arg
/// triples for the grouped gate/down grids. Prefill-width MoE has expert multiplicity
/// ncols·K/E (≈6 at 192 cols) — each tile streams its expert's weights ONCE for all its pairs,
/// the guaranteed dedup the decode-width sort couldn't buy (L2 already covered multiplicity 2).
const MOE_SEG_BUILD: &str = r#"
@group(0) @binding(0) var<storage, read> sel: array<u32>; // [n_pairs] expert ids
@group(0) @binding(1) var<storage, read_write> texp: array<u32>; // [n_pairs] tile→expert
@group(0) @binding(2) var<storage, read_write> tpair: array<u32>; // [n_pairs, 8] tile pairs
@group(0) @binding(3) var<storage, read_write> args: array<u32>; // [8] two dispatch triples
@group(0) @binding(4) var<uniform> dims: vec4<u32>; // (n_pairs, E, gx_gate, gx_down)
const EMAX = 256u;
var<workgroup> cnt: array<atomic<u32>, 256>;
var<workgroup> off: array<u32, 257>;
var<workgroup> curs: array<atomic<u32>, 256>;
var<workgroup> perm: array<u32, 1536>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
let np = dims.x;
let e_n = dims.y;
let t = lid.x;
atomicStore(&cnt[t], 0u);
workgroupBarrier();
for (var i = t; i < np; i = i + 256u) { atomicAdd(&cnt[min(sel[i], EMAX - 1u)], 1u); }
workgroupBarrier();
if (t == 0u) {
var run = 0u;
for (var e = 0u; e < e_n; e = e + 1u) {
off[e] = run;
run = run + atomicLoad(&cnt[e]);
}
off[e_n] = run;
}
workgroupBarrier();
atomicStore(&curs[t], off[min(t, e_n - 1u)]);
workgroupBarrier();
for (var i = t; i < np; i = i + 256u) {
let pos = atomicAdd(&curs[min(sel[i], EMAX - 1u)], 1u);
perm[pos] = i;
}
workgroupBarrier();
if (t == 0u) {
var nt = 0u;
for (var e = 0u; e < e_n; e = e + 1u) {
var i = off[e];
while (i < off[e + 1u]) {
texp[nt] = e;
for (var j = 0u; j < 8u; j = j + 1u) {
if (i + j < off[e + 1u]) { tpair[nt * 8u + j] = perm[i + j]; }
else { tpair[nt * 8u + j] = 0xFFFFFFFFu; }
}
nt = nt + 1u;
i = i + 8u;
}
}
args[0] = dims.z; args[1] = nt; args[2] = 1u; args[3] = 0u;
args[4] = dims.w; args[5] = nt; args[6] = 1u; args[7] = 0u;
}
}
"#;
/// Grouped-MoE GATE (indirect grid `(mi/64, n_tiles)`): one WG per (row band, tile) streams
/// the tile's expert w1/w3 rows ONCE and dots them against the ≤8 staged pair columns —
/// per-pair weight traffic ÷ tile occupancy. Inner product fully UNROLLED over the 8 pairs
/// (static registers — dynamically-indexed privates land in DRAM scratch, the measured 4.2×
/// GEMM lesson). Output layout identical to the per-pair gate (`y[pair·mi + row]`), so the
/// downstream down kernels are unaffected. Pure baseline WGSL (no subgroups) — every adapter.
fn moe_gate_grouped_src() -> String {
let mut inner = String::new();
for w4 in 0..4 {
let comp = ["x", "y", "z", "w"][w4];
inner.push_str(&format!(
" let ql1_{w4} = d1 * q4_lo(q1v.{comp}); let qh1_{w4} = d1 * q4_hi(q1v.{comp});\n\
\x20 let ql3_{w4} = d3 * q4_lo(q3v.{comp}); let qh3_{w4} = d3 * q4_hi(q3v.{comp});\n"
));
for p in 0..8 {
inner.push_str(&format!(
" g{p} = g{p} + dot(ql1_{w4}, xs[{lo}u + {p}u]) + dot(qh1_{w4}, xs[{hi}u + {p}u]);\n\
\x20 u{p} = u{p} + dot(ql3_{w4}, xs[{lo}u + {p}u]) + dot(qh3_{w4}, xs[{hi}u + {p}u]);\n",
lo = w4 * 8,
hi = (4 + w4) * 8,
));
}
}
let mut decl = String::new();
for p in 0..8 {
decl.push_str(&format!(" var g{p} = 0.0;\n var u{p} = 0.0;\n"));
}
let mut epi = String::new();
for p in 0..8 {
epi.push_str(&format!(
" {{\n let pair = tpair[t * 8u + {p}u];\n if (pair != 0xFFFFFFFFu) {{\n var act: f32;\n if (epsm.y != 0.0) {{\n let g3 = g{p} * g{p} * g{p};\n let targ = clamp(0.7978845608028654 * (g{p} + 0.044715 * g3), -20.0, 20.0);\n act = 0.5 * g{p} * (1.0 + tanh(targ));\n }} else {{\n act = g{p} / (1.0 + exp(-g{p}));\n }}\n y[pair * mi + row] = act * u{p};\n }}\n }}\n"
));
}
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> s1: array<f16>;
@group(0) @binding(1) var<storage, read> q1: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> s3: array<f16>;
@group(0) @binding(3) var<storage, read> q3: array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> x: array<vec4<f32>>; // [ncols, h/4] normed
@group(0) @binding(5) var<storage, read> texp: array<u32>;
@group(0) @binding(6) var<storage, read> tpair: array<u32>;
@group(0) @binding(7) var<storage, read_write> y: array<f32>; // [n_pairs, mi]
@group(0) @binding(8) var<uniform> dims: vec4<u32>; // (mi, h, _, K)
@group(0) @binding(9) var<uniform> epsm: vec4<f32>;
var<workgroup> xs: array<vec4<f32>, 64>; // [8 k-vec4s][8 pairs] block tile of x
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let mi = dims.x; let h = dims.y; let kx = dims.w;
let t = wid.y;
let e = texp[t];
let row = wid.x * 64u + lid.x;
let nblk = h / 32u;
let wrow = (e * mi + min(row, mi - 1u)) * nblk;
{decl}
for (var b = 0u; b < nblk; b = b + 1u) {{
// Stage the block's x for the 8 pairs: xs[e4·8 + p] = vec4 e4 of pair p's column.
{{
let p = lid.x % 8u;
let e4 = lid.x / 8u;
let pair = tpair[t * 8u + p];
var v = vec4<f32>(0.0);
if (pair != 0xFFFFFFFFu) {{
let col = pair / kx;
v = x[col * (h / 4u) + b * 8u + e4];
}}
xs[e4 * 8u + p] = v;
}}
workgroupBarrier();
{{
let d1 = f32(s1[wrow + b]);
let q1v = q1[wrow + b];
let d3 = f32(s3[wrow + b]);
let q3v = q3[wrow + b];
{inner}
}}
workgroupBarrier();
}}
if (row >= mi) {{ return; }}
{epi}
}}
"#
)
}
/// Grouped-MoE DOWN (indirect grid `(h/64, n_tiles)`): streams the tile's expert w2 rows once,
/// dots against ≤8 staged pair activations, and writes `wsel`-scaled per-pair PARTIALS
/// (`partial[pair·h + row]`) — the cross-pair accumulation into `cur` would race across tiles
/// (no f32 atomics in core WGSL), so [`MOE_DOWN_REDUCE`] folds the K partials + shared-expert
/// epilogue per column afterwards. Same unroll discipline as the gate.
fn moe_down_grouped_src() -> String {
let mut inner = String::new();
for w4 in 0..4 {
let comp = ["x", "y", "z", "w"][w4];
inner.push_str(&format!(
" let ql_{w4} = d * q4_lo(qv.{comp}); let qh_{w4} = d * q4_hi(qv.{comp});\n"
));
for p in 0..8 {
inner.push_str(&format!(
" a{p} = a{p} + dot(ql_{w4}, gs[{lo}u + {p}u]) + dot(qh_{w4}, gs[{hi}u + {p}u]);\n",
lo = w4 * 8,
hi = (4 + w4) * 8,
));
}
}
let mut decl = String::new();
for p in 0..8 {
decl.push_str(&format!(" var a{p} = 0.0;\n"));
}
let mut epi = String::new();
for p in 0..8 {
epi.push_str(&format!(
" {{\n let pair = tpair[t * 8u + {p}u];\n if (pair != 0xFFFFFFFFu) {{\n partial[pair * h + row] = wsel[pair] * a{p};\n }}\n }}\n"
));
}
let q4fn = crate::Q4_FN;
format!(
r#"{q4fn}
@group(0) @binding(0) var<storage, read> scales: array<f16>;
@group(0) @binding(1) var<storage, read> quants: array<vec4<u32>>;
@group(0) @binding(2) var<storage, read> g: array<vec4<f32>>; // [n_pairs, mi/4]
@group(0) @binding(3) var<storage, read> texp: array<u32>;
@group(0) @binding(4) var<storage, read> tpair: array<u32>;
@group(0) @binding(5) var<storage, read> wsel: array<f32>; // [n_pairs]
@group(0) @binding(6) var<storage, read_write> partial: array<f32>; // [n_pairs, h]
@group(0) @binding(7) var<uniform> dims: vec4<u32>; // (h, mi, _, K)
var<workgroup> gs: array<vec4<f32>, 64>; // [8 k-vec4s][8 pairs] block tile of g
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {{
let h = dims.x; let mi = dims.y;
let t = wid.y;
let e = texp[t];
let row = wid.x * 64u + lid.x;
let nblk = mi / 32u;
let wrow = (e * h + min(row, h - 1u)) * nblk;
{decl}
for (var b = 0u; b < nblk; b = b + 1u) {{
{{
let p = lid.x % 8u;
let e4 = lid.x / 8u;
let pair = tpair[t * 8u + p];
var v = vec4<f32>(0.0);
if (pair != 0xFFFFFFFFu) {{
v = g[pair * (mi / 4u) + b * 8u + e4];
}}
gs[e4 * 8u + p] = v;
}}
workgroupBarrier();
{{
let d = f32(scales[wrow + b]);
let qv = quants[wrow + b];
{inner}
}}
workgroupBarrier();
}}
if (row >= h) {{ return; }}
{epi}
}}
"#
)
}
/// Grouped-MoE epilogue: `cur[col] += mlpout[col]·sg[col] + Σ_slot partial[col·K+slot]` —
/// the slot sum runs in the SAME 0..K-1 order as the fused down's in-kernel loop. Grid
/// `(h/64, ncols)`.
const MOE_DOWN_REDUCE: &str = r#"
@group(0) @binding(0) var<storage, read> partial: array<f32>; // [n_pairs, h]
@group(0) @binding(1) var<storage, read> yshared: array<f32>; // [ncols, h] UNSCALED
@group(0) @binding(2) var<storage, read> sg: array<f32>; // [ncols]
@group(0) @binding(3) var<storage, read_write> cur: array<f32>; // [ncols, h]
@group(0) @binding(4) var<uniform> dims: vec4<u32>; // (h, _, _, K)
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let h = dims.x; let kx = dims.w;
let col = wid.y;
let row = wid.x * 64u + lid.x;
if (row >= h) { return; }
let o = col * h + row;
var s = yshared[o] * sg[col];
for (var slot = 0u; slot < kx; slot = slot + 1u) {
s = s + partial[(col * kx + slot) * h + row];
}
cur[o] = cur[o] + s;
}
"#;
/// Test-only re-exports: the kernel benches build a pipeline/bind-group by hand so they
/// measure EXACTLY what the engine dispatches (see tests/moshi_kernel_bw.rs).
pub fn pipeline_pub(ctx: &GpuCtx, label: &str, src: &str) -> wgpu::ComputePipeline {
pipeline(ctx, label, src)
}
pub fn uni_pub(ctx: &GpuCtx, data: &[u8]) -> wgpu::Buffer {
uni(ctx, data)
}
pub fn make_bg_pub(
ctx: &GpuCtx,
pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer],
dims: &wgpu::Buffer,
) -> wgpu::BindGroup {
make_bg(ctx, pl, bufs, dims)
}
/// Shader modules compile UNCHECKED by default: naga's checked mode rides an array-bounds
/// clamp on every buffer access — measured +9.5% e2e on the OCR engine just from dropping it
/// in the plan kernels (numerics byte-identical; in-bounds clamps never fire). The kernel
/// generators' asserts + the gate suite (parity, spec, e2e strings) are the OOB protection.
/// `OSFKB_UNCHECKED_SHADERS=0` restores the clamps for debugging a suspected OOB.
pub(crate) fn shader_checks() -> wgpu::ShaderRuntimeChecks {
static UNCHECKED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
if *UNCHECKED
.get_or_init(|| std::env::var("OSFKB_UNCHECKED_SHADERS").ok().as_deref() != Some("0"))
{
wgpu::ShaderRuntimeChecks::unchecked()
} else {
wgpu::ShaderRuntimeChecks::checked()
}
}
/// [`shader_checks`]-honoring replacement for `Device::create_shader_module` — every runtime
/// shader site routes through this so ALL inference (towers, DeltaNet, encoder, grammar)
/// gets the unchecked default, not just the batch-plan kernels.
pub(crate) trait ShaderModuleTuned {
fn shader_module_tuned(&self, desc: wgpu::ShaderModuleDescriptor<'_>) -> wgpu::ShaderModule;
}
impl ShaderModuleTuned for wgpu::Device {
fn shader_module_tuned(&self, desc: wgpu::ShaderModuleDescriptor<'_>) -> wgpu::ShaderModule {
// SAFETY: unchecked() promises in-bounds indexing — held by the generators and gates.
unsafe { self.create_shader_module_trusted(desc, shader_checks()) }
}
}
pub(crate) fn pipeline(ctx: &GpuCtx, label: &str, src: &str) -> wgpu::ComputePipeline {
// SAFETY: with unchecked() this promises the WGSL never indexes out of bounds — held by
// the kernel generators' asserts and the gate suite; the flag defaults off.
let m = unsafe {
ctx.device.create_shader_module_trusted(
wgpu::ShaderModuleDescriptor {
label: Some(label),
source: wgpu::ShaderSource::Wgsl(src.into()),
},
shader_checks(),
)
};
ctx.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(label),
layout: None,
module: &m,
entry_point: Some("main"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
})
}
fn u32x4(a: u32, b: u32, c: u32, d: u32) -> [u32; 4] {
[a, b, c, d]
}
/// The GPU LFM2 engine: a **pre-built dispatch plan** (all bind groups + uniforms baked once) that is
/// replayed every token, so the per-token CPU cost is just recording + a few small `write_buffer`s —
/// not ~108 bind-group/uniform allocations. The plan's bind groups hold Arc refs to the buffers, so
/// they stay valid after `new` moves everything into `Self`.
/// Key of the pre-encoded command-buffer cache: (plan serial, ncols, need_logit mask).
type CmdbufKey = (u64, u32, u64);
/// Pre-encoded command buffer for the next batch step, keyed by [`CmdbufKey`].
type CmdbufCache = std::cell::RefCell<Option<(CmdbufKey, wgpu::CommandBuffer)>>;
/// llama.cpp-family eligibility (WG=32 `subgroupAdd` kernels): adapters with VALIDATED 32-wide
/// subgroup semantics — a reported 32..32 guarantee, or `GpuCtx::subgroups32_effective`'s
/// runtime probe of the actual lcpp pipeline (integer-exact). The old vendor-shaped guard
/// ("Metal reports variable subgroup sizes and produced garbage under force") is superseded by
/// the probe: garbage semantics now FAIL the probe and keep the tree family, healthy semantics
/// qualify, and the force flag stays a no-op anywhere unhealthy. Shared by the batch plans AND
/// the M=1 plan so solo and batched columns pick the same family on a given adapter
/// (single-family FP order is what keeps `batched == solo` bitwise).
///
/// DEFAULT = every validated adapter, measured on both backends: V100/Vulkan 3-7.4× per kernel
/// at 1-16 columns; Metal/M4 Max +24% end-to-end on the 4B dense decode (65.7 → 81.6 tok/s,
/// interleaved min-of-3, 2026-07-15 — the old `v3-sg` family default predates the probe, when
/// Metal could not run lcpp at all). `OSFKB_GEMV_LCPP=0` pins the old per-backend family.
fn lcpp_family(ctx: &GpuCtx) -> bool {
let ok = ctx.subgroups32_effective();
match std::env::var("OSFKB_GEMV_LCPP").ok().as_deref() {
Some("0") => false,
_ => ok,
}
}
/// lcpp-family MLP: split rmsnorm + the two-stream lcpp gate/up kernel instead of the fused
/// `mlp_k` (which was the largest per-layer kernel left on the old family once the plain GEMVs
/// moved). Decode-width plans only — wide prefill keeps the tiled GEMM/fused shapes.
/// `OSFKB_MLP_LCPP=0` pins the fused kernel (A/B kill switch).
fn mlp_lcpp_enabled() -> bool {
std::env::var("OSFKB_MLP_LCPP").ok().as_deref() != Some("0")
}
/// Command-buffer reuse (P2) kill switch: `OSFKB_CMDBUF=0` forces a fresh encode every round.
///
/// Read ONCE per engine (stored in `Lfm2Gpu::cmdbuf_reuse`) rather than per round: `reusable` sits
/// on the per-round encode path, and an `env::var` there would put a lookup — and an allocation —
/// inside the very loop this optimization exists to shorten. Deliberately NOT a process-wide
/// `OnceLock`: that would latch on first read and make the toggle unexercisable in-process, so the
/// bitwise gate below could not construct an engine each way.
///
/// Reuse is purely an encode-time economy — the dispatches and their floating-point order are
/// identical either way — so flipping it must not move a single output bit.
fn cmdbuf_reuse_from_env() -> bool {
!matches!(
std::env::var("OSFKB_CMDBUF").ok().as_deref(),
Some("0") | Some("off")
)
}
/// Construction knobs for the serving engine. `pool_tokens` sizes the paged KV buffers and
/// `max_seq_tokens` is the per-sequence (prompt + generated) cap the scheduler enforces — BOTH are
/// fixed at construction (the KV buffers cannot be resized). [`Lfm2Gpu::new`] derives them from the
/// checkpoint + `OSFKB_KV_POOL_TOKENS`; the serving binary overrides `pool_tokens` from a VRAM
/// budget via [`EngineOpts::from_vram`], giving the precedence `--vram-gb` flag > env > default.
#[derive(Clone, Copy, Debug)]
pub struct EngineOpts {
pub pool_tokens: usize,
pub max_seq_tokens: usize,
}
impl EngineOpts {
/// f16 KV bytes per token: K and V (×2), over EVERY layer's cache at 2 bytes each. The pool
/// allocates one KV buffer per layer index (conv layers included — their buffer is unused but
/// still allocated), so this uses `n_layers`, matching the real VRAM footprint rather than just
/// the attention layers.
pub fn kv_bytes_per_token(cfg: &crate::weights::Lfm2Config) -> usize {
// Per-layer widths, KV-OWNING layers only (Gemma-4 edge consumers share their donor's
// pool); uniform models reduce to the old 2·L·nkv·hd·2.
(0..cfg.n_layers)
.filter(|&li| cfg.owns_kv(li))
.map(|li| 2 * cfg.n_kv_heads * cfg.head_dim_at(li) * 2)
.sum()
}
/// The default [`Lfm2Gpu::new`] sizing: pool from `OSFKB_KV_POOL_TOKENS` (floored at `MAX_T` so
/// the dense/solo layout is always representable), per-sequence cap = `MAX_T`.
pub fn sized_from_env() -> Self {
let pool_tokens = std::env::var("OSFKB_KV_POOL_TOKENS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(MAX_T)
.max(MAX_T)
.next_multiple_of(PAGE_BLK);
Self {
pool_tokens,
max_seq_tokens: MAX_T,
}
}
/// Size the KV pool from a VRAM budget: `floor(vram_bytes · kv_fraction / bytes_per_token)`
/// tokens, rounded DOWN to a `PAGE_BLK` multiple (stay within budget) and floored at 2 blocks
/// (trash + one usable). The per-sequence cap stays at the shipped `MAX_T` geometry (the pool
/// limit is enforced separately by the scheduler's never-fits reject).
pub fn from_vram(cfg: &crate::weights::Lfm2Config, vram_bytes: u64, kv_fraction: f64) -> Self {
let bpt = Self::kv_bytes_per_token(cfg).max(1) as u64;
let budget = (vram_bytes as f64 * kv_fraction).max(0.0) as u64;
let toks = (budget / bpt) as usize;
let pool_tokens = (toks / PAGE_BLK * PAGE_BLK).max(2 * PAGE_BLK);
Self {
pool_tokens,
max_seq_tokens: MAX_T,
}
}
}
pub struct Lfm2Gpu {
pub w: Weights,
conv_state: Vec<wgpu::Buffer>, // [hidden*conv_l] per layer (conv layers)
k_cache: Vec<wgpu::Buffer>, // paged KV pool per layer: [pool_tokens, n_kv, head_dim]
v_cache: Vec<wgpu::Buffer>,
/// Block tables `[1 + MAX_SLOTS rows × PROW]`: row 0 is the IDENTITY mapping (block i → slot i)
/// used by every single-sequence path — dense layout as a special case of paging, so the same
/// compiled kernels serve solo, speculative and multi-sequence batches (the bitwise contract).
/// Rows 1.. belong to server sequence slots (rewritten on block allocation).
pub(crate) btab: wgpu::Buffer,
/// M=1 per-column metadata `[(pos, btab row)]` — row 0 (identity) for all solo paths.
cmeta1: wgpu::Buffer,
/// Image-embedding arena `[img_rows, hidden]`: the vision tower's output, substituted into the
/// residual stream at image-placeholder columns by [`GATHER_K_MM`]. `None` on a text-only model,
/// which then binds and dispatches exactly what it always did.
///
/// Sized once at construction rather than per request: a request-sized buffer would have to be
/// created on the hot path (and `batch_step` takes `&self`). `OSFKB_IMG_TOKENS` overrides.
img_embed: Option<wgpu::Buffer>,
img_rows: usize,
/// KV pool capacity in tokens (`pool_tokens / PAGE_BLK` blocks; the last block is the trash
/// block idle batch columns write into).
pub(crate) pool_tokens: usize,
/// Per-sequence (prompt + generated) token cap the scheduler enforces (≤ MAX_T for the shipped
/// geometry). Set at construction from [`EngineOpts`].
pub(crate) max_seq_tokens: usize,
s: Scratch,
/// Shared attention meta uniform `(n_heads, n_kv, head_dim, T)`; only `T=pos+1` changes per token.
attn_t: wgpu::Buffer,
/// Gemma-4 edge: the wide-head (nh, nkv, HD_FULL, T) uniform twin, stepped with attn_t.
attn_t2: Option<wgpu::Buffer>,
/// Per-layer `[q_norm(hd) | k_norm(hd)]` — ONE binding for what used to be two. The qkrc
/// kernel bound 11 storage buffers; WebGPU's floor (and Chrome's actual grant, on the very
/// Metal adapter that natively offers 31) is 10, so the browser REJECTED the pipeline and
/// silently ran no KV writes at all — attention then read an all-zero cache and the model
/// emitted one constant token. Concatenating is not a browser branch: identical values in
/// identical order, one descriptor fewer, bitwise-identical output on every backend.
/// `None` on layers without attention. See `tests/webgpu_limits.rs`.
qk_norm_cat: Vec<Option<wgpu::Buffer>>,
plan: Vec<Step>,
/// The k-column pipelines shared by the M=1 plan (gy=1) and every [`SpecPlan`] — sharing the
/// compiled module is what makes batched columns bitwise-equal to the sequential loop.
spec_pls: SpecPls,
/// Qwen2 q/k/v-bias add pass (`QKV_BIAS_ADD`); built for every model but only DISPATCHED on
/// arches that carry `Op::Attn::qkv_bias` (Qwen2), so it is inert elsewhere.
qkv_bias_pl: wgpu::ComputePipeline,
/// TurboQuant KV quantization (opt-in via `OSFKB_TQ_KV=1` at engine construction): the M=1
/// plan quantizes each attention layer's freshly-written K/V to 4-bit rotated coordinates and
/// attends from the quantized cache. `None` = full-precision KV (default).
tq: Option<TqResources>,
/// ROWPACK4 twins of the Q1 weight matrices, keyed by plan-build order (deterministic:
/// every band plan visits sites in the same order). Built lazily by the first
/// [`Self::build_batch_plan`] with the f16 path active — one GPU repack pass per matrix
/// (~+3.7 GB VRAM alongside the legacy layout, which the wide/prefill kernels keep using).
q1_rp4: std::cell::RefCell<std::collections::HashMap<u32, (wgpu::Buffer, wgpu::Buffer)>>,
/// Lazily-built resources for [`Self::argmax_masked`] (the constrained-decode fast path); only the
/// constrained-OSFQL caller touches it, so plain greedy decode never pays for it.
masked_am: std::sync::OnceLock<MaskedArgmax>,
/// Cached decode pipelines (gather + masked-argmax block/merge) for [`Self::decode_constrained`],
/// built on first use so a query never recompiles their WGSL shaders (~450ms/query saved).
decode_pls: std::sync::OnceLock<DecodePls>,
/// The shared tiled lm_head (one compilation for the M=1 plan and every batch plan).
q4lm: crate::Q4LmHead,
/// KC=16 pipelines for the wide prefill plan (None on adapters without subgroups).
wide_pls: Option<WidePls>,
/// Gated-DeltaNet pipelines + slotted per-layer state (Qwen3.5 hybrid layers).
dn_pls: Option<DnPls>,
/// DN slot count actually allocated (see `OSFKB_DN_SLOTS`); trash slot = dn_slots-1.
dn_slots: usize,
dn_ring: Vec<Option<wgpu::Buffer>>,
dn_state: Vec<Option<wgpu::Buffer>>,
/// Mamba per-layer recurrence state (SSM `[d_inner·d_state]` + conv `[d_inner·(d_conv-1)]`),
/// persisting across decode tokens; `None` on non-Mamba layers. Reset with the KV/DN state.
mamba_ssm_state: Vec<Option<wgpu::Buffer>>,
mamba_conv_state: Vec<Option<wgpu::Buffer>>,
/// RWKV per-layer recurrence state, persisting across decode tokens; `None` on non-RWKV layers.
/// `rwkv_wkv_state` is `[num | den | max]` (3·hidden, max init −1e38); the token-shift states
/// are the previous-token hidden for the time-mix and channel-mix sub-layers. Reset with KV/DN.
rwkv_wkv_state: Vec<Option<wgpu::Buffer>>,
rwkv_att_shift: Vec<Option<wgpu::Buffer>>,
rwkv_ffn_shift: Vec<Option<wgpu::Buffer>>,
/// Zero step-counter for the unchained paths (chained batches own a bumped counter).
cnt0: wgpu::Buffer,
/// Rows per workgroup of the GEMV family (8 subgroup / 16 fallback) — grid math must match.
gemv_rows_per_wg: u32,
/// Which GEMV family the M=1 plan carries (`"lcpp-split"` | `"v3-sg"` | `"tree"`) —
/// introspection for the family gates; the batch plans pick per-plan (see `lcpp_family`).
m1_family: &'static str,
/// Trace cell: (prep_s, encode_s, read_start) of the LAST batch step (see
/// [`Self::take_step_breakdown`]) — small-width overhead attribution.
step_breakdown: std::cell::RefCell<(f64, f64, wclock::Instant)>,
/// Command-buffer cache hits across all plans (see `BatchPlan::cmdbuf`) — observability +
/// the reuse gate's hook.
cmdbuf_hits: std::cell::Cell<u64>,
/// P2 command-buffer reuse enabled (`OSFKB_CMDBUF=0` disables) — see [`cmdbuf_reuse_from_env`].
cmdbuf_reuse: bool,
/// Lazily-built argmax pipelines for [`Self::argmax_logits`] (pipeline-stage sampling).
argmax_pls: std::sync::OnceLock<ArgmaxPls>,
/// Prefix-KV cache (for [`Self::prefill_cached`]): the last prefilled prompt, and per-position
/// snapshots of every layer's conv ring-buffer. KV reuse is free (the cache is positionally
/// overwritten, never cleared); only the conv state needs restoring to the divergence point.
cached_tokens: Vec<u32>,
conv_snap: Vec<wgpu::Buffer>, // per layer: [CACHE_MAX * hidden * conv_l]
}
/// On-GPU grammar-masked argmax: the masked parallel-argmax pipelines + the reused mask/partials/chosen
/// buffers and their bind groups (bound once to the engine's `s.logits`).
struct MaskedArgmax {
block_pl: wgpu::ComputePipeline,
merge_pl: wgpu::ComputePipeline,
mask_buf: wgpu::Buffer, // ceil(vocab/32) u32 — 1 bit/token
chosen: wgpu::Buffer, // 1 u32 — the picked token id
block_bg: wgpu::BindGroup,
merge_bg: wgpu::BindGroup,
_pv: wgpu::Buffer,
_pi: wgpu::Buffer,
_posu: wgpu::Buffer,
}
/// The pipelines [`Lfm2Gpu::decode_constrained`] dispatches per token: token gather + the
/// grammar-sparse head pick (mask compact → indirect-args → sparse lm_head → sparse argmax).
/// Compiling WGSL is ~150ms/shader on Metal; building these once and cloning the (Arc-backed)
/// handles per call removes the per-query shader recompilation.
struct DecodePls {
gather: wgpu::ComputePipeline,
compact: wgpu::ComputePipeline,
sargs: wgpu::ComputePipeline,
shead: crate::Q4LmHeadSparse,
sargmax: wgpu::ComputePipeline,
sargmax_k: wgpu::ComputePipeline,
}
/// Per-call resources for the grammar-sparse head pick: the compacted allowed-id list, the
/// slot-aligned sparse logits, and the on-GPU count/indirect buffer (`smeta = [count, ix, iy, iz,
/// frozen]`), bound over the caller's grammar-mask and chosen buffers. Encoded in two halves around
/// the forward: [`Self::encode_compact`] (mask → ids/count/indirect) before it, [`Self::encode_pick`]
/// (indirect sparse head + argmax) after the final norm — so the 65 536-row tied head reads ONLY the
/// grammar-allowed rows, and the loop still never syncs with the CPU.
struct SparsePick {
ids: wgpu::Buffer,
svals: wgpu::Buffer,
smeta: wgpu::Buffer,
compact_bg: wgpu::BindGroup,
sargs_bg: wgpu::BindGroup,
shead_bg: wgpu::BindGroup,
sargmax_bg: wgpu::BindGroup,
}
/// The batched-verify resources for [`Lfm2Gpu::decode_constrained_spec`] — built once per
/// (engine, grammar, k) by [`Lfm2Gpu::make_spec_plan`]. Holds the k-column dispatch plan (bind
/// groups keep the scratch/weight buffers alive), the per-batch upload targets, the per-conv-layer
/// ring snapshots, and one sparse pick per column. `drafts` is public so the grammar side can bind
/// its snapshot-advance kernel to it.
pub struct SpecPlan {
/// Draft columns per batch (2..=16).
pub k: usize,
plan: Vec<Step>,
fed: wgpu::Buffer,
/// The k draft token ids for the current batch (uploaded per batch; the grammar's
/// snapshot-advance kernel reads `drafts[col]`).
pub drafts: wgpu::Buffer,
/// The batched grammar mask, one u32 per (column, token): the grammar's `mask_k` kernel writes
/// `maskk[col·vocab + id]`; each column's sparse pick compacts its own slice.
pub maskk: wgpu::Buffer,
csk: wgpu::Buffer,
snk: wgpu::Buffer,
csk_l: wgpu::Buffer,
snk_l: wgpu::Buffer,
attn_tk: wgpu::Buffer,
/// Per-column paged metadata `[(pos, btab row)] × k` — row 0 (identity) for the spec path.
cmeta_k: wgpu::Buffer,
hnorm: wgpu::Buffer,
chosen_k: wgpu::Buffer,
conv_snapk: Vec<(usize, wgpu::Buffer)>,
picks: Vec<SparsePick>,
gather_pl: wgpu::ComputePipeline,
gather_bg: wgpu::BindGroup,
/// The k-column residual stream (kept for the test bisection hook).
spec_cur: wgpu::Buffer,
/// MAP_READ staging for `chosen_k`, copied inside the batch's own submit — a separate
/// readback encoder/submit costs ~a millisecond per batch on Metal.
chosen_staging: wgpu::Buffer,
}
/// Continuous-batching plan (see [`Lfm2Gpu::make_batch_plan`]): the k-column stack over paged KV
/// plus per-column tiled-head/argmax resources. One instance serves the whole engine lifetime.
pub struct BatchPlan {
/// Per-DN-layer (state, conv-ring) per-column snapshots — Some only on plans built
/// by [`Lfm2Gpu::make_batch_plan_spec`]; consumed by [`Lfm2Gpu::dn_restore`].
dn_snaps: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>>,
/// Max concurrent columns per step.
pub k: usize,
/// Unique id for the pre-encoded command-buffer cache (part of the key so a plan swap can
/// never replay another plan's buffer).
serial: u64,
/// Pre-encoded command buffer for this plan's next step (see `stage_step_submit`): the
/// cache lives ON the plan so per-group plan rotation in serving keeps every group's hot
/// buffer instead of evicting through one engine-wide slot.
cmdbuf: CmdbufCache,
plan: Vec<Step>,
fed: wgpu::Buffer,
/// Per-column embedding SOURCE: `-1` = gather from the token table, `>= 0` = that row of the
/// engine's image arena. All `-1` on a text model (and the text gather never binds it).
isrc: wgpu::Buffer,
csk: wgpu::Buffer,
snk: wgpu::Buffer,
csk_l: wgpu::Buffer,
snk_l: wgpu::Buffer,
cmeta_k: wgpu::Buffer,
/// k-column residual stream (pipeline stages read it back / write it in).
cur: wgpu::Buffer,
/// Dev bisection handles (see read_batch_qn/attn_o_for_tests); harmless clones.
qn: wgpu::Buffer,
attn_o: wgpu::Buffer,
qkv: wgpu::Buffer,
hidden_staging: wgpu::Buffer,
/// k-column final norm (kept for tests / future spec integration).
pub hnorm: wgpu::Buffer,
logits_k: wgpu::Buffer,
chosen_k: wgpu::Buffer,
chosen_staging: wgpu::Buffer,
/// On-GPU token gather (chained decode only) — absent when the f32 embedding exceeds the
/// device's storage-binding limit (e.g. a 3.1 GB 32B embed vs Vulkan's 2 GB cap). The
/// unchained stage path copies embedding rows instead (exact bytes, no binding).
gather: Option<(wgpu::ComputePipeline, wgpu::BindGroup)>,
lm_pl: wgpu::ComputePipeline,
/// Head bind group + grid — only on the last pipeline stage (full models: always).
lm: Option<(wgpu::BindGroup, u32, u32)>,
block_pl: wgpu::ComputePipeline,
merge_pl: wgpu::ComputePipeline,
heads: Vec<BatchHead>,
cnt: wgpu::Buffer,
bump_pl: wgpu::ComputePipeline,
bump_bg: wgpu::BindGroup,
/// Per conv layer (LFM2): `[1 + MAX_SLOTS, hidden, conv_l]` ring state, slot-addressed.
conv_slots: Vec<Option<wgpu::Buffer>>,
/// DeltaNet slot-owner kernel active (batch plans on hybrid archs) — informational.
#[allow(dead_code)]
pub(crate) dn_parallel: bool,
}
/// A pipeline stage's output: the residual stream to ship to the next stage, or (last stage)
/// the greedy token to send back to the coordinator.
pub enum StageOut {
Hidden(Vec<f32>),
Token(u32),
}
/// One profiled dispatch: `(plan index, gx, gy, microseconds)`.
pub type ProfiledStep = (usize, u32, u32, f64, &'static str);
/// A BATCHED pipeline stage's output (see [`Lfm2Gpu::batch_stage_step`]): the k-column residual
/// stream, or (last stage) the k chosen tokens.
pub enum StageBatchOut {
Hidden(Vec<f32>),
Tokens(Vec<u32>),
}
/// Lazily-built argmax-over-logits pipelines (pipeline-stage sampling).
struct ArgmaxPls {
block: wgpu::ComputePipeline,
merge: wgpu::ComputePipeline,
block_bg: wgpu::BindGroup,
merge_bg: wgpu::BindGroup,
chosen: wgpu::Buffer,
_keep: Vec<wgpu::Buffer>,
}
/// Per-column argmax bind groups of a [`BatchPlan`] (the head itself is one k-column dispatch).
struct BatchHead {
block_bg: wgpu::BindGroup,
merge_bg: wgpu::BindGroup,
}
/// One column of a continuous-batching step: feed `token` at `pos` for the sequence whose KV
/// block-table row is `btrow`; `need_logit` selects the columns that pay for head + argmax
/// (decode columns and the final column of a prefill chunk).
#[derive(Clone, Copy, Debug)]
pub struct BatchCol {
pub token: u32,
/// LINEAR position: the KV slot and the causal mask. Always the token's index in its sequence —
/// an image does NOT get to renumber this, or attention would read the wrong keys.
pub pos: u32,
pub btrow: u32,
pub need_logit: bool,
/// 3-D M-RoPE position `(t, h, w)` — the ROPE ANGLE only, never the KV slot.
///
/// For text this is `[pos; 3]`, which makes the rope bit-for-bit identical to the 1-D path (all
/// three axes carry the same number, so the axis map is a no-op). Image tokens share one `t`
/// while `h`/`w` walk the merged grid, which is the entire reason M-RoPE exists.
pub mpos: [u32; 3],
/// Row of the engine's image-embedding arena that REPLACES this column's token embedding, or
/// `-1` to gather from the token table as usual.
///
/// An image placeholder token id carries no information — every image expands to the same
/// repeated id — so the vision tower's output has to be substituted into the residual stream
/// before layer 0. This is that substitution.
pub embed_row: i32,
}
impl BatchCol {
/// An ordinary TEXT column: uniform 3-D position, embedding gathered from the token table.
/// Byte-for-byte the behaviour every caller had before M-RoPE existed.
pub fn text(token: u32, pos: u32, btrow: u32, need_logit: bool) -> Self {
Self {
token,
pos,
btrow,
need_logit,
mpos: [pos; 3],
embed_row: -1,
}
}
}
/// TurboQuant per-engine resources: fixed rotation + centroid table + per-layer quantized caches.
struct TqResources {
rot: wgpu::Buffer, // [hd, hd] orthogonal
cents: wgpu::Buffer, // [16] Lloyd–Max (4-bit)
sketch: wgpu::Buffer, // [hd, hd] QJL sign sketch (stage-2 keys)
/// Per attention layer: (packed 4-bit KV `[MAX_T, nkv, 2, hd/8]`, scales+γ `[MAX_T, nkv, 3]`,
/// key sign bits `[MAX_T, nkv, hd/32]`).
caches: Vec<Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer)>>,
quant_pl: wgpu::ComputePipeline,
attn_pl: wgpu::ComputePipeline,
}
/// KC=16 GEMV-family pipelines for the WIDE PREFILL plan (subgroup adapters only): dequant once,
/// dot 16 positions — 4× the weight amortization of the decode tiles. Per-column math is
/// bitwise-identical to the KC4 kernels (gated in `tests/prefill_kc16.rs`).
struct WidePls {
gemv: wgpu::ComputePipeline,
gn32: wgpu::ComputePipeline,
mlp: wgpu::ComputePipeline,
}
/// Gated-DeltaNet + Qwen3.5 gating pipelines (shared by the M=1 and batch plans).
struct DnPls {
conv: wgpu::ComputePipeline,
step: wgpu::ComputePipeline,
step_par: wgpu::ComputePipeline,
/// Chunk-parallel recurrence for WIDE (prefill) plans (dv-split; see [`dn_chunk_pk_src`]).
chunk_pk: wgpu::ComputePipeline,
/// Deferred gated-RMSNorm epilogue of the dv-split chunk kernel.
chunk_onorm: wgpu::ComputePipeline,
gate_mul: wgpu::ComputePipeline,
/// Compiled, not yet dispatched — staged for the Qwen3.5 gating path.
#[allow(dead_code)]
sgate: wgpu::ComputePipeline,
/// Compiled, not yet dispatched — staged for the Qwen3.5 gating path.
#[allow(dead_code)]
add: wgpu::ComputePipeline,
}
/// The k-column pipelines (see [`Lfm2Gpu`]::spec_pls). One compilation serves the M=1 plan (gy=1)
/// and all spec plans, so per-column FP behavior cannot drift between them.
struct SpecPls {
rmsnorm: wgpu::ComputePipeline,
rmsnorm_add: wgpu::ComputePipeline,
moe_router: wgpu::ComputePipeline,
moe_gate: wgpu::ComputePipeline,
moe_down: wgpu::ComputePipeline,
/// Qwen3.5 dispatch-fused variants (router absorbs norm+sgate; down absorbs the residual add).
/// Compiled, not yet dispatched — the split router is still the live path.
#[allow(dead_code)]
moe_router_fused: wgpu::ComputePipeline,
moe_down_fused: wgpu::ComputePipeline,
/// Small-batch row-strip expert kernels (k ≤ 8): the expert re-reads dominate the verify.
moe_gate_nbar: wgpu::ComputePipeline,
moe_down_nbar_fused: wgpu::ComputePipeline,
/// llama.cpp-shaped expert kernels (subgroup adapters; supersede the row-strips at k ≤ 8).
moe_gate_lcpp: wgpu::ComputePipeline,
moe_down_lcpp_fused: wgpu::ComputePipeline,
/// Router split (norm+sgate / coalesced logits / pick) — the fused router's serial logit
/// loop was the single largest small-batch kernel (~500 µs/layer).
router_norm_sgate: wgpu::ComputePipeline,
router_logits: wgpu::ComputePipeline,
router_pick: wgpu::ComputePipeline,
qkrc_k: wgpu::ComputePipeline,
attn_k: wgpu::ComputePipeline,
/// Gemma-4 edge wide-head (256 < hd ≤ 512) twins; `None` on every other model.
qkrc_k_wide: Option<wgpu::ComputePipeline>,
attn_k_wide: Option<wgpu::ComputePipeline>,
/// Gemma-4 edge PLE pipelines (gather, prep, batched act); `None` without PLE.
ple_gather: Option<wgpu::ComputePipeline>,
ple_prep: Option<wgpu::ComputePipeline>,
ple_act: Option<wgpu::ComputePipeline>,
/// Split-KV flash-decode pair (`None` when `OSFKB_ATTN_SPLITKV` disables it): the
/// partial twin of `attn_k` plus the LSE merge; `attn_skv_z` is the split count.
attn_ksp: Option<wgpu::ComputePipeline>,
attn_kmerge: Option<wgpu::ComputePipeline>,
attn_skv_z: u32,
gemv_k: wgpu::ComputePipeline,
/// Barrier-free small-batch plain GEMV (lab: 1.9–3.6× at ncols ≤ 8 on V100) — bitwise-equal
/// per column to `gemv_k`, selected by the plan builders when the batch width is ≤ 8.
gemv_nbar: wgpu::ComputePipeline,
/// llama.cpp-shaped 4-row subgroup GEMV (Vulkan serving path; 3-7× — see the lab).
gemv_lcpp: wgpu::ComputePipeline,
/// Tiled Q4 GEMM for prefill widths (see [`crate::gemm_q4_src`]) — portable baseline WGSL.
gemm: wgpu::ComputePipeline,
/// Grouped-MoE prefill set (tile builder + gate/down + reduce) — see [`MOE_SEG_BUILD`].
moe_seg_build: wgpu::ComputePipeline,
moe_gate_grouped: wgpu::ComputePipeline,
moe_down_grouped: wgpu::ComputePipeline,
moe_down_reduce: wgpu::ComputePipeline,
/// Tiled prefill attention (see [`ATTN_TILED`]) — wide plans, full-attention layers.
attn_tiled: wgpu::ComputePipeline,
gn32_k: wgpu::ComputePipeline,
mlp_k: wgpu::ComputePipeline,
/// llama.cpp-shaped dense MLP gate+up (lcpp family: split rmsnorm feeds it, like the qkv
/// site). See [`crate::mlp_gate_q4_lcpp_src`].
mlp_lcpp: wgpu::ComputePipeline,
/// BINARY-weight (Q1) twins of the lcpp GEMV + fused MLP (`OSFKB_Q1_WEIGHTS=1`, Bonsai/
/// BitNet). Same bindings as the Q4 kernels — the M=1 plan swaps ONLY the pipeline, the
/// Q1 data rides the existing `Q4` weight containers. `None` off the flag / without lcpp.
gemv_q1: Option<wgpu::ComputePipeline>,
/// Q1 twin of the tiled wide GEMM (weight-SHARED: BN=16 columns per weight stream) —
/// what keeps Q1 WIDE prefill off the per-column re-read path. Env-gated like gemv_q1;
/// no subgroup requirement (portable baseline WGSL, same as `gemm`).
gemm_q1: Option<wgpu::ComputePipeline>,
mlp_q1: Option<wgpu::ComputePipeline>,
convdot_k: wgpu::ComputePipeline,
convmix_k: wgpu::ComputePipeline,
}
/// The grammar-side pieces [`Lfm2Gpu::decode_constrained_spec`] dispatches — all opaque wgpu
/// primitives so the engine stays grammar-agnostic. `snap_adv_bgs[p]` must bind the grammar's
/// snapshot-advance entry for column `p` (snapshot the FSM state into slot `p` of `snap_buf`, then
/// advance by `drafts[p]`); `state_bytes` is the FSM state size (the snapshot stride).
pub struct SpecGrammar<'a> {
/// Batched mask kernel: one dispatch of grid (mask_wg, k) fills all k columns of `maskk` from
/// the FSM snapshots laid down by the snap_advance chain.
pub mask_k_pl: &'a wgpu::ComputePipeline,
pub mask_k_bg: &'a wgpu::BindGroup,
pub mask_wg: u32,
pub snap_adv_pl: &'a wgpu::ComputePipeline,
pub snap_adv_bgs: &'a [wgpu::BindGroup],
pub advance_pl: &'a wgpu::ComputePipeline,
pub advance_bg: &'a wgpu::BindGroup,
pub state_buf: &'a wgpu::Buffer,
pub snap_buf: &'a wgpu::Buffer,
pub chosen_in: &'a wgpu::Buffer,
pub state_bytes: u64,
}
impl SparsePick {
/// Column-slot variant for the batched-verify path: the head reads `hnorm` at byte offset
/// `col·h·4` (column `col` of the batched final norm) and the argmax writes `chosen[col]`.
#[allow(clippy::too_many_arguments)]
fn new_col(
ctx: &GpuCtx,
dp: &DecodePls,
w: &Weights,
hnorm_k: &wgpu::Buffer,
col: u32,
mask_buf: &wgpu::Buffer,
chosen_k: &wgpu::Buffer,
posu: &wgpu::Buffer,
) -> Self {
let vocab = w.cfg.vocab;
let h = w.cfg.hidden;
let ids = ctx.empty(vocab);
let svals = ctx.empty(vocab);
let smeta = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sparse_meta"),
size: 32,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::INDIRECT
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
ctx.queue.write_buffer(&smeta, 0, &[0u8; 32]);
// This column's slice of the batched mask (`maskk[col·vocab ..]`); vocab·4 B is far above
// the 256 B storage-offset alignment.
let vocab_bytes = (vocab as u64) * 4;
let compact_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("mask_compact_col"),
layout: &dp.compact.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: mask_buf,
offset: (col as u64) * vocab_bytes,
size: std::num::NonZeroU64::new(vocab_bytes),
}),
},
wgpu::BindGroupEntry {
binding: 1,
resource: ids.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: smeta.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: posu.as_entire_binding(),
},
],
});
let sargs_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sparse_args"),
layout: &dp.sargs.get_bind_group_layout(0),
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: smeta.as_entire_binding(),
}],
});
let shead_bg = dp.shead.make_at(
ctx,
&w.embed_q4,
hnorm_k,
(col as u64) * (h as u64) * 4,
(h as u64) * 4,
&svals,
&smeta,
&ids,
vocab as u32,
h as u32,
);
let colu = uni(ctx, bytemuck::cast_slice(&u32x4(col, 0, 0, 0)));
let sargmax_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sparse_argmax_k"),
layout: &dp.sargmax_k.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: svals.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: ids.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: smeta.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: chosen_k.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: colu.as_entire_binding(),
},
],
});
Self {
ids,
svals,
smeta,
compact_bg,
sargs_bg,
shead_bg,
sargmax_bg,
}
}
/// Sparse head + column-slot argmax (the batched-verify pick; pipelines differ only in the
/// argmax variant).
fn encode_pick_col(&self, p: &mut wgpu::ComputePass<'_>, dp: &DecodePls) {
p.set_pipeline(dp.shead.pipeline());
p.set_bind_group(0, &self.shead_bg, &[]);
p.dispatch_workgroups_indirect(&self.smeta, 4);
p.set_pipeline(&dp.sargmax_k);
p.set_bind_group(0, &self.sargmax_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
fn new(
ctx: &GpuCtx,
dp: &DecodePls,
w: &Weights,
hnorm: &wgpu::Buffer,
mask_buf: &wgpu::Buffer,
chosen_buf: &wgpu::Buffer,
posu: &wgpu::Buffer,
) -> Self {
let vocab = w.cfg.vocab;
let ids = ctx.empty(vocab);
let svals = ctx.empty(vocab);
// [count, ix, iy, iz, frozen] — INDIRECT so the sparse head's grid is set on-GPU; zeroed
// once here (the SPARSE_ARGS kernel re-zeroes the running count after each freeze).
let smeta = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sparse_meta"),
size: 32,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::INDIRECT
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
ctx.queue.write_buffer(&smeta, 0, &[0u8; 32]);
let compact_bg = make_bg(ctx, &dp.compact, &[mask_buf, &ids, &smeta], posu);
let sargs_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sparse_args"),
layout: &dp.sargs.get_bind_group_layout(0),
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: smeta.as_entire_binding(),
}],
});
let shead_bg = dp.shead.make(
ctx,
&w.embed_q4,
hnorm,
&svals,
&smeta,
&ids,
vocab as u32,
w.cfg.hidden as u32,
);
let sargmax_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sparse_argmax"),
layout: &dp.sargmax.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: svals.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: ids.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: smeta.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: chosen_buf.as_entire_binding(),
},
],
});
Self {
ids,
svals,
smeta,
compact_bg,
sargs_bg,
shead_bg,
sargmax_bg,
}
}
/// Compact the (already-written) grammar mask into the allowed-id list and freeze the count +
/// indirect grid. Runs BEFORE the forward (the mask depends only on the FSM state).
fn encode_compact(&self, p: &mut wgpu::ComputePass<'_>, dp: &DecodePls, vocab: u32) {
p.set_pipeline(&dp.compact);
p.set_bind_group(0, &self.compact_bg, &[]);
p.dispatch_workgroups(vocab.div_ceil(256), 1, 1);
p.set_pipeline(&dp.sargs);
p.set_bind_group(0, &self.sargs_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
/// Sparse head over the allowed rows (indirect grid from `smeta[1..4]`) + lowest-id argmax →
/// `chosen[0]`. Runs AFTER the final norm (`hnorm`) is written.
fn encode_pick(&self, p: &mut wgpu::ComputePass<'_>, dp: &DecodePls) {
p.set_pipeline(dp.shead.pipeline());
p.set_bind_group(0, &self.shead_bg, &[]);
p.dispatch_workgroups_indirect(&self.smeta, 4);
p.set_pipeline(&dp.sargmax);
p.set_bind_group(0, &self.sargmax_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
}
impl MaskedArgmax {
fn new(ctx: &GpuCtx, vocab: usize, logits: &wgpu::Buffer) -> Self {
let block_pl = pipeline(ctx, "argmax_block_masked", ARGMAX_BLOCK_MASKED);
let merge_pl = pipeline(ctx, "argmax_merge_chosen", ARGMAX_MERGE_CHOSEN);
let mask_buf = ctx.empty(vocab.div_ceil(32));
let pv = ctx.empty(256);
let pi = ctx.empty(256);
let chosen = ctx.empty(1);
let posu = uni(ctx, bytemuck::cast_slice(&u32x4(0, 0, vocab as u32, 0)));
let block_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("argmax_block_masked"),
layout: &block_pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: logits.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: pv.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: pi.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: posu.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: mask_buf.as_entire_binding(),
},
],
});
let merge_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("argmax_merge_chosen"),
layout: &merge_pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: pv.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: pi.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: chosen.as_entire_binding(),
},
],
});
Self {
block_pl,
merge_pl,
mask_buf,
chosen,
block_bg,
merge_bg,
_pv: pv,
_pi: pi,
_posu: posu,
}
}
}
/// One step of the cached decode plan.
enum Step {
/// A compute dispatch with a pre-built bind group.
D {
pl: wgpu::ComputePipeline,
bg: wgpu::BindGroup,
gx: u32,
gy: u32,
/// Kernel tag for profiling (derived from the pipeline variable at the build site).
tag: &'static str,
},
/// An INDIRECT compute dispatch: the workgroup counts live in `args` at `offset` (written
/// on-GPU by an earlier step — the grouped-MoE tile count is data-dependent).
Di {
pl: wgpu::ComputePipeline,
bg: wgpu::BindGroup,
args: wgpu::Buffer,
offset: u64,
/// Kernel tag for profiling; carried for parity with the direct-dispatch step.
#[allow(dead_code)]
tag: &'static str,
},
/// Copy `embed[token] → cur` (token-dependent source offset).
Embed,
}
/// Per-token scratch buffers, allocated once and reused (content overwritten each forward).
struct Scratch {
cur: wgpu::Buffer,
normed: wgpu::Buffer,
bcx: wgpu::Buffer,
convy: wgpu::Buffer,
opout: wgpu::Buffer,
qkv: wgpu::Buffer, // fused q|k|v output [(nh+2*nkv)*hd]
qn: wgpu::Buffer,
attn: wgpu::Buffer,
gate: wgpu::Buffer,
mlpout: wgpu::Buffer,
hnorm: wgpu::Buffer,
logits: wgpu::Buffer,
cosb: wgpu::Buffer,
sinb: wgpu::Buffer,
/// Local-theta RoPE tables (Gemma-3 sliding layers; unused by LFM2 — global theta only).
cosb_l: wgpu::Buffer,
sinb_l: wgpu::Buffer,
/// MoE per-token scratch: selected expert ids/weights (top_k) and the expert GLU intermediate.
moe_sel: wgpu::Buffer,
/// Qwen3.5 shared-expert gate scalar (fused router output), [n_cols].
moe_sg: wgpu::Buffer,
/// Router logits staging [n_cols, E] (router split stage 2 → 3).
moe_logits: wgpu::Buffer,
moe_wsel: wgpu::Buffer,
moe_gate: wgpu::Buffer,
/// Gated-DeltaNet per-token scratch (Qwen3.5): fused projection out, conved qkv, head core.
dn_mixed: wgpu::Buffer,
dn_conved: wgpu::Buffer,
dn_core: wgpu::Buffer,
/// Attention-output gate rows (Qwen3.5 `attn_output_gate`).
agate: wgpu::Buffer,
/// DeepSeek MLA per-token scratch: q low-rank, compressed KV (latent+k_pe), normed latent,
/// and the kv_b expansion. Sized to the model's MLA dims; 1-element on non-MLA models.
mla_qc: wgpu::Buffer,
mla_qln: wgpu::Buffer,
mla_kvc: wgpu::Buffer,
mla_kvln: wgpu::Buffer,
mla_kvb: wgpu::Buffer,
/// Gemma-4 edge PLE: merged per-layer inputs [n_layers·ple_dim], the context-projection
/// staging, the per-layer gate/activation scratch [ple_dim], and the 1-element token-id
/// buffer the gather kernel reads (written queue-ordered per step).
ple: wgpu::Buffer,
ple_ctx: wgpu::Buffer,
ple_g: wgpu::Buffer,
ple_a: wgpu::Buffer,
ple_tok: wgpu::Buffer,
}
impl Lfm2Gpu {
/// Rows available in the image-embedding arena (`0` on a text-only model).
pub fn image_rows(&self) -> usize {
self.img_rows
}
/// Load the vision tower's output into the arena at `row0`, so that columns carrying
/// `embed_row = row0 + i` receive `rows[i]` instead of a token embedding.
///
/// `rows` is `[n, hidden]`, exactly what [`crate::vision::VisionTower::forward`] returns.
/// Refuses rather than truncating: a silently short upload would leave stale embeddings from a
/// previous request in the arena, and the model would read someone else's image.
pub fn upload_image_embeds(&self, ctx: &GpuCtx, row0: usize, rows: &[f32]) -> Result<()> {
let h = self.w.cfg.hidden;
let img = anyhow::Context::context(
self.img_embed.as_ref(),
"this model has no vision arena (text-only checkpoint)",
)?;
anyhow::ensure!(
rows.len().is_multiple_of(h),
"image embeddings are {} floats, not a multiple of hidden ({h})",
rows.len()
);
let n = rows.len() / h;
anyhow::ensure!(
row0 + n <= self.img_rows,
"image embeddings [{row0}, {}) overflow the {}-row arena — raise OSFKB_IMG_TOKENS",
row0 + n,
self.img_rows
);
// Chunked write + flush: `write_buffer` allocates its staging in the HOST-VISIBLE heap
// (~256 MB on Vulkan/NVIDIA, NOT the 32 GB device heap), and staging is only reclaimed
// after a submit is polled. A burst of admissions staged several multi-MB embed sets
// between engine polls — on a queue also carrying tower uploads — and write_buffer
// itself panicked "Out of Memory" with 28 GB of VRAM free (backtrace: this call, via
// Scheduler::step/admit; the cause of every "mystery OOM" at high OCR concurrency).
// Submitting after each chunk and polling completed work bounds outstanding staging to
// one chunk per upload. Admission-path only — never per-token.
const CHUNK_FLOATS: usize = 512 * 1024; // 2 MB
for (i, chunk) in rows.chunks(CHUNK_FLOATS).enumerate() {
ctx.queue.write_buffer(
img,
(row0 * h * 4 + i * CHUNK_FLOATS * 4) as u64,
bytemuck::cast_slice(chunk),
);
ctx.queue.submit(std::iter::empty());
let _ = ctx.device.poll(wgpu::PollType::Poll);
}
Ok(())
}
/// Build the engine with the default KV sizing (`OSFKB_KV_POOL_TOKENS` or `MAX_T`); see
/// [`Self::new_with_opts`] for the serving path that sizes the pool from a VRAM budget.
pub fn new(ctx: &GpuCtx, w: Weights) -> Self {
let opts = EngineOpts::sized_from_env();
Self::new_with_opts(ctx, w, opts)
}
/// [`Self::new`] with explicit [`EngineOpts`] (KV `pool_tokens` + per-sequence `max_seq_tokens`).
/// The pool size is fixed here — the KV buffers are allocated to exactly `pool_tokens`.
pub fn new_with_opts(ctx: &GpuCtx, w: Weights, opts: EngineOpts) -> Self {
let nl = w.cfg.n_layers;
let conv_state: Vec<wgpu::Buffer> = (0..nl)
.map(|_| ctx.empty(w.cfg.hidden * w.cfg.conv_l))
.collect();
let conv_snap: Vec<wgpu::Buffer> = (0..nl)
.map(|_| ctx.empty(CACHE_MAX * w.cfg.hidden * w.cfg.conv_l))
.collect();
// KV pool capacity (tokens) is fixed by the caller's opts; rounded to a PAGE_BLK multiple
// and floored at 2 blocks (one usable + trash). The default path (`new`) keeps the old
// env/MAX_T behavior; the serving path sizes it from a VRAM budget.
let pool_tokens = opts
.pool_tokens
.next_multiple_of(PAGE_BLK)
.max(2 * PAGE_BLK);
let max_seq_tokens = opts.max_seq_tokens.clamp(PAGE_BLK, MAX_T);
// f16 pool: halves the attention-side KV traffic (both caches are only ever touched by
// qkrc (f16 write) / ATTN_K + TQ stage-1 (f16 read) — DN/conv state stays f32, matching
// the checkpoints' mamba_ssm_dtype and the reference recurrence).
// KV pools ONLY for attention layers: the 48 DeltaNet layers never touch K/V (their
// state is fixed-size; every k_cache[li] bind site sits under an Op::Attn arm, and
// the TQ caches below already gate on layer_is_attn). Full-size buffers for all 64
// layers cost 6.4 GB of which 4.8 GB was dead — on a 32 GB card that was the
// difference between "full" and "room for a second model". Non-attn layers get a
// one-page dummy so per-layer indexing (and any incidental bind) stays valid.
// Per-layer width (Gemma-4 edge: full-attention layers run wider heads), and KV-shared
// CONSUMERS get a CLONE of their donor's pool (wgpu buffers are ref-counted, a clone
// shares the resource) — every `k_cache[li]` bind site then reads the donor's cache with
// no per-site indirection, and consumers never write (their qkrc runs Q-only).
let alloc_kv = |li: usize| {
let kv_l = w.cfg.n_kv_heads * w.cfg.head_dim_at(li);
if w.cfg.layer_is_attn[li] {
ctx.empty_f16(pool_tokens * kv_l)
} else {
ctx.empty_f16(PAGE_BLK * kv_l)
}
};
let mut k_cache: Vec<wgpu::Buffer> = Vec::with_capacity(nl);
let mut v_cache: Vec<wgpu::Buffer> = Vec::with_capacity(nl);
for li in 0..nl {
let src = w.cfg.kv_src_at(li);
if src != li {
// Donors always precede consumers (checked at config parse).
k_cache.push(k_cache[src].clone());
v_cache.push(v_cache[src].clone());
} else {
k_cache.push(alloc_kv(li));
v_cache.push(alloc_kv(li));
}
}
// Block tables: row 0 = identity (solo paths), rows 1..=MAX_SLOTS = server sequence
// slots, row MAX_SLOTS+1 = trash (every logical block → the sacrificial last pool block,
// where idle batch columns park their KV writes).
let mut btab_init = vec![0u32; (2 + MAX_SLOTS) * PAGE_ROW];
for (i, e) in btab_init.iter_mut().take(PAGE_ROW).enumerate() {
*e = i as u32;
}
let trash_block = (pool_tokens / PAGE_BLK - 1) as u32;
let dn_slots_btab = std::env::var("OSFKB_DN_SLOTS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.map(|v| v.clamp(3, 2 + MAX_SLOTS))
.unwrap_or(2 + MAX_SLOTS);
// The shrunken-DN trash row (see `dn_trash_row`) must also park KV writes safely.
for e in btab_init[(dn_slots_btab - 1) * PAGE_ROW..dn_slots_btab * PAGE_ROW].iter_mut() {
*e = trash_block;
}
for e in btab_init[(1 + MAX_SLOTS) * PAGE_ROW..].iter_mut() {
*e = trash_block;
}
let btab = ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("btab"),
contents: bytemuck::cast_slice(&btab_init),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
});
let cmeta1 = ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("cmeta1"),
contents: bytemuck::cast_slice(&[0u32, 0u32]),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
});
let cnt0 = ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("cnt0"),
contents: bytemuck::cast_slice(&[0u32]),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
});
let (h, im, hd) = (w.cfg.hidden, w.cfg.intermediate, w.cfg.head_dim);
let qd = w.cfg.n_heads * hd;
// Gemma-4 edge geometry: scratch sizes take the MAXIMA; per-layer dispatches read the
// per-layer widths inside the loop. On every other model the maxima equal the uniform
// scalars and nothing changes.
let hd_x = w.cfg.max_head_dim();
let im_x = w.cfg.max_intermediate();
let qd_x = w.cfg.n_heads * hd_x;
let kd_x = w.cfg.n_kv_heads * hd_x;
let g4e = hd_x != hd
|| !w.cfg.edge.layer_kv_src.is_empty()
|| !w.cfg.edge.layer_intermediate.is_empty()
|| w.cfg.edge.ple_dim > 0;
// RoPE table widths: the GLOBAL-theta pair serves the full-attention layers, the
// LOCAL-theta pair the sliding ones — on an edge model those widths differ (512/256).
let hd_glob = w.cfg.head_dim_full();
let hd_slid = w.cfg.head_dim_sliding();
// Gated-DeltaNet geometry + SLOTTED persistent state (slot 0 = solo, 1..=MAX_SLOTS =
// server sequences, last = trash): conv ring [slots, conv_dim, K] and recurrent
// [slots, nv, dk, dv] per DN layer.
let (dn_nk, dn_nv, dn_dk, dn_dv, dn_kn) = (
w.cfg.dn_nk,
w.cfg.dn_nv,
w.cfg.dn_dk,
w.cfg.dn_dv,
w.cfg.dn_kernel,
);
let dn_conv_dim = 2 * dn_nk * dn_dk + dn_nv * dn_dv;
let dn_in_rows = dn_conv_dim + dn_nv * dn_dv + 2 * dn_nv;
// `OSFKB_DN_SLOTS`: shrink the DN slot slabs for single-stream deployments — the
// default 194 slots cost nv·dk·dv·4 ≈ 2 MB × slots × DN-layers (12.2 GB on the 35B),
// which is what kept the full model off one 32 GB V100. The last slot is the TRASH
// slot (idle batch columns park their state writes there — see `dn_trash_row`).
let dn_slots = std::env::var("OSFKB_DN_SLOTS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.map(|v| v.clamp(3, 2 + MAX_SLOTS))
.unwrap_or(2 + MAX_SLOTS);
let dn_ring: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::DeltaNet { .. })
.then(|| ctx.storage(&vec![0f32; dn_slots * dn_conv_dim * dn_kn]))
})
.collect();
let dn_state: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::DeltaNet { .. })
.then(|| ctx.storage(&vec![0f32; dn_slots * dn_nv * dn_dk * dn_dv]))
})
.collect();
let mamba_ssm_state: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| match &l.op {
Op::Mamba {
d_inner, d_state, ..
} => Some(ctx.storage(&vec![0f32; d_inner * d_state])),
_ => None,
})
.collect();
let mamba_conv_state: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| match &l.op {
Op::Mamba {
d_inner, d_conv, ..
} => Some(ctx.storage(&vec![0f32; d_inner * (d_conv - 1)])),
_ => None,
})
.collect();
let rwkv_wkv_state: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::Rwkv { .. }).then(|| {
// [num | den | max]; the running log-sum-exp pivot `max` starts at -1e38.
let mut s = vec![0f32; 3 * w.cfg.hidden];
s[2 * w.cfg.hidden..].fill(-1e38);
ctx.storage(&s)
})
})
.collect();
let rwkv_att_shift: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::Rwkv { .. }).then(|| ctx.storage(&vec![0f32; w.cfg.hidden]))
})
.collect();
let rwkv_ffn_shift: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::Rwkv { .. }).then(|| ctx.storage(&vec![0f32; w.cfg.hidden]))
})
.collect();
// MLA scratch sizing from the first Mla layer (1-element on non-MLA models).
let (mla_qc_n, mla_kvc_n, mla_kvln_n, mla_kvb_n) = w
.layers
.iter()
.find_map(|l| match &l.op {
Op::Mla {
qk_nope,
qk_rope,
v_head_dim,
kv_lora,
q_lora,
..
} => Some((
(*q_lora).max(1),
kv_lora + qk_rope,
(*kv_lora).max(1),
w.cfg.n_heads * (qk_nope + v_head_dim),
)),
_ => None,
})
.unwrap_or((1, 1, 1, 1));
let s = Scratch {
cur: ctx.empty(h),
normed: ctx.empty(h),
bcx: ctx.empty(3 * h),
convy: ctx.empty(h),
opout: ctx.empty(h.max(qd_x)),
qkv: ctx.empty(qd_x + 2 * kd_x),
qn: ctx.empty(qd_x),
attn: ctx.empty(qd_x),
gate: ctx.empty(im_x),
mlpout: ctx.empty(h),
hnorm: ctx.empty(h),
logits: ctx.empty(w.cfg.vocab),
cosb: ctx.empty(hd_glob),
sinb: ctx.empty(hd_glob),
cosb_l: ctx.empty(hd_slid),
sinb_l: ctx.empty(hd_slid),
moe_sel: ctx.empty(w.cfg.top_k_experts.max(1)),
moe_sg: ctx.empty(1),
moe_logits: ctx.empty(w.cfg.num_experts.max(1)),
moe_wsel: ctx.empty(w.cfg.top_k_experts.max(1)),
moe_gate: ctx.empty((w.cfg.moe_intermediate * w.cfg.top_k_experts).max(1)),
dn_mixed: ctx.empty(dn_in_rows.max(1)),
dn_conved: ctx.empty(dn_conv_dim.max(1)),
dn_core: ctx.empty((w.cfg.dn_nv * w.cfg.dn_dv).max(1)),
agate: ctx.empty(if w.cfg.arch == crate::weights::Arch::Qwen35 {
qd
} else {
1
}),
mla_qc: ctx.empty(mla_qc_n),
mla_qln: ctx.empty(mla_qc_n),
mla_kvc: ctx.empty(mla_kvc_n),
mla_kvln: ctx.empty(mla_kvln_n),
mla_kvb: ctx.empty(mla_kvb_n),
// Gemma-4 edge PLE scratch (1-element on every other model).
ple: ctx.empty((nl * w.cfg.edge.ple_dim).max(1)),
ple_ctx: ctx.empty((nl * w.cfg.edge.ple_dim).max(1)),
ple_g: ctx.empty(w.cfg.edge.ple_dim.max(1)),
ple_a: ctx.empty(w.cfg.edge.ple_dim.max(1)),
ple_tok: ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("ple_tok"),
contents: bytemuck::cast_slice(&[0u32]),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
}),
};
let (nh, nkv) = (w.cfg.n_heads, w.cfg.n_kv_heads);
let (eps, conv_l, vocab) = (w.cfg.eps, w.cfg.conv_l, w.cfg.vocab);
let attn_t = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, 1)),
);
// Gemma-4 edge: the wide-head layers' (nh, nkv, HD_GLOBAL, T) twin — updated alongside
// attn_t at every step (see upload_rope).
let attn_t2 = (hd_glob != hd).then(|| {
uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd_glob as u32, 1)),
)
});
// `[q_norm | k_norm]` per layer, once (a GPU-side copy of two ≤512-float vectors).
let qk_norm_cat: Vec<Option<wgpu::Buffer>> = w
.layers
.iter()
.enumerate()
.map(|(li, layer)| {
let Op::Attn { q_norm, k_norm, .. } = &layer.op else {
return None;
};
let hd_l = w.cfg.head_dim_at(li);
let bytes = (hd_l * 4) as u64;
debug_assert!(
q_norm.size() >= bytes && k_norm.size() >= bytes,
"layer {li}: qk-norm vector shorter than its head_dim"
);
let cat = ctx.empty(2 * hd_l);
let mut enc = ctx.device.create_command_encoder(&Default::default());
enc.copy_buffer_to_buffer(q_norm, 0, &cat, 0, bytes);
enc.copy_buffer_to_buffer(k_norm, 0, &cat, bytes, bytes);
ctx.queue.submit([enc.finish()]);
Some(cat)
})
.collect();
// 1-slot ring snapshot for the M=1 conv-mix dispatches (written, never read).
let conv_mix_snap1 = ctx.empty(h * conv_l);
// The M=1 plan uses the SAME compiled k-column pipelines as the speculative batched plan,
// dispatched with gy = 1 (column 0). This is what makes the speculative path's per-column
// results BITWISE-identical to the sequential loop: not "source-identical kernels" (two
// compilations may contract FMAs differently — measured: last-ulp drift on the conv split)
// but the SAME pipeline object, whose column-0 FP instruction stream is the same by
// construction. The one structural change vs the old fused plan: the conv layer runs as
// projection (parallel) + ring mix (sequential tail) — one extra small dispatch per conv
// layer, paid by every path equally.
let wg1k = ctx.device.limits().max_compute_invocations_per_workgroup >= 1024;
let rmsnorm = if ctx.subgroups && wg1k {
pipeline(ctx, "rmsnorm", RMSNORM_SG)
} else {
pipeline(ctx, "rmsnorm", RMSNORM)
};
let center = pipeline(ctx, "center", CENTER);
let ln_affine = pipeline(ctx, "layernorm_affine", LAYERNORM_AFFINE);
let mlp_act_mul = pipeline(ctx, "mlp_act_mul", MLP_ACT_MUL);
let bias_add = pipeline(ctx, "bias_add", BIAS_ADD);
let qk_fullnorm = pipeline(ctx, "qk_fullnorm", QK_FULLNORM);
let mamba_conv = pipeline(ctx, "mamba_conv", MAMBA_CONV);
let mamba_ssm = pipeline(ctx, "mamba_ssm", MAMBA_SSM);
let rwkv_tokenshift = pipeline(ctx, "rwkv_tokenshift", RWKV_TOKENSHIFT);
let rwkv_wkv = pipeline(ctx, "rwkv_wkv", RWKV_WKV);
let rwkv_sigmul = pipeline(ctx, "rwkv_sigmul", RWKV_SIGMUL);
let copy_vec = pipeline(ctx, "copy_vec", COPY_VEC);
let qkrc_k = pipeline(
ctx,
"qkrc_k",
// GLM-OCR and Llama/Mistral/Qwen2 have NO decoder qk-norm — the no-norm kernel variant.
&qk_norm_rope_cache_k_src(eps, arch_has_qk_norm(w.cfg.arch)),
);
let qkv_bias_pl = pipeline(ctx, "qkv_bias_add", QKV_BIAS_ADD);
// Gemma-4 edge wide twins (256 < head_dim ≤ 512) — only built when a layer needs them.
let qkrc_k_wide = (hd_x > 256).then(|| {
pipeline(
ctx,
"qkrc_k_wide",
&qk_norm_rope_cache_k_wide_src(eps, arch_has_qk_norm(w.cfg.arch)),
)
});
let attn_k_wide =
(hd_x > 256).then(|| pipeline(ctx, "attn_k_wide", &page_geom(&attn_k_wide_src())));
// PLE pipelines (Gemma-4 edge).
let ple_pls = (w.cfg.edge.ple_dim > 0).then(|| {
assert!(
w.cfg.edge.ple_dim <= 256,
"PLE_PREP assumes ple_dim ≤ 256 (one workgroup per layer-group)"
);
(
pipeline(ctx, "ple_gather", PLE_GATHER),
pipeline(ctx, "ple_prep", PLE_PREP),
pipeline(ctx, "ple_act_k", PLE_ACT_K),
)
});
// Kernel-family choice is BACKEND-TUNED (measured, 35B stage, 16 cols): subgroup GEMVs
// win 2.1× on Metal but LOSE 16% on Vulkan/V100 (tree fallback 42.0ms vs 48.6ms) —
// per-lane f16 loads and the subgroup reduction lower poorly there. Env overrides win.
let sg = match std::env::var("OSFKB_SUBGROUP_GEMV").ok().as_deref() {
Some("1") => ctx.subgroups,
Some("0") => false,
_ => ctx.subgroups && ctx.backend.starts_with("Metal"),
} && std::env::var_os("OSFKB_NO_SUBGROUP_GEMV").is_none();
// vec4 attention (`ATTN_K_V4`) — OPT-IN (`OSFKB_ATTN_V4=1`) until a quiet window
// proves it: at landing time the box ran a VM + rustc jobs and the 4B A/B scattered
// 22.5–31.2 tok/s across rounds (the same evidence bar that kept the staged encoder
// read opt-in). Gates are green on both arms; one pipeline serves the M=1, spec and
// batch plans alike, so batch==solo stays bitwise whichever is chosen.
let attn_v4 = w.cfg.head_dim.is_multiple_of(4)
&& std::env::var("OSFKB_ATTN_V4").ok().as_deref() == Some("1");
// Parallel-V upgrade of the v4 body (16 dim-slots x 16 key-lanes): needs hd<=64 so the
// dim slots cover the head. Applies to BOTH the plain and split-KV arms via v4_body.
let attn_v4par = attn_v4
&& w.cfg.head_dim <= 64
&& std::env::var("OSFKB_ATTN_V4PAR").ok().as_deref() == Some("1");
let v4_body: std::borrow::Cow<'static, str> = if attn_v4par {
attn_k_v4_par_src().into()
} else {
ATTN_K_V4.into()
};
let attn_k = if attn_v4 {
pipeline(ctx, "attn_k_v4", &page_geom(&v4_body))
} else {
pipeline(ctx, "attn_k", &page_geom(ATTN_K))
};
// DeepSeek MLA assembly (built for every model; only dispatched on Op::Mla layers).
let mla_assemble = pipeline(
ctx,
"mla_assemble",
&page_geom(&mla_assemble_src(w.cfg.rope_theta)),
);
// Split-KV twins (FlashDecoding): patched from WHICHEVER base kernel was selected, so
// the split arm inherits the same scalar/vec4 body and batch==solo stays bitwise —
// both the M=1 and batched plans take the same pair at the same z.
let attn_skv_z = attn_splitkv_z();
let attn_ksp = (attn_skv_z > 1).then(|| {
let base: &str = if attn_v4 { &v4_body } else { ATTN_K };
pipeline(
ctx,
if attn_v4 { "attn_ksp_v4" } else { "attn_ksp" },
&page_geom(&attn_split_patch(base, attn_v4)),
)
});
let attn_kmerge = (attn_skv_z > 1).then(|| pipeline(ctx, "attn_kmerge", ATTN_K_MERGE));
// GEMV family: subgroup kernels (one subgroup per row, zero barriers) where the adapter
// supports them — the workgroup-tree fallback keeps GL-class adapters working. The row
// tiling differs: v3 packs 256/subgroup_size rows per workgroup (8 at ssz=32).
let gemv_rows_per_wg: u32 = if sg { 8 } else { 16 };
let gemv_k = if sg {
pipeline(ctx, "gemv_q4_k_sg", &crate::gemv_q4_k_sg_src())
} else {
pipeline(ctx, "gemv_q4_k", &crate::gemv_q4_k_src())
};
let gn32_k = if sg {
pipeline(
ctx,
"q4_gemv_norm32_k_sg",
&crate::q4_gemv_norm32_k_sg_src(),
)
} else {
pipeline(ctx, "q4_gemv_norm32_k", &crate::q4_gemv_norm32_k_src())
};
// Moshi: the FFN weights are Q8_0 (its gating matrices are q4-fragile — see
// `Arch::Moshi` and the moshi_lm q4sim sweep); the fused-MLP and down-proj steps get
// q8 kernels, everything else stays Q4. Subgroup-only for now (loud refusal below);
// the portable fallback kernel is follow-up work.
let moshi_q8 = w.cfg.arch == crate::weights::Arch::Moshi;
// Moshi's FFN is Q8_0 (its gating matrices are q4-fragile — see `Arch::Moshi` and the
// moshi_lm q4sim sweep). Subgroup Q8 kernels where available, the portable
// workgroup-tree twins otherwise — so Moshi runs on any adapter (verified via
// `OSFKB_NO_SUBGROUPS`). The Q8 GEMV/MLP tiling matches its Q4 sibling, so the q8 steps
// ride the same `gemv_rows_per_wg` dispatch; only the down-proj needs an explicit count.
let mlp_k = if moshi_q8 {
if sg {
pipeline(ctx, "mlp_gate_q8_k_sg", &crate::mlp_gate_q8_k_sg_src())
} else {
pipeline(ctx, "mlp_gate_q8_k", &crate::mlp_gate_q8_k_src())
}
} else if sg {
pipeline(ctx, "mlp_gate_q4_k_sg", &crate::mlp_gate_q4_k_sg_src())
} else {
pipeline(ctx, "mlp_gate_q4_k", &crate::mlp_gate_q4_k_src())
};
let gemv_q8_ffn = moshi_q8.then(|| {
if sg && std::env::var("OSFKB_NO_SOLO_GEMV").is_err() {
pipeline(ctx, "gemv_q8_k_sg_solo", &crate::gemv_q8_k_sg_solo_src())
} else if sg {
pipeline(ctx, "gemv_q8_k_sg", &crate::gemv_q8_k_sg_src())
} else {
pipeline(ctx, "gemv_q8_k", &crate::gemv_q8_k_src())
}
});
// Rows-per-workgroup for the Q8 GEMV: matches the kernel (subgroup = 8 rows/wg,
// portable = 16), independent of the general `m1_rpw` (which the lcpp/Q1 paths retune).
let moshi_q8_rpw: u32 = if sg { 8 } else { 16 };
let convdot_k = pipeline(ctx, "conv_dot_k", &crate::conv_dot_k_src(eps));
let convmix_k = pipeline(ctx, "conv_mix_k", crate::CONV_MIX_K);
let rmsnorm_add = if ctx.subgroups && wg1k {
pipeline(ctx, "rmsnorm_add", RMSNORM_ADD_SG)
} else {
pipeline(ctx, "rmsnorm_add", RMSNORM_ADD)
};
// TurboQuant (opt-in): 4-bit rotated-domain KV. Rotation/centroids are deterministic and
// data-free (arXiv 2504.19874 stage 1); per-attn-layer quantized caches replace the f32
// reads in attention while qkrc's f32 writes become the quantizer's staging.
assert!(
!(g4e && std::env::var_os("OSFKB_TQ_KV").is_some()),
"OSFKB_TQ_KV is not wired for Gemma-4 edge geometry (per-layer head widths)"
);
let tq = if std::env::var_os("OSFKB_TQ_KV").is_some() {
let rotm = crate::turboquant::orthogonal_rotation(hd, 0x7157_0002);
let cents = crate::turboquant::lloyd_max_centroids(4);
let caches: Vec<Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer)>> = (0..nl)
.map(|li| {
w.cfg.layer_is_attn[li].then(|| {
(
ctx.empty(MAX_T * w.cfg.n_kv_heads * 2 * (hd / 8)),
ctx.empty(MAX_T * w.cfg.n_kv_heads * 3),
ctx.empty(MAX_T * w.cfg.n_kv_heads * (hd / 32)),
)
})
})
.collect();
Some(TqResources {
rot: ctx.storage(&rotm),
cents: ctx.storage(¢s),
sketch: ctx.storage(&crate::turboquant::gaussian_sketch(hd, 0x7157_0003)),
caches,
quant_pl: pipeline(ctx, "tq_quant", TQ_QUANT),
attn_pl: pipeline(ctx, "attn_tq", ATTN_TQ),
})
} else {
None
};
// The ONE tiled lm_head compilation every path shares (M=1 plan column 0, batched serving
// columns): sharing the compiled pipeline object is the bitwise batch==solo contract.
let q4lm = crate::Q4LmHead::new(ctx, w.cfg.head_kind());
let moe_router_pl = pipeline(ctx, "moe_router", MOE_ROUTER);
let moe_router_v3_pl = pipeline(ctx, "moe_router_v3", MOE_ROUTER_V3);
let moe_router_fused_pl = pipeline(ctx, "moe_router_fused", MOE_ROUTER_FUSED);
let moe_down_fused_pl = pipeline(ctx, "moe_down_fused", &moe_down_q4_fused_src());
let router_norm_sgate_pl = pipeline(ctx, "router_norm_sgate", ROUTER_NORM_SGATE);
let router_logits_pl = pipeline(ctx, "router_logits", ROUTER_LOGITS);
let router_pick_pl = if ctx.subgroups && ctx.subgroups32 {
pipeline(ctx, "router_pick_sg32", ROUTER_PICK_SG32)
} else {
pipeline(ctx, "router_pick", ROUTER_PICK)
};
let moe_gate_pl = pipeline(ctx, "moe_gate", &moe_gate_q4_src());
let moe_down_pl = pipeline(ctx, "moe_down", &moe_down_q4_src());
// Wide-prefill pipelines need only HARDWARE subgroup support — independent of the
// decode-family choice (KC16 subgroup kernels beat the tree at prefill widths even on
// Vulkan: measured +30% end-to-end on the 35B prompt phase).
let wide_pls = ctx.subgroups.then(|| WidePls {
gemv: pipeline(ctx, "gemv_q4_k_sg16", &crate::gemv_q4_k_sg16_src()),
gn32: pipeline(
ctx,
"q4_gemv_norm32_k_sg16",
&crate::q4_gemv_norm32_k_sg16_src(),
),
mlp: pipeline(ctx, "mlp_gate_q4_k_sg16", &crate::mlp_gate_q4_k_sg16_src()),
});
let dn_pls = (w.cfg.arch == crate::weights::Arch::Qwen35).then(|| DnPls {
chunk_pk: pipeline(ctx, "dn_chunk_pk", &dn_chunk_pk_src()),
chunk_onorm: pipeline(ctx, "dn_chunk_onorm", DN_CHUNK_ONORM),
conv: pipeline(ctx, "dn_conv_k", DN_CONV_K),
step: pipeline(ctx, "dn_step_k", DN_STEP_K),
step_par: {
// Sub-lane count is bounded by what the DEVICE granted (browsers stay at
// 256 invocations => ic<=2); default ic=4 (WG 512) — the band=1 profile
// measured dn_step 13x over its state-I/O roofline from occupancy
// starvation (48 WGs x 4 warps on 80 SMs).
let cap = (ctx.device.limits().max_compute_invocations_per_workgroup / 128).max(1);
match std::env::var("OSFKB_DN_PK2").ok().as_deref() {
Some("0") => pipeline(ctx, "dn_step_pk", DN_STEP_PK),
Some("2") => pipeline(ctx, "dn_step_pk2", &dn_step_pk2_src(2.min(cap))),
Some("8") => pipeline(ctx, "dn_step_pk2", &dn_step_pk2_src(8.min(cap))),
_ if cap == 1 => pipeline(ctx, "dn_step_pk", DN_STEP_PK),
_ => pipeline(ctx, "dn_step_pk2", &dn_step_pk2_src(8.min(cap))),
}
},
gate_mul: pipeline(ctx, "gate_mul_k", GATE_MUL_K),
sgate: pipeline(ctx, "sgate_scale_k", SGATE_SCALE_K),
add: pipeline(ctx, "add_k", ADD_K),
});
let spec_pls = SpecPls {
rmsnorm: rmsnorm.clone(),
rmsnorm_add: rmsnorm_add.clone(),
moe_router: moe_router_pl.clone(),
moe_gate: moe_gate_pl.clone(),
moe_down: moe_down_pl.clone(),
moe_router_fused: moe_router_fused_pl.clone(),
moe_down_fused: moe_down_fused_pl.clone(),
moe_gate_nbar: pipeline(ctx, "moe_gate_nbar", &moe_gate_q4_nbar_src()),
moe_down_nbar_fused: pipeline(
ctx,
"moe_down_nbar_fused",
&moe_down_q4_nbar_fused_src(),
),
moe_gate_lcpp: pipeline(ctx, "moe_gate_lcpp", &moe_gate_q4_lcpp_src()),
moe_down_lcpp_fused: pipeline(
ctx,
"moe_down_lcpp_fused",
&moe_down_q4_lcpp_fused_src(),
),
router_norm_sgate: pipeline(ctx, "router_norm_sgate", ROUTER_NORM_SGATE),
router_logits: pipeline(ctx, "router_logits", ROUTER_LOGITS),
router_pick: if ctx.subgroups && ctx.subgroups32 {
pipeline(ctx, "router_pick_sg32", ROUTER_PICK_SG32)
} else {
pipeline(ctx, "router_pick", ROUTER_PICK)
},
qkrc_k: qkrc_k.clone(),
attn_k: attn_k.clone(),
qkrc_k_wide: qkrc_k_wide.clone(),
attn_k_wide: attn_k_wide.clone(),
ple_gather: ple_pls.as_ref().map(|p| p.0.clone()),
ple_prep: ple_pls.as_ref().map(|p| p.1.clone()),
ple_act: ple_pls.as_ref().map(|p| p.2.clone()),
attn_ksp: attn_ksp.clone(),
attn_kmerge: attn_kmerge.clone(),
attn_skv_z,
gemv_k: gemv_k.clone(),
gemv_nbar: pipeline(ctx, "gemv_q4_k_nbar", &crate::gemv_q4_k_nbar_src()),
gemv_lcpp: pipeline(ctx, "gemv_q4_k_lcpp", &crate::gemv_q4_k_lcpp_src()),
gemm: pipeline(ctx, "gemm_q4", &crate::gemm_q4_src()),
moe_seg_build: pipeline(ctx, "moe_seg_build", MOE_SEG_BUILD),
moe_gate_grouped: pipeline(ctx, "moe_gate_grouped", &moe_gate_grouped_src()),
moe_down_grouped: pipeline(ctx, "moe_down_grouped", &moe_down_grouped_src()),
moe_down_reduce: pipeline(ctx, "moe_down_reduce", MOE_DOWN_REDUCE),
// Codegen the tiled attention for THIS head_dim (32-multiple, ≤256) so hd=256
// models get the 8-column K/V sharing; the const hd=128 body is otherwise identical.
attn_tiled: if hd % 32 == 0 && hd <= 256 {
pipeline(ctx, "attn_tiled", &page_geom(&attn_tiled_src(hd)))
} else {
pipeline(ctx, "attn_tiled", &page_geom(ATTN_TILED))
},
gn32_k: gn32_k.clone(),
mlp_k: mlp_k.clone(),
mlp_lcpp: pipeline(ctx, "mlp_gate_q4_lcpp", &crate::mlp_gate_q4_lcpp_src()),
gemv_q1: (ctx.subgroups32_effective()
&& std::env::var("OSFKB_Q1_WEIGHTS").ok().as_deref() == Some("1"))
.then(|| pipeline(ctx, "gemv_q1_k_lcpp", &crate::gemv_q1_k_lcpp_src())),
gemm_q1: (std::env::var("OSFKB_Q1_WEIGHTS").ok().as_deref() == Some("1"))
.then(|| pipeline(ctx, "gemm_q1", &crate::gemm_q1_src())),
mlp_q1: (ctx.subgroups32_effective()
&& std::env::var("OSFKB_Q1_WEIGHTS").ok().as_deref() == Some("1"))
.then(|| pipeline(ctx, "mlp_gate_q1_lcpp", &crate::mlp_gate_q1_lcpp_src())),
convdot_k: convdot_k.clone(),
convmix_k: convmix_k.clone(),
};
// ── Build the dispatch plan ONCE (bind groups + uniforms baked; replayed every token) ──
// Split-KV partials: per (head, split) — `hd` accumulator floats + (m, l) at a
// vec4-aligned stride. One buffer serves every layer (the plan is serial).
let attn_part = ctx.storage(&vec![0f32; nh * attn_skv_z as usize * (hd + 4)]);
let attn_merge_u = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, hd as u32, attn_skv_z, 0)),
);
let mf = |a: f32, b: f32| uni(ctx, bytemuck::cast_slice(&[a, b, eps, 0.0]));
// RMSNORM_ADD with a layer_scalar in kp.w (0 = 1.0 — Gemma-4 edge only).
let mf_ls = |a: f32, b: f32, s: f32| uni(ctx, bytemuck::cast_slice(&[a, b, eps, s]));
let epsm = uni(ctx, bytemuck::cast_slice(&[eps, 0.0f32, 0.0, 0.0]));
let udims =
|a: u32, b: u32, acc: u32, d: u32| uni(ctx, bytemuck::cast_slice(&u32x4(a, b, acc, d)));
let bg2 = |pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer],
dims: &wgpu::Buffer,
eu: &wgpu::Buffer| {
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32,
resource: dims.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32 + 1,
resource: eu.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &entries,
})
};
// Gemma-3: the gate activation flag rides in epsm.y (0 = SiLU, 1 = gelu_pytorch_tanh).
let epsm_act = uni(
ctx,
bytemuck::cast_slice(&[eps, if w.cfg.act_gelu { 1.0f32 } else { 0.0 }, 0.0, 0.0]),
);
// The M=1 plan takes the SAME GEMV family the batch plans pick on this adapter: on
// guaranteed-32-subgroup Vulkan that is the llama.cpp-shaped family (lab-measured
// 3-7.4× over the tree at these widths), with fused-norm sites SPLIT into
// rmsnorm + plain lcpp GEMV — the split measured better than the inline-norm fusion
// (which re-read each WG's column twice; the fused variant was deleted). Sharing the
// family also restores single-FP-order `solo == batched-column` on Vulkan; Metal keeps
// the v3 subgroup family and GL-class adapters the tree, byte-identical to before.
let m1_lcpp = lcpp_family(ctx);
// BINARY-weight (Q1) plan: every projection's `Q4` container actually holds Q1 sign
// weights, so the M=1 plan runs the Q1 lcpp GEMV/MLP twins. The bind groups are
// unchanged (same two buffers, same 5/8 bindings) — only these pipelines swap. Q1
// REQUIRES the lcpp family (its subgroup shape); `SpecPls` only builds the twins when
// the flag is set AND the probe validated 32-wide subgroups, so this is an all-or-error
// gate rather than a silent Q4-kernel-on-Q1-bytes fallback.
assert!(
spec_pls.gemv_q1.is_none() || m1_lcpp,
"OSFKB_Q1_WEIGHTS needs the lcpp GEMV family (validated 32-wide subgroups)"
);
// SOLO GEMV for the M=1 decode plan: every dispatch here is single-column, and the
// KC-loop kernels carry an `array<f32, KC>` accumulator that Metal spills to scratch
// (dynamic indexing) — 1.22-1.26x on all three Moshi shapes, measured interleaved in
// tests/moshi_kernel_bw.rs. Batch/prefill plans keep the KC kernels unchanged.
let (m1_gemv, m1_rpw) = if let Some(q1) = &spec_pls.gemv_q1 {
(q1.clone(), 4u32)
} else if m1_lcpp {
(spec_pls.gemv_lcpp.clone(), 4u32)
} else if sg && std::env::var("OSFKB_NO_SOLO_GEMV").is_err() {
(
pipeline(ctx, "gemv_q4_k_sg_solo", &crate::gemv_q4_k_sg_solo_src()),
gemv_rows_per_wg,
)
} else {
(gemv_k.clone(), gemv_rows_per_wg)
};
// ── f16-STORAGE weights (`OSFKB_DECODE_PRECISION=f16`, incl. per-layer `blk:` mixes) ──
// The f16 kernels take FOUR bindings where the Q4 kernels take five (there is no scale
// table), so the pipeline and its bind group must be chosen together — and since the
// dtype is PER LAYER (`cfg.wdtype_at`), that choice happens at the top of the layer loop
// below (`wf16`/`gv_pl` are loop locals; every GEMV site routes through them). f16
// REQUIRES the lcpp family: the twins are lcpp-shaped (WG == subgroup == 32, 4 rows/WG,
// so a Q4-lcpp layer and an f16 layer share `m1_rpw` = 4 and the grid math). All-or-error
// asserts rather than a silent Q4-kernel-on-f16-bytes fallback, matching the Q1 contract.
let any_f16 = w.cfg.any_f16_body();
let any_nonq4 = w.cfg.any_nonq4_body();
assert!(
!any_nonq4 || m1_lcpp,
"non-Q4_0 weight containers need the lcpp GEMV family (validated 32-wide subgroups)"
);
assert!(
!(any_nonq4 && spec_pls.gemv_q1.is_some()),
"f16/native layers cannot mix with the Q1 (Bonsai) plan — its pipelines replace the \
base family globally"
);
let m1_gemv_f16 =
any_f16.then(|| pipeline(ctx, "gemv_f16_k_lcpp", &crate::gemv_f16_k_lcpp_src()));
let mlp_f16 =
any_f16.then(|| pipeline(ctx, "mlp_gate_f16_lcpp", &crate::mlp_gate_f16_lcpp_src()));
// Native-GGUF pipelines, built only for the container kinds this model actually uses
// (`load_gguf_native` fills cfg.layer_kinds; everything else leaves it empty).
use crate::weights::WKind;
let kind_used = |k: WKind| {
(0..w.cfg.n_layers).any(|li| {
let lk = w.cfg.kinds_at(li);
[lk.qkv, lk.o, lk.gate_up, lk.down].contains(&k)
})
};
let q4k_pl =
kind_used(WKind::Q4K).then(|| pipeline(ctx, "gemv_q4k", &crate::gemv_q4k_k_lcpp_src()));
let q6k_pl =
kind_used(WKind::Q6K).then(|| pipeline(ctx, "gemv_q6k", &crate::gemv_q6k_k_lcpp_src()));
let q5k_pl =
kind_used(WKind::Q5K).then(|| pipeline(ctx, "gemv_q5k", &crate::gemv_q5k_k_lcpp_src()));
let iq4nl_pl = kind_used(WKind::Iq4Nl)
.then(|| pipeline(ctx, "gemv_iq4_nl", &crate::gemv_iq4_nl_k_lcpp_src()));
let iq4xs_pl = kind_used(WKind::Iq4Xs)
.then(|| pipeline(ctx, "gemv_iq4_xs", &crate::gemv_iq4_xs_k_lcpp_src()));
// Grid-codebook IQ family: one pipeline per used kind, indexed by iqg_idx.
let iqg_pls: Vec<Option<wgpu::ComputePipeline>> = IQG_KINDS
.iter()
.map(|&(kk, ty)| {
kind_used(kk).then(|| pipeline(ctx, "gemv_iq_grid", &crate::gemv_iq_grid_src(ty)))
})
.collect();
let q5n_pl = kind_used(WKind::Q5_0N)
.then(|| pipeline(ctx, "gemv_q5_0n", &crate::gemv_q5_0n_k_lcpp_src()));
let q8n_pl = kind_used(WKind::Q8_0N)
.then(|| pipeline(ctx, "gemv_q8_0n", &crate::gemv_q8_0n_k_lcpp_src()));
let mlp_q8n = (0..w.cfg.n_layers)
.any(|li| w.cfg.kinds_at(li).gate_up == WKind::Q8_0N)
.then(|| pipeline(ctx, "mlp_gate_q8_0n", &crate::mlp_gate_q8_0n_lcpp_src()));
let m1_family = if any_f16 {
"lcpp-f16"
} else if m1_lcpp {
"lcpp-split"
} else if sg {
"v3-sg"
} else {
"tree"
};
// lcpp-split scratch: the rmsnorm output the split projections read ([hidden] f32; bind
// groups hold it alive). Allocated tiny on every arm to keep construction unconditional.
let normed_x = ctx.empty(h);
// LayerNorm arches (OLMo1: non-parametric LN) center the norm input here, then feed it to
// the RMSNorm-fused kernels with ones weights — RMSNorm(x-mean)·1 = LayerNorm(x). The raw
// residual `s.cur` is never centered; only this scratch is. Tiny + allocated unconditionally.
let is_ln = crate::weights::arch_is_layernorm(w.cfg.arch);
let is_affine_ln = crate::weights::arch_layernorm_affine(w.cfg.arch);
// PARALLEL arches (Falcon): ONE LayerNorm(cur) feeds BOTH attn and mlp, both accumulating
// into the same residual (no intermediate norm). The shared norm is hoisted to the top of
// the layer (into normed_x) BEFORE attn touches cur, so the mlp reads the pre-attn norm.
let is_parallel = crate::weights::arch_is_parallel(w.cfg.arch);
// NON-GATED MLP (Falcon): fc2(gelu(fc1·x)) — no gate/silu·mul. Realised as gemv(w3=fc1) →
// act·mul with a ones "up" (act(g)·1 = act(g)) → gemv(w2=fc2). `ones_im` is that ones vector.
let mlp_nogate = crate::weights::arch_mlp_nogate(w.cfg.arch);
// Phi-2 biases: fc1 + (parallel-combined o_proj+fc2) packed into each layer's `qkv_bias`
// AFTER the [q|k|v] section (offsets qkv_rows, qkv_rows+im); lm_head bias packed into
// `embedding_norm` after [weight|bias] (offset 2·hidden). BIAS_ADD reads the slices.
let has_biases = crate::weights::arch_has_biases(w.cfg.arch);
let qkv_rows_b = (nh * hd + 2 * nkv * hd) as u32; // where the FFN biases start in qkv_bias
// OLMo2 POST-norm: the attn/mlp project the RAW residual (no pre-norm; RMSNorm weight=1 ≠
// identity, so the norm must be genuinely SKIPPED), a FULL-VECTOR qk-norm runs on the qkv,
// and the sub-layer OUTPUT is RMSNorm'd before the residual add (the existing sandwich
// post_op_norm/post_ffn_norm). `qd`/`kd` are the q/k projection widths for QK_FULLNORM.
let skip_prenorm = crate::weights::arch_skip_prenorm(w.cfg.arch);
let (qd_fn, kd_fn) = ((nh * hd) as u32, (nkv * hd) as u32);
// MLP_ACT_MUL activation mode: 3 = ReLU (Nemotron squared-relu), 2 = exact-erf GELU
// (Falcon), 1 = tanh GELU, 0 = SiLU.
let act_mode = if crate::weights::arch_act_relu(w.cfg.arch) {
3u32
} else if crate::weights::arch_gelu_exact(w.cfg.arch) {
2u32
} else {
u32::from(w.cfg.act_gelu)
};
let centered = ctx.empty(h);
// Affine-LayerNorm scratch: gate/up projections of the fully-normed input (fed to the
// standalone act·mul), so the FFN needs no fused-norm kernel — backend-agnostic.
let ln_gate_raw = ctx.empty(im_x);
let ln_up_raw = ctx.empty(im_x);
let ones_im = ctx.storage(&vec![1.0f32; im_x]);
// Mamba mixer scratch (sized from the first Op::Mamba layer; 1-elem on non-Mamba models).
let (mb_di, mb_ds, mb_dtrk) = w
.layers
.iter()
.find_map(|l| match &l.op {
Op::Mamba {
d_inner,
d_state,
dt_rank,
..
} => Some((*d_inner, *d_state, *dt_rank)),
_ => None,
})
.unwrap_or((1, 1, 1));
let mb_xz = ctx.empty(2 * mb_di);
let mb_xconv = ctx.empty(mb_di);
// Sized for the dt_proj gemv's padded read too (pad cols multiply the B/C tail by zero).
let mb_dbl = ctx.empty((mb_dtrk + 2 * mb_ds).max(mb_dtrk.div_ceil(32) * 32));
let mb_dtraw = ctx.empty(mb_di);
let mb_y = ctx.empty(mb_di);
// RWKV mixer scratch (sized from the first Op::Rwkv layer; 1-elem on non-RWKV models). The
// key projection is `ffn_dim`-wide in the channel-mix, `hidden` in the time-mix → size to
// the max; `rk_relu` holds the channel-mix relu² (ffn_dim wide).
let rk_ffn = w
.layers
.iter()
.find_map(|l| match &l.op {
Op::Rwkv { ffn_dim, .. } => Some(*ffn_dim),
_ => None,
})
.unwrap_or(1);
let rk_x = ctx.empty(h); // norm(cur)
let rk_xk = ctx.empty(h);
let rk_xv = ctx.empty(h);
let rk_xr = ctx.empty(h);
let rk_key = ctx.empty(rk_ffn.max(h)); // key proj (ffn_dim in channel-mix, hidden in time-mix)
let rk_relu = ctx.empty(rk_ffn); // relu²(key) for the channel-mix
let rk_v = ctx.empty(h); // value proj (hidden either way)
let rk_r = ctx.empty(h); // receptance proj
let rk_wkv = ctx.empty(h); // WKV output
let rk_rwkv = ctx.empty(h); // sigmoid(r)·wkv (time-mix)
let mut plan: Vec<Step> = Vec::new();
// Push a CENTER step (cur → centered) and return the buffer the norm kernels should read:
// `centered` on NON-parametric LayerNorm arches (OLMo), raw `cur` otherwise. (Affine
// LayerNorm arches take the `is_affine_ln` branch instead and never call this.) Defined
// after `plan` so the macro body's references resolve at this scope.
macro_rules! norm_in {
() => {{
if is_ln && !is_affine_ln {
plan.push(Step::D {
tag: "center",
pl: center.clone(),
bg: make_bg(
ctx,
¢er,
&[&s.cur, ¢ered],
&udims(h as u32, 0, 0, 0),
),
gx: 1,
gy: 1,
});
¢ered
} else {
&s.cur
}
}};
}
if w.cfg.stage_first {
plan.push(Step::Embed);
}
// Gemma-4 edge PLE model-level stage (runs on the scaled embedding in `s.cur`):
// ple ← dequant(table[tok]) (table pre-scaled ×sqrt(pd) at load)
// ple_ctx ← per_layer_model_projection · cur
// ple ← (groupnorm(ple_ctx · h^-0.5)·w + ple) · 2^-0.5 (PLE_PREP, per group)
if let (Some((gpl, ppl, _)), Some(pw)) = (&ple_pls, &w.ple) {
let pd = w.cfg.edge.ple_dim;
let cols = (nl * pd) as u32;
plan.push(Step::D {
tag: "ple_gather",
pl: gpl.clone(),
bg: make_bg(
ctx,
gpl,
&[
&pw.table.scales,
&pw.table.quants,
&s.ple_tok,
&s.ple,
&cnt0,
],
&udims(cols, 1, 0, 0),
),
gx: cols.div_ceil(256),
gy: 1,
});
let ple_f16 = w.cfg.wdtype == crate::weights::WDtype::F16;
let mp_pl: &wgpu::ComputePipeline = if ple_f16 {
m1_gemv_f16.as_ref().expect("f16 pipelines built when used")
} else {
&m1_gemv
};
let mp_bufs: Vec<&wgpu::Buffer> = if ple_f16 {
vec![&pw.model_proj.quants, &s.cur, &s.ple_ctx]
} else {
vec![
&pw.model_proj.scales,
&pw.model_proj.quants,
&s.cur,
&s.ple_ctx,
]
};
plan.push(Step::D {
tag: "gemv_k",
pl: mp_pl.clone(),
bg: make_bg(ctx, mp_pl, &mp_bufs, &udims(cols, h as u32, 0, 1)),
gx: cols.div_ceil(m1_rpw),
gy: 1,
});
plan.push(Step::D {
tag: "ple_prep",
pl: ppl.clone(),
bg: make_bg(
ctx,
ppl,
&[&s.ple_ctx, &pw.proj_norm, &s.ple],
&mf(pd as f32, 1.0 / (h as f32).sqrt()),
),
gx: nl as u32,
gy: 1,
});
}
for li in 0..nl {
let layer = &w.layers[li];
// Gemma-4 edge per-layer geometry — on uniform models these equal the captured
// scalars and every dispatch below is unchanged.
let hd_l = w.cfg.head_dim_at(li);
let (qd_l, kd_l) = (nh * hd_l, nkv * hd_l);
let im_l = w.cfg.intermediate_at(li);
let owns_kv = w.cfg.owns_kv(li);
let wide_l = hd_l > 256;
let attn_t_l: &wgpu::Buffer = if hd_l == hd {
&attn_t
} else {
attn_t2
.as_ref()
.expect("wide attn_t twin built for wide layers")
};
// THIS layer's per-SITE container kinds pick each site's kernel family + bind-group
// shape; a Q4_0 site in a mixed plan must NOT take a no-scale-table pipeline. `wf16`
// remains the layer-level policy signal for the sites that have no per-site kind
// (DeltaNet/Conv archs never carry layer_kinds).
let wf16 = w.cfg.wdtype_at(li) == crate::weights::WDtype::F16;
let kinds = w.cfg.kinds_at(li);
let pl_of = |k: crate::weights::WKind| -> &wgpu::ComputePipeline {
use crate::weights::WKind;
match k {
WKind::Q4_0 => &m1_gemv,
WKind::F16 => m1_gemv_f16.as_ref().expect("f16 pipelines built when used"),
WKind::Q4K => q4k_pl.as_ref().expect("q4k pipeline built when used"),
WKind::Q6K => q6k_pl.as_ref().expect("q6k pipeline built when used"),
WKind::Q5K => q5k_pl.as_ref().expect("q5k pipeline built when used"),
WKind::Q5_0N => q5n_pl.as_ref().expect("q5_0n pipeline built when used"),
WKind::Q8_0N => q8n_pl.as_ref().expect("q8_0n pipeline built when used"),
WKind::Iq4Nl => iq4nl_pl.as_ref().expect("iq4_nl pipeline built when used"),
WKind::Iq4Xs => iq4xs_pl.as_ref().expect("iq4_xs pipeline built when used"),
k if k.is_iq_grid() => iqg_pls[iqg_idx(k)]
.as_ref()
.expect("iq-grid pipeline built when used"),
_ => unreachable!(),
}
};
let gv_pl = pl_of(if wf16 {
crate::weights::WKind::F16
} else {
crate::weights::WKind::Q4_0
});
let qkv_pl = pl_of(kinds.qkv);
let o_pl = pl_of(kinds.o);
// Macros, not closures (the binding lists borrow from their arguments) — and defined
// HERE, inside the loop, because macro_rules is hygienic: the bodies can only see
// definition-site locals, and they must see THIS layer's `wf16`/`kinds`.
macro_rules! gv_bufs {
($scales:expr, $quants:expr, $a:expr, $b:expr) => {
if wf16 {
vec![$quants, $a, $b]
} else {
vec![$scales, $quants, $a, $b]
}
};
}
// Kind-aware twin: every non-Q4_0 container is a no-scale-table layout.
macro_rules! k_bufs {
($k:expr, $scales:expr, $quants:expr, $a:expr, $b:expr) => {
if $k == crate::weights::WKind::Q4_0 || $k.is_iq_grid() {
// Q4_0: real scale table. IQ grid kinds: the codebook rides the
// scales slot (loader stores it there) — same 5-binding shape.
vec![$scales, $quants, $a, $b]
} else {
vec![$quants, $a, $b]
}
};
}
let onorm = &layer.operator_norm;
// Phi-2: the layer's `qkv_bias` also carries the FFN biases after the [q|k|v] section.
let layer_bias = if has_biases {
match &layer.op {
crate::weights::Op::Attn { qkv_bias, .. } => qkv_bias.as_ref(),
_ => None,
}
} else {
None
};
// Parallel arches (Falcon): hoist the single shared LayerNorm(cur) → normed_x here,
// BEFORE attn mutates cur. Both the attn and the mlp read this same normed_x below.
if is_parallel {
plan.push(Step::D {
tag: "layernorm_affine",
pl: ln_affine.clone(),
bg: make_bg(
ctx,
&ln_affine,
&[&s.cur, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
}
match &layer.op {
Op::Conv {
in_proj,
conv_w,
out_proj,
} => {
// norm + in_proj(B/C/x) projection (parallel), then the ring mix (k=1: the
// 1-slot snapshot is written and unused), then out_proj (+residual).
plan.push(Step::D {
tag: "convdot_k",
pl: convdot_k.clone(),
bg: make_bg(
ctx,
&convdot_k,
&[&in_proj.scales, &in_proj.quants, &s.cur, onorm, &s.bcx],
&udims(h as u32, conv_l as u32, 0, 0),
),
gx: h as u32,
gy: 1,
});
plan.push(Step::D {
tag: "convmix_k",
pl: convmix_k.clone(),
bg: make_bg(
ctx,
&convmix_k,
&[
&s.bcx,
conv_w,
&conv_state[li],
&s.convy,
&conv_mix_snap1,
&cmeta1,
&cnt0,
],
&udims(h as u32, conv_l as u32, 1, 0),
),
gx: (h as u32).div_ceil(64),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&out_proj.scales, &out_proj.quants, &s.convy, &s.cur),
&udims(h as u32, h as u32, 1, 1),
),
gx: (h as u32).div_ceil(m1_rpw),
gy: 1,
});
}
Op::DeltaNet {
in_proj,
conv_w,
gpar,
out_proj,
} => {
let dn = dn_pls.as_ref().expect("DeltaNet layers imply dn_pls");
// Norm + [qkv|z|b|a] projection: SPLIT rmsnorm + lcpp GEMV on the lcpp
// family (mirrors the batch plans), fused gn32 otherwise.
if m1_lcpp {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.cur, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&in_proj.scales, &in_proj.quants, &normed_x, &s.dn_mixed),
&udims(dn_in_rows as u32, h as u32, 0, 1),
),
gx: (dn_in_rows as u32).div_ceil(m1_rpw),
gy: 1,
});
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&in_proj.scales, &in_proj.quants, &s.cur, onorm, &s.dn_mixed],
&udims(dn_in_rows as u32, h as u32, 0, 1),
&epsm,
),
gx: (dn_in_rows as u32).div_ceil(gemv_rows_per_wg),
gy: 1,
});
}
// Depthwise conv over the qkv span (slot 0 via cmeta1), then the recurrence.
plan.push(Step::D {
tag: "dn.conv",
pl: dn.conv.clone(),
bg: bg2(
&dn.conv,
&[
&s.dn_mixed,
conv_w,
dn_ring[li].as_ref().expect("dn ring"),
&s.dn_conved,
&cmeta1,
&cnt0,
],
&udims(dn_conv_dim as u32, dn_kn as u32, 1, 0),
&udims(dn_in_rows as u32, 0, 0, 0),
),
gx: (dn_conv_dim as u32).div_ceil(64),
gy: 1,
});
plan.push(Step::D {
tag: "dn.step",
pl: dn.step.clone(),
bg: {
let bufs: [&wgpu::Buffer; 7] = [
&s.dn_conved,
&s.dn_mixed,
gpar,
dn_state[li].as_ref().expect("dn state"),
&s.dn_core,
&cmeta1,
&cnt0,
];
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
let d1 = udims(dn_nv as u32, dn_nk as u32, dn_dk as u32, dn_dv as u32);
let d2 = udims(dn_in_rows as u32, dn_conv_dim as u32, 1, 0);
entries.push(wgpu::BindGroupEntry {
binding: 7,
resource: d1.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 8,
resource: d2.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 9,
resource: epsm.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &dn.step.get_bind_group_layout(0),
entries: &entries,
})
},
gx: dn_nv as u32,
gy: 1,
});
// out_proj (+ residual add fused).
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&out_proj.scales, &out_proj.quants, &s.dn_core, &s.cur),
&udims(h as u32, (dn_nv * dn_dv) as u32, 1, 1),
),
gx: (h as u32).div_ceil(m1_rpw),
gy: 1,
});
}
Op::Attn {
qkv,
o,
q_norm: _,
k_norm: _,
window,
local_rope,
attn_gate,
qkv_bias,
} => {
let qk_norm_cat = qk_norm_cat[li]
.as_ref()
.expect("attention layer implies a qk-norm concat");
// Norm + q|k|v projection (concatenated): split on the lcpp family (the
// attn_gate below reuses this site's normed_x — same layer input), fused
// gn32 otherwise.
// KV-shared consumers project Q ONLY (their loader qkv holds just q rows).
let qkv_rows = if !owns_kv {
qd_l
} else if w.cfg.layer_k_eq_v[li] {
qd_l + kd_l
} else {
qd_l + 2 * kd_l
};
if skip_prenorm {
// OLMo2 post-norm: project the RAW residual (no pre-norm).
plan.push(Step::D {
tag: "gemv_k",
pl: qkv_pl.clone(),
bg: make_bg(
ctx,
qkv_pl,
&k_bufs!(kinds.qkv, &qkv.scales, &qkv.quants, &s.cur, &s.qkv),
&udims(qkv_rows as u32, h as u32, 0, 1),
),
gx: (qkv_rows as u32).div_ceil(m1_rpw),
gy: 1,
});
} else if is_affine_ln {
// Affine LayerNorm → normed_x (fully normed: (x-mean)/std·w + b), then a
// PLAIN projection (the norm is not fused, so it's backend-agnostic). On
// PARALLEL arches normed_x was already computed (hoisted) above — reuse it.
if !is_parallel {
plan.push(Step::D {
tag: "layernorm_affine",
pl: ln_affine.clone(),
bg: make_bg(
ctx,
&ln_affine,
&[&s.cur, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
}
plan.push(Step::D {
tag: "gemv_k",
pl: qkv_pl.clone(),
bg: make_bg(
ctx,
qkv_pl,
&k_bufs!(kinds.qkv, &qkv.scales, &qkv.quants, &normed_x, &s.qkv),
&udims(qkv_rows as u32, h as u32, 0, 1),
),
gx: (qkv_rows as u32).div_ceil(m1_rpw),
gy: 1,
});
} else {
let ni = norm_in!();
if m1_lcpp {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[ni, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: qkv_pl.clone(),
bg: make_bg(
ctx,
qkv_pl,
&k_bufs!(
kinds.qkv,
&qkv.scales,
&qkv.quants,
&normed_x,
&s.qkv
),
&udims(qkv_rows as u32, h as u32, 0, 1),
),
gx: (qkv_rows as u32).div_ceil(m1_rpw),
gy: 1,
});
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&qkv.scales, &qkv.quants, ni, onorm, &s.qkv],
&udims(qkv_rows as u32, h as u32, 0, 1),
&epsm,
),
gx: (qkv_rows as u32).div_ceil(gemv_rows_per_wg),
gy: 1,
});
}
}
// ONE fused kernel: qk-RMSNorm + RoPE + KV-cache write. At gy=1 the k-variant
// reads cs/sn[0..hd]. Sliding layers (Gemma-3) bind the LOCAL-theta tables.
let (csb, snb) = if *local_rope {
(&s.cosb_l, &s.sinb_l)
} else {
(&s.cosb, &s.sinb)
};
// Gemma-4 attention flags: weightless V-norm on every layer, K=V on global
// layers, scaling fixed at 1.0. Zero for every other architecture.
let g4 = w.cfg.arch == crate::weights::Arch::Gemma4;
let rd_af = if w.cfg.rotary_dim > 0 && w.cfg.rotary_dim < hd {
w.cfg.rotary_dim as u32
} else {
0
};
// af.y: 0 plain, 1 K=V, 2 Q-only (KV-shared consumer).
let af_mode = if !owns_kv {
2
} else {
u32::from(w.cfg.layer_k_eq_v[li])
};
let af = udims(u32::from(g4), af_mode, 0, rd_af);
// Qwen2 q/k/v bias: add to the fused qkv (1 column, M=1) before qkrc.
if let Some(bias) = qkv_bias {
let total = qkv_rows as u32;
plan.push(Step::D {
tag: "qkv_bias",
pl: qkv_bias_pl.clone(),
bg: make_bg(
ctx,
&qkv_bias_pl,
&[&s.qkv, bias],
&udims(qkv_rows as u32, total, 0, 0),
),
gx: total.div_ceil(256),
gy: 1,
});
}
// OLMo2: full-vector qk-norm on the fused qkv BEFORE rope (qkrc then does rope
// with no per-head norm, since arch_has_qk_norm=false). q_norm/k_norm here.
if skip_prenorm {
plan.push(Step::D {
tag: "qk_fullnorm",
pl: qk_fullnorm.clone(),
bg: bg2(
&qk_fullnorm,
&[&s.qkv, onorm], // operator_norm packs [q_norm | k_norm]
&udims(qd_fn, kd_fn, 0, 0),
&epsm,
),
gx: 2,
gy: 1,
});
}
let qp: &wgpu::ComputePipeline = if wide_l {
qkrc_k_wide.as_ref().expect("wide qkrc built")
} else {
&qkrc_k
};
plan.push(Step::D {
tag: "qkrc_k",
pl: qp.clone(),
bg: bg2(
qp,
&[
&s.qkv,
qk_norm_cat,
csb,
snb,
&s.qn,
&k_cache[li],
&v_cache[li],
&btab,
&cmeta1,
&cnt0,
],
attn_t_l,
&af,
),
// Q-only consumers dispatch the q heads alone: the k/v branches (and
// the donor-cache writes) never run.
gx: if owns_kv {
(nh + 2 * nkv) as u32
} else {
nh as u32
},
gy: 1,
});
if let Some((t, (kvq, kvs, ksig))) = tq
.as_ref()
.and_then(|t| t.caches[li].as_ref().map(|c| (t, c)))
{
// Quantize the position just written by qkrc (f32 cache = staging: 4-bit
// MSE for V, 4-bit MSE + 1-bit QJL + γ for K), then attend from the
// quantized cache with the unbiased key-score correction.
plan.push(Step::D {
tag: "t.quant_pl",
pl: t.quant_pl.clone(),
bg: make_bg(
ctx,
&t.quant_pl,
&[
&k_cache[li],
&v_cache[li],
&t.rot,
&t.cents,
&t.sketch,
kvq,
kvs,
ksig,
],
attn_t_l,
),
gx: nkv as u32,
gy: 2,
});
// Default 0 = MSE-only scores: MEASURED per-step argmax agreement on
// Qwen3-0.6B teacher-forced is 18/32 MSE-only vs 8/32 with the unbiased
// QJL correction (its zero-mean noise ≈ the 4-bit MSE error magnitude, and
// near-tie argmax prefers small bias over equal-size noise at short
// context). The unbiased path (OSFKB_TQ_CORR=1) remains for long-context /
// aggregate-quality use, where the paper's guarantees apply.
let corr_mode = std::env::var("OSFKB_TQ_CORR")
.ok()
.and_then(|v| v.parse::<u32>().ok())
.unwrap_or(0);
plan.push(Step::D {
tag: "t.attn_pl",
pl: t.attn_pl.clone(),
bg: bg2(
&t.attn_pl,
&[&s.qn, kvq, kvs, &t.rot, &t.cents, &t.sketch, ksig, &s.attn],
attn_t_l,
&udims(*window, u32::from(g4), corr_mode, 0),
),
gx: nh as u32,
gy: 1,
});
} else if let (Some(ksp), Some(kmg), false) = (&attn_ksp, &attn_kmerge, wide_l)
{
// Split-KV pair: z splits ride grid.y (no Step/replay change; the
// count sits in lw.w), the LSE merge folds them — the KV-length
// parallelism a single long-context stream was missing.
plan.push(Step::D {
tag: "attn_k",
pl: ksp.clone(),
bg: bg2(
ksp,
&[
&s.qn,
&k_cache[li],
&v_cache[li],
&attn_part,
&btab,
&cmeta1,
&cnt0,
],
attn_t_l,
&udims(*window, u32::from(g4), 0, attn_skv_z),
),
gx: nh as u32,
gy: attn_skv_z,
});
plan.push(Step::D {
tag: "attn_merge",
pl: kmg.clone(),
bg: make_bg(ctx, kmg, &[&attn_part, &s.attn], &attn_merge_u),
gx: nh as u32,
gy: 1,
});
} else {
let ak: &wgpu::ComputePipeline = if wide_l {
attn_k_wide.as_ref().expect("wide attn_k built")
} else {
&attn_k
};
plan.push(Step::D {
tag: "attn_k",
pl: ak.clone(),
bg: bg2(
ak,
&[
&s.qn,
&k_cache[li],
&v_cache[li],
&s.attn,
&btab,
&cmeta1,
&cnt0,
],
attn_t_l,
&udims(*window, u32::from(g4), 0, 0),
),
gx: nh as u32,
gy: 1,
});
}
if let Some(ag) = attn_gate {
// Qwen3.5 attention output gate: gate = ag·norm(x) (same norm+GEMV as
// qkv), then attn ← attn · sigmoid(gate), pre-o_proj. The lcpp arm
// reuses the qkv site's normed_x (same layer input, still untouched).
let dn = dn_pls.as_ref().expect("attn_gate implies dn_pls");
if m1_lcpp {
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&ag.scales, &ag.quants, &normed_x, &s.agate),
&udims(qd as u32, h as u32, 0, 1),
),
gx: (qd as u32).div_ceil(m1_rpw),
gy: 1,
});
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&ag.scales, &ag.quants, &s.cur, onorm, &s.agate],
&udims(qd as u32, h as u32, 0, 1),
&epsm,
),
gx: (qd as u32).div_ceil(gemv_rows_per_wg),
gy: 1,
});
}
plan.push(Step::D {
tag: "dn.gate_mul",
pl: dn.gate_mul.clone(),
bg: make_bg(
ctx,
&dn.gate_mul,
&[&s.agate, &s.attn],
&udims(qd as u32, 0, 0, 0),
),
gx: (qd as u32).div_ceil(64),
gy: 1,
});
}
// o-proj: LFM2 adds the residual in the GEMV (acc=1); Gemma-3 post-norms the
// block output first (sandwich norm), so acc=0 → opout, then RMSNORM_ADD.
let acc = u32::from(layer.post_op_norm.is_none());
let o_target = if layer.post_op_norm.is_some() {
&s.opout
} else {
&s.cur
};
plan.push(Step::D {
tag: "gemv_k",
pl: o_pl.clone(),
bg: make_bg(
ctx,
o_pl,
&k_bufs!(kinds.o, &o.scales, &o.quants, &s.attn, o_target),
&udims(h as u32, qd_l as u32, acc, 1),
),
gx: (h as u32).div_ceil(m1_rpw),
gy: 1,
});
if let Some(pn) = &layer.post_op_norm {
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(
ctx,
&rmsnorm_add,
&[&s.opout, pn, &s.cur],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
}
}
Op::Mla {
q_a,
q_a_norm,
q_or_qb,
kv_a,
kv_a_norm,
kv_b,
o,
qk_nope,
qk_rope,
v_head_dim,
kv_lora,
q_lora,
} => {
let qkh = qk_nope + qk_rope;
macro_rules! gemv {
($w:expr, $inp:expr, $out:expr, $rows:expr, $cols:expr, $acc:expr) => {{
let w: &crate::weights::Q4 = $w;
let rows: usize = $rows;
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&w.scales, &w.quants, $inp, $out),
&udims(rows as u32, ($cols) as u32, $acc, 1),
),
gx: (rows as u32).div_ceil(m1_rpw),
gy: 1,
});
}};
}
// 1. layer norm
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.cur, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
// 2. q: q_a → q_a_layernorm → q_b, or direct q_proj → s.qkv (raw q [nh·qkh])
if *q_lora > 0 {
gemv!(q_a.as_ref().unwrap(), &normed_x, &s.mla_qc, *q_lora, h, 0);
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.mla_qc, q_a_norm.as_ref().unwrap(), &s.mla_qln],
&mf(*q_lora as f32, 1.0),
),
gx: 1,
gy: 1,
});
gemv!(q_or_qb, &s.mla_qln, &s.qkv, nh * qkh, *q_lora, 0);
} else {
gemv!(q_or_qb, &normed_x, &s.qkv, nh * qkh, h, 0);
}
// 3. kv_a → s.mla_kvc [kv_lora | k_pe]; 4. layernorm the latent → s.mla_kvln
gemv!(kv_a, &normed_x, &s.mla_kvc, kv_lora + qk_rope, h, 0);
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.mla_kvc, kv_a_norm, &s.mla_kvln],
&mf(*kv_lora as f32, 1.0),
),
gx: 1,
gy: 1,
});
// 5. kv_b(latent) → s.mla_kvb [nh·(qk_nope+v_head_dim)]
gemv!(
kv_b,
&s.mla_kvln,
&s.mla_kvb,
nh * (qk_nope + v_head_dim),
*kv_lora,
0
);
// 6. assemble: rope pe, broadcast k_pe, pad v → s.qn + KV cache
let mkp = udims(
nh as u32,
*qk_nope as u32,
*qk_rope as u32,
*v_head_dim as u32,
);
let msp = udims(1, *kv_lora as u32, 0, 0);
plan.push(Step::D {
tag: "mla_assemble",
pl: mla_assemble.clone(),
bg: bg2(
&mla_assemble,
&[
&s.qkv,
&s.mla_kvb,
&s.mla_kvc,
&s.qn,
&k_cache[li],
&v_cache[li],
&btab,
&cmeta1,
&cnt0,
],
&mkp,
&msp,
),
gx: nh as u32,
gy: 1,
});
// 7. attention (reuse ATTN_K; hd = qk_head_dim; v is zero-padded so dims ≥ v_head_dim vanish)
plan.push(Step::D {
tag: "attn_k",
pl: attn_k.clone(),
bg: bg2(
&attn_k,
&[
&s.qn,
&k_cache[li],
&v_cache[li],
&s.attn,
&btab,
&cmeta1,
&cnt0,
],
&attn_t,
&udims(0, 0, 0, 0),
),
gx: nh as u32,
gy: 1,
});
// 8. o_proj (padded to [hidden, nh·qkh]) with residual add (acc = 1)
gemv!(o, &s.attn, &s.cur, h, nh * qkh, 1);
}
Op::Mamba {
in_proj,
conv_w,
conv_b,
x_proj,
dt_proj,
ssm_par,
out_proj,
d_inner,
d_state,
d_conv,
dt_rank,
} => {
let (di, ds, dc, dtrk) = (*d_inner, *d_state, *d_conv, *dt_rank);
let cst = mamba_conv_state[li].as_ref().expect("mamba conv state");
let sst = mamba_ssm_state[li].as_ref().expect("mamba ssm state");
macro_rules! mgemv {
($w:expr, $inp:expr, $out:expr, $rows:expr, $cols:expr, $acc:expr) => {{
let w: &crate::weights::Q4 = $w;
let rows: usize = $rows;
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&w.scales, &w.quants, $inp, $out),
&udims(rows as u32, ($cols) as u32, $acc, 1),
),
gx: (rows as u32).div_ceil(m1_rpw),
gy: 1,
});
}};
}
// 1. RMSNorm pre-norm.
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.cur, onorm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
// 2. in_proj → [x_ssm(di) | z(di)] in mb_xz.
mgemv!(in_proj, &normed_x, &mb_xz, 2 * di, h, 0);
// 3. depthwise conv1d step (+ SiLU) → mb_xconv (reads x_ssm = mb_xz[0..di]).
plan.push(Step::D {
tag: "mamba_conv",
pl: mamba_conv.clone(),
bg: make_bg(
ctx,
&mamba_conv,
&[&mb_xz, conv_w, conv_b, cst, &mb_xconv],
&udims(di as u32, dc as u32, 0, 0),
),
gx: (di as u32).div_ceil(64),
gy: 1,
});
// 4. x_proj → [dt | B | C] in mb_dbl.
mgemv!(x_proj, &mb_xconv, &mb_dbl, dtrk + 2 * ds, di, 0);
// 5. dt_proj reads mb_dbl → mb_dtraw. Cols are padded to a multiple of 32 (the
// pad columns' weights are zero, so the B/C tail of mb_dbl contributes nothing).
mgemv!(dt_proj, &mb_dbl, &mb_dtraw, di, dtrk.div_ceil(32) * 32, 0);
// 6. selective-scan step (updates ssm state) → mb_y (gated by SiLU(z)).
plan.push(Step::D {
tag: "mamba_ssm",
pl: mamba_ssm.clone(),
bg: make_bg(
ctx,
&mamba_ssm,
&[&mb_xconv, &mb_dtraw, &mb_dbl, &mb_xz, ssm_par, sst, &mb_y],
&udims(di as u32, ds as u32, dtrk as u32, di as u32),
),
gx: (di as u32).div_ceil(64),
gy: 1,
});
// 7. out_proj (+ residual).
mgemv!(out_proj, &mb_y, &s.cur, h, di, 1);
}
Op::Rwkv {
ln0,
att_mix,
att_k,
att_v,
att_r,
att_o,
time_decay,
time_first,
ffn_mix,
ffn_k,
ffn_v,
ffn_r,
ffn_dim,
} => {
let fd = *ffn_dim;
let wst = rwkv_wkv_state[li].as_ref().expect("rwkv wkv state");
let ash = rwkv_att_shift[li].as_ref().expect("rwkv att shift");
let fsh = rwkv_ffn_shift[li].as_ref().expect("rwkv ffn shift");
macro_rules! rgemv {
($w:expr, $inp:expr, $out:expr, $rows:expr, $cols:expr, $acc:expr) => {{
let w: &crate::weights::Q4 = $w;
let rows: usize = $rows;
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&w.scales, &w.quants, $inp, $out),
&udims(rows as u32, ($cols) as u32, $acc, 1),
),
gx: (rows as u32).div_ceil(m1_rpw),
gy: 1,
});
}};
}
// Affine LayerNorm dispatch (kp = (n, _, eps, _)); single workgroup.
macro_rules! ln_af {
($src:expr, $norm:expr, $dst:expr) => {
plan.push(Step::D {
tag: "layernorm_affine",
pl: ln_affine.clone(),
bg: make_bg(
ctx,
&ln_affine,
&[$src, $norm, $dst],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
};
}
// A one-thread-per-channel elementwise dispatch (token-shift / WKV / sig-mul).
macro_rules! chan {
($tag:literal, $pl:expr, $bufs:expr, $u:expr) => {
plan.push(Step::D {
tag: $tag,
pl: $pl.clone(),
bg: make_bg(ctx, &$pl, $bufs, $u),
gx: (h as u32).div_ceil(64),
gy: 1,
});
};
}
// ln0 (block 0 only): affine LayerNorm on the embeddings, which permanently
// replaces the residual stream (RWKV's `pre_ln`). LAYERNORM_AFFINE can't alias
// src=dst in one dispatch, so it lands in rk_x and is copied back into cur.
if let Some(ln0) = ln0 {
ln_af!(&s.cur, ln0, &rk_x);
plan.push(Step::D {
tag: "copy_vec",
pl: copy_vec.clone(),
bg: make_bg(
ctx,
©_vec,
&[&rk_x, &s.cur],
&udims(h as u32, 0, 0, 0),
),
gx: (h as u32).div_ceil(64),
gy: 1,
});
}
// ===== TIME-MIX (WKV linear attention) =====
// 1. ln1 → rk_x.
ln_af!(&s.cur, onorm, &rk_x);
// 2. token-shift mix → [xk|xv|xr]; att-shift state updated to this token (rk_x).
chan!(
"rwkv_tokenshift",
rwkv_tokenshift,
&[&rk_x, ash, att_mix, &rk_xk, &rk_xv, &rk_xr],
&udims(h as u32, 0, 0, 0)
);
// 3. k/v/r projections.
rgemv!(att_k, &rk_xk, &rk_key, h, h, 0);
rgemv!(att_v, &rk_xv, &rk_v, h, h, 0);
rgemv!(att_r, &rk_xr, &rk_r, h, h, 0);
// 4. WKV recurrence → rk_wkv; updates the [num|den|max] state.
chan!(
"rwkv_wkv",
rwkv_wkv,
&[&rk_key, &rk_v, time_decay, time_first, wst, &rk_wkv],
&udims(h as u32, 0, 0, 0)
);
// 5. gate by sigmoid(receptance): rk_rwkv = sigmoid(r)·wkv (overwrite, add=0).
chan!(
"rwkv_sigmul",
rwkv_sigmul,
&[&rk_r, &rk_wkv, &rk_rwkv],
&udims(h as u32, 0, 0, 0)
);
// 6. output projection (+ residual).
rgemv!(att_o, &rk_rwkv, &s.cur, h, h, 1);
// ===== CHANNEL-MIX =====
// 7. ln2 → rk_x.
ln_af!(&s.cur, &layer.ffn_norm, &rk_x);
// 8. token-shift mix → [xk | (unused) | xr]; ffn-shift state updated to rk_x.
chan!(
"rwkv_tokenshift",
rwkv_tokenshift,
&[&rk_x, fsh, ffn_mix, &rk_xk, &rk_xv, &rk_xr],
&udims(h as u32, 0, 0, 0)
);
// 9. key projection (→ ffn_dim), then relu² via MLP_ACT_MUL mode 3 (relu(k)·k).
rgemv!(ffn_k, &rk_xk, &rk_key, fd, h, 0);
plan.push(Step::D {
tag: "mlp_act_mul",
pl: mlp_act_mul.clone(),
bg: make_bg(
ctx,
&mlp_act_mul,
&[&rk_key, &rk_key, &rk_relu],
&udims(fd as u32, 3, 0, 0),
),
gx: (fd as u32).div_ceil(256),
gy: 1,
});
// 10. value projection (ffn_dim → hidden) and receptance projection.
rgemv!(ffn_v, &rk_relu, &rk_v, h, fd, 0);
rgemv!(ffn_r, &rk_xr, &rk_r, h, h, 0);
// 11. gate by sigmoid(receptance) and ADD into the residual (add=1).
chan!(
"rwkv_sigmul",
rwkv_sigmul,
&[&rk_r, &rk_v, &s.cur],
&udims(h as u32, 1, 0, 0)
);
}
}
// ffn-norm + w1 + w3 + act·mul → gate, then w2 (residual or sandwich-normed add).
// On the lcpp family the norm is SPLIT out (rmsnorm → normed_x, the same split the
// qkv site makes) and gate+up run the two-stream lcpp kernel — the fused mlp_k was
// the largest per-layer kernel left on the old family. `OSFKB_MLP_LCPP=0` pins the
// fused kernel for A/B.
if is_affine_ln || skip_prenorm {
// Affine LayerNorm → normed_x, then the projections as PLAIN gemvs + the standalone
// act·mul → s.gate. No fused-norm MLP kernel, so it runs on every backend. PARALLEL
// arches reuse the hoisted normed_x; NON-GATED arches (Falcon) skip the w1 gate and
// multiply the activation by ones (act(fc1·x)·1 = act(fc1·x)). OLMo2 (skip_prenorm)
// projects the RAW residual (no LN) and post-norms the output below.
let ffn_src = if skip_prenorm { &s.cur } else { &normed_x };
if !is_parallel && !skip_prenorm {
plan.push(Step::D {
tag: "layernorm_affine",
pl: ln_affine.clone(),
bg: make_bg(
ctx,
&ln_affine,
&[&s.cur, &layer.ffn_norm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
}
if !mlp_nogate {
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&layer.w1.scales, &layer.w1.quants, ffn_src, &ln_gate_raw),
&udims(im_l as u32, h as u32, 0, 1),
),
gx: (im_l as u32).div_ceil(m1_rpw),
gy: 1,
});
}
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&layer.w3.scales, &layer.w3.quants, ffn_src, &ln_up_raw),
&udims(im_l as u32, h as u32, 0, 1),
),
gx: (im_l as u32).div_ceil(m1_rpw),
gy: 1,
});
// Phi-2: add the fc1 bias to the up-projection (ln_up_raw) BEFORE the activation.
if let Some(b) = layer_bias {
plan.push(Step::D {
tag: "bias_add",
pl: bias_add.clone(),
bg: make_bg(
ctx,
&bias_add,
&[b, &ln_up_raw],
&udims(im as u32, qkv_rows_b, 0, 0),
),
gx: (im as u32).div_ceil(256),
gy: 1,
});
}
// gate stream: fc1 output for non-gated (act·1), the w1 gate otherwise.
let gate_src = if mlp_nogate { &ln_up_raw } else { &ln_gate_raw };
let up_src = if mlp_nogate { &ones_im } else { &ln_up_raw };
plan.push(Step::D {
tag: "mlp_act_mul",
pl: mlp_act_mul.clone(),
bg: make_bg(
ctx,
&mlp_act_mul,
&[gate_src, up_src, &s.gate],
&udims(im_l as u32, act_mode, 0, 0),
),
gx: (im_l as u32).div_ceil(256),
gy: 1,
});
} else {
let ni_ffn = norm_in!();
// Moshi's FFN is Q8_0 — the lcpp family only has Q4 MLP kernels, so Moshi takes
// the mlp_k route (swapped to the q8 kernel above).
if m1_lcpp && mlp_lcpp_enabled() && !moshi_q8 {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[ni_ffn, &layer.ffn_norm, &normed_x],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
// Q1 swaps the fused-MLP pipeline for its sign-decoding twin (same 8 bindings).
// f16 swaps it for a SIX-binding twin (no scale tables), so the bind group is
// built to match the chosen pipeline rather than shared.
let mlp_pl = match kinds.gate_up {
crate::weights::WKind::Q8_0N => mlp_q8n
.as_ref()
.expect("q8_0n MLP built when a gate_up site uses it"),
crate::weights::WKind::F16 => mlp_f16
.as_ref()
.expect("f16 pipelines are built whenever any layer is f16"),
_ if wf16 => mlp_f16
.as_ref()
.expect("f16 pipelines are built whenever any layer is f16"),
_ => spec_pls.mlp_q1.as_ref().unwrap_or(&spec_pls.mlp_lcpp),
};
let mlp_bufs: Vec<&wgpu::Buffer> =
if wf16 || kinds.gate_up != crate::weights::WKind::Q4_0 {
vec![&layer.w1.quants, &layer.w3.quants, &normed_x, &s.gate]
} else {
vec![
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
&normed_x,
&s.gate,
]
};
plan.push(Step::D {
tag: "mlp_k",
pl: mlp_pl.clone(),
bg: bg2(
mlp_pl,
&mlp_bufs,
&udims(im_l as u32, h as u32, 0, 1),
&epsm_act,
),
gx: (im_l as u32).div_ceil(4),
gy: 1,
});
} else {
plan.push(Step::D {
tag: "mlp_k",
pl: mlp_k.clone(),
bg: bg2(
&mlp_k,
&[
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
ni_ffn,
&layer.ffn_norm,
&s.gate,
],
&udims(im_l as u32, h as u32, 0, 1),
&epsm_act,
),
gx: (im_l as u32).div_ceil(gemv_rows_per_wg),
gy: 1,
});
}
}
// Qwen3.5: the "dense" GLU is the SHARED EXPERT — its output is gated by
// sigmoid(shared_gate·xn) BEFORE the routed experts join, and the whole MoE block
// adds to the residual at the end (never directly into cur).
// ...but ONLY on the MoE flavor. On a DENSE Qwen3.5 (`moe: None`) there is no shared
// expert and no MoE block to add `mlpout` back, so the GLU must accumulate straight
// into the residual like every other dense arch — otherwise its output is computed and
// silently dropped, and the model decodes garbage.
// The shared-gated MoE path (dense GLU is the sigmoid-gated shared expert, added back
// via the fused epilogue). Keyed off `shared_gate` — true for Qwen3.5 AND Qwen2-MoE,
// the only arches with a gated shared expert. (`q35` kept as the local name.)
let q35 = layer.moe.as_ref().is_some_and(|m| m.shared_gate.is_some());
let acc = u32::from(layer.post_ffn_norm.is_none() && !q35);
let d_target = if layer.post_ffn_norm.is_some() || q35 {
&s.mlpout
} else {
&s.cur
};
// Moshi's down-proj is Q8_0 like the rest of its FFN; every other arch keeps Q4.
// The Q8 kernel has its own row-tiling (`moshi_q8_rpw`), so dispatch it explicitly
// rather than with `m1_rpw` (whose lcpp/Q1 retune would over/under-cover the rows).
// NOTE: this site does NOT go through `gv_pl` — it has its own Moshi-Q8 arm — so the
// f16 selection has to be repeated here. Missing it is exactly how the first f16
// bring-up produced a silently DEAD down-projection (the MLP computed, then its
// result was multiplied by garbage weights and contributed ~nothing).
let (down_pl, down_rpw) = if kinds.down != crate::weights::WKind::Q4_0 {
(pl_of(kinds.down), 4u32)
} else if wf16 {
(pl_of(crate::weights::WKind::F16), 4u32)
} else {
match gemv_q8_ffn.as_ref() {
Some(p) => (p, moshi_q8_rpw),
None => (&m1_gemv, m1_rpw),
}
};
plan.push(Step::D {
tag: "gemv_k",
pl: down_pl.clone(),
bg: make_bg(
ctx,
down_pl,
&k_bufs!(
kinds.down,
&layer.w2.scales,
&layer.w2.quants,
&s.gate,
d_target
),
&udims(h as u32, im_l as u32, acc, 1),
),
gx: (h as u32).div_ceil(down_rpw),
gy: 1,
});
// MoE (Gemma-4): route the ffn-normed input, then ACCUMULATE the top-k weighted expert
// outputs onto the dense-MLP output (the dense path is the architecture's shared
// expert), all before the sandwich norm.
if let Some(moe) = &layer.moe {
let (e_n, kx, mi) = (
w.cfg.num_experts as u32,
w.cfg.top_k_experts as u32,
w.cfg.moe_intermediate as u32,
);
if let Some(sg_w) = &moe.shared_gate {
// Router split: norm+sgate → coalesced logits → pick (see kernel docs).
plan.push(Step::D {
tag: "router_norm_sgate_pl",
pl: router_norm_sgate_pl.clone(),
bg: bg2(
&router_norm_sgate_pl,
&[&s.cur, &layer.ffn_norm, sg_w, &s.normed, &s.moe_sg],
&udims(e_n, h as u32, kx, 0),
&epsm,
),
gx: 1,
gy: 1,
});
plan.push(Step::D {
tag: "router_logits_pl",
pl: router_logits_pl.clone(),
bg: make_bg(
ctx,
&router_logits_pl,
&[&moe.router, &s.normed, &s.moe_logits],
&udims(e_n, h as u32, kx, 0),
),
gx: e_n.div_ceil(16),
gy: 1,
});
plan.push(Step::D {
tag: "router_pick_pl",
pl: router_pick_pl.clone(),
bg: make_bg(
ctx,
&router_pick_pl,
&[
&s.moe_logits,
&moe.per_expert_scale,
&s.moe_sel,
&s.moe_wsel,
],
&udims(e_n, h as u32, kx, 0),
),
gx: 1,
gy: 1,
});
} else {
// ffn-normed input for router + experts (the dense mlp norms internally).
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&s.cur, &layer.ffn_norm, &s.normed],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
});
if let Some(v3) = &moe.router_v3 {
// DeepSeek-V3 aux-loss-free router: sigmoid + bias + group-limited top-k.
// Emits sel/wsel directly (weights already renorm'd × scaling) — the
// downstream moe_gate/moe_down are shared with the softmax path.
plan.push(Step::D {
tag: "moe_router_v3_pl",
pl: moe_router_v3_pl.clone(),
bg: bg2(
&moe_router_v3_pl,
&[&moe.router, &v3.bias, &s.normed, &s.moe_sel, &s.moe_wsel],
&udims(e_n, h as u32, kx, v3.n_group),
&uni(
ctx,
bytemuck::cast_slice(&[
v3.topk_group,
v3.scaling.to_bits(),
0u32,
0u32,
]),
),
),
gx: 1,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "moe_router_pl",
pl: moe_router_pl.clone(),
bg: make_bg(
ctx,
&moe_router_pl,
&[
&moe.router,
&moe.per_expert_scale,
&s.normed,
&s.moe_sel,
&s.moe_wsel,
],
&udims(e_n, h as u32, kx, 0),
),
gx: 1,
gy: 1,
});
}
}
plan.push(Step::D {
tag: "moe_gate_pl",
pl: moe_gate_pl.clone(),
bg: bg2(
&moe_gate_pl,
&[
&moe.w1.scales,
&moe.w1.quants,
&moe.w3.scales,
&moe.w3.quants,
&s.normed,
&s.moe_sel,
&s.moe_gate,
],
&udims(mi, h as u32, 0, kx),
&epsm_act,
),
gx: mi,
gy: kx,
});
if q35 {
// FUSED down: `cur += shared·g + routed` in the epilogue (see the derivation).
plan.push(Step::D {
tag: "moe_down_fused_pl",
pl: moe_down_fused_pl.clone(),
bg: make_bg_down_fused(
ctx,
&moe_down_fused_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&s.moe_gate,
&s.moe_sel,
&s.moe_wsel,
&s.mlpout,
&s.moe_sg,
&s.cur,
],
&udims(h as u32, mi, 0, kx),
),
gx: h as u32,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "moe_down_pl",
pl: moe_down_pl.clone(),
bg: make_bg(
ctx,
&moe_down_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&s.moe_gate,
&s.moe_sel,
&s.moe_wsel,
d_target,
],
&udims(h as u32, mi, 0, kx),
),
gx: h as u32,
gy: 1,
});
}
}
if let Some(pn) = &layer.post_ffn_norm {
// Gemma-4 edge: layer_scalar rides THIS add only when no PLE block follows
// (HF order: FFN add → PLE sub-block → ×layer_scalar).
let ls_here = layer
.edge
.as_ref()
.filter(|e| e.ple.is_none())
.map_or(0.0, |e| e.layer_scalar);
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(
ctx,
&rmsnorm_add,
&[&s.mlpout, pn, &s.cur],
&mf_ls(h as f32, 1.0, ls_here),
),
gx: 1,
gy: 1,
});
}
// Gemma-4 edge PLE injection — a THIRD sandwich sub-block on the post-FFN stream:
// x += post_norm(proj(gelu(gate(x)) · ple[li])), then x ×= layer_scalar.
if let Some(ledge) = &layer.edge {
if let Some(lp) = &ledge.ple {
let pd = w.cfg.edge.ple_dim;
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&lp.gate.scales, &lp.gate.quants, &s.cur, &s.ple_g),
&udims(pd as u32, h as u32, 0, 1),
),
gx: (pd as u32).div_ceil(m1_rpw),
gy: 1,
});
// gelu(gate) · the layer's slice of the merged per-layer inputs — the MLP
// act kernel with its up-stream offset (dims.z) selecting the slice.
plan.push(Step::D {
tag: "mlp_act_mul",
pl: mlp_act_mul.clone(),
bg: make_bg(
ctx,
&mlp_act_mul,
&[&s.ple_g, &s.ple, &s.ple_a],
&udims(pd as u32, act_mode, (li * pd) as u32, 0),
),
gx: (pd as u32).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gv_pl.clone(),
bg: make_bg(
ctx,
gv_pl,
&gv_bufs!(&lp.proj.scales, &lp.proj.quants, &s.ple_a, &s.opout),
&udims(h as u32, pd as u32, 0, 1),
),
gx: (h as u32).div_ceil(m1_rpw),
gy: 1,
});
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(
ctx,
&rmsnorm_add,
&[&s.opout, &lp.post_norm, &s.cur],
&mf_ls(h as f32, 1.0, ledge.layer_scalar),
),
gx: 1,
gy: 1,
});
}
}
// Phi-2: add the combined o_proj+fc2 bias to the residual once, after both the attn
// and mlp outputs have accumulated into cur (parallel → they share this add).
if let Some(b) = layer_bias {
plan.push(Step::D {
tag: "bias_add",
pl: bias_add.clone(),
bg: make_bg(
ctx,
&bias_add,
&[b, &s.cur],
&udims(h as u32, qkv_rows_b + im as u32, 0, 0),
),
gx: (h as u32).div_ceil(256),
gy: 1,
});
}
}
if w.cfg.stage_last {
// Build the final-norm step first (norm_in! itself pushes CENTER for OLMo, so it must
// not run inside a `plan.push(...)` argument — that would double-borrow `plan`).
let final_norm = if is_affine_ln {
Step::D {
tag: "layernorm_affine",
pl: ln_affine.clone(),
bg: make_bg(
ctx,
&ln_affine,
&[&s.cur, &w.embedding_norm, &s.hnorm],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
}
} else {
let ni_final = norm_in!();
Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[ni_final, &w.embedding_norm, &s.hnorm],
&mf(h as f32, 1.0),
),
gx: 1,
gy: 1,
}
};
plan.push(final_norm);
let (lm_bg, lm_gx, lm_gy) = q4lm.make(
ctx,
&w.embed_q4,
&s.hnorm,
&s.logits,
vocab as u32,
h as u32,
1,
);
plan.push(Step::D {
tag: "q4lm.pipeline_for(1)",
pl: q4lm.pipeline_for(1).clone(),
bg: lm_bg,
gx: lm_gx,
gy: lm_gy,
});
// Phi-2: untied lm_head bias, packed into `embedding_norm` after [weight|bias].
if has_biases {
plan.push(Step::D {
tag: "bias_add",
pl: bias_add.clone(),
bg: make_bg(
ctx,
&bias_add,
&[&w.embedding_norm, &s.logits],
&udims(vocab as u32, 2 * h as u32, 0, 0),
),
gx: (vocab as u32).div_ceil(256),
gy: 1,
});
}
}
// A multimodal checkpoint (one carrying an M-RoPE section) gets the arena; a text model gets
// nothing and behaves exactly as before. EXCEPTION: a model whose oversized embedding forces
// the split `GATHER_K_F16` (e.g. NuExtract3) runs text-only — the vision arena and the split
// gather are not co-wired yet (image rows would need the same lo/hi split), and the gather
// path asserts `img_embed.is_none()`. Such a model is served for text extraction anyway, so
// disable the arena rather than panic. (Q1/GGUF also sets `embed_f16` but is mrope-less.)
let img_rows = if w.cfg.mrope_section.is_some() && w.embed_f16.is_none() {
std::env::var("OSFKB_IMG_TOKENS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(DEFAULT_IMG_TOKENS)
} else {
0
};
let img_embed = (img_rows > 0).then(|| ctx.empty(img_rows * w.cfg.hidden));
Self {
q1_rp4: std::cell::RefCell::new(std::collections::HashMap::new()),
w,
conv_state,
k_cache,
v_cache,
btab,
cmeta1,
img_embed,
img_rows,
pool_tokens,
max_seq_tokens,
s,
attn_t,
attn_t2,
qk_norm_cat,
plan,
spec_pls,
qkv_bias_pl,
tq,
q4lm,
wide_pls,
dn_pls,
dn_slots,
dn_ring,
dn_state,
mamba_ssm_state,
mamba_conv_state,
rwkv_wkv_state,
rwkv_att_shift,
rwkv_ffn_shift,
cnt0,
gemv_rows_per_wg,
m1_family,
step_breakdown: std::cell::RefCell::new((0.0, 0.0, wclock::Instant::now())),
cmdbuf_hits: std::cell::Cell::new(0),
cmdbuf_reuse: cmdbuf_reuse_from_env(),
masked_am: std::sync::OnceLock::new(),
decode_pls: std::sync::OnceLock::new(),
argmax_pls: std::sync::OnceLock::new(),
cached_tokens: Vec::new(),
conv_snap,
}
}
/// Upload the RoPE cos/sin tables for `pos` — the global-theta pair always, plus the
/// local-theta pair when the architecture has sliding layers (Gemma-3).
fn upload_rope(&self, ctx: &GpuCtx, pos: usize) {
let c = &self.w.cfg;
// The GLOBAL-theta pair serves the full-attention layers: on Gemma-4 edge that width is
// the WIDE head, with the proportional zero-frequency identity tail.
let (coss, sins) = if c.edge.rope_prop_pairs > 0 {
rope_cs_prop(c.head_dim_full(), c.rope_theta, c.edge.rope_prop_pairs, pos)
} else {
rope_cs_padded(c.head_dim_full(), c.rotary_dim, c.rope_theta, pos)
};
ctx.queue
.write_buffer(&self.s.cosb, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&self.s.sinb, 0, bytemuck::cast_slice(&sins));
if matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
) {
let (cl, sl) = rope_cs(c.head_dim_sliding(), c.rope_local_theta, pos);
ctx.queue
.write_buffer(&self.s.cosb_l, 0, bytemuck::cast_slice(&cl));
ctx.queue
.write_buffer(&self.s.sinb_l, 0, bytemuck::cast_slice(&sl));
}
// The wide-head uniform twin steps in lockstep with attn_t (same pos+1).
if let Some(t2) = &self.attn_t2 {
ctx.queue.write_buffer(
t2,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.max_head_dim() as u32,
(pos + 1) as u32,
)),
);
}
}
/// Reset decode caches (conv ring buffers + KV cache) for a fresh sequence.
pub fn reset(&self, ctx: &GpuCtx) {
for b in &self.conv_state {
ctx.queue.write_buffer(
b,
0,
&vec![0u8; (self.w.cfg.hidden * self.w.cfg.conv_l) * 4],
);
}
for (li, ring) in self.dn_ring.iter().enumerate() {
if let (Some(ring), Some(st)) = (ring.as_ref(), self.dn_state[li].as_ref()) {
let zr = vec![0f32; (ring.size() / 4) as usize];
ctx.queue.write_buffer(ring, 0, bytemuck::cast_slice(&zr));
let zs = vec![0f32; (st.size() / 4) as usize];
ctx.queue.write_buffer(st, 0, bytemuck::cast_slice(&zs));
}
}
// Mamba recurrence state (SSM + conv) + RWKV token-shift state reset to zero, like KV/DN.
for st in self
.mamba_ssm_state
.iter()
.chain(self.mamba_conv_state.iter())
.chain(self.rwkv_att_shift.iter())
.chain(self.rwkv_ffn_shift.iter())
.flatten()
{
let z = vec![0f32; (st.size() / 4) as usize];
ctx.queue.write_buffer(st, 0, bytemuck::cast_slice(&z));
}
// RWKV WKV state is [num | den | max]; num/den reset to 0 but the log-sum-exp pivot `max`
// resets to -1e38 (the running-max identity), matching the constructor.
for st in self.rwkv_wkv_state.iter().flatten() {
let n = (st.size() / 4) as usize;
let mut z = vec![0f32; n];
z[2 * (n / 3)..].fill(-1e38);
ctx.queue.write_buffer(st, 0, bytemuck::cast_slice(&z));
}
}
/// Byte size of one layer's conv ring buffer.
fn conv_bytes(&self) -> u64 {
(self.w.cfg.hidden * self.w.cfg.conv_l * 4) as u64
}
/// Save every layer's current conv ring-buffer into `conv_snap[..][pos]` (the state *for*
/// position `pos`, i.e. before forwarding `tokens[pos]`). One submit; no readback.
fn snapshot_conv(&self, ctx: &GpuCtx, pos: usize) {
let stride = self.conv_bytes();
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (cs, snap) in self.conv_state.iter().zip(&self.conv_snap) {
enc.copy_buffer_to_buffer(cs, 0, snap, pos as u64 * stride, stride);
}
ctx.queue.submit([enc.finish()]);
}
/// Restore every layer's conv ring-buffer from `conv_snap[..][pos]` (the state for position `pos`).
fn restore_conv(&self, ctx: &GpuCtx, pos: usize) {
let stride = self.conv_bytes();
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (cs, snap) in self.conv_state.iter().zip(&self.conv_snap) {
enc.copy_buffer_to_buffer(snap, pos as u64 * stride, cs, 0, stride);
}
ctx.queue.submit([enc.finish()]);
}
/// Prefix-cached prefill: reuse the KV + conv state for the longest prompt prefix shared with the
/// previously prefilled prompt, and only run the forward over the diverging suffix. The KV cache is
/// positionally overwritten (never cleared), so a matching prefix's KV is already valid; the conv
/// ring buffer is restored from the snapshot at the divergence point. Invariant: after each call,
/// `conv_snap[p]` holds the conv state for `cached_tokens[0..p]` (`p < CACHE_MAX`), maintained
/// transitively — so the restore at the next call's LCP is always valid.
pub fn prefill_cached(&mut self, ctx: &GpuCtx, tokens: &[u32]) {
let mut lcp = 0usize;
while lcp < self.cached_tokens.len()
&& lcp < tokens.len()
&& self.cached_tokens[lcp] == tokens[lcp]
{
lcp += 1;
}
lcp = lcp.min(CACHE_MAX - 1); // clamp to a snapshot slot we actually hold
if lcp == 0 {
self.reset(ctx);
} else {
self.restore_conv(ctx, lcp);
}
for (pos, &tok) in tokens.iter().enumerate().skip(lcp) {
if pos < CACHE_MAX {
self.snapshot_conv(ctx, pos);
}
self.forward_only(ctx, tok, pos);
}
// Snapshot the FINAL conv state (for position tokens.len()) so an identical / extending next
// prompt can restore at LCP == len — the loop above only snapshots up to len-1.
if tokens.len() < CACHE_MAX {
self.snapshot_conv(ctx, tokens.len());
}
self.cached_tokens = tokens.to_vec();
}
}
pub(crate) fn uni(ctx: &GpuCtx, data: &[u8]) -> wgpu::Buffer {
ctx.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("kp"),
contents: data,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
})
}
/// Build a reusable bind group: `bufs` on bindings 0.. and the uniform `meta` on the last binding.
fn make_bg1(ctx: &GpuCtx, pl: &wgpu::ComputePipeline, buf: &wgpu::Buffer) -> wgpu::BindGroup {
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: buf.as_entire_binding(),
}],
})
}
/// Bind group for MOE_DOWN_FUSED: the base kernel's uniform stays at binding 6, so the two
/// appended buffers (shared-gate scalar, residual) land at 7/8 — explicit indices, not append.
fn make_bg_down_fused(
ctx: &GpuCtx,
pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer; 8],
dims: &wgpu::Buffer,
) -> wgpu::BindGroup {
let bindings = [0u32, 1, 2, 3, 4, 5, 7, 8];
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.zip(bindings)
.map(|(b, i)| wgpu::BindGroupEntry {
binding: i,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: 6,
resource: dims.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &entries,
})
}
/// Grid-codebook IQ kinds in a fixed order: (WKind, ggml type id). Pipelines for the family
/// are built and dispatched by this index.
const IQG_KINDS: [(crate::weights::WKind, u32); 7] = [
(crate::weights::WKind::Iq2Xxs, 16),
(crate::weights::WKind::Iq2Xs, 17),
(crate::weights::WKind::Iq2S, 22),
(crate::weights::WKind::Iq3Xxs, 18),
(crate::weights::WKind::Iq3S, 21),
(crate::weights::WKind::Iq1S, 19),
(crate::weights::WKind::Iq1M, 29),
];
fn iqg_idx(k: crate::weights::WKind) -> usize {
IQG_KINDS
.iter()
.position(|&(kk, _)| kk == k)
.expect("iqg_idx on a non-grid kind")
}
pub(crate) fn make_bg(
ctx: &GpuCtx,
pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer],
meta: &wgpu::Buffer,
) -> wgpu::BindGroup {
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32,
resource: meta.as_entire_binding(),
});
if std::env::var("OSFKB_BG_TRACE").is_ok() {
if bufs.iter().any(|b| b.size() == 0) {
let sizes: Vec<u64> = bufs.iter().map(|b| b.size()).collect();
eprintln!("BG_TRACE: zero-size binding, all sizes {sizes:?}");
}
}
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &entries,
})
}
/// RoPE cos/sin for `head_dim` at position `pos` (CPU; cheap, uploaded per token).
/// [`rope_cs`] with partial-rotary support: tables built at `rd` dims (frequency exponents over
/// `rd`) and zero-padded to `hd` — the qkrc kernel never reads past `rd` when `af.w = rd`.
fn rope_cs_padded(hd: usize, rd: usize, theta: f32, pos: usize) -> (Vec<f32>, Vec<f32>) {
let d = if rd > 0 && rd < hd { rd } else { hd };
let (mut c, mut s) = rope_cs(d, theta, pos);
c.resize(hd, 0.0);
s.resize(hd, 0.0);
(c, s)
}
fn rope_cs(hd: usize, theta: f32, pos: usize) -> (Vec<f32>, Vec<f32>) {
let half = hd / 2;
let (mut cos, mut sin) = (vec![0f32; hd], vec![0f32; hd]);
for i in 0..half {
let f = (pos as f32) * theta.powf(-2.0 * i as f32 / hd as f32);
let (s, c) = f.sin_cos();
cos[i] = c;
cos[i + half] = c;
sin[i] = s;
sin[i + half] = s;
}
(cos, sin)
}
/// Gemma-4 "proportional" rope tables: only the first `pairs` frequency pairs rotate — their
/// exponent denominator stays the FULL head width — and the rest get zero frequency
/// (cos=1/sin=0), which the standard full-width rope kernel applies as identity. That is
/// `transformers` `_compute_proportional_rope_parameters`, table-side.
fn rope_cs_prop(hd: usize, theta: f32, pairs: usize, pos: usize) -> (Vec<f32>, Vec<f32>) {
let half = hd / 2;
let (mut cos, mut sin) = (vec![0f32; hd], vec![0f32; hd]);
for i in 0..half {
let (s, c) = if i < pairs {
((pos as f32) * theta.powf(-2.0 * i as f32 / hd as f32)).sin_cos()
} else {
(0.0, 1.0)
};
cos[i] = c;
cos[i + half] = c;
sin[i] = s;
sin[i + half] = s;
}
(cos, sin)
}
/// Which axis (0 = t, 1 = h, 2 = w) each rotary frequency rotates by, under **interleaved** M-RoPE.
///
/// HF builds the full `(3, seq, half)` frequency block and then overwrites slices of the T row:
/// axis `d ∈ {h, w}` claims indices `offset, offset+3, offset+6, …` up to `section[d] · 3`, with
/// `offset = d`. Everything not claimed stays T. For the common `[11, 11, 10]` over a 32-wide rotary
/// half this reduces exactly to `axis = i % 3` — T H W T H W … — but it is derived here rather than
/// hardcoded, so a checkpoint with a different section still gets the right map (and the pinning
/// test asserts the reduction actually holds for [11,11,10]).
fn mrope_axis_map(section: [usize; 3], half: usize, chunked: bool) -> Vec<u8> {
let mut axes = vec![0u8; half];
if chunked {
// GLM-OCR: CONTIGUOUS sections — lanes [0, s0) → t, [s0, s0+s1) → h, rest → w.
let mut i = 0usize;
for (d, &len) in section.iter().enumerate() {
for _ in 0..len {
if i >= half {
break;
}
axes[i] = d as u8;
i += 1;
}
}
return axes;
}
for (d, &len) in section.iter().enumerate().skip(1) {
let mut i = d;
while i < (len * 3).min(half) {
axes[i] = d as u8;
i += 3;
}
}
axes
}
/// Builds one column's rope cos/sin, dispatching on whether the model carries an M-RoPE section.
///
/// The axis map depends only on the config, so it is computed ONCE per staging call rather than per
/// column. When the model has no section this is exactly the old 1-D call — and even when it does,
/// a text column (uniform `mpos`) produces bit-identical tables, so turning M-RoPE on cannot perturb
/// a text-only workload.
struct RopeStager {
prop_pairs: usize,
axes: Option<Vec<u8>>,
hd: usize,
rotary_dim: usize,
theta: f32,
}
impl RopeStager {
fn new(c: &crate::weights::Lfm2Config, hd: usize) -> Self {
let rd = if c.rotary_dim > 0 && c.rotary_dim < hd {
c.rotary_dim
} else {
hd
};
Self {
axes: c
.mrope_section
.map(|s| mrope_axis_map(s, rd / 2, c.arch == crate::weights::Arch::GlmOcr)),
hd,
rotary_dim: c.rotary_dim,
theta: c.rope_theta,
prop_pairs: c.edge.rope_prop_pairs,
}
}
fn col(&self, mpos: [u32; 3], pos: u32) -> (Vec<f32>, Vec<f32>) {
// Gemma-4 edge "proportional" rope: zero-frequency identity tail (see rope_cs_prop).
if self.prop_pairs > 0 {
return rope_cs_prop(self.hd, self.theta, self.prop_pairs, pos as usize);
}
match &self.axes {
Some(a) => mrope_cs_padded(self.hd, self.rotary_dim, self.theta, a, mpos),
None => rope_cs_padded(self.hd, self.rotary_dim, self.theta, pos as usize),
}
}
}
/// [`rope_cs_padded`] with a 3-D position: frequency `i` rotates by `pos3[axes[i]]`.
///
/// With a uniform position (`t == h == w`, i.e. every text token) this is **bitwise identical** to
/// the 1-D path — the axis map only selects which of three equal numbers to use. That identity is
/// the whole safety argument for turning M-RoPE on for a model that previously ran text-only, and
/// `mrope_is_bitwise_identical_to_1d_rope_on_text` pins it.
fn mrope_cs_padded(
hd: usize,
rd: usize,
theta: f32,
axes: &[u8],
pos3: [u32; 3],
) -> (Vec<f32>, Vec<f32>) {
let d = if rd > 0 && rd < hd { rd } else { hd };
let half = d / 2;
debug_assert_eq!(axes.len(), half, "axis map must cover the rotary half");
let (mut cos, mut sin) = (vec![0f32; hd], vec![0f32; hd]);
for i in 0..half {
let pos = pos3[axes[i] as usize];
// Same expression as `rope_cs`, so the uniform-position case is bit-for-bit identical.
let f = (pos as f32) * theta.powf(-2.0 * i as f32 / d as f32);
let (s, c) = f.sin_cos();
cos[i] = c;
cos[i + half] = c;
sin[i] = s;
sin[i + half] = s;
}
(cos, sin)
}
impl Lfm2Gpu {
/// Single-token forward at absolute position `pos`; returns logits `[vocab]`. Replays the cached
/// plan: the only per-token CPU work is 3 small `write_buffer`s (cos/sin + the attn `T`) + the
/// embed/KV copies + recording the dispatches — no bind-group/uniform allocation.
/// Wall time of `bp.plan` encoded `inner` times back-to-back in ONE command buffer and ONE
/// compute pass — the instrument that separates per-dispatch cost from submit/poll latency.
///
/// Dispatch count scales with `inner`; the submit, the poll and the readback happen ONCE
/// regardless. So the DIFFERENCE cancels every fixed cost:
///
/// ```text
/// wall(inner=2) − wall(inner=1) = kernel time + inter-dispatch cost, for one plan pass
/// ```
///
/// Subtract the kernel-time sum (which [`Self::profile_batch_step`] measures soundly) and what
/// remains is the inter-dispatch dead time, alone. That is the one question the P7 accounting
/// could only answer by elimination: ~47% of production's GPU-side verify is not kernel
/// execution, and the survivors were "barriers between dispatches" vs "submit/poll latency".
/// Elimination is how I mis-attributed this cost twice already — first to the profiler's own
/// per-pass overhead, then to the spec plan's snapshot kernels. This measures it instead.
/// See `bench/P7_profile_the_verify_is_gap_bound.md`.
///
/// TIMING ONLY. Repeating the plan re-runs its KV/DN state writes, so the model state and any
/// output after a repeated call are meaningless — nothing is read back and nothing may be. Call
/// it only on a plan already warmed by a real step, and throw the engine away afterwards.
/// `with_head` also dispatches the LM head each pass, mirroring what `batch_stage_step` does on
/// the last stage when any column wants a logit. The head is NOT in `bp.plan` — it is dispatched
/// separately — so [`Self::profile_batch_step`] never times it, and on a 248 320-token vocab
/// that omission is not small. Running the probe both ways isolates the head's cost by
/// difference, using the same fixed-cost cancellation as the `inner` sweep.
pub fn time_batch_step_repeated(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
inner: usize,
with_head: bool,
) -> Result<f64> {
anyhow::ensure!(inner >= 1, "inner must be >= 1");
anyhow::ensure!(
!with_head || bp.lm.is_some(),
"--with-head needs a stage that carries the LM head (stage_last)"
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
for _ in 0..inner {
for step in &bp.plan {
match step {
Step::D { pl, bg, gx, gy, .. } => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
Step::Di {
pl,
bg,
args,
offset,
..
} => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
// Skipped, exactly as `batch_stage_step`'s pass skips it: the token→hidden
// embed is an explicit buffer copy encoded BEFORE the pass, not a dispatch.
// It is fixed per submit, so it neither belongs in the repeated body nor
// survives the 2×−1× subtraction.
Step::Embed => {}
}
}
if with_head {
let (lm_bg, lm_gx, lm_gy) = bp.lm.as_ref().expect("checked above");
p.set_pipeline(&bp.lm_pl);
p.set_bind_group(0, lm_bg, &[]);
p.dispatch_workgroups(*lm_gx, *lm_gy, 1);
}
}
}
let t0 = wclock::Instant::now();
ctx.queue.submit([enc.finish()]);
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
Ok(t0.elapsed().as_secs_f64())
}
/// Profile one M=1 forward: replay the plan with ONE COMPUTE PASS PER STEP, each pass
/// timestamped, and return `(step_index, gx, gy, micros)` per dispatch plus the wall time.
///
/// **The per-dispatch times are sound; `wall − busy` is NOT.** One compute pass per step means
/// this replay pays a full barrier plus a timestamp resolve at every dispatch — serialization
/// that production (which issues the whole plan inside a SINGLE pass, see
/// [`Self::batch_stage_step`]) never pays. So the gap between wall and GPU-busy here is mostly
/// this instrument's own cost, and reading it as engine dead time is an error I actually made
/// (`bench/P7_profile_the_verify_is_gap_bound.md`). Use the RELATIVE split — which kernels burn
/// the time — and use [`Self::time_batch_step_repeated`] for anything about gaps.
///
/// Note also that the returned steps cover `bp.plan` only: the LM head, which
/// `batch_stage_step` dispatches when `stage_last && need_logit`, is NOT among them, so the
/// kernel-time sum omits it.
pub fn profile_batch_step(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
) -> Result<(Vec<ProfiledStep>, f64)> {
anyhow::ensure!(ctx.timestamps, "adapter lacks TIMESTAMP_QUERY");
let c = &self.w.cfg;
let hd = c.head_dim;
let trash_row = (self.dn_slots - 1) as u32;
let mut cm = vec![0u32; bp.k * 2];
let mut coss = Vec::with_capacity(bp.k * hd);
let mut sins = Vec::with_capacity(bp.k * hd);
// Global-theta tables serve the full-attention layers — the WIDE heads on an edge model.
let rope = RopeStager::new(c, c.head_dim_full());
for p in 0..bp.k {
let (pos, row, mpos) = match cols.get(p) {
Some(col) => (col.pos, col.btrow, col.mpos),
None => (0, trash_row, [0; 3]),
};
cm[p * 2] = pos;
cm[p * 2 + 1] = row;
let (cs, sn) = rope.col(mpos, pos);
coss.extend_from_slice(&cs);
sins.extend_from_slice(&sn);
}
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
ctx.queue
.write_buffer(&bp.cmeta_k, 0, bytemuck::cast_slice(&cm));
ctx.queue
.write_buffer(&bp.csk, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&bp.snk, 0, bytemuck::cast_slice(&sins));
let steps: Vec<(usize, &Step)> = bp
.plan
.iter()
.enumerate()
.filter(|(_, st)| matches!(st, Step::D { .. }))
.collect();
let nq = (steps.len() * 2) as u32;
let qs = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
label: Some("prof_b"),
ty: wgpu::QueryType::Timestamp,
count: nq,
});
let qbuf = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("prof_b_resolve"),
size: u64::from(nq) * 8,
usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let qstage = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("prof_b_stage"),
size: u64::from(nq) * 8,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if c.stage_first {
for (p, col) in cols.iter().enumerate() {
let (esrc, eoff) = self.w.embed_row_src(col.token);
enc.copy_buffer_to_buffer(
esrc,
eoff,
&bp.cur,
(p * c.hidden * 4) as u64,
(c.hidden * 4) as u64,
);
}
}
for (qi, (_, st)) in steps.iter().enumerate() {
if let Step::D { pl, bg, gx, gy, .. } = st {
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: None,
timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
query_set: &qs,
beginning_of_pass_write_index: Some((qi * 2) as u32),
end_of_pass_write_index: Some((qi * 2 + 1) as u32),
}),
});
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
enc.resolve_query_set(&qs, 0..nq, &qbuf, 0);
enc.copy_buffer_to_buffer(&qbuf, 0, &qstage, 0, u64::from(nq) * 8);
let t0 = wclock::Instant::now();
ctx.queue.submit([enc.finish()]);
let slice = qstage.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
let wall = t0.elapsed().as_secs_f64() * 1e6;
rx.recv()
.map_err(|e| anyhow::anyhow!("prof staging: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let raw: Vec<u64> =
bytemuck::cast_slice(&slice.get_mapped_range().expect("mapped range")).to_vec();
qstage.unmap();
let out = steps
.iter()
.enumerate()
.map(|(qi, (si, st))| {
let (gx, gy, tag) = match st {
Step::D { gx, gy, tag, .. } => (*gx, *gy, *tag),
_ => (0, 0, ""),
};
let dt = raw[qi * 2 + 1].saturating_sub(raw[qi * 2]) as f64
* f64::from(ctx.ts_period)
/ 1e3;
(*si, gx, gy, dt, tag)
})
.collect();
Ok((out, wall))
}
pub fn profile_forward(
&self,
ctx: &GpuCtx,
token: u32,
pos: usize,
) -> Result<(Vec<ProfiledStep>, f64)> {
anyhow::ensure!(ctx.timestamps, "adapter lacks TIMESTAMP_QUERY");
let c = &self.w.cfg;
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
let steps: Vec<(usize, &Step)> = self
.plan
.iter()
.enumerate()
.filter(|(_, st)| matches!(st, Step::D { .. }))
.collect();
let nq = (steps.len() * 2) as u32;
let qs = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
label: Some("prof"),
ty: wgpu::QueryType::Timestamp,
count: nq,
});
let qbuf = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("prof_resolve"),
size: u64::from(nq) * 8,
usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let qstage = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("prof_stage"),
size: u64::from(nq) * 8,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if c.stage_first {
let (esrc, eoff) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(esrc, eoff, &self.s.cur, 0, (c.hidden * 4) as u64);
}
for (qi, (_, st)) in steps.iter().enumerate() {
if let Step::D { pl, bg, gx, gy, .. } = st {
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: None,
timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
query_set: &qs,
beginning_of_pass_write_index: Some((qi * 2) as u32),
end_of_pass_write_index: Some((qi * 2 + 1) as u32),
}),
});
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
enc.resolve_query_set(&qs, 0..nq, &qbuf, 0);
enc.copy_buffer_to_buffer(&qbuf, 0, &qstage, 0, u64::from(nq) * 8);
let t0 = wclock::Instant::now();
ctx.queue.submit([enc.finish()]);
let slice = qstage.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
let wall = t0.elapsed().as_secs_f64() * 1e6;
rx.recv()
.map_err(|e| anyhow::anyhow!("prof staging: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let raw: Vec<u64> =
bytemuck::cast_slice(&slice.get_mapped_range().expect("mapped range")).to_vec();
qstage.unmap();
let out = steps
.iter()
.enumerate()
.map(|(qi, (si, st))| {
let (gx, gy, tag) = match st {
Step::D { gx, gy, tag, .. } => (*gx, *gy, *tag),
_ => (0, 0, ""),
};
let dt = raw[qi * 2 + 1].saturating_sub(raw[qi * 2]) as f64
* f64::from(ctx.ts_period)
/ 1e3;
(*si, gx, gy, dt, tag)
})
.collect();
Ok((out, wall))
}
pub fn forward(&self, ctx: &GpuCtx, token: u32, pos: usize) -> Result<Vec<f32>> {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
let _ = h;
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if self.w.cfg.edge.ple_dim > 0 {
// PLE gather input for this step (queue-ordered before the submit below).
ctx.queue
.write_buffer(&self.s.ple_tok, 0, bytemuck::cast_slice(&[token]));
}
self.record(&mut enc, token, pos);
ctx.queue.submit([enc.finish()]);
ctx.read(&self.s.logits, c.vocab)
}
/// The Moshi entry: the CALLER supplies the input embedding (the 17-table sum) instead of a
/// token id. Returns the logits (for `Arch::Moshi`: the TEXT head) and leaves the
/// post-final-norm hidden in `s.hnorm` (read with [`Self::read_hnorm_for_tests`]) — the
/// depformer's conditioning input.
pub fn forward_from_embeds(&self, ctx: &GpuCtx, emb: &[f32], pos: usize) -> Result<Vec<f32>> {
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
self.record_from_embeds(ctx, &mut enc, emb, pos)?;
ctx.queue.submit([enc.finish()]);
ctx.read(&self.s.logits, self.w.cfg.vocab)
}
/// The recording half of [`Self::forward_from_embeds`]: queue the preamble writes and the
/// full plan into the CALLER's encoder — the Moshi full-GPU step chains the text argmax
/// and the depformer cycle into the same submit.
pub(crate) fn record_from_embeds(
&self,
ctx: &GpuCtx,
enc: &mut wgpu::CommandEncoder,
emb: &[f32],
pos: usize,
) -> Result<()> {
let c = &self.w.cfg;
anyhow::ensure!(
emb.len() == c.hidden,
"embeds len {} != hidden {}",
emb.len(),
c.hidden
);
let (hd, nh, nkv) = (c.head_dim, c.n_heads, c.n_kv_heads);
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
// Queue-ordered: lands in `s.cur` before the submitted plan runs.
ctx.queue
.write_buffer(&self.s.cur, 0, bytemuck::cast_slice(emb));
self.record_steps(enc, None, pos);
Ok(())
}
/// [`Self::forward_from_embeds`] without the logits readback: leaves logits in `s.logits`
/// (combine with [`Self::argmax_logits`] — a 4-byte token read instead of the full row) and
/// the post-final-norm hidden in `s.hnorm`. The Moshi fast path.
pub fn forward_from_embeds_only(&self, ctx: &GpuCtx, emb: &[f32], pos: usize) -> Result<()> {
let c = &self.w.cfg;
anyhow::ensure!(
emb.len() == c.hidden,
"embeds len {} != hidden {}",
emb.len(),
c.hidden
);
let (hd, nh, nkv) = (c.head_dim, c.n_heads, c.n_kv_heads);
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
ctx.queue
.write_buffer(&self.s.cur, 0, bytemuck::cast_slice(emb));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
self.record_steps(&mut enc, None, pos);
ctx.queue.submit([enc.finish()]);
Ok(())
}
/// Submit the forward for `token` at `pos` but leave the 65536-wide logits in `s.logits` — no
/// CPU readback. The submit is non-blocking, so the caller can build the grammar mask on the CPU
/// while the GPU runs the forward, then call [`Self::argmax_masked`]. This is the constrained-decode
/// fast path: it removes the 256 KB/token logit copy and overlaps the grammar scan with the forward.
pub fn forward_only(&self, ctx: &GpuCtx, token: u32, pos: usize) {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
let _ = h;
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if self.w.cfg.edge.ple_dim > 0 {
// PLE gather input for this step (queue-ordered before the submit below).
ctx.queue
.write_buffer(&self.s.ple_tok, 0, bytemuck::cast_slice(&[token]));
}
self.record(&mut enc, token, pos);
ctx.queue.submit([enc.finish()]);
}
/// On-GPU grammar-masked argmax over the logits left in `s.logits` by [`Self::forward_only`].
/// `mask` is `ceil(vocab/32)` u32 words, bit `i` set ⇒ token `i` is grammar-allowed; the GPU picks
/// the max-logit allowed token and only that single id is read back. The caller MUST ensure at
/// least one bit is set (a well-formed grammar always permits one continuation); an all-zero mask
/// yields an arbitrary id.
pub fn argmax_masked(&self, ctx: &GpuCtx, mask: &[u32]) -> Result<u32> {
let m = self
.masked_am
.get_or_init(|| MaskedArgmax::new(ctx, self.w.cfg.vocab, &self.s.logits));
ctx.queue
.write_buffer(&m.mask_buf, 0, bytemuck::cast_slice(mask));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(&m.block_pl);
p.set_bind_group(0, &m.block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&m.merge_pl);
p.set_bind_group(0, &m.merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
ctx.queue.submit([enc.finish()]);
Ok(ctx.read_u32(&m.chosen, 1)?[0])
}
/// Record the plan into `enc`: a **single compute pass per run of dispatches**, ended/reopened
/// only at the embed/KV copies (one Metal encoder per run, not per dispatch — the key lever from
/// the WebGPU 1400 tok/s driver).
fn record(&self, enc: &mut wgpu::CommandEncoder, token: u32, pos: usize) {
self.record_steps(enc, Some(token), pos)
}
/// `token: None` skips the `Step::Embed` table copy — the caller has already written the
/// input embedding into `s.cur` (the Moshi 17-table-sum entry, `forward_from_embeds`).
pub(crate) fn hnorm_buf(&self) -> &wgpu::Buffer {
&self.s.hnorm
}
/// The lm_head logits buffer + its vocab — the Moshi sampler adds Gumbel noise onto the
/// text logits between the temporal forward and `record_argmax`.
/// SESSION SNAPSHOT: copy the first `tokens` KV entries of every layer into fresh buffers.
/// PersonaPlex priming (voice replay + role prompt) is deterministic, so a server can pay
/// it ONCE at startup and restore it per session instead of re-running ~100 forwards
/// (measured 13 s of dead air at every connect).
pub fn kv_snapshot(&self, ctx: &GpuCtx, tokens: usize) -> Vec<(wgpu::Buffer, wgpu::Buffer)> {
let c = &self.w.cfg;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
// Per-layer widths (Gemma-4 edge); consumers snapshot their donor's pool AGAIN — the
// restore then writes identical bytes twice, which keeps the zip pairing trivial.
let out: Vec<(wgpu::Buffer, wgpu::Buffer)> = self
.k_cache
.iter()
.zip(&self.v_cache)
.enumerate()
.map(|(li, (k, v))| {
let kv_l = tokens * c.n_kv_heads * c.head_dim_at(li);
let bytes = (kv_l * 2) as u64; // f16 pool
let ks = ctx.empty_f16(kv_l);
let vs = ctx.empty_f16(kv_l);
enc.copy_buffer_to_buffer(k, 0, &ks, 0, bytes);
enc.copy_buffer_to_buffer(v, 0, &vs, 0, bytes);
(ks, vs)
})
.collect();
ctx.queue.submit([enc.finish()]);
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
out
}
/// Restore a [`Self::kv_snapshot`] — one submit of 2×n_layers buffer copies (µs).
pub fn kv_restore(&self, ctx: &GpuCtx, snap: &[(wgpu::Buffer, wgpu::Buffer)], tokens: usize) {
let c = &self.w.cfg;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (li, ((k, v), (ks, vs))) in self.k_cache.iter().zip(&self.v_cache).zip(snap).enumerate()
{
let bytes = (tokens * c.n_kv_heads * c.head_dim_at(li) * 2) as u64;
enc.copy_buffer_to_buffer(ks, 0, k, 0, bytes);
enc.copy_buffer_to_buffer(vs, 0, v, 0, bytes);
}
ctx.queue.submit([enc.finish()]);
}
pub(crate) fn logits_buf(&self) -> &wgpu::Buffer {
&self.s.logits
}
pub(crate) fn text_vocab(&self) -> usize {
self.w.cfg.vocab
}
pub(crate) fn record_steps(
&self,
enc: &mut wgpu::CommandEncoder,
token: Option<u32>,
_pos: usize,
) {
let c = &self.w.cfg;
let h = c.hidden;
let mut i = 0;
while i < self.plan.len() {
match &self.plan[i] {
Step::Embed => {
if let Some(token) = token {
// caller may have pre-written `s.cur` (Moshi embeds-sum entry)
let (esrc, eoff) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(esrc, eoff, &self.s.cur, 0, (h * 4) as u64);
}
i += 1;
}
// The M=1 plan never carries Di steps (grouped MoE is wide-plan-only); the
// arm keeps the match exhaustive and correct if that ever changes.
Step::D { .. } | Step::Di { .. } => {
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
loop {
match self.plan.get(i) {
Some(Step::D { pl, bg, gx, gy, .. }) => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
i += 1;
}
Some(Step::Di {
pl,
bg,
args,
offset,
..
}) => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
i += 1;
}
_ => break,
}
}
}
}
}
}
/// Fully on-GPU greedy decode: prefill `prompt`, generate `gen` tokens, return the `gen` outputs.
/// The token lives in a GPU buffer (gather-by-GPU-token + on-GPU argmax → `tokens[pos+1]`), and
/// each step is submitted WITHOUT a per-token poll, so the CPU recording of step t+1 overlaps the
/// GPU execution of step t — removing both the per-token sync and the logits readback.
pub fn decode(
&self,
ctx: &GpuCtx,
prompt: &[u32],
ngen: usize,
) -> Result<(Vec<u32>, f64, f64)> {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
let _ = h;
self.reset(ctx);
let total = prompt.len() + ngen;
// Q1 (Bonsai-class) stores the gather table f16 (5 GB f32 > Metal's 4 GB/buffer cap).
let q1_embed = std::env::var("OSFKB_Q1_WEIGHTS").ok().as_deref() == Some("1");
let gather_pl = if q1_embed {
pipeline(ctx, "gather_f16", GATHER_F16)
} else {
pipeline(ctx, "gather", GATHER)
};
let argmax_pl = pipeline(ctx, "argmax", ARGMAX);
let tokens = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("tokens"),
size: (total * 4) as u64,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut init = vec![0u32; total];
init[..prompt.len()].copy_from_slice(prompt);
ctx.queue
.write_buffer(&tokens, 0, bytemuck::cast_slice(&init));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, prompt.len() as u32, c.vocab as u32, h as u32)),
);
let gather_bg = if q1_embed {
// The f16 table exceeds wgpu-Vulkan's 2 GiB-4 per-BINDING clamp (the BUFFER is
// fine), so bind it as two sub-ranges split at vocab/2 — the row stride h*2 B is
// 256-aligned for every production h, and the kernel picks lo/hi by token id.
// Same data, zero extra VRAM. (Metal binds it whole; this split is for Vulkan.)
let e16 = self
.w
.embed_f16
.as_ref()
.expect("Q1 decode: loader provides the f16 gather twin");
let lo_size = (c.vocab as u64 / 2) * (h as u64 * 2);
let hi_size = e16.size() - lo_size;
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("gather_f16_split"),
layout: &gather_pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: e16,
offset: 0,
size: std::num::NonZeroU64::new(lo_size),
}),
},
wgpu::BindGroupEntry {
binding: 1,
resource: tokens.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: self.s.cur.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: posu.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: e16,
offset: lo_size,
size: std::num::NonZeroU64::new(hi_size),
}),
},
],
})
} else {
make_bg(
ctx,
&gather_pl,
&[&self.w.embed, &tokens, &self.s.cur],
&posu,
)
};
let _ = &argmax_pl;
// Parallel argmax: 256-block scan → 1-workgroup merge (was 1 underutilized workgroup).
let argmax_block_pl = pipeline(ctx, "argmax_block", ARGMAX_BLOCK);
let argmax_merge_pl = pipeline(ctx, "argmax_merge", ARGMAX_MERGE);
let pv = ctx.empty(256);
let pi = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("pi"),
size: 256 * 4,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let block_bg = make_bg(ctx, &argmax_block_pl, &[&self.s.logits, &pv, &pi], &posu);
let merge_bg = make_bg(ctx, &argmax_merge_pl, &[&pv, &pi, &tokens], &posu);
// Split prefill vs decode wall-clock so the raw (unconstrained) path can report a decode-only
// rate + the prompt-prefill ms, like the constrained path.
let gen_start = prompt.len().max(1);
let t = wclock::Instant::now();
let mut prefill_secs = -1.0f64;
for pos in 0..total - 1 {
if pos + 1 >= gen_start && prefill_secs < 0.0 {
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
prefill_secs = t.elapsed().as_secs_f64();
}
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
ctx.queue.write_buffer(
&posu,
0,
bytemuck::cast_slice(&u32x4(
pos as u32,
prompt.len() as u32,
c.vocab as u32,
h as u32,
)),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if self.w.cfg.edge.ple_dim > 0 {
// PLE gather input: this step's token id, straight from the GPU token ring
// (generated tokens never round-trip to the host in this loop).
enc.copy_buffer_to_buffer(&tokens, (pos * 4) as u64, &self.s.ple_tok, 0, 4);
}
{
// ONE compute pass: gather + the whole plan (Embed skipped, no Kv) + argmax. Metal
// serializes dispatches within a pass, so the data deps hold and we pay zero
// pass-break overhead (the lever that took the qkrc fusion from 240→306).
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(&gather_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
for step in &self.plan {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
} else if let Step::Di {
pl,
bg,
args,
offset,
..
} = step
{
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
}
p.set_pipeline(&argmax_block_pl);
p.set_bind_group(0, &block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&argmax_merge_pl);
p.set_bind_group(0, &merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
ctx.queue.submit([enc.finish()]); // NO poll — pipeline CPU/GPU across tokens
}
let toks = ctx.read_u32(&tokens, total)?;
let secs = t.elapsed().as_secs_f64();
if prefill_secs < 0.0 {
prefill_secs = secs;
}
let decode_secs = (secs - prefill_secs).max(1e-9);
Ok((
toks[prompt.len()..].to_vec(),
ngen as f64 / decode_secs,
prefill_secs * 1000.0,
))
}
/// Throughput (tok/s) of the on-GPU **constrained** loop: the decode loop (`gather → plan →
/// argmax`, one submit/token, no per-token readback — the same pipelined path that hits ~907) plus
/// `extras` — the caller's grammar `mask` + `advance` dispatches — run as extra passes per token.
/// Each extra is `(pipeline, bind_group, workgroups_x)` and operates on the caller's own buffers, so
/// this measures exactly what the grammar adds to the loop. KV is skipped (matching [`Self::decode`]),
/// keeping it apples-to-apples with the pure-decode baseline. The generated ids are not grammar-
/// correct here (the mask isn't wired into this argmax); this method answers throughput only.
pub fn bench_constrained_decode(
&self,
ctx: &GpuCtx,
prompt: &[u32],
ngen: usize,
extras: &[(&wgpu::ComputePipeline, &wgpu::BindGroup, u32)],
) -> Result<f64> {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
self.reset(ctx);
let total = prompt.len() + ngen;
let gather_pl = pipeline(ctx, "gather", GATHER);
let tokens = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("tokens"),
size: (total * 4) as u64,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut init = vec![0u32; total];
init[..prompt.len()].copy_from_slice(prompt);
ctx.queue
.write_buffer(&tokens, 0, bytemuck::cast_slice(&init));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, prompt.len() as u32, c.vocab as u32, h as u32)),
);
let gather_bg = make_bg(
ctx,
&gather_pl,
&[&self.w.embed, &tokens, &self.s.cur],
&posu,
);
let argmax_block_pl = pipeline(ctx, "argmax_block", ARGMAX_BLOCK);
let argmax_merge_pl = pipeline(ctx, "argmax_merge", ARGMAX_MERGE);
let pv = ctx.empty(256);
let pi = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("pi"),
size: 256 * 4,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let block_bg = make_bg(ctx, &argmax_block_pl, &[&self.s.logits, &pv, &pi], &posu);
let merge_bg = make_bg(ctx, &argmax_merge_pl, &[&pv, &pi, &tokens], &posu);
let t = wclock::Instant::now();
for pos in 0..total - 1 {
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
ctx.queue.write_buffer(
&posu,
0,
bytemuck::cast_slice(&u32x4(
pos as u32,
prompt.len() as u32,
c.vocab as u32,
h as u32,
)),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(&gather_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
for step in &self.plan {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
} else if let Step::Di {
pl,
bg,
args,
offset,
..
} = step
{
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
}
p.set_pipeline(&argmax_block_pl);
p.set_bind_group(0, &block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&argmax_merge_pl);
p.set_bind_group(0, &merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
for &(pl, bg, wg) in extras {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(wg, 1, 1);
}
}
ctx.queue.submit([enc.finish()]);
}
ctx.read_u32(&tokens, 1)?; // single poll forces all ngen steps to complete
Ok(ngen as f64 / t.elapsed().as_secs_f64())
}
/// **Correct** on-GPU constrained decode: like [`Self::bench_constrained_decode`] but the masked
/// argmax actually reads the grammar `mask` (so the chosen token is grammar-valid) and `advance`
/// consumes that chosen token (so the FSM state threads correctly across tokens) — all on-device,
/// one submit/token, no per-token readback. The caller supplies the grammar's `mask_buf` (one u32
/// per token, written by `mask_bg`) and `chosen_buf` (1 u32, written here by the merge, read by
/// `advance_bg`); `set_state` on the grammar must seed the initial state first. Returns the `ngen`
/// generated token ids and the achieved tok/s.
///
/// Prompt-aware: the first `prompt.len()` positions build the prompt's KV (gather + plan only, no
/// masking) and the prompt tokens are preserved; the grammar engages only for the generated suffix
/// (`tokens[prompt.len()..]`). A 1-token prompt (the throughput bench) generates from pos 0, so its
/// behavior is unchanged. Stopping at the grammar's `Done` is the caller's job: it walks the returned
/// tokens through the CPU grammar and truncates at completion (any post-Done all-zero-mask argmax
/// emits a sentinel that the truncation discards).
///
/// Returns `(generated_tokens, decode_only_tok_s, prefill_ms)`. The tok/s excludes the prompt
/// prefill — for a short query behind a long few-shot prompt, the prefill otherwise dominates the
/// wall-clock and crushes the apparent rate, so callers can report the steady-state decode speed.
#[allow(clippy::too_many_arguments)]
/// The decode-loop pipelines (gather + sparse-pick stack), compiled once per engine — each WGSL
/// module is ~150ms on Metal, so per-call compilation would dominate short queries.
fn decode_pipelines(&self, ctx: &GpuCtx) -> &DecodePls {
self.decode_pls.get_or_init(|| DecodePls {
gather: pipeline(ctx, "gather", GATHER),
compact: pipeline(ctx, "mask_compact", MASK_COMPACT),
sargs: pipeline(ctx, "sparse_args", SPARSE_ARGS),
shead: crate::Q4LmHeadSparse::new(ctx),
sargmax: pipeline(ctx, "sparse_argmax", SPARSE_ARGMAX),
sargmax_k: pipeline(ctx, "sparse_argmax_k", SPARSE_ARGMAX_K),
})
}
/// Test-only: run the sparse pick (compact → indirect sparse head → lowest-id argmax) against
/// the CURRENT `hnorm` (from a prior forward) and an explicit one-u32-per-token mask. Returns
/// `(chosen, count, ids[0..count], svals[0..count])` so a test can assert the sparse dots are
/// bitwise-identical to the full head's logits.
#[doc(hidden)]
pub fn sparse_pick_for_tests(
&self,
ctx: &GpuCtx,
mask: &[u32],
) -> Result<(u32, u32, Vec<u32>, Vec<f32>)> {
self.assert_head_not_f16("sparse_pick_for_tests");
let (vocab, h) = (self.w.cfg.vocab, self.w.cfg.hidden);
anyhow::ensure!(mask.len() == vocab, "mask must be one u32 per vocab token");
let dp = self.decode_pipelines(ctx);
let mask_buf = ctx.storage_bytes(bytemuck::cast_slice(mask));
let chosen = ctx.empty(1);
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, vocab as u32, h as u32)),
);
let sp = SparsePick::new(ctx, dp, &self.w, &self.s.hnorm, &mask_buf, &chosen, &posu);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
sp.encode_compact(&mut p, dp, vocab as u32);
sp.encode_pick(&mut p, dp);
}
ctx.queue.submit([enc.finish()]);
let cnt = ctx.read_u32(&sp.smeta, 5)?[4];
let ids = if cnt == 0 {
Vec::new()
} else {
ctx.read_u32(&sp.ids, cnt as usize)?
};
let svals = if cnt == 0 {
Vec::new()
} else {
ctx.read(&sp.svals, cnt as usize)?
};
let chosen_v = ctx.read_u32(&chosen, 1)?[0];
Ok((chosen_v, cnt, ids, svals))
}
/// Build the batched-verify plan for speculative constrained decode: k-column scratch, the
/// k-column dispatch plan (bitwise-per-column-identical kernel variants; NO lm_head — the
/// per-column grammar-sparse pick replaces it), per-conv-layer ring snapshots for rollback, and
/// one [`SparsePick`] per column over the caller's grammar-mask buffer. Build once per
/// (engine, grammar, k) and reuse across queries.
pub fn make_spec_plan(&self, ctx: &GpuCtx, k: usize) -> SpecPlan {
assert!(
!self.w.cfg.is_edge(),
"gemma4-edge: the spec/verify plan is not wired for per-layer geometry yet — \
the solo (M=1) path serves these models"
);
self.assert_head_not_f16("the grammar-constrained spec plan");
assert!((2..=16).contains(&k), "spec k must be in 2..=16");
assert!(
self.w.cfg.arch != crate::weights::Arch::Qwen35,
"speculative OSFQL decoding is not wired for the Qwen3.5 hybrid architecture"
);
assert!(
!crate::weights::arch_is_layernorm(self.w.cfg.arch),
"LayerNorm arches (OLMo) are wired for M=1 decode only — no spec plan yet"
);
let w = &self.w;
let c = &w.cfg;
let (h, im, hd) = (c.hidden, c.intermediate, c.head_dim);
let (nh, nkv, nl) = (c.n_heads, c.n_kv_heads, c.n_layers);
let (qd, kd) = (nh * hd, nkv * hd);
let (eps, conv_l, vocab) = (c.eps, c.conv_l, c.vocab);
let dp = self.decode_pipelines(ctx);
let fed = ctx.empty(k);
let drafts = ctx.empty(k);
let csk = ctx.empty(k * hd);
let snk = ctx.empty(k * hd);
let csk_l = ctx.empty(k * hd);
let snk_l = ctx.empty(k * hd);
let attn_tk = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, 1)),
);
let cmeta_k = ctx.empty(k * 2);
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, vocab as u32, h as u32)),
);
let maskk = ctx.empty(k * vocab);
let cur = ctx.empty(k * h);
let bcx = ctx.empty(k * 3 * h);
let convy = ctx.empty(k * h);
let qkv = ctx.empty(k * (qd + 2 * kd));
let qn = ctx.empty(k * qd);
let attn_o = ctx.empty(k * qd);
let gate = ctx.empty(k * im);
let opout_k = ctx.empty(k * h);
let mlpout_k = ctx.empty(k * h);
let normed_k = ctx.empty(k * h);
let moe_sel_k = ctx.empty(k * w.cfg.top_k_experts.max(1));
let moe_wsel_k = ctx.empty(k * w.cfg.top_k_experts.max(1));
let moe_gate_k = ctx.empty((k * w.cfg.moe_intermediate * w.cfg.top_k_experts).max(1));
let hnorm = ctx.empty(k * h);
let chosen_k = ctx.empty(k);
let chosen_staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("chosen_staging"),
size: (k * 4) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
// The SAME pipeline objects the M=1 plan replays (the bitwise contract — see `new`):
// on lcpp-family adapters the M=1 plan runs norm-SPLIT + lcpp GEMV, so the spec plan
// must pick the identical family or `spec == sequential` byte-identity breaks there.
let lcpp = lcpp_family(ctx);
let gemv_k = if lcpp {
self.spec_pls.gemv_lcpp.clone()
} else {
self.spec_pls.gemv_k.clone()
};
let gn32_k = self.spec_pls.gn32_k.clone();
let mlp_k = self.spec_pls.mlp_k.clone();
let convdot_k = self.spec_pls.convdot_k.clone();
let convmix_k = self.spec_pls.convmix_k.clone();
let qkrc_k = self.spec_pls.qkrc_k.clone();
let attn_k = self.spec_pls.attn_k.clone();
let attn_ksp = self.spec_pls.attn_ksp.clone();
let attn_kmerge = self.spec_pls.attn_kmerge.clone();
let attn_skv_z = self.spec_pls.attn_skv_z;
let rmsnorm = self.spec_pls.rmsnorm.clone();
let rmsnorm_add = self.spec_pls.rmsnorm_add.clone();
let moe_router_pl = self.spec_pls.moe_router.clone();
let moe_gate_pl = self.spec_pls.moe_gate.clone();
let moe_down_pl = self.spec_pls.moe_down.clone();
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, k as f32, eps, 0.0f32]));
let epsm_act = uni(
ctx,
bytemuck::cast_slice(&[eps, if w.cfg.act_gelu { 1.0f32 } else { 0.0 }, 0.0, 0.0]),
);
let gather_pl = pipeline(ctx, "gather_k", GATHER_K);
let gather_bg = make_bg(ctx, &gather_pl, &[&w.embed, &fed, &cur, &self.cnt0], &posu);
// Bind-group builder for the two-uniform kernels (dims at n-2, epsm at n-1).
let bg2 = |pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer],
dims: &wgpu::Buffer,
epsm: &wgpu::Buffer| {
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32,
resource: dims.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32 + 1,
resource: epsm.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &entries,
})
};
let epsm = uni(ctx, bytemuck::cast_slice(&[eps, 0.0f32, 0.0, 0.0]));
let udims =
|a: u32, b: u32, acc: u32, d: u32| uni(ctx, bytemuck::cast_slice(&u32x4(a, b, acc, d)));
let mut plan: Vec<Step> = Vec::new();
let mut conv_snapk: Vec<(usize, wgpu::Buffer)> = Vec::new();
let kk = k as u32;
// Split-KV partials for this plan's column capacity (shared across layers).
let attn_part = ctx.storage(&vec![0f32; k * nh * attn_skv_z as usize * (hd + 4)]);
let attn_merge_u = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, hd as u32, attn_skv_z, 0)),
);
let kt = kk.div_ceil(4); // column tiles (KC = 4) for the weight-shared GEMV kernels
// GEMV dispatch geometry per family: lcpp = 4 rows/WG, one real column per gy slice;
// tree/sg = the weight-shared KC4 column tiles.
let (gemv_rpw, gemv_gy) = if lcpp {
(4u32, kk)
} else {
(self.gemv_rows_per_wg, kt)
};
// lcpp-split scratch: the rmsnorm output the split qkv projection reads ([k, hidden]).
let normed_x_k = ctx.empty(k * h);
for li in 0..nl {
let layer = &w.layers[li];
let onorm = &layer.operator_norm;
match &layer.op {
// MLA is M=1-only for now (the `forward()` path); the spec/MTP plan isn't wired.
Op::Mla { .. } => unimplemented!("MLA has no spec/MTP plan yet"),
Op::Mamba { .. } => unimplemented!("Mamba is M=1 decode only; no spec plan"),
Op::Rwkv { .. } => unimplemented!("RWKV is M=1 decode only; no spec plan"),
Op::DeltaNet { .. } => {
unreachable!("spec plan refuses Qwen3.5 at entry")
}
Op::Conv {
in_proj,
conv_w,
out_proj,
} => {
let snap = ctx.empty(k * h * conv_l);
plan.push(Step::D {
tag: "convdot_k",
pl: convdot_k.clone(),
bg: make_bg(
ctx,
&convdot_k,
&[&in_proj.scales, &in_proj.quants, &cur, onorm, &bcx],
&udims(h as u32, conv_l as u32, 0, 0),
),
gx: h as u32,
gy: kk,
});
plan.push(Step::D {
tag: "convmix_k",
pl: convmix_k.clone(),
bg: make_bg(
ctx,
&convmix_k,
&[
&bcx,
conv_w,
&self.conv_state[li],
&convy,
&snap,
&cmeta_k,
&self.cnt0,
],
&udims(h as u32, conv_l as u32, kk, 0),
),
gx: (h as u32).div_ceil(64),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&out_proj.scales, &out_proj.quants, &convy, &cur],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
conv_snapk.push((li, snap));
}
Op::Attn {
qkv: qkv_w,
o,
q_norm: _,
k_norm: _,
window,
local_rope,
attn_gate: _,
qkv_bias,
} => {
let qk_norm_cat = self.qk_norm_cat[li]
.as_ref()
.expect("attention layer implies a qk-norm concat");
let qkv_rows = if w.cfg.layer_k_eq_v[li] {
qd + kd
} else {
qd + 2 * kd
};
if lcpp {
// Norm SPLIT + lcpp GEMV — the M=1 plan's family on this adapter.
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(ctx, &rmsnorm, &[&cur, onorm, &normed_x_k], &mfk(h as f32)),
gx: kk,
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&qkv_w.scales, &qkv_w.quants, &normed_x_k, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: (qkv_rows as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&qkv_w.scales, &qkv_w.quants, &cur, onorm, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
&epsm,
),
gx: (qkv_rows as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
let (csb, snb) = if *local_rope {
(&csk_l, &snk_l)
} else {
(&csk, &snk)
};
let g4 = w.cfg.arch == crate::weights::Arch::Gemma4;
let af = udims(u32::from(g4), u32::from(w.cfg.layer_k_eq_v[li]), 0, 0);
// Qwen2 q/k/v bias over the k-column fused qkv, before qkrc.
if let Some(bias) = qkv_bias {
let total = (qkv_rows * k) as u32;
plan.push(Step::D {
tag: "qkv_bias",
pl: self.qkv_bias_pl.clone(),
bg: make_bg(
ctx,
&self.qkv_bias_pl,
&[&qkv, bias],
&udims(qkv_rows as u32, total, 0, 0),
),
gx: total.div_ceil(256),
gy: 1,
});
}
plan.push(Step::D {
tag: "qkrc_k",
pl: qkrc_k.clone(),
bg: bg2(
&qkrc_k,
&[
&qkv,
qk_norm_cat,
csb,
snb,
&qn,
&self.k_cache[li],
&self.v_cache[li],
&self.btab,
&cmeta_k,
&self.cnt0,
],
&attn_tk,
&af,
),
gx: (nh + 2 * nkv) as u32,
gy: kk,
});
if let (Some(ksp), Some(kmg)) = (&attn_ksp, &attn_kmerge) {
// Same split pair as the M=1 plan at the same z — batch==solo stays
// bitwise by construction (identical per-(col,head,split) walks and
// an identical fixed-order merge).
plan.push(Step::D {
tag: "attn_k",
pl: ksp.clone(),
bg: bg2(
ksp,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_part,
&self.btab,
&cmeta_k,
&self.cnt0,
],
&attn_tk,
&udims(*window, u32::from(g4), 0, attn_skv_z),
),
gx: nh as u32,
gy: kk * attn_skv_z,
});
plan.push(Step::D {
tag: "attn_merge",
pl: kmg.clone(),
bg: make_bg(ctx, kmg, &[&attn_part, &attn_o], &attn_merge_u),
gx: nh as u32,
gy: kk,
});
} else {
plan.push(Step::D {
tag: "attn_k",
pl: attn_k.clone(),
bg: bg2(
&attn_k,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_o,
&self.btab,
&cmeta_k,
&self.cnt0,
],
&attn_tk,
&udims(*window, u32::from(g4), 0, 0),
),
gx: nh as u32,
gy: kk,
});
}
let acc = u32::from(layer.post_op_norm.is_none());
let o_target = if layer.post_op_norm.is_some() {
&opout_k
} else {
&cur
};
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&o.scales, &o.quants, &attn_o, o_target],
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some(pn) = &layer.post_op_norm {
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(ctx, &rmsnorm_add, &[&opout_k, pn, &cur], &mfk(h as f32)),
gx: kk,
gy: 1,
});
}
}
}
// ffn-norm + gate/up: SPLIT rmsnorm + two-stream lcpp kernel on the lcpp family
// (mirrors the qkv split; `OSFKB_MLP_LCPP=0` pins the fused kernel), fused mlp_k
// otherwise.
if lcpp && mlp_lcpp_enabled() {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &layer.ffn_norm, &normed_x_k],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
// Q1 swaps the fused-MLP pipeline for its sign-decoding twin (same 8 bindings).
let mlp_split_pl = self
.spec_pls
.mlp_q1
.as_ref()
.unwrap_or(&self.spec_pls.mlp_lcpp)
.clone();
plan.push(Step::D {
tag: "mlp_k",
pl: mlp_split_pl.clone(),
bg: bg2(
&mlp_split_pl,
&[
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
&normed_x_k,
&gate,
],
&udims(im as u32, h as u32, 0, k as u32),
&epsm_act,
),
gx: (im as u32).div_ceil(4),
gy: kk,
});
} else {
plan.push(Step::D {
tag: "mlp_k",
pl: mlp_k.clone(),
bg: bg2(
&mlp_k,
&[
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
&cur,
&layer.ffn_norm,
&gate,
],
&udims(im as u32, h as u32, 0, k as u32),
&epsm_act,
),
gx: (im as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
let acc = u32::from(layer.post_ffn_norm.is_none());
let d_target = if layer.post_ffn_norm.is_some() {
&mlpout_k
} else {
&cur
};
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&layer.w2.scales, &layer.w2.quants, &gate, d_target],
&udims(h as u32, im as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
// MoE (Gemma-4): per-column routing + weighted expert accumulate (see the M=1 plan).
if let Some(moe) = &layer.moe {
let (e_n, kx, mi) = (
w.cfg.num_experts as u32,
w.cfg.top_k_experts as u32,
w.cfg.moe_intermediate as u32,
);
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &layer.ffn_norm, &normed_k],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
plan.push(Step::D {
tag: "moe_router_pl",
pl: moe_router_pl.clone(),
bg: make_bg(
ctx,
&moe_router_pl,
&[
&moe.router,
&moe.per_expert_scale,
&normed_k,
&moe_sel_k,
&moe_wsel_k,
],
&udims(e_n, h as u32, kx, 0),
),
gx: kk,
gy: 1,
});
plan.push(Step::D {
tag: "moe_gate_pl",
pl: moe_gate_pl.clone(),
bg: bg2(
&moe_gate_pl,
&[
&moe.w1.scales,
&moe.w1.quants,
&moe.w3.scales,
&moe.w3.quants,
&normed_k,
&moe_sel_k,
&moe_gate_k,
],
&udims(mi, h as u32, 0, kx),
&epsm_act,
),
gx: mi,
gy: kk * kx,
});
plan.push(Step::D {
tag: "moe_down_pl",
pl: moe_down_pl.clone(),
bg: make_bg(
ctx,
&moe_down_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&moe_gate_k,
&moe_sel_k,
&moe_wsel_k,
d_target,
],
&udims(h as u32, mi, 0, kx),
),
gx: h as u32,
gy: kk,
});
}
if let Some(pn) = &layer.post_ffn_norm {
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(ctx, &rmsnorm_add, &[&mlpout_k, pn, &cur], &mfk(h as f32)),
gx: kk,
gy: 1,
});
}
}
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &w.embedding_norm, &hnorm],
&uni(
ctx,
bytemuck::cast_slice(&[h as f32, k as f32, eps, 0.0f32]),
),
),
gx: kk,
gy: 1,
});
let picks = (0..k)
.map(|p| SparsePick::new_col(ctx, dp, w, &hnorm, p as u32, &maskk, &chosen_k, &posu))
.collect();
SpecPlan {
k,
plan,
fed,
drafts,
maskk,
csk,
snk,
csk_l,
snk_l,
attn_tk,
cmeta_k,
hnorm,
chosen_k,
conv_snapk,
picks,
gather_pl,
gather_bg,
spec_cur: cur,
chosen_staging,
}
}
/// Pipeline-stage step (see [`crate::shard`]): run this shard's layers at `pos`. Stage 0 feeds
/// `token`; later stages feed `hidden_in` (the previous stage's residual stream, written
/// straight into `s.cur`). Returns the residual stream — or, on the LAST stage, the greedy
/// token (4-byte hop instead of a 600 KB logits ship).
pub fn stage_forward(
&self,
ctx: &GpuCtx,
pos: usize,
token: u32,
hidden_in: Option<&[f32]>,
) -> Result<StageOut> {
let c = &self.w.cfg;
assert_eq!(
hidden_in.is_none(),
c.stage_first,
"stage 0 takes a token; later stages take hidden"
);
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
if let Some(hidden) = hidden_in {
anyhow::ensure!(hidden.len() == c.hidden, "hidden hop size");
ctx.queue
.write_buffer(&self.s.cur, 0, bytemuck::cast_slice(hidden));
}
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if self.w.cfg.edge.ple_dim > 0 {
// PLE gather input for this step (queue-ordered before the submit below).
ctx.queue
.write_buffer(&self.s.ple_tok, 0, bytemuck::cast_slice(&[token]));
}
self.record(&mut enc, token, pos);
ctx.queue.submit([enc.finish()]);
if c.stage_last {
Ok(StageOut::Token(self.argmax_logits(ctx)?))
} else {
Ok(StageOut::Hidden(ctx.read(&self.s.cur, c.hidden)?))
}
}
/// Read the first `ncols` columns of a batch plan's residual stream (pre-final-norm) —
/// the MTP-training data dump path.
pub fn read_cur_for_tests(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
ncols: usize,
) -> Result<Vec<f32>> {
ctx.read(&bp.cur, ncols * self.w.cfg.hidden)
}
/// Dev bisection: read a batch plan's qn / attn_o scratch (first `n` floats).
pub fn read_batch_qn_for_tests(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
n: usize,
) -> Result<Vec<f32>> {
ctx.read(&bp.qn, n)
}
pub fn read_batch_attn_o_for_tests(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
n: usize,
) -> Result<Vec<f32>> {
ctx.read(&bp.attn_o, n)
}
/// (prep, encode, poll+read) seconds of the last batch step — the poll+read component is
/// measured from the post-submit instant to THIS call (call right after the step returns).
pub fn take_step_breakdown(&self) -> (f64, f64, f64) {
let (p, e, r0) = *self.step_breakdown.borrow();
(p, e, r0.elapsed().as_secs_f64())
}
/// Greedy argmax over `s.logits` on-GPU (block scan + merge — the same kernel pair the solo
/// decode loop dispatches, so a sharded pick equals the full-model pick).
pub fn argmax_logits(&self, ctx: &GpuCtx) -> Result<u32> {
let pls = self.argmax_pls.get_or_init(|| {
let block = pipeline(ctx, "argmax_block", ARGMAX_BLOCK);
let merge = pipeline(ctx, "argmax_merge_col", ARGMAX_MERGE_COL);
let pv = ctx.empty(256);
let pi = ctx.empty(256);
let chosen = ctx.empty(1);
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, self.w.cfg.vocab as u32, 0)),
);
let colu = uni(ctx, bytemuck::cast_slice(&u32x4(0, 0, 0, 0)));
let block_bg = make_bg(ctx, &block, &[&self.s.logits, &pv, &pi], &posu);
let merge_bg = make_bg(ctx, &merge, &[&pv, &pi, &chosen, &self.cnt0], &colu);
ArgmaxPls {
block,
merge,
block_bg,
merge_bg,
chosen,
_keep: vec![pv, pi],
}
});
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
Self::record_argmax_passes(pls, &mut enc);
ctx.queue.submit([enc.finish()]);
Ok(ctx.read_u32(&pls.chosen, 1)?[0])
}
fn record_argmax_passes(pls: &ArgmaxPls, enc: &mut wgpu::CommandEncoder) {
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(&pls.block);
p.set_bind_group(0, &pls.block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&pls.merge);
p.set_bind_group(0, &pls.merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
/// Record the greedy text-argmax into the CALLER's encoder and return the `chosen` token
/// buffer (1 × u32) — the Moshi full-GPU step copies it into the depformer's token chain.
pub(crate) fn record_argmax(
&self,
ctx: &GpuCtx,
enc: &mut wgpu::CommandEncoder,
) -> &wgpu::Buffer {
let pls = self.argmax_pls.get_or_init(|| {
let block = pipeline(ctx, "argmax_block", ARGMAX_BLOCK);
let merge = pipeline(ctx, "argmax_merge_col", ARGMAX_MERGE_COL);
let pv = ctx.empty(256);
let pi = ctx.empty(256);
let chosen = ctx.empty(1);
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, self.w.cfg.vocab as u32, 0)),
);
let colu = uni(ctx, bytemuck::cast_slice(&u32x4(0, 0, 0, 0)));
let block_bg = make_bg(ctx, &block, &[&self.s.logits, &pv, &pi], &posu);
let merge_bg = make_bg(ctx, &merge, &[&pv, &pi, &chosen, &self.cnt0], &colu);
ArgmaxPls {
block,
merge,
block_bg,
merge_bg,
chosen,
_keep: vec![pv, pi],
}
});
Self::record_argmax_passes(pls, enc);
&pls.chosen
}
/// Build the continuous-batching plan: the k-column layer stack (THE SAME compiled pipelines
/// as the M=1 and speculative plans — the bitwise contract) where each column is an independent
/// (sequence, position) pair addressed through its own block-table row, plus a per-column tiled
/// lm_head + argmax. Attention-only architectures (Qwen3, Gemma-3/4): LFM2's conv layers need
/// per-slot ring state, which [`crate::CONV_MIX_K`] supports but this builder does not wire yet.
/// f16 weight storage is wired for the M=1 decode plan ONLY. Every other plan still builds
/// Q4 bind groups (five buffers) and Q4 kernels, which would read f16 bytes as nibbles and
/// return confident nonsense — so they refuse loudly instead. Wiring them is the next step,
/// exactly as the Q1 family sequenced its own arms.
fn assert_not_f16(&self, what: &str) {
assert!(
!self.w.cfg.any_nonq4_body(),
"f16 weight storage is not wired for {what} yet (M=1 decode only) — refusing rather \
than decoding f16 bytes with the Q4 kernels"
);
}
/// The grammar-SPARSE head reads `w.embed_q4` with a Q4-assuming kernel whose sparse dots must
/// be bitwise-equal to the full head's logits (the constrained-pick soundness contract). An
/// f16 head has no sparse twin, so the constrained paths refuse it rather than read f16 bytes
/// as nibbles. (The plain full head IS wired for f16 — see `Q4LmHead::pipeline_for`.)
fn assert_head_not_f16(&self, what: &str) {
assert!(
self.w.cfg.head_kind() == crate::weights::WKind::Q4_0,
"an f16 LM head is not wired for {what} yet (the grammar-sparse head has no f16 twin) \
— refusing rather than decoding f16 bytes with the Q4 sparse kernel"
);
}
pub fn make_batch_plan(&self, ctx: &GpuCtx, k: usize) -> BatchPlan {
self.build_batch_plan(ctx, k, false, false, false)
}
/// SPEC-VERIFY plan: the wide (KC=16) GEMV/MLP families — weight reads amortized over the
/// span columns, which is what makes a (1+k)-column verify cheaper than k+1 decode steps —
/// but the DECODE attention kernels, NOT wide's prefill-local vec4/RB attention. The verify
/// soundness contract is that every column is bitwise-equal to the narrow decode plan (so
/// speculation can never change an emitted token), and the vec4 attention's different
/// summation order breaks exactly that (measured: spec-vs-plain divergence after ~13 tokens
/// on a thin-margin model). Requires subgroups, like the wide plan.
pub fn make_batch_plan_verify(&self, ctx: &GpuCtx, k: usize) -> BatchPlan {
assert!(
self.wide_pls.is_some(),
"the verify plan's KC16 family needs subgroup support on this adapter"
);
self.build_batch_plan(ctx, k, true, false, true)
}
/// WIDE PREFILL plan: same structure as the decode plan but the GEMV family runs the KC=16
/// kernels (16-position weight sharing). Requires subgroup support. Columns are
/// bitwise-equal to the decode plan's, so a prompt prefilled wide continues into decode with
/// identical KV/state.
pub fn make_batch_plan_wide(&self, ctx: &GpuCtx, k: usize) -> BatchPlan {
assert!(
self.wide_pls.is_some(),
"wide prefill needs subgroup support on this adapter"
);
self.build_batch_plan(ctx, k, true, false, false)
}
/// Copy micro-batch column `col`'s residual stream (`bp.cur`, PRE-final-norm) into `dst`
/// — seeds the MTP draft chain from the last accepted position without a host round-trip.
pub fn copy_cur_col_to(&self, ctx: &GpuCtx, bp: &BatchPlan, col: usize, dst: &wgpu::Buffer) {
let h = self.w.cfg.hidden as u64;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
enc.copy_buffer_to_buffer(&bp.cur, col as u64 * h * 4, dst, 0, h * 4);
ctx.queue.submit([enc.finish()]);
}
/// Roll the DN recurrent state (and conv ring) of `slot` back to its value AFTER column
/// `col` of the LAST micro-batch run on `bp` — the speculative-verify rejection path. KV is
/// rolled back positionally for free; only the DeltaNet recurrence needs explicit snapshots.
pub fn dn_restore(&self, ctx: &GpuCtx, bp: &BatchPlan, slot: u32, col: usize) {
let c = &self.w.cfg;
let slab = (c.dn_nv * c.dn_dk * c.dn_dv * 4) as u64;
let cd = 2 * c.dn_nk * c.dn_dk + c.dn_nv * c.dn_dv;
let rslab = (cd * c.dn_kernel * 4) as u64;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (li, snaps) in bp.dn_snaps.iter().enumerate() {
let Some((snap, csnap)) = snaps else { continue };
let st = self.dn_state[li].as_ref().expect("dn state");
let ring = self.dn_ring[li].as_ref().expect("dn ring");
enc.copy_buffer_to_buffer(snap, col as u64 * slab, st, u64::from(slot) * slab, slab);
enc.copy_buffer_to_buffer(
csnap,
col as u64 * rslab,
ring,
u64::from(slot) * rslab,
rslab,
);
}
ctx.queue.submit([enc.finish()]);
}
/// TEST HOOK: read `slot`'s DN recurrent state slab of layer `li`.
pub fn read_dn_state_for_tests(
&self,
ctx: &GpuCtx,
li: usize,
slot: usize,
) -> Result<Vec<f32>> {
let c = &self.w.cfg;
let slab = c.dn_nv * c.dn_dk * c.dn_dv;
let st = self.dn_state[li].as_ref().expect("dn state");
let all = ctx.read(st, (slot + 1) * slab)?;
Ok(all[slot * slab..].to_vec())
}
/// Speculative-verify variant: DN conv/state kernels also SNAPSHOT per column, enabling
/// [`Self::dn_restore`] rollback on partial draft acceptance (MTP self-speculation).
pub fn make_batch_plan_spec(&self, ctx: &GpuCtx, k: usize) -> BatchPlan {
self.build_batch_plan(ctx, k, false, true, false)
}
/// llama.cpp-shaped GEMV eligibility: Vulkan + subgroups (the shape LOSES on Metal's tile
/// caches at these widths — backend-aware like the earlier sg-vs-tree finding). Kill switch:
/// `OSFKB_GEMV_LCPP=0`.
fn lcpp_gemv(&self, ctx: &GpuCtx) -> bool {
lcpp_family(ctx)
}
/// Get-or-build the ROWPACK4 twin of one Q1 weight (see the `q1_rp4` field). `key` is the
/// deterministic plan-build visit order; the repack dispatch runs ONCE per key and the
/// buffers live for the engine's lifetime (shared by every band plan).
#[allow(clippy::too_many_arguments)]
fn rp4_twin(
&self,
ctx: &GpuCtx,
repack: &wgpu::ComputePipeline,
key: u32,
scales: &wgpu::Buffer,
quants: &wgpu::Buffer,
m: usize,
n: usize,
) -> (wgpu::Buffer, wgpu::Buffer) {
if let Some(t) = self.q1_rp4.borrow().get(&key) {
return (t.0.clone(), t.1.clone());
}
let g4 = m.div_ceil(4);
let nblk = n / 32;
let nsc = n / 128;
let scales4 = ctx.empty(g4 * nsc * 4);
let bits4 = ctx.empty(g4 * nblk * 4);
let d = uni(
ctx,
bytemuck::cast_slice(&u32x4(m as u32, nblk as u32, nsc as u32, 0)),
);
let bg = make_bg(ctx, repack, &[scales, quants, &scales4, &bits4], &d);
let mut enc = ctx.device.create_command_encoder(&Default::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(repack);
p.set_bind_group(0, &bg, &[]);
p.dispatch_workgroups(((g4 * nblk.max(nsc)) as u32).div_ceil(256), 1, 1);
}
ctx.queue.submit([enc.finish()]);
self.q1_rp4
.borrow_mut()
.insert(key, (scales4.clone(), bits4.clone()));
(scales4, bits4)
}
fn build_batch_plan(
&self,
ctx: &GpuCtx,
k: usize,
wide: bool,
dn_snap: bool,
// Spec-verify: keep every ATTENTION choice identical to the narrow decode plan (bitwise
// contract) while still taking the wide GEMV families. See `make_batch_plan_verify`.
verify: bool,
) -> BatchPlan {
// Gemma-4 edge: per-layer geometry rides the per-layer SHADOWED dims below; the KC16 /
// wide-GEMM / split-K speed families assume uniform widths across layers, so edge models
// force the plain per-column family (a later, measured lever — not a correctness need).
let edge = self.w.cfg.is_edge();
assert!(
(1..=MAX_SLOTS).contains(&k),
"batch k must be in 1..=MAX_SLOTS"
);
assert!(
!crate::weights::arch_is_layernorm(self.w.cfg.arch),
"LayerNorm arches (OLMo) are wired for M=1 decode only — no batched plan yet"
);
assert!(
self.w.cfg.arch != crate::weights::Arch::Mamba,
"Mamba is wired for M=1 decode only — no batched plan yet"
);
// The DECODE plan runs the wide (KC=16, subgroup) GEMV family whenever the batch fills at
// least one 16-column tile (k >= 16): KC4 tiling reads every weight-shared matrix
// ceil(k/4)× per step, KC16 reads it ceil(k/16)× — 4× fewer weight re-reads. Measured on
// Apple M4 Max / claim-extractor-4B (interleaved A/B, thermal-robust): +15% tok/s at k=16,
// +25% at k=32, +80% at k=64 (KC16 also removes KC4's large-batch throughput regression —
// its per-step weight traffic grew ceil(k/4)×). BELOW one full tile KC16 pads to 16 columns
// and LOSES (measured 0.52× at k=8), so the crossover is exactly the tile width. Requires
// subgroups (`wide_pls`; falls back to KC4 otherwise); per-column results stay bitwise-equal
// to the KC4 subgroup family (prefill_kc16 gate). `OSFKB_DECODE_KC16=0/1` forces off/on.
let kc16 = !edge
&& (wide
|| (self.wide_pls.is_some()
&& match std::env::var("OSFKB_DECODE_KC16").ok().as_deref() {
Some("0") => false,
Some("1") => true,
_ => k >= 16,
}));
// Q1 (Bonsai/BitNet, `OSFKB_Q1_WEIGHTS`): sign-block weights exist ONLY in the lcpp
// GEMV/MLP family — every other family (tiled GEMM, the KC16 wide shapes, the gn32
// fused-norm GEMV, the tree/nbar families) decodes the sign bits as Q4 nibbles, which
// is not noise but constant-token garbage (measured on Bonsai-27B: every prompt argmaxed
// to the same id). Batch plans therefore (a) force the lcpp wide arm on, (b) keep the
// Q4-only shapes off, and (c) swap the GEMV/MLP pipelines for their Q1 twins below —
// both twins are bind-compatible by design ("the plan swaps ONLY the pipeline").
let q1 = self.spec_pls.gemv_q1.is_some();
if std::env::var("OSFKB_PLAN_TRACE").ok().as_deref() == Some("1") {
eprintln!(
"[plan] k={k} wide={wide} kc16={kc16} q1={q1} lcpp_gemv={} wide_lcpp={} wide_gemm={}",
self.lcpp_gemv(ctx),
kc16 && self.lcpp_gemv(ctx)
&& (q1 || std::env::var("OSFKB_WIDE_LCPP").ok().as_deref() != Some("0")),
kc16 && !q1 && std::env::var("OSFKB_WIDE_GEMM").ok().as_deref() != Some("0"),
);
}
assert!(
!q1 || self.lcpp_gemv(ctx),
"OSFKB_Q1_WEIGHTS batch plans need the lcpp GEMV family (validated 32-wide subgroups)"
);
// WIDE plans on guaranteed-32-subgroup adapters take the LCPP family (see the family
// selection below); their remaining tree-family sites (attention qkv gn32) then need
// KC=4 tile geometry, not KC=16.
let wide_lcpp = kc16
&& self.lcpp_gemv(ctx)
&& (q1 || std::env::var("OSFKB_WIDE_LCPP").ok().as_deref() != Some("0"));
// WIDE plans route the projections through the tiled Q4 GEMM (portable baseline WGSL,
// every adapter): BN=32 columns share one weight stream — weights read ⌈k/32⌉× per
// step instead of k× (the flat-in-width profile proved the per-column shapes scale
// linearly). `OSFKB_WIDE_GEMM=0` reverts to the previous per-column family.
// The tiled wide-GEMM family has Q4_0/F16/Q8_0N members (gemm_q4 + the generated
// twins); kinds outside that set (K-quant/Q5_0N natives) fall back to the lcpp family's
// per-column reads — correct for every kind, slower at prefill widths.
let gemm_kinds_ok = (0..self.w.layers.len()).all(|li| {
let sk = self.w.cfg.kinds_at(li);
[sk.qkv, sk.o, sk.gate_up, sk.down].iter().all(|x| {
matches!(
x,
crate::weights::WKind::Q4_0
| crate::weights::WKind::F16
| crate::weights::WKind::Q8_0N
)
})
});
let wide_gemm =
kc16 && gemm_kinds_ok && std::env::var("OSFKB_WIDE_GEMM").ok().as_deref() != Some("0");
let kc_tile: u32 = if kc16 && !wide_lcpp && !wide_gemm {
16
} else {
4
};
// Conv layers (LFM2) get per-slot ring state: [1 + MAX_SLOTS, hidden, conv_l] per conv
// layer, column-addressed through cmeta's slot index (same trick as the KV block tables).
assert!(
self.tq.is_none(),
"batched serving is incompatible with OSFKB_TQ_KV"
);
let w = &self.w;
let c = &w.cfg;
let (h, im, hd) = (c.hidden, c.intermediate, c.head_dim);
let (nh, nkv, nl) = (c.n_heads, c.n_kv_heads, c.n_layers);
let (qd, _kd) = (nh * hd, nkv * hd);
// Edge maxima for buffer sizing; the layer loop below SHADOWS (hd, qd, kd, im) per layer.
let hd_x = c.max_head_dim();
let im_x = c.max_intermediate();
let (qd_x, kd_x) = (nh * hd_x, nkv * hd_x);
let (hd_glob, hd_slid) = (c.head_dim_full(), c.head_dim_sliding());
let (eps, vocab) = (c.eps, c.vocab);
let kk = k as u32;
let kt = kk.div_ceil(kc_tile); // column tiles for the weight-shared GEMV kernels
// Chain-capable buffers: step s of a chained submit reads slice [s·k ..] via the on-GPU
// step counter; the unchained path (cnt = 0) uses slice 0 — same buffers, same kernels.
let fed = ctx.empty((CHAIN_MAX + 1) * k);
// Rope tables per LAYER TYPE: the global pair serves the full-attention layers (the
// wide heads on an edge model), the local pair the sliding ones.
let csk = ctx.empty(CHAIN_MAX * k * hd_glob);
let snk = ctx.empty(CHAIN_MAX * k * hd_glob);
let csk_l = ctx.empty(CHAIN_MAX * k * hd_slid);
let snk_l = ctx.empty(CHAIN_MAX * k * hd_slid);
let attn_tk = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, 0)),
);
// The wide-head layers' twin (Gemma-4 edge full-attention: hd 512-class). Static per
// plan, exactly like attn_tk — positions ride cmeta.
let attn_tk2 = (hd_glob != hd).then(|| {
uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd_glob as u32, 0)),
)
});
// ATTN_TILED reads kp.w = kcols (the plan width — parked columns sit on the trash
// btab row exactly as under ATTN_K). The shared `attn_tk` carries 0 there, which
// masked EVERY column (cvalid = col < 0): the tiled kernel wrote nothing and wide
// full-attention µbatches consumed stale attn_o. Dedicated uniform, correct width.
let attn_tiled_tk = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, k as u32)),
);
let cmeta_k = ctx.empty(CHAIN_MAX * k * 2);
let cnt = ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("batch_cnt"),
contents: bytemuck::cast_slice(&[0u32]),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
});
let cur = ctx.empty(k * h);
let qkv = ctx.empty(k * (qd_x + 2 * kd_x));
let qn = ctx.empty(k * qd_x);
let attn_o = ctx.empty(k * qd_x);
let gate = ctx.empty(k * im_x);
let opout_k = ctx.empty(k * h);
let mlpout_k = ctx.empty(k * h);
let normed_k = ctx.empty(k * h);
let moe_sel_k = ctx.empty(k * c.top_k_experts.max(1));
let moe_sg_k = ctx.empty(k);
let moe_logits_k = ctx.empty(k * c.num_experts.max(1));
let normed_x_k = ctx.empty(k * h);
// Tiled-GEMM MLP (`OSFKB_WIDE_MLP_GEMM=0` reverts): the wide plan routes the PROJECTIONS
// through the tiled Q4 GEMM but left gate+up on the per-column fused `mlp_k` — 62% of a
// short-position prefill step by the profile, at ~1 TF/s vs the tiled GEMM's ~3.3. Run
// gate and up as two tiled GEMMs into these scratch buffers, then one act·mul fold.
let wide_mlp_gemm =
wide_gemm && std::env::var("OSFKB_WIDE_MLP_GEMM").ok().as_deref() != Some("0");
let (gate_raw, up_raw, mlp_actmul_pl) = if wide_mlp_gemm {
(
ctx.empty(k * im),
ctx.empty(k * im),
Some(pipeline(ctx, "mlp_act_mul", MLP_ACT_MUL)),
)
} else {
(ctx.empty(1), ctx.empty(1), None)
};
let moe_wsel_k = ctx.empty(k * c.top_k_experts.max(1));
let moe_gate_k = ctx.empty((k * c.moe_intermediate * c.top_k_experts).max(1));
// GROUPED MoE (wide prefill plans only): expert-tiled dispatch (see MOE_SEG_BUILD).
// Guards: ≤256 experts (the builder's shared histogram), ≤8 top-k (the unrolled pair
// tile), ≤1536 pairs (the shared scatter). `OSFKB_MOE_GROUPED=0` kills.
let n_pairs = k * c.top_k_experts.max(1);
let moe_grouped = kc16
&& c.arch == crate::weights::Arch::Qwen35
&& c.num_experts <= 256
&& c.top_k_experts <= 8
&& n_pairs <= 1536
&& std::env::var("OSFKB_MOE_GROUPED").ok().as_deref() != Some("0");
let moe_grp = moe_grouped.then(|| {
let texp = ctx.empty(n_pairs);
let tpair = ctx.empty(n_pairs * 8);
let dpartial = ctx.empty(n_pairs * h);
let args = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("moe_grouped_args"),
size: 32,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::INDIRECT
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
(texp, tpair, dpartial, args)
});
let hnorm = ctx.empty(k * h);
let conv_l = c.conv_l;
let conv_slots: Vec<Option<wgpu::Buffer>> = self
.w
.layers
.iter()
.map(|l| {
matches!(l.op, Op::Conv { .. })
.then(|| ctx.storage(&vec![0f32; (2 + MAX_SLOTS) * h * conv_l]))
})
.collect();
let conv_snap_scratch = ctx.empty(k * h * conv_l);
let bcx_k = ctx.empty(k * 3 * h);
let c_dn = &self.w.cfg;
let dn_cd = 2 * c_dn.dn_nk * c_dn.dn_dk + c_dn.dn_nv * c_dn.dn_dv;
let dn_inr = dn_cd + c_dn.dn_nv * c_dn.dn_dv + 2 * c_dn.dn_nv;
let dn_mixed_k = ctx.empty((k * dn_inr).max(1));
let dn_conved_k = ctx.empty((k * dn_cd).max(1));
let dn_core_k = ctx.empty((k * c_dn.dn_nv * c_dn.dn_dv).max(1));
// dv-split chunk kernel: per-(col, head, half) Σo² partials for the deferred norm.
let dn_opart_k = ctx.empty((k * c_dn.dn_nv * 2).max(1));
let agate_k = ctx.empty(if c_dn.arch == crate::weights::Arch::Qwen35 {
k * qd
} else {
1
});
let logits_k = ctx.empty(k * vocab);
let pv_k = ctx.empty(k * 256);
let pi_k = ctx.empty(k * 256);
let chosen_k = ctx.empty(CHAIN_MAX * k);
let chosen_staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("batch_chosen_staging"),
size: (CHAIN_MAX * k * 4) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
// WIDE plans on guaranteed-32-subgroup adapters take the LCPP family too (`wide_lcpp`
// decided above, before the tile geometry): the KC16 weight-shared shapes ran at
// single-digit GB/s effective on V100 (wide-plan profile 2026-07-05: gn32_k
// 2.9 ms/dispatch = 37% of a 64-col prefill step), while lcpp's per-column reads
// stream at full bandwidth — weight re-reads are cheaper than the weight-shared
// kernels' latency, the same measured lesson as the decode widths. `OSFKB_WIDE_LCPP=0`
// reverts (and Metal keeps the KC16 family — no subgroups32).
let (gemv_k, gn32_k, mlp_k, gemv_rpw, gemv_gy) = if wide_gemm {
(
self.spec_pls.gemm.clone(),
self.spec_pls.gn32_k.clone(),
self.spec_pls.mlp_k.clone(),
32,
kk.div_ceil(16),
)
} else if wide_lcpp {
(
self.spec_pls.gemv_lcpp.clone(),
self.spec_pls.gn32_k.clone(),
self.spec_pls.mlp_k.clone(),
4,
kk,
)
} else if kc16 {
let wp = self.wide_pls.as_ref().expect("kc16 requires wide_pls");
(
wp.gemv.clone(),
wp.gn32.clone(),
wp.mlp.clone(),
self.gemv_rows_per_wg,
kt,
)
} else if self.lcpp_gemv(ctx) {
// llama.cpp-shaped GEMV (WG=32, 4 rows, subgroupAdd): lab-measured 3-7.4× on
// V100/Vulkan at every batch width. Different FP order than the tree family —
// batch==solo stays exact (both sides run this kernel); Metal keeps the old family
// (the OSFQL constrained product path's bitwise anchors live there).
(
self.spec_pls.gemv_lcpp.clone(),
self.spec_pls.gn32_k.clone(),
self.spec_pls.mlp_k.clone(),
4,
kk,
)
} else if k <= 8 {
// Small-batch: barrier-free plain GEMV (bitwise-equal; lab-measured 1.9–3.6×). The
// fused-norm/mlp kernels keep their staged shape (their norm rides the staging).
(
self.spec_pls.gemv_nbar.clone(),
self.spec_pls.gn32_k.clone(),
self.spec_pls.mlp_k.clone(),
16,
kt,
)
} else {
(
self.spec_pls.gemv_k.clone(),
self.spec_pls.gn32_k.clone(),
self.spec_pls.mlp_k.clone(),
self.gemv_rows_per_wg,
kt,
)
};
// SPLIT-K wide GEMM family (`OSFKB_WIDE_SPLITK=<S>`, S >= 2; wide_gemm Q4 plans only):
// the serial-k gemm_q4 runs 48-144 workgroups of 32 threads on an 80-SM V100 — measured
// ~1.2-3.4 TF/s per projection. The split-K twin partitions the block walk over S
// k-slices (gy carries ncol_tiles x S) into an f32 partials scratch, and a fixed-order
// reduce step folds them + applies the acc flag: V100-raced 2-3.7x per projection,
// per-layer projection sum 1855 -> 951 us (tests/gemm_volta_lab.rs, q4split section).
// NOT bitwise vs the serial walk (summation grouping) -> strictly opt-in; the
// wide-bitwise gates run with it OFF.
let splitk_s: Option<usize> = (wide_gemm && !q1 && !self.w.cfg.any_nonq4_body())
.then(|| {
std::env::var("OSFKB_WIDE_SPLITK")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|&s| s >= 2)
})
.flatten();
let (gemv_k, gemv_gy, splitk) = match splitk_s {
Some(s) => {
let c = &self.w.cfg;
let qkv_rows = (c.n_heads + 2 * c.n_kv_heads) * c.head_dim;
let max_m = [
c.hidden,
c.intermediate,
qkv_rows,
2 * c.moe_intermediate,
c.dn_nv * c.dn_dv + 2 * c.dn_nk * c.dn_dk + c.dn_nv * c.dn_dv,
]
.into_iter()
.max()
.expect("non-empty");
(
pipeline(ctx, "gemm_q4_splitk", &crate::gemm_q4_splitk_src(s)),
kk.div_ceil(16) * s as u32,
Some((
pipeline(
ctx,
"gemm_q4_splitk_red",
&crate::gemm_q4_splitk_reduce_src(s),
),
ctx.empty(s * k * max_m),
)),
)
}
None => (gemv_k, gemv_gy, None),
};
// Q1 swaps ONLY the pipeline: the GEMV twin is bind-compatible with gemv_q4_k_lcpp,
// and the GEMM twin with gemm_q4 (same dims uniform, same tile geometry) — the wide
// arm keeps its weight-sharing (ceil(k/16) weight streams/step, NOT k).
let gemv_k = if q1 {
if wide_gemm {
self.spec_pls
.gemm_q1
.clone()
.expect("gemm_q1 built whenever OSFKB_Q1_WEIGHTS selects the Q1 plan")
} else {
self.spec_pls
.gemv_q1
.clone()
.expect("q1 asserted lcpp above")
}
} else {
gemv_k
};
// ── Non-Q4_0 containers in BATCH plans: per-site kind dispatch, lcpp family only ──
// Mirrors the M=1 plan: each site picks its pipeline + bind-group shape from the
// layer's per-site kind. Geometry is untouched — every non-Q4_0 kernel is lcpp-shaped
// (4 rows/WG, col = wid.y), the same shape the lcpp arms already dispatch.
let b_any_nonq4 = self.w.cfg.any_nonq4_body();
assert!(
!b_any_nonq4 || wide_gemm || wide_lcpp || (self.lcpp_gemv(ctx) && !kc16),
"non-Q4_0 batch plans need the lcpp GEMV or wide-GEMM family (validated 32-wide \
subgroups)"
);
let b_nl = self.w.layers.len();
let bkind_used = |kk: crate::weights::WKind| {
(0..b_nl).any(|li| {
let sk = self.w.cfg.kinds_at(li);
[sk.qkv, sk.o, sk.gate_up, sk.down].contains(&kk)
})
};
let bf16 = bkind_used(crate::weights::WKind::F16).then(|| {
if wide_gemm {
pipeline(ctx, "gemm_f16", &crate::gemm_f16_src())
} else {
pipeline(ctx, "gemv_f16_k_lcpp", &crate::gemv_f16_k_lcpp_src())
}
});
let bq4k = bkind_used(crate::weights::WKind::Q4K)
.then(|| pipeline(ctx, "gemv_q4k", &crate::gemv_q4k_k_lcpp_src()));
let bq5k = bkind_used(crate::weights::WKind::Q5K)
.then(|| pipeline(ctx, "gemv_q5k", &crate::gemv_q5k_k_lcpp_src()));
let biq4nl = bkind_used(crate::weights::WKind::Iq4Nl)
.then(|| pipeline(ctx, "gemv_iq4_nl", &crate::gemv_iq4_nl_k_lcpp_src()));
let biq4xs = bkind_used(crate::weights::WKind::Iq4Xs)
.then(|| pipeline(ctx, "gemv_iq4_xs", &crate::gemv_iq4_xs_k_lcpp_src()));
let biqg: Vec<Option<wgpu::ComputePipeline>> = IQG_KINDS
.iter()
.map(|&(kk, ty)| {
bkind_used(kk).then(|| pipeline(ctx, "gemv_iq_grid", &crate::gemv_iq_grid_src(ty)))
})
.collect();
let bq6k = bkind_used(crate::weights::WKind::Q6K)
.then(|| pipeline(ctx, "gemv_q6k", &crate::gemv_q6k_k_lcpp_src()));
let bq5n = bkind_used(crate::weights::WKind::Q5_0N)
.then(|| pipeline(ctx, "gemv_q5_0n", &crate::gemv_q5_0n_k_lcpp_src()));
let bq8n = bkind_used(crate::weights::WKind::Q8_0N).then(|| {
if wide_gemm {
pipeline(ctx, "gemm_q8_0n", &crate::gemm_q8_0n_src())
} else {
pipeline(ctx, "gemv_q8_0n", &crate::gemv_q8_0n_k_lcpp_src())
}
});
let bmlp_f16 = (0..b_nl)
.any(|li| self.w.cfg.kinds_at(li).gate_up == crate::weights::WKind::F16)
.then(|| pipeline(ctx, "mlp_gate_f16_lcpp", &crate::mlp_gate_f16_lcpp_src()));
let bmlp_q8n = (0..b_nl)
.any(|li| self.w.cfg.kinds_at(li).gate_up == crate::weights::WKind::Q8_0N)
.then(|| pipeline(ctx, "mlp_gate_q8_0n", &crate::mlp_gate_q8_0n_lcpp_src()));
let bpl =
|kk: crate::weights::WKind, dflt: &wgpu::ComputePipeline| -> wgpu::ComputePipeline {
use crate::weights::WKind;
match kk {
WKind::Q4_0 => dflt.clone(),
WKind::F16 => bf16.clone().expect("f16 batch pipeline built when used"),
WKind::Q4K => bq4k.clone().expect("q4k batch pipeline built when used"),
WKind::Q5K => bq5k.clone().expect("q5k batch pipeline built when used"),
WKind::Q6K => bq6k.clone().expect("q6k batch pipeline built when used"),
WKind::Q5_0N => bq5n.clone().expect("q5_0n batch pipeline built when used"),
WKind::Q8_0N => bq8n.clone().expect("q8_0n batch pipeline built when used"),
WKind::Iq4Nl => biq4nl
.clone()
.expect("iq4_nl batch pipeline built when used"),
WKind::Iq4Xs => biq4xs
.clone()
.expect("iq4_xs batch pipeline built when used"),
k if k.is_iq_grid() => biqg[iqg_idx(k)]
.clone()
.expect("iq-grid batch pipeline built when used"),
_ => unreachable!(),
}
};
// Pre-transposed-x Q1 wide GEMM (`OSFKB_Q1_XT=0` kills): the ablation put 52% of
// gemm_q1 in its staging skeleton; reading activations already in the shared-tile
// layout removes the in-kernel transpose (and its dynamic-[comp] select chains) —
// measured 3.83 -> 5.07 TFLOPS same-run at the MLP shape, FP order unchanged
// (bitwise-equal results). Each GEMM site gains a ~20 us transpose dispatch of its
// input; one scratch sized for the widest input serves every site (plans are serial).
let q1_xt: Option<(wgpu::ComputePipeline, wgpu::ComputePipeline, wgpu::Buffer)> = (q1
&& wide_gemm
&& std::env::var("OSFKB_Q1_XT").ok().as_deref() != Some("0"))
.then(|| {
let w_max = h.max(im_x).max(qd_x).max(c.dn_nv * c.dn_dv);
(
pipeline(ctx, "transpose_cols", &crate::transpose_cols_src()),
pipeline(ctx, "gemm_q1_xt", &crate::gemm_q1_xt_src()),
ctx.empty(w_max * (k + 4)),
)
});
// Packed-f16 activations for the NON-wide Q1 GEMV (`OSFKB_Q1_XF16=0` kills): the
// decode ablation put x-loads at 41% of gemv_q1; reading x as packed f16 pairs
// (vec4<u32>, unpack2x16float) halves the x-load instructions — measured
// 1.30-1.50x on all five Bonsai decode shapes. Each site gains a ~3 us pack
// dispatch of its input; one scratch (u32s, f16 pairs) serves every site (plans
// are serial). Wide prefill keeps the f32 xt path. Numerics: activations round
// to f16 at the GEMV boundary (accumulation stays f32) — batch==solo bitwise
// holds because every plan width shares this same path.
#[allow(clippy::type_complexity)]
#[allow(clippy::type_complexity)]
let q1_f16: Option<(
wgpu::ComputePipeline,
wgpu::ComputePipeline,
wgpu::Buffer,
wgpu::ComputePipeline,
wgpu::ComputePipeline,
wgpu::ComputePipeline,
wgpu::ComputePipeline,
wgpu::ComputePipeline,
wgpu::Buffer,
wgpu::Buffer,
wgpu::Buffer,
)> = (q1 && !wide_gemm && std::env::var("OSFKB_Q1_XF16").ok().as_deref() != Some("0"))
.then(|| {
let w_max = h.max(im_x).max(qd_x).max(c.dn_nv * c.dn_dv);
(
pipeline(ctx, "pack_f16", &crate::pack_f16_src()),
pipeline(ctx, "gemv_q1_f16x", &crate::gemv_q1_f16x_src()),
// u32 words: w_max*k f16 elems / 2, +pad; ctx.empty counts f32-sized slots
ctx.empty(w_max * k / 2 + 16),
pipeline(ctx, "mlp_q1_f16x", &crate::mlp_gate_q1_f16x_src()),
pipeline(ctx, "q1_rp4_repack", &crate::q1_rp4_repack_src()),
pipeline(ctx, "gemv_q1_rp4_f16x", &crate::gemv_q1_rp4_f16x_src()),
pipeline(ctx, "mlp_q1_rp4_f16x", &crate::mlp_gate_q1_rp4_f16x_src()),
pipeline(ctx, "rmsnorm_sg_pk", RMSNORM_SG_PK),
// dedicated packed twin of normed_x_k, written by rmsnorm_sg_pk — the
// five consumer sites bind it directly and skip their pack dispatch.
ctx.empty(h * k / 2 + 16),
// y16 trash for gemv/mlp sites whose packed twin has no consumer
// (DEDICATED — never the shared pack scratch: plans are serial and a
// dummy write between a pack and its consumer would corrupt it).
ctx.empty(w_max * k / 2 + 16),
// gate16: the fused MLP's packed output, consumed by the w2 gemv
ctx.empty(im_x * k / 2 + 16),
)
});
// ROWPACK4 (`OSFKB_Q1_RP4=0` falls back to the flat-layout f16 kernels): 1.55-1.66x
// measured COLD over the flat path (tests/gemv_cold_bench.rs) — the 4 rows' sign
// words load as ONE vec4<u32> and their scales as ONE vec4<f32>.
let rp4_on = std::env::var("OSFKB_Q1_RP4").ok().as_deref() != Some("0");
let mut rp4_key = 0u32;
let qkrc_k = self.spec_pls.qkrc_k.clone();
let attn_k = self.spec_pls.attn_k.clone();
let attn_k_wide = self.spec_pls.attn_k_wide.clone();
let qkrc_k_wide = self.spec_pls.qkrc_k_wide.clone();
let attn_ksp = self.spec_pls.attn_ksp.clone();
let attn_kmerge = self.spec_pls.attn_kmerge.clone();
let attn_skv_z = self.spec_pls.attn_skv_z;
let attn_tiled_pl = self.spec_pls.attn_tiled.clone();
// VEC4 prefill attention (`OSFKB_PREFILL_V4=0` reverts): the WIDE plan's hd=256
// attention ran the SCALAR ATTN_K (the tiled kernel needs hd≤128); ATTN_K_V4's hd/4
// vec4 QK+V is +15% at long context (3841-tok prefill 151 → 173 tok/s). Gated to
// `wide` ONLY — the chained/verify plan must stay bitwise-equal to the M=1 decode, and
// that runs whatever the global OSFKB_ATTN_V4 selects, not this prefill-local vec4.
let prefill_v4 = wide
&& !verify
&& hd % 4 == 0
&& std::env::var("OSFKB_PREFILL_V4").ok().as_deref() != Some("0");
let attn_pfx_v4 =
prefill_v4.then(|| pipeline(ctx, "attn_k_v4_prefill", &page_geom(ATTN_K_V4)));
// RB query-column register blocking (`OSFKB_PREFILL_RB2=0` reverts to plain vec4):
// cuts the K/V DRAM re-reads on full-attention prefill layers by RB×. Built alongside
// the plain vec4 (used for windowed layers, where the shared-key-stream assumption
// fails). `OSFKB_PREFILL_RB=N` picks the block (default 2 = the M4 sweet spot).
let rb_n = std::env::var("OSFKB_PREFILL_RB")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|n| (2..=16).contains(n))
.unwrap_or(2);
let attn_pfx_rb2 = (prefill_v4
&& std::env::var("OSFKB_PREFILL_RB2").ok().as_deref() != Some("0"))
.then(|| pipeline(ctx, "attn_k_v4_rb", &page_geom(&attn_k_v4_rb_src(rb_n))));
// Built lazily per plan (cheap — one small module) rather than in SpecPls: the θ
// rides the config and the variant is opt-in.
// The GPU-rope qkrc kernel drops the CPU-staged cos/sin buffers (11 → 9 storage
// buffers), which is REQUIRED under browser WebGPU's 10-storage-buffers-per-stage cap
// (native Metal/Vulkan allow ≥16, so the table path's 11 is fine there). Both Qwen3
// and Qwen3.5-MoE are single-θ Qwen-family, which the GPU-rope kernel handles; Gemma's
// dual-theta local rope keeps the table kernel. On wasm it's the default for Qwen; on
// native it stays opt-in (`OSFKB_ROPE_GPU=1`) so the bitwise table path is unchanged.
let qwen_family = matches!(
c.arch,
crate::weights::Arch::Qwen3 | crate::weights::Arch::Qwen35
);
let rope_gpu = qwen_family
&& (cfg!(target_arch = "wasm32")
|| std::env::var("OSFKB_ROPE_GPU").ok().as_deref() == Some("1"));
let qkrc_gpu_pl = pipeline(
ctx,
"qkrc_gpu_k",
&qk_norm_rope_gpu_k_src(eps, c.rope_theta, arch_has_qk_norm(c.arch)),
);
let rmsnorm = self.spec_pls.rmsnorm.clone();
let rmsnorm_add = self.spec_pls.rmsnorm_add.clone();
let moe_router_pl = self.spec_pls.moe_router.clone();
let moe_down_fused_pl = self.spec_pls.moe_down_fused.clone();
let moe_gate_nbar_pl = self.spec_pls.moe_gate_nbar.clone();
let moe_down_nbar_fused_pl = self.spec_pls.moe_down_nbar_fused.clone();
let moe_gate_lcpp_pl = self.spec_pls.moe_gate_lcpp.clone();
let moe_down_lcpp_fused_pl = self.spec_pls.moe_down_lcpp_fused.clone();
let moe_seg_build_pl = self.spec_pls.moe_seg_build.clone();
let moe_gate_grouped_pl = self.spec_pls.moe_gate_grouped.clone();
let moe_down_grouped_pl = self.spec_pls.moe_down_grouped.clone();
let moe_down_reduce_pl = self.spec_pls.moe_down_reduce.clone();
let router_norm_sgate_pl = self.spec_pls.router_norm_sgate.clone();
let router_logits_pl = self.spec_pls.router_logits.clone();
let router_pick_pl = self.spec_pls.router_pick.clone();
let moe_gate_pl = self.spec_pls.moe_gate.clone();
let moe_down_pl = self.spec_pls.moe_down.clone();
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, k as f32, eps, 0.0f32]));
// RMSNORM_ADD with a layer_scalar in kp.w, and the PLE_PREP uniform (pd, hscale, eps,
// row stride) — Gemma-4 edge only; 0 in kp.w is bitwise-neutral everywhere else.
let mfk_ls = |a: f32, s: f32| uni(ctx, bytemuck::cast_slice(&[a, k as f32, eps, s]));
let mfp = |a: f32, b: f32, s: f32| uni(ctx, bytemuck::cast_slice(&[a, b, eps, s]));
let epsm = uni(ctx, bytemuck::cast_slice(&[eps, 0.0f32, 0.0, 0.0]));
let epsm_act = uni(
ctx,
bytemuck::cast_slice(&[eps, if c.act_gelu { 1.0f32 } else { 0.0 }, 0.0, 0.0]),
);
let udims =
|a: u32, b: u32, acc: u32, d: u32| uni(ctx, bytemuck::cast_slice(&u32x4(a, b, acc, d)));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(k as u32, 0, vocab as u32, h as u32)),
);
let embed_bytes = (vocab as u64) * (h as u64) * 4;
// Per-column embedding source: -1 = the token table, otherwise a row of the image arena.
// Same shape as `fed`, and (like `fed`) uploaded as DATA, so the command buffer that reads
// it stays request-independent and cacheable.
let isrc = ctx.empty((CHAIN_MAX + 1) * k);
let gather = if c.stage_first && self.w.embed_f16.is_some() {
// Q1/GGUF: the f32 table (5 GB) can never bind; the f16 twin binds as two halves.
assert!(
self.img_embed.is_none(),
"vision arena + Q1 f16 embed is not wired yet (image rows would need the same split)"
);
let e16 = self.w.embed_f16.as_ref().unwrap();
let pl = pipeline(ctx, "gather_k_f16", GATHER_K_F16);
let lo_size = (vocab as u64 / 2) * (h as u64) * 2;
let hi_size = e16.size() - lo_size;
let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("gather_k_f16_split"),
layout: &pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: e16,
offset: 0,
size: std::num::NonZeroU64::new(lo_size),
}),
},
wgpu::BindGroupEntry {
binding: 1,
resource: fed.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: cur.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: cnt.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: posu.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 5,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: e16,
offset: lo_size,
size: std::num::NonZeroU64::new(hi_size),
}),
},
],
});
Some((pl, bg))
} else {
(c.stage_first && embed_bytes <= ctx.device.limits().max_storage_buffer_binding_size)
.then(|| match &self.img_embed {
// Multimodal: the gather picks each column's source from `isrc`.
Some(img) => {
let pl = pipeline(ctx, "gather_k_mm", GATHER_K_MM);
let bg =
make_bg(ctx, &pl, &[&w.embed, &fed, &cur, &cnt, img, &isrc], &posu);
(pl, bg)
}
// Text-only: byte-for-byte the pipeline and bind group this always had.
None => {
let pl = pipeline(ctx, "gather_k", GATHER_K);
let bg = make_bg(ctx, &pl, &[&w.embed, &fed, &cur, &cnt], &posu);
(pl, bg)
}
})
};
let bg2 = |pl: &wgpu::ComputePipeline,
bufs: &[&wgpu::Buffer],
dims: &wgpu::Buffer,
eu: &wgpu::Buffer| {
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32,
resource: dims.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: bufs.len() as u32 + 1,
resource: eu.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &pl.get_bind_group_layout(0),
entries: &entries,
})
};
let dn_snap_pls = dn_snap.then(|| {
(
pipeline(ctx, "dn_conv_k_snap", &dn_conv_k_snap_src()),
pipeline(ctx, "dn_step_pk_snap", &dn_step_pk_snap_src()),
)
});
// State-resident DN step (DN_STEP_RES): MEASURED WORSE than the global-state walk at
// 5 cols (147 → 189 µs/layer) — the 32 KB shared residency caps occupancy at one
// 64-thread WG per SM, and small-width decode is latency-bound, not traffic-bound.
// Kept behind OSFKB_DN_RES=1 for re-evaluation on other widths/hardware.
let dn_res = k <= 8 && std::env::var("OSFKB_DN_RES").ok().as_deref() == Some("1");
let dn_step_res_pl = dn_res.then(|| {
if dn_snap {
pipeline(ctx, "dn_step_res_snap", &dn_step_res_snap_src())
} else {
pipeline(ctx, "dn_step_res", DN_STEP_RES)
}
});
let dn_onorm_pl = dn_res.then(|| pipeline(ctx, "dn_onorm", DN_ONORM));
let mut dn_snaps: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>> =
(0..nl).map(|_| None).collect();
// Split-KV partials for this plan's column capacity (the split pair fires only on
// the non-kc16 decode fallback below).
let attn_part = ctx.storage(&vec![0f32; k * nh * attn_skv_z as usize * (hd_x + 4)]);
let attn_merge_u = uni(
ctx,
bytemuck::cast_slice(&u32x4(nh as u32, hd as u32, attn_skv_z, 0)),
);
let mut plan: Vec<Step> = Vec::new();
// ── Gemma-4 edge PLE model-level stage ── `cur` holds the k gathered embeddings when
// the plan runs (the drivers dispatch the embed gather first, every chained round), and
// the gather below reads the SAME per-round `fed` slice, so chaining stays correct.
let ple_bufs = (c.edge.ple_dim > 0).then(|| {
let pd = c.edge.ple_dim;
(
ctx.empty(k * nl * pd), // merged per-layer inputs
ctx.empty(k * nl * pd), // context-projection staging
ctx.empty(k * pd), // per-layer gate scratch
ctx.empty(k * pd), // per-layer activation scratch
)
});
if let (Some((ple_b, ple_ctx_b, _, _)), Some(pw)) = (&ple_bufs, &w.ple) {
// The PLE mats are packed at the BASE dtype; the plain batched GEMV below is the
// Q4_0 shape (the M=1 plan handles the f16 policy — wire the twin here if a mixed
// policy ever meets an edge model).
assert!(
c.wdtype == crate::weights::WDtype::Q4,
"gemma4-edge batched: Q4 body policy only (f16 PLE twin not wired)"
);
debug_assert!(splitk.is_none(), "edge forces the plain gemv family");
let pd = c.edge.ple_dim;
let cols = (nl * pd) as u32;
let gp = self
.spec_pls
.ple_gather
.as_ref()
.expect("ple pipelines built");
let pp = self
.spec_pls
.ple_prep
.as_ref()
.expect("ple pipelines built");
plan.push(Step::D {
tag: "ple_gather",
pl: gp.clone(),
bg: make_bg(
ctx,
gp,
&[&pw.table.scales, &pw.table.quants, &fed, ple_b, &cnt],
&udims(cols, k as u32, 0, 0),
),
gx: cols.div_ceil(256),
gy: kk,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[
&pw.model_proj.scales,
&pw.model_proj.quants,
&cur,
ple_ctx_b,
],
&udims(cols, h as u32, 0, k as u32),
),
gx: cols.div_ceil(gemv_rpw),
gy: gemv_gy,
});
plan.push(Step::D {
tag: "ple_prep",
pl: pp.clone(),
bg: make_bg(
ctx,
pp,
&[ple_ctx_b, &pw.proj_norm, ple_b],
&mfp(pd as f32, 1.0 / (h as f32).sqrt(), cols as f32),
),
gx: nl as u32,
gy: kk,
});
}
#[allow(clippy::needless_range_loop)] // li also indexes cfg.layer_k_eq_v + conv_slots
for li in 0..nl {
let layer = &w.layers[li];
let onorm = &layer.operator_norm;
// Gemma-4 edge per-layer geometry SHADOWS the uniform dims: every dispatch in this
// loop reads these, so uniform models see identical values and edge models get the
// per-layer widths with no per-site changes. Scratch was sized at the maxima above.
let hd = c.head_dim_at(li);
let (qd, kd) = (nh * hd, nkv * hd);
let im = c.intermediate_at(li);
let owns_kv = c.owns_kv(li);
let wide_l = hd > 256;
let attn_tk_l: &wgpu::Buffer = if hd == c.head_dim {
&attn_tk
} else {
attn_tk2
.as_ref()
.expect("wide attn_tk twin built for wide layers")
};
let _ = (im, owns_kv, wide_l, attn_tk_l);
if let Op::DeltaNet {
in_proj,
conv_w,
gpar,
out_proj,
} = &layer.op
{
let dn = self.dn_pls.as_ref().expect("DeltaNet layers imply dn_pls");
if (self.lcpp_gemv(ctx) && !kc16) || wide_lcpp || wide_gemm {
// Norm SPLIT (measured better than the inline-norm fusion, which re-reads
// each WG's column twice: +1.9 ms/round): rmsnorm + 4-row subgroup GEMV.
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "rmsnorm",
pl: xf.7.clone(),
bg: make_bg(
ctx,
&xf.7,
&[&cur, onorm, &normed_x_k, &xf.8],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(ctx, &rmsnorm, &[&cur, onorm, &normed_x_k], &mfk(h as f32)),
gx: kk,
gy: 1,
});
}
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&normed_x_k, &xtp.2],
&udims(h as u32, k as u32, 0, 0),
),
gx: ((h as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&in_proj.scales, &in_proj.quants, &xtp.2, &dn_mixed_k],
&udims(dn_inr as u32, h as u32, 0, k as u32),
),
gx: (dn_inr as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&in_proj.scales,
&in_proj.quants,
(dn_inr as u32) as usize,
(h as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.8, &dn_mixed_k, &xf.9],
&udims(dn_inr as u32, h as u32, 0, k as u32),
),
gx: (dn_inr as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&in_proj.scales, &in_proj.quants, &xf.8, &dn_mixed_k],
&udims(dn_inr as u32, h as u32, 0, k as u32),
),
gx: (dn_inr as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[
&in_proj.scales,
&in_proj.quants,
&normed_x_k,
splitk.as_ref().map_or(&dn_mixed_k, |(_, p)| p),
],
&udims(dn_inr as u32, h as u32, 0, k as u32),
),
gx: (dn_inr as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, &dn_mixed_k],
&udims(dn_inr as u32, h as u32, 0, k as u32),
),
gx: ((dn_inr) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&in_proj.scales, &in_proj.quants, &cur, onorm, &dn_mixed_k],
&udims(dn_inr as u32, h as u32, 0, k as u32),
&epsm,
),
gx: (dn_inr as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
if let Some((conv_snap_pl, _)) = &dn_snap_pls {
// Snapshot buffers: state after EVERY column (rollback granularity).
let snap = ctx.empty(k * c_dn.dn_nv * c_dn.dn_dk * c_dn.dn_dv);
let csnap = ctx.empty(k * dn_cd * c_dn.dn_kernel);
let bufs: [&wgpu::Buffer; 6] = [
&dn_mixed_k,
conv_w,
self.dn_ring[li].as_ref().expect("dn ring"),
&dn_conved_k,
&cmeta_k,
&cnt,
];
let d1 = udims(dn_cd as u32, c_dn.dn_kernel as u32, kk, k as u32);
let d2 = udims(dn_inr as u32, 0, 0, 0);
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: 6,
resource: d1.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 7,
resource: d2.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 8,
resource: csnap.as_entire_binding(),
});
plan.push(Step::D {
tag: "conv_snap_pl",
pl: conv_snap_pl.clone(),
bg: ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &conv_snap_pl.get_bind_group_layout(0),
entries: &entries,
}),
gx: (dn_cd as u32).div_ceil(64),
gy: 1,
});
dn_snaps[li] = Some((snap, csnap));
} else {
plan.push(Step::D {
tag: "dn.conv",
pl: dn.conv.clone(),
bg: bg2(
&dn.conv,
&[
&dn_mixed_k,
conv_w,
self.dn_ring[li].as_ref().expect("dn ring"),
&dn_conved_k,
&cmeta_k,
&cnt,
],
&udims(dn_cd as u32, c_dn.dn_kernel as u32, kk, k as u32),
&udims(dn_inr as u32, 0, 0, 0),
),
gx: (dn_cd as u32).div_ceil(64),
gy: 1,
});
}
// WIDE (prefill) plans take the CHUNK-PARALLEL recurrence (state touched
// 3×/chunk instead of 3×/position — the prefill wall; `OSFKB_DN_CHUNK=0`
// reverts to the sequential kernel). Decode/spec plans keep the sequential
// family (snapshots + small widths live there).
let dn_chunk = wide
&& dn_step_res_pl.is_none()
&& dn_snap_pls.is_none()
&& std::env::var("OSFKB_DN_CHUNK").ok().as_deref() != Some("0");
let dn_step_pl = if dn_chunk {
&dn.chunk_pk
} else if let Some(res_pl) = &dn_step_res_pl {
res_pl
} else if let Some((_, step_snap_pl)) = &dn_snap_pls {
step_snap_pl
} else {
&dn.step_par
};
plan.push(Step::D {
tag: "dn_step_pl",
pl: dn_step_pl.clone(),
bg: {
let bufs: [&wgpu::Buffer; 7] = [
&dn_conved_k,
&dn_mixed_k,
gpar,
self.dn_state[li].as_ref().expect("dn state"),
&dn_core_k,
&cmeta_k,
&cnt,
];
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
let d1 = udims(
c_dn.dn_nv as u32,
c_dn.dn_nk as u32,
c_dn.dn_dk as u32,
c_dn.dn_dv as u32,
);
let d2 = udims(dn_inr as u32, dn_cd as u32, kk, k as u32);
entries.push(wgpu::BindGroupEntry {
binding: 7,
resource: d1.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 8,
resource: d2.as_entire_binding(),
});
// The RES variants moved the norm to DN_ONORM and never read epsu —
// Naga's auto-layout OMITS unused bindings, so the entry must too. The
// dv-split CHUNK variant replaces binding 9 with the Σo² partials.
if dn_chunk {
entries.push(wgpu::BindGroupEntry {
binding: 9,
resource: dn_opart_k.as_entire_binding(),
});
} else if dn_step_res_pl.is_none() {
entries.push(wgpu::BindGroupEntry {
binding: 9,
resource: epsm.as_entire_binding(),
});
}
if dn_snap_pls.is_some() {
let (snap, _) = dn_snaps[li].as_ref().expect("snap allocated by conv");
entries.push(wgpu::BindGroupEntry {
binding: 10,
resource: snap.as_entire_binding(),
});
}
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &dn_step_pl.get_bind_group_layout(0),
entries: &entries,
})
},
gx: if dn_step_res_pl.is_some() || dn_chunk {
c_dn.dn_nv as u32 * 2
} else {
c_dn.dn_nv as u32
},
gy: kk,
});
if dn_chunk {
// Deferred gated RMSNorm + z-gate over the RAW chunk outputs.
let d1 = udims(
c_dn.dn_nv as u32,
c_dn.dn_nk as u32,
c_dn.dn_dk as u32,
c_dn.dn_dv as u32,
);
let d2 = udims(dn_inr as u32, dn_cd as u32, kk, k as u32);
plan.push(Step::D {
tag: "dn_chunk_onorm",
pl: dn.chunk_onorm.clone(),
bg: {
let bufs: [&wgpu::Buffer; 4] =
[&dn_core_k, &dn_mixed_k, gpar, &dn_opart_k];
let mut entries: Vec<wgpu::BindGroupEntry> = bufs
.iter()
.enumerate()
.map(|(i, b)| wgpu::BindGroupEntry {
binding: i as u32,
resource: b.as_entire_binding(),
})
.collect();
entries.push(wgpu::BindGroupEntry {
binding: 4,
resource: d1.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 5,
resource: d2.as_entire_binding(),
});
entries.push(wgpu::BindGroupEntry {
binding: 6,
resource: epsm.as_entire_binding(),
});
ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &dn.chunk_onorm.get_bind_group_layout(0),
entries: &entries,
})
},
gx: c_dn.dn_nv as u32,
gy: kk,
});
}
if let Some(onorm_pl) = &dn_onorm_pl {
plan.push(Step::D {
tag: "dn_onorm_pl",
pl: onorm_pl.clone(),
bg: bg2(
onorm_pl,
&[&dn_mixed_k, gpar, &dn_core_k],
&udims(
c_dn.dn_nv as u32,
dn_inr as u32,
dn_cd as u32,
c_dn.dn_dv as u32,
),
&epsm,
),
gx: c_dn.dn_nv as u32,
gy: kk,
});
}
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&dn_core_k, &xtp.2],
&udims((c_dn.dn_nv * c_dn.dn_dv) as u32, k as u32, 0, 0),
),
gx: (((c_dn.dn_nv * c_dn.dn_dv) as u32) * (k as u32).div_ceil(4))
.div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&out_proj.scales, &out_proj.quants, &xtp.2, &cur],
&udims(h as u32, (c_dn.dn_nv * c_dn.dn_dv) as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "packf16",
pl: xf.0.clone(),
bg: make_bg(
ctx,
&xf.0,
&[&dn_core_k, &xf.2],
&udims(((c_dn.dn_nv * c_dn.dn_dv) as u32) * k as u32, 0, 0, 0),
),
gx: (((c_dn.dn_nv * c_dn.dn_dv) as u32) * (k as u32) / 2).div_ceil(256),
gy: 1,
});
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&out_proj.scales,
&out_proj.quants,
(h as u32) as usize,
((c_dn.dn_nv * c_dn.dn_dv) as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.2, &cur, &xf.9],
&udims(h as u32, (c_dn.dn_nv * c_dn.dn_dv) as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&out_proj.scales, &out_proj.quants, &xf.2, &cur],
&udims(h as u32, (c_dn.dn_nv * c_dn.dn_dv) as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[
&out_proj.scales,
&out_proj.quants,
&dn_core_k,
splitk.as_ref().map_or(&cur, |(_, p)| p),
],
&udims(h as u32, (c_dn.dn_nv * c_dn.dn_dv) as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, &cur],
&udims(h as u32, (c_dn.dn_nv * c_dn.dn_dv) as u32, 1, k as u32),
),
gx: ((h) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
}
if let Op::Conv {
in_proj,
conv_w,
out_proj,
} = &layer.op
{
let slot_state = conv_slots[li].as_ref().expect("conv slot state");
plan.push(Step::D {
tag: "self.spec_pls.convdot_k",
pl: self.spec_pls.convdot_k.clone(),
bg: make_bg(
ctx,
&self.spec_pls.convdot_k,
&[&in_proj.scales, &in_proj.quants, &cur, onorm, &bcx_k],
&udims(h as u32, conv_l as u32, 0, 0),
),
gx: h as u32,
gy: kk,
});
plan.push(Step::D {
tag: "self.spec_pls.convmix_k",
pl: self.spec_pls.convmix_k.clone(),
bg: make_bg(
ctx,
&self.spec_pls.convmix_k,
&[
&bcx_k,
conv_w,
slot_state,
&opout_k,
&conv_snap_scratch,
&cmeta_k,
&cnt,
],
&udims(h as u32, conv_l as u32, kk, k as u32),
),
gx: (h as u32).div_ceil(64),
gy: 1,
});
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&opout_k, &xtp.2],
&udims(h as u32, k as u32, 0, 0),
),
gx: ((h as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&out_proj.scales, &out_proj.quants, &xtp.2, &cur],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "packf16",
pl: xf.0.clone(),
bg: make_bg(
ctx,
&xf.0,
&[&opout_k, &xf.2],
&udims((h as u32) * k as u32, 0, 0, 0),
),
gx: ((h as u32) * (k as u32) / 2).div_ceil(256),
gy: 1,
});
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&out_proj.scales,
&out_proj.quants,
(h as u32) as usize,
(h as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.2, &cur, &xf.9],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&out_proj.scales, &out_proj.quants, &xf.2, &cur],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[
&out_proj.scales,
&out_proj.quants,
&opout_k,
splitk.as_ref().map_or(&cur, |(_, p)| p),
],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, &cur],
&udims(h as u32, h as u32, 1, k as u32),
),
gx: ((h) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
}
if let Op::Attn {
qkv: qkv_w,
o,
q_norm: _,
k_norm: _,
window,
local_rope,
attn_gate,
qkv_bias,
} = &layer.op
{
let qk_norm_cat = self.qk_norm_cat[li]
.as_ref()
.expect("attention layer implies a qk-norm concat");
// KV-shared consumers project Q ONLY (their loader qkv holds just q rows).
let qkv_rows = if !owns_kv {
qd
} else if c.layer_k_eq_v[li] {
qd + kd
} else {
qd + 2 * kd
};
if (self.lcpp_gemv(ctx) && !kc16) || wide_lcpp || wide_gemm {
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "rmsnorm",
pl: xf.7.clone(),
bg: make_bg(
ctx,
&xf.7,
&[&cur, onorm, &normed_x_k, &xf.8],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(ctx, &rmsnorm, &[&cur, onorm, &normed_x_k], &mfk(h as f32)),
gx: kk,
gy: 1,
});
}
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&normed_x_k, &xtp.2],
&udims(h as u32, k as u32, 0, 0),
),
gx: ((h as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&qkv_w.scales, &qkv_w.quants, &xtp.2, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: (qkv_rows as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&qkv_w.scales,
&qkv_w.quants,
(qkv_rows as u32) as usize,
(h as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.8, &qkv, &xf.9],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: (qkv_rows as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&qkv_w.scales, &qkv_w.quants, &xf.8, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: (qkv_rows as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
let kq = self.w.cfg.kinds_at(li).qkv;
let qp = bpl(kq, &gemv_k);
// split-K (Q4-only, gated at construction) writes partials first.
let dstq = splitk.as_ref().map_or(&qkv, |(_, p)| p);
let bufs: Vec<&wgpu::Buffer> = if kq == crate::weights::WKind::Q4_0 {
vec![&qkv_w.scales, &qkv_w.quants, &normed_x_k, dstq]
} else {
vec![&qkv_w.quants, &normed_x_k, dstq]
};
plan.push(Step::D {
tag: "gemv_k",
pl: qp.clone(),
bg: make_bg(
ctx,
&qp,
&bufs,
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: (qkv_rows as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
),
gx: ((qkv_rows) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&qkv_w.scales, &qkv_w.quants, &cur, onorm, &qkv],
&udims(qkv_rows as u32, h as u32, 0, k as u32),
&epsm,
),
gx: (qkv_rows as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
let (csb, snb) = if *local_rope {
(&csk_l, &snk_l)
} else {
(&csk, &snk)
};
let g4 = c.arch == crate::weights::Arch::Gemma4;
let rd_af = if c.rotary_dim > 0 && c.rotary_dim < hd {
c.rotary_dim as u32
} else {
0
};
// af.y: 0 plain, 1 K=V, 2 Q-only (KV-shared consumer).
let af_mode = if !owns_kv {
2
} else {
u32::from(c.layer_k_eq_v[li])
};
let af = udims(u32::from(g4), af_mode, k as u32, rd_af);
// Qwen2 q/k/v bias over the kk-column fused qkv (covers every column qkrc reads),
// before either qkrc variant.
if let Some(bias) = qkv_bias {
let total = (qkv_rows * kk as usize) as u32;
plan.push(Step::D {
tag: "qkv_bias",
pl: self.qkv_bias_pl.clone(),
bg: make_bg(
ctx,
&self.qkv_bias_pl,
&[&qkv, bias],
&udims(qkv_rows as u32, total, 0, 0),
),
gx: total.div_ceil(256),
gy: 1,
});
}
// ROPE-FROM-POS (`OSFKB_ROPE_GPU=1`, Qwen single-θ only): cos/sin computed
// from the column's cmeta position in-kernel — no CPU-staged tables. The
// GPU-resident MTP round enabler (positions are data-dependent there).
if rope_gpu {
plan.push(Step::D {
tag: "qkrc_k",
pl: qkrc_gpu_pl.clone(),
bg: bg2(
&qkrc_gpu_pl,
&[
&qkv,
qk_norm_cat,
&qn,
&self.k_cache[li],
&self.v_cache[li],
&self.btab,
&cmeta_k,
&cnt,
],
attn_tk_l,
&af,
),
gx: if owns_kv {
(nh + 2 * nkv) as u32
} else {
nh as u32
},
gy: kk,
});
} else {
let qp: &wgpu::ComputePipeline = if wide_l {
qkrc_k_wide.as_ref().expect("wide qkrc built")
} else {
&qkrc_k
};
plan.push(Step::D {
tag: "qkrc_k",
pl: qp.clone(),
bg: bg2(
qp,
&[
&qkv,
qk_norm_cat,
csb,
snb,
&qn,
&self.k_cache[li],
&self.v_cache[li],
&self.btab,
&cmeta_k,
&cnt,
],
attn_tk_l,
&af,
),
gx: if owns_kv {
(nh + 2 * nkv) as u32
} else {
nh as u32
},
gy: kk,
});
}
// WIDE full-attention layers take the TILED kernel (KV chunk staged once for
// 8 columns — the per-column span re-read was 30.8% of the prefill stage);
// windowed layers and decode plans keep ATTN_K. `OSFKB_ATTN_TILED=0` kills.
// hd ≤ 128 stays on the tiled kernel; hd=256 (Qwen3.5/Bonsai) took the
// codegen'd tiled path for a while but MEASURED SLOWER than vec4 ATTN_K_V4 at
// long context (CB=16 forces 2× the barriers of hd=128; +15% went the wrong
// way — see the `prefill_v4` win below). `OSFKB_ATTN_TILED256=1` re-enables the
// codegen path for adapters where K/V bandwidth, not barriers, is the wall.
let hd256_tiled =
hd <= 256 && std::env::var("OSFKB_ATTN_TILED256").ok().as_deref() == Some("1");
// VERIFY plans make the NARROW plan's attention choices (tiled off, split pair
// exactly when narrow takes it): the spec-verify contract is per-column bitwise
// equality with plain decode, and the tiled/vec4 prefill kernels change the
// accumulation order.
let kc16_attn = kc16 && !verify;
let attn_tiled_on = !wide_l
&& kc16_attn
&& *window == 0
&& hd.is_multiple_of(32)
&& (hd <= 128 || hd256_tiled)
&& std::env::var("OSFKB_ATTN_TILED").ok().as_deref() != Some("0");
if attn_tiled_on {
plan.push(Step::D {
tag: "attn_k",
pl: attn_tiled_pl.clone(),
bg: bg2(
&attn_tiled_pl,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_o,
&self.btab,
&cmeta_k,
&cnt,
],
&attn_tiled_tk,
&udims(0, u32::from(g4), k as u32, 0),
),
gx: nh as u32,
gy: kk.div_ceil(8),
});
} else if let (Some(ksp), Some(kmg), false) =
(&attn_ksp, &attn_kmerge, kc16_attn || wide_l)
{
// CHAINED-decode split pair: lw.z stays the chained step stride (the
// split patch reads its cmeta row via cnt[0]·lw.z + col), lw.w carries z.
plan.push(Step::D {
tag: "attn_k",
pl: ksp.clone(),
bg: bg2(
ksp,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_part,
&self.btab,
&cmeta_k,
&cnt,
],
attn_tk_l,
&udims(*window, u32::from(g4), k as u32, attn_skv_z),
),
gx: nh as u32,
gy: kk * attn_skv_z,
});
plan.push(Step::D {
tag: "attn_merge",
pl: kmg.clone(),
bg: make_bg(ctx, kmg, &[&attn_part, &attn_o], &attn_merge_u),
gx: nh as u32,
gy: kk,
});
} else if let (Some(rb2), 0, false) = (&attn_pfx_rb2, *window, wide_l) {
// Full-attention prefill: RB columns/workgroup, RB× fewer K/V reads.
plan.push(Step::D {
tag: "attn_k",
pl: rb2.clone(),
bg: bg2(
rb2,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_o,
&self.btab,
&cmeta_k,
&cnt,
],
attn_tk_l,
&udims(*window, u32::from(g4), k as u32, 0),
),
gx: nh as u32,
gy: kk.div_ceil(rb_n as u32),
});
} else {
// Prefill takes the vec4 kernel where built (wide + hd%4); same bindings.
// Wide-head layers (Gemma-4 edge) route to the wide ATTN_K twin — the vec4 /
// rb2 / split / tiled families all assume thread-owns-dim ≤ 256.
let ap = if wide_l {
attn_k_wide.as_ref().expect("wide attn_k built")
} else {
attn_pfx_v4.as_ref().unwrap_or(&attn_k)
};
plan.push(Step::D {
tag: "attn_k",
pl: ap.clone(),
bg: bg2(
ap,
&[
&qn,
&self.k_cache[li],
&self.v_cache[li],
&attn_o,
&self.btab,
&cmeta_k,
&cnt,
],
attn_tk_l,
&udims(*window, u32::from(g4), k as u32, 0),
),
gx: nh as u32,
gy: kk,
});
}
if let Some(ag) = attn_gate {
let dn = self.dn_pls.as_ref().expect("attn_gate implies dn_pls");
if (self.lcpp_gemv(ctx) && !kc16) || wide_lcpp || wide_gemm {
// Reuses the qkv site's normed_x_k (same layer input, already computed).
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&normed_x_k, &xtp.2],
&udims(h as u32, k as u32, 0, 0),
),
gx: ((h as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&ag.scales, &ag.quants, &xtp.2, &agate_k],
&udims(qd as u32, h as u32, 0, k as u32),
),
gx: (qd as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&ag.scales,
&ag.quants,
(qd as u32) as usize,
(h as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.8, &agate_k, &xf.9],
&udims(qd as u32, h as u32, 0, k as u32),
),
gx: (qd as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&ag.scales, &ag.quants, &xf.8, &agate_k],
&udims(qd as u32, h as u32, 0, k as u32),
),
gx: (qd as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[
&ag.scales,
&ag.quants,
&normed_x_k,
splitk.as_ref().map_or(&agate_k, |(_, p)| p),
],
&udims(qd as u32, h as u32, 0, k as u32),
),
gx: (qd as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, &agate_k],
&udims(qd as u32, h as u32, 0, k as u32),
),
gx: ((qd) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
} else {
plan.push(Step::D {
tag: "gn32_k",
pl: gn32_k.clone(),
bg: bg2(
&gn32_k,
&[&ag.scales, &ag.quants, &cur, onorm, &agate_k],
&udims(qd as u32, h as u32, 0, k as u32),
&epsm,
),
gx: (qd as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
plan.push(Step::D {
tag: "dn.gate_mul",
pl: dn.gate_mul.clone(),
bg: make_bg(
ctx,
&dn.gate_mul,
&[&agate_k, &attn_o],
&udims(qd as u32, 0, 0, 0),
),
gx: (qd as u32).div_ceil(64),
gy: kk,
});
}
let acc = u32::from(layer.post_op_norm.is_none());
let o_target = if layer.post_op_norm.is_some() {
&opout_k
} else {
&cur
};
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&attn_o, &xtp.2],
&udims(qd as u32, k as u32, 0, 0),
),
gx: ((qd as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&o.scales, &o.quants, &xtp.2, o_target],
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "packf16",
pl: xf.0.clone(),
bg: make_bg(
ctx,
&xf.0,
&[&attn_o, &xf.2],
&udims((qd as u32) * k as u32, 0, 0, 0),
),
gx: ((qd as u32) * (k as u32) / 2).div_ceil(256),
gy: 1,
});
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&o.scales,
&o.quants,
(h as u32) as usize,
(qd as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.2, o_target, &xf.9],
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&o.scales, &o.quants, &xf.2, o_target],
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
let ko = self.w.cfg.kinds_at(li).o;
let op_pl = bpl(ko, &gemv_k);
let dsto = splitk.as_ref().map_or(o_target, |(_, p)| p);
let bufs: Vec<&wgpu::Buffer> = if ko == crate::weights::WKind::Q4_0 {
vec![&o.scales, &o.quants, &attn_o, dsto]
} else {
vec![&o.quants, &attn_o, dsto]
};
plan.push(Step::D {
tag: "gemv_k",
pl: op_pl.clone(),
bg: make_bg(
ctx,
&op_pl,
&bufs,
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, o_target],
&udims(h as u32, qd as u32, acc, k as u32),
),
gx: ((h) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
if let Some(pn) = &layer.post_op_norm {
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(ctx, &rmsnorm_add, &[&opout_k, pn, &cur], &mfk(h as f32)),
gx: kk,
gy: 1,
});
}
}
// ffn-norm + gate/up: SPLIT rmsnorm + two-stream lcpp kernel on the lcpp DECODE
// arms (wide prefill keeps the tiled-GEMM/fused shapes — the lcpp per-column
// weight re-read is a decode-width trade). `OSFKB_MLP_LCPP=0` pins the fused
// kernel.
if let Some(actmul) = &mlp_actmul_pl {
// Tiled-GEMM MLP: split rmsnorm, then gate and up as two tiled Q4 GEMMs, then
// the act·mul fold. The two GEMMs run at the projections' ~3.3 TF/s instead of
// the fused per-column kernel's ~1; the fold is memory-bound over k·im.
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "rmsnorm",
pl: xf.7.clone(),
bg: make_bg(
ctx,
&xf.7,
&[&cur, &layer.ffn_norm, &normed_x_k, &xf.8],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &layer.ffn_norm, &normed_x_k],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
}
for (w, dst) in [(&layer.w1, &gate_raw), (&layer.w3, &up_raw)] {
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&normed_x_k, &xtp.2],
&udims(h as u32, k as u32, 0, 0),
),
gx: ((h as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(), // = the tiled `gemm` under wide_gemm
bg: make_bg(
ctx,
&xtp.1,
&[&w.scales, &w.quants, &xtp.2, dst],
&udims(im as u32, h as u32, 0, k as u32),
),
gx: (im as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&w.scales,
&w.quants,
(im as u32) as usize,
(h as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(), // = the tiled `gemm` under wide_gemm
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.8, dst, &xf.9],
&udims(im as u32, h as u32, 0, k as u32),
),
gx: (im as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(), // = the tiled `gemm` under wide_gemm
bg: make_bg(
ctx,
&xf.1,
&[&w.scales, &w.quants, &xf.8, dst],
&udims(im as u32, h as u32, 0, k as u32),
),
gx: (im as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
// Kind-aware: f16/Q8_0N gate/up run their GEMM twins here — the
// full wide-MLP speedup, not just correctness.
let kg = self.w.cfg.kinds_at(li).gate_up;
let wp = bpl(kg, &gemv_k);
let dstm = splitk.as_ref().map_or(dst, |(_, p)| p);
let bufs: Vec<&wgpu::Buffer> = if kg == crate::weights::WKind::Q4_0 {
vec![&w.scales, &w.quants, &normed_x_k, dstm]
} else {
vec![&w.quants, &normed_x_k, dstm]
};
plan.push(Step::D {
tag: "gemv_k",
pl: wp.clone(), // = the tiled `gemm` family under wide_gemm
bg: make_bg(
ctx,
&wp,
&bufs,
&udims(im as u32, h as u32, 0, k as u32),
),
gx: (im as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, dst],
&udims(im as u32, h as u32, 0, k as u32),
),
gx: ((im) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
}
plan.push(Step::D {
tag: "mlp_k",
pl: actmul.clone(),
bg: make_bg(
ctx,
actmul,
&[&gate_raw, &up_raw, &gate],
&udims((k * im) as u32, u32::from(c.act_gelu), 0, 0),
),
gx: ((k * im) as u32).div_ceil(256),
gy: 1,
});
} else if self.lcpp_gemv(ctx)
&& (q1 || b_any_nonq4 || (!wide && !kc16 && mlp_lcpp_enabled()))
{
if let Some(xf) = &q1_f16 {
plan.push(Step::D {
tag: "rmsnorm",
pl: xf.7.clone(),
bg: make_bg(
ctx,
&xf.7,
&[&cur, &layer.ffn_norm, &normed_x_k, &xf.8],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &layer.ffn_norm, &normed_x_k],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
}
// Q1 (Bonsai) swaps the fused-MLP pipeline for its sign-decoding twin — same 8
// bindings, same pre-normed input. Without this the Q4 kernel decodes the
// sign-packed weights as nibbles (wrong layout AND stride) → the residual grows
// unbounded (the DN/attn projections were already on gemv_q1; only this was missed).
let mlp_split_pl = self
.spec_pls
.mlp_q1
.as_ref()
.unwrap_or(&self.spec_pls.mlp_lcpp)
.clone();
// f16 x for the fused MLP too: one pack of the (mlp-)normed x, both
// weight streams then share the halved x-load count.
let (mlp_pl_sel, mlp_x): (&wgpu::ComputePipeline, &wgpu::Buffer) =
if let (Some(xf), true) = (&q1_f16, self.spec_pls.mlp_q1.is_some()) {
(&xf.3, &xf.8)
} else {
(&mlp_split_pl, &normed_x_k)
};
// ROWPACK4 twins for BOTH mlp streams (gate w1 + up w3) when active.
let mlp_rp4: Option<((wgpu::Buffer, wgpu::Buffer), (wgpu::Buffer, wgpu::Buffer))> =
match (&q1_f16, rp4_on, self.spec_pls.mlp_q1.is_some()) {
(Some(xf), true, true) => {
let t1 = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&layer.w1.scales,
&layer.w1.quants,
im,
h,
);
let t3 = self.rp4_twin(
ctx,
&xf.4,
rp4_key + 1,
&layer.w3.scales,
&layer.w3.quants,
im,
h,
);
Some((t1, t3))
}
_ => None,
};
rp4_key += 2;
if let (Some(xf), Some((t1, t3))) = (&q1_f16, &mlp_rp4) {
plan.push(Step::D {
tag: "mlp_k",
pl: xf.6.clone(),
bg: bg2(
&xf.6,
&[&t1.0, &t1.1, &t3.0, &t3.1, mlp_x, &gate, &xf.10],
&udims(im as u32, h as u32, 0, k as u32),
&epsm_act,
),
gx: (im as u32).div_ceil(4),
gy: kk,
});
} else {
let kg = self.w.cfg.kinds_at(li).gate_up;
let mpl: wgpu::ComputePipeline = match kg {
crate::weights::WKind::Q8_0N => {
bmlp_q8n.clone().expect("q8_0n batch MLP built when used")
}
crate::weights::WKind::F16 => {
bmlp_f16.clone().expect("f16 batch MLP built when used")
}
_ => mlp_pl_sel.clone(),
};
let bufs: Vec<&wgpu::Buffer> = if kg == crate::weights::WKind::Q4_0 {
vec![
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
mlp_x,
&gate,
]
} else {
vec![&layer.w1.quants, &layer.w3.quants, mlp_x, &gate]
};
plan.push(Step::D {
tag: "mlp_k",
pl: mpl.clone(),
bg: bg2(
&mpl,
&bufs,
&udims(im as u32, h as u32, 0, k as u32),
&epsm_act,
),
gx: (im as u32).div_ceil(4),
gy: kk,
});
}
} else {
plan.push(Step::D {
tag: "mlp_k",
pl: mlp_k.clone(),
bg: bg2(
&mlp_k,
&[
&layer.w1.scales,
&layer.w1.quants,
&layer.w3.scales,
&layer.w3.quants,
&cur,
&layer.ffn_norm,
&gate,
],
&udims(im as u32, h as u32, 0, k as u32),
&epsm_act,
),
gx: (im as u32).div_ceil(self.gemv_rows_per_wg),
gy: kt,
});
}
// Same as the decode plan: the `mlpout` detour is the shared-gated MoE path (Qwen3.5 or
// Qwen2-MoE). Keyed off `shared_gate`; a DENSE Qwen3.5 (moe=None) has no block to fold it
// back in, so it accumulates into the residual.
let q35 = layer.moe.as_ref().is_some_and(|m| m.shared_gate.is_some());
let acc = u32::from(layer.post_ffn_norm.is_none() && !q35);
let d_target = if layer.post_ffn_norm.is_some() || q35 {
&mlpout_k
} else {
&cur
};
if let Some(xtp) = &q1_xt {
plan.push(Step::D {
tag: "xt",
pl: xtp.0.clone(),
bg: make_bg(
ctx,
&xtp.0,
&[&gate, &xtp.2],
&udims(im as u32, k as u32, 0, 0),
),
gx: ((im as u32) * (k as u32).div_ceil(4)).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: xtp.1.clone(),
bg: make_bg(
ctx,
&xtp.1,
&[&layer.w2.scales, &layer.w2.quants, &xtp.2, d_target],
&udims(h as u32, im as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
if let Some(xf) = &q1_f16 {
if !rp4_on {
plan.push(Step::D {
tag: "packf16",
pl: xf.0.clone(),
bg: make_bg(
ctx,
&xf.0,
&[&gate, &xf.2],
&udims((im as u32) * k as u32, 0, 0, 0),
),
gx: ((im as u32) * (k as u32) / 2).div_ceil(256),
gy: 1,
});
}
if rp4_on {
let rp4t = self.rp4_twin(
ctx,
&xf.4,
rp4_key,
&layer.w2.scales,
&layer.w2.quants,
(h as u32) as usize,
(im as u32) as usize,
);
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.5.clone(),
bg: make_bg(
ctx,
&xf.5,
&[&rp4t.0, &rp4t.1, &xf.10, d_target, &xf.9],
&udims(h as u32, im as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
} else {
rp4_key += 1;
plan.push(Step::D {
tag: "gemv_k",
pl: xf.1.clone(),
bg: make_bg(
ctx,
&xf.1,
&[&layer.w2.scales, &layer.w2.quants, &xf.2, d_target],
&udims(h as u32, im as u32, acc, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
}
} else {
let kdn = self.w.cfg.kinds_at(li).down;
let dp = bpl(kdn, &gemv_k);
let bufs: Vec<&wgpu::Buffer> =
if kdn == crate::weights::WKind::Q4_0 || kdn.is_iq_grid() {
vec![&layer.w2.scales, &layer.w2.quants, &gate, d_target]
} else {
vec![&layer.w2.quants, &gate, d_target]
};
plan.push(Step::D {
tag: "gemv_k",
pl: dp.clone(),
bg: make_bg(ctx, &dp, &bufs, &udims(h as u32, im as u32, acc, k as u32)),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
if let Some((rpl, part)) = splitk.as_ref() {
plan.push(Step::D {
tag: "gemv_red",
pl: rpl.clone(),
bg: make_bg(
ctx,
rpl,
&[part, d_target],
&udims(h as u32, im as u32, acc, k as u32),
),
gx: ((h) as u32 * kk).div_ceil(256),
gy: 1,
});
}
}
}
if let Some(moe) = &layer.moe {
let (e_n, kx, mi) = (
c.num_experts as u32,
c.top_k_experts as u32,
c.moe_intermediate as u32,
);
if let Some(sg_w) = &moe.shared_gate {
// Router split: norm+sgate → coalesced logits → pick (see kernel docs).
plan.push(Step::D {
tag: "router_norm_sgate_pl",
pl: router_norm_sgate_pl.clone(),
bg: bg2(
&router_norm_sgate_pl,
&[&cur, &layer.ffn_norm, sg_w, &normed_k, &moe_sg_k],
&udims(e_n, h as u32, kx, 0),
&epsm,
),
gx: kk,
gy: 1,
});
plan.push(Step::D {
tag: "router_logits_pl",
pl: router_logits_pl.clone(),
bg: make_bg(
ctx,
&router_logits_pl,
&[&moe.router, &normed_k, &moe_logits_k],
&udims(e_n, h as u32, kx, 0),
),
gx: e_n.div_ceil(16),
gy: kk,
});
plan.push(Step::D {
tag: "router_pick_pl",
pl: router_pick_pl.clone(),
bg: make_bg(
ctx,
&router_pick_pl,
&[
&moe_logits_k,
&moe.per_expert_scale,
&moe_sel_k,
&moe_wsel_k,
],
&udims(e_n, h as u32, kx, 0),
),
gx: kk,
gy: 1,
});
} else {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &layer.ffn_norm, &normed_k],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
plan.push(Step::D {
tag: "moe_router_pl",
pl: moe_router_pl.clone(),
bg: make_bg(
ctx,
&moe_router_pl,
&[
&moe.router,
&moe.per_expert_scale,
&normed_k,
&moe_sel_k,
&moe_wsel_k,
],
&udims(e_n, h as u32, kx, 0),
),
gx: kk,
gy: 1,
});
}
// wide_lcpp extends the lcpp expert kernels to prefill widths too: the tree
// down re-reads each column's activations once PER ROW (L2-bound — profiled
// 2.1 ms/layer at 64 cols); the lcpp row-strip shape reads them once per 4
// rows. Decode widths keep their measured selections.
let lcpp_experts = (k <= 8 || wide_lcpp) && self.lcpp_gemv(ctx);
if let Some((texp, tpair, dpartial, args)) = moe_grp.as_ref().filter(|_| q35) {
// GROUPED prefill MoE: tile build → indirect gate/down → reduce epilogue.
plan.push(Step::D {
tag: "moe_seg_build",
pl: moe_seg_build_pl.clone(),
bg: make_bg(
ctx,
&moe_seg_build_pl,
&[&moe_sel_k, texp, tpair, args],
&udims(kk * kx, e_n, mi.div_ceil(64), (h as u32).div_ceil(64)),
),
gx: 1,
gy: 1,
});
plan.push(Step::Di {
tag: "gate_pl",
pl: moe_gate_grouped_pl.clone(),
bg: bg2(
&moe_gate_grouped_pl,
&[
&moe.w1.scales,
&moe.w1.quants,
&moe.w3.scales,
&moe.w3.quants,
&normed_k,
texp,
tpair,
&moe_gate_k,
],
&udims(mi, h as u32, 0, kx),
&epsm_act,
),
args: args.clone(),
offset: 0,
});
plan.push(Step::Di {
tag: "down_pl",
pl: moe_down_grouped_pl.clone(),
bg: make_bg(
ctx,
&moe_down_grouped_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&moe_gate_k,
texp,
tpair,
&moe_wsel_k,
dpartial,
],
&udims(h as u32, mi, 0, kx),
),
args: args.clone(),
offset: 16,
});
plan.push(Step::D {
tag: "moe_down_reduce",
pl: moe_down_reduce_pl.clone(),
bg: make_bg(
ctx,
&moe_down_reduce_pl,
&[dpartial, &mlpout_k, &moe_sg_k, &cur],
&udims(h as u32, 0, 0, kx),
),
gx: (h as u32).div_ceil(64),
gy: kk,
});
} else {
let (gate_pl, gate_gx) = if lcpp_experts {
(&moe_gate_lcpp_pl, mi.div_ceil(4))
} else if k <= 8 {
(&moe_gate_nbar_pl, mi.div_ceil(16))
} else {
(&moe_gate_pl, mi)
};
plan.push(Step::D {
tag: "gate_pl",
pl: gate_pl.clone(),
bg: bg2(
gate_pl,
&[
&moe.w1.scales,
&moe.w1.quants,
&moe.w3.scales,
&moe.w3.quants,
&normed_k,
&moe_sel_k,
&moe_gate_k,
],
&udims(mi, h as u32, 0, kx),
&epsm_act,
),
gx: gate_gx,
gy: kk * kx,
});
if q35 {
// FUSED down: `cur += shared·g + routed` in the epilogue.
let (down_pl, down_gx) = if lcpp_experts {
(&moe_down_lcpp_fused_pl, (h as u32).div_ceil(4))
} else if k <= 8 {
(&moe_down_nbar_fused_pl, (h as u32).div_ceil(16))
} else {
(&moe_down_fused_pl, h as u32)
};
plan.push(Step::D {
tag: "down_pl",
pl: down_pl.clone(),
bg: make_bg_down_fused(
ctx,
down_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&moe_gate_k,
&moe_sel_k,
&moe_wsel_k,
&mlpout_k,
&moe_sg_k,
&cur,
],
&udims(h as u32, mi, 0, kx),
),
gx: down_gx,
gy: kk,
});
} else {
plan.push(Step::D {
tag: "moe_down_pl",
pl: moe_down_pl.clone(),
bg: make_bg(
ctx,
&moe_down_pl,
&[
&moe.w2.scales,
&moe.w2.quants,
&moe_gate_k,
&moe_sel_k,
&moe_wsel_k,
d_target,
],
&udims(h as u32, mi, 0, kx),
),
gx: h as u32,
gy: kk,
});
}
}
}
if let Some(pn) = &layer.post_ffn_norm {
// Gemma-4 edge: layer_scalar rides THIS add only when no PLE block follows
// (HF order: FFN add → PLE sub-block → ×layer_scalar).
let ls_here = layer
.edge
.as_ref()
.filter(|e| e.ple.is_none())
.map_or(0.0, |e| e.layer_scalar);
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(
ctx,
&rmsnorm_add,
&[&mlpout_k, pn, &cur],
&mfk_ls(h as f32, ls_here),
),
gx: kk,
gy: 1,
});
}
// Gemma-4 edge PLE injection — the third sandwich sub-block, then ×layer_scalar.
// (nested if-let kept: collapsing into a let-chain buries the PLE gate condition)
#[allow(clippy::collapsible_if)]
if let (Some(ledge), Some((ple_b, _, ple_g_b, ple_a_b))) = (&layer.edge, &ple_bufs) {
if let Some(lp) = &ledge.ple {
let pd = c.edge.ple_dim;
let pa = self.spec_pls.ple_act.as_ref().expect("ple pipelines built");
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&lp.gate.scales, &lp.gate.quants, &cur, ple_g_b],
&udims(pd as u32, h as u32, 0, k as u32),
),
gx: (pd as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
plan.push(Step::D {
tag: "ple_act_k",
pl: pa.clone(),
bg: make_bg(
ctx,
pa,
&[ple_g_b, ple_b, ple_a_b],
&udims(
(k * pd) as u32,
(li * pd) as u32,
pd as u32,
(nl * pd) as u32,
),
),
gx: ((k * pd) as u32).div_ceil(256),
gy: 1,
});
plan.push(Step::D {
tag: "gemv_k",
pl: gemv_k.clone(),
bg: make_bg(
ctx,
&gemv_k,
&[&lp.proj.scales, &lp.proj.quants, ple_a_b, &opout_k],
&udims(h as u32, pd as u32, 0, k as u32),
),
gx: (h as u32).div_ceil(gemv_rpw),
gy: gemv_gy,
});
plan.push(Step::D {
tag: "rmsnorm_add",
pl: rmsnorm_add.clone(),
bg: make_bg(
ctx,
&rmsnorm_add,
&[&opout_k, &lp.post_norm, &cur],
&mfk_ls(h as f32, ledge.layer_scalar),
),
gx: kk,
gy: 1,
});
}
}
}
if c.stage_last {
plan.push(Step::D {
tag: "rmsnorm",
pl: rmsnorm.clone(),
bg: make_bg(
ctx,
&rmsnorm,
&[&cur, &w.embedding_norm, &hnorm],
&mfk(h as f32),
),
gx: kk,
gy: 1,
});
}
// ONE k-column head dispatch: all columns' logits from a single pass over the head
// weights (the batch's dominant bandwidth). THE ENGINE'S OWN pipeline object — batched
// column c is bitwise-equal to a solo forward by construction, not by compiler luck.
// Non-last pipeline stages hold no head (dummy weights) — the plan carries no lm.
let lm = c.stage_last.then(|| {
self.q4lm.make(
ctx,
&w.embed_q4,
&hnorm,
&logits_k,
vocab as u32,
h as u32,
k as u32,
)
});
let hidden_staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("batch_hidden_staging"),
size: (k * h * 4) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let block_pl = pipeline(ctx, "argmax_block", ARGMAX_BLOCK);
let merge_pl = pipeline(ctx, "argmax_merge_col", ARGMAX_MERGE_COL);
let bump_pl = pipeline(ctx, "cnt_bump", CNT_BUMP);
let bump_bg = make_bg1(ctx, &bump_pl, &cnt);
let mut heads = Vec::with_capacity(k);
for p in 0..k {
let colu = uni(ctx, bytemuck::cast_slice(&u32x4(p as u32, k as u32, 0, 0)));
fn sub(b: &wgpu::Buffer, off: u64, size: u64) -> wgpu::BindingResource<'_> {
wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer: b,
offset: off,
size: std::num::NonZeroU64::new(size),
})
}
let block_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &block_pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: sub(&logits_k, (p * vocab * 4) as u64, (vocab * 4) as u64),
},
wgpu::BindGroupEntry {
binding: 1,
resource: sub(&pv_k, (p * 256 * 4) as u64, 256 * 4),
},
wgpu::BindGroupEntry {
binding: 2,
resource: sub(&pi_k, (p * 256 * 4) as u64, 256 * 4),
},
wgpu::BindGroupEntry {
binding: 3,
resource: posu.as_entire_binding(),
},
],
});
let merge_bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &merge_pl.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: sub(&pv_k, (p * 256 * 4) as u64, 256 * 4),
},
wgpu::BindGroupEntry {
binding: 1,
resource: sub(&pi_k, (p * 256 * 4) as u64, 256 * 4),
},
wgpu::BindGroupEntry {
binding: 2,
resource: chosen_k.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: cnt.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: colu.as_entire_binding(),
},
],
});
heads.push(BatchHead { block_bg, merge_bg });
}
// Dev bisection instrument: truncate the plan to the first N steps
// (`OSFKB_PLAN_TRUNC=N`) — with `read_cur_for_tests` this binary-searches the first
// diverging step between two plans. Dev-only; unset = untouched.
if let Ok(n) = std::env::var("OSFKB_PLAN_TRUNC") {
if let Ok(n) = n.parse::<usize>() {
let tags: Vec<&str> = plan
.iter()
.map(|s| match s {
Step::D { tag, .. } | Step::Di { tag, .. } => *tag,
Step::Embed => "embed",
})
.collect();
eprintln!(
"[plan] TRUNCATED to {n} of {} steps; tags: {tags:?}",
plan.len()
);
plan.truncate(n);
}
}
BatchPlan {
k,
plan,
fed,
csk,
snk,
csk_l,
snk_l,
cmeta_k,
cur,
qn: qn.clone(),
attn_o: attn_o.clone(),
qkv: qkv.clone(),
hidden_staging,
hnorm,
logits_k,
chosen_k,
chosen_staging,
gather,
isrc,
lm_pl: self.q4lm.pipeline_for(k as u32).clone(),
lm,
block_pl,
merge_pl,
heads,
cnt,
bump_pl,
bump_bg,
conv_slots,
dn_snaps,
dn_parallel: !wide && self.dn_pls.is_some(),
cmdbuf: std::cell::RefCell::new(None),
serial: {
use std::sync::atomic::{AtomicU64, Ordering};
static SERIAL: AtomicU64 = AtomicU64::new(1);
SERIAL.fetch_add(1, Ordering::Relaxed)
},
}
}
/// Zero one sequence slot's DeltaNet state (conv rings + recurrent S) — a freshly admitted
/// sequence must not inherit the previous occupant's recurrence.
pub fn zero_dn_slot(&self, ctx: &GpuCtx, slot: usize) {
let c = &self.w.cfg;
let cd = 2 * c.dn_nk * c.dn_dk + c.dn_nv * c.dn_dv;
let zr = vec![0f32; cd * c.dn_kernel.max(1)];
let zs = vec![0f32; c.dn_nv * c.dn_dk * c.dn_dv];
for (ring, st) in self.dn_ring.iter().zip(&self.dn_state) {
if let (Some(ring), Some(st)) = (ring.as_ref(), st.as_ref()) {
ctx.queue.write_buffer(
ring,
(slot * zr.len() * 4) as u64,
bytemuck::cast_slice(&zr),
);
ctx.queue
.write_buffer(st, (slot * zs.len() * 4) as u64, bytemuck::cast_slice(&zs));
}
}
}
/// Zero one sequence slot's conv rings (fresh sequence; LFM2 batch serving).
pub fn zero_conv_slot(&self, ctx: &GpuCtx, bp: &BatchPlan, slot: usize) {
let c = &self.w.cfg;
let z = vec![0f32; c.hidden * c.conv_l];
for buf in bp.conv_slots.iter().flatten() {
ctx.queue.write_buffer(
buf,
((1 + slot) * c.hidden * c.conv_l * 4) as u64,
bytemuck::cast_slice(&z),
);
}
}
/// One continuous-batching step: run every column's (token, position) through the stack —
/// columns address their own sequence's KV blocks via `btrow` — then compute logits + greedy
/// argmax for the columns that need a token. Returns `chosen` for all k columns (only
/// `need_logit` columns are meaningful). Columns beyond `cols.len()` are parked on the trash
/// block-table row. Per-column results are bitwise-equal to a solo run of the same sequence.
pub fn batch_step(&self, ctx: &GpuCtx, bp: &BatchPlan, cols: &[BatchCol]) -> Result<Vec<u32>> {
match self.batch_stage_step(ctx, bp, cols, None)? {
StageBatchOut::Tokens(t) => Ok(t),
StageBatchOut::Hidden(_) => unreachable!("full model ends with the head"),
}
}
/// [`Self::batch_step`] generalized to a PIPELINE STAGE: non-first stages feed the previous
/// stage's residual stream (`hidden_in`, `cols.len()·hidden` floats) instead of tokens;
/// non-last stages return the residual stream instead of tokens. Column c's compute is the
/// same compiled pipelines as a solo run of that sequence — batch invariance holds per stage,
/// and stages compose it across devices.
pub fn batch_stage_step(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
hidden_in: Option<&[f32]>,
) -> Result<StageBatchOut> {
self.batch_stage_step_inner(ctx, bp, cols, hidden_in, false)
}
/// [`Self::batch_stage_step`] with `bp.cur` ALREADY holding the residual stream on-GPU
/// (the MTP draft path: the fc-merge dispatches fill `cur` before the layer replays).
pub(crate) fn batch_stage_step_cur_ready(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
) -> Result<StageBatchOut> {
self.batch_stage_step_inner(ctx, bp, cols, None, true)
}
/// Encode one batch step's command buffer (embed input + plan pass + head + readback
/// copies). Except for the stage-first copy fallback (token offsets baked into buffer
/// copies), the buffer has no data-dependent content and is reusable across rounds of
/// identical (plan, ncols, need_logit) shape — see `BatchPlan::cmdbuf`.
fn encode_step_cmdbuf(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
use_gather: bool,
) -> wgpu::CommandBuffer {
let c = &self.w.cfg;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if c.stage_first {
if use_gather {
// Tokens flow through `fed` + the gather dispatch (same bytes as row copies):
// no data-dependent offsets in the command buffer ⇒ it is reusable across
// rounds with the same (plan, ncols, need_logit) shape.
let (pl, bg) = bp.gather.as_ref().expect("checked by caller");
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups((c.hidden as u32).div_ceil(64), cols.len() as u32, 1);
} else {
// Exact-byte row copies — fallback when the f32 embed exceeds the device's
// storage-binding cap (no gather bind group). NuExtract-3 lands here: its embedding
// table is 248320 × 2560 × 4 = 2.54 GB.
//
// This path must honour `embed_row` too. It originally did not, and the failure was
// exactly the one this migration keeps producing: the model ran, emitted well-formed
// JSON, and had never seen the image — every image token silently got the
// placeholder id's embedding, which is the same vector for every image ever sent.
for (p, col) in cols.iter().enumerate() {
let (src, off) = match self.img_embed.as_ref() {
Some(img) if col.embed_row >= 0 => {
(img, (col.embed_row as u64) * (c.hidden as u64) * 4)
}
_ => {
debug_assert!(
col.embed_row < 0,
"column wants image row {} but the engine has no vision arena — \
gathering the placeholder token's embedding would be silently wrong",
col.embed_row
);
self.w.embed_row_src(col.token)
}
};
enc.copy_buffer_to_buffer(
src,
off,
&bp.cur,
(p * c.hidden * 4) as u64,
(c.hidden * 4) as u64,
);
}
}
}
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
for step in &bp.plan {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
} else if let Step::Di {
pl,
bg,
args,
offset,
..
} = step
{
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
}
if c.stage_last && cols.iter().any(|c| c.need_logit) {
// One k-column head pass (weights read once for ALL columns); idle/trash
// columns' logits are computed and ignored — with the shared weight stream the
// marginal column is bandwidth-free.
let (lm_bg, lm_gx, lm_gy) = bp.lm.as_ref().expect("stage_last has a head");
p.set_pipeline(&bp.lm_pl);
p.set_bind_group(0, lm_bg, &[]);
p.dispatch_workgroups(*lm_gx, *lm_gy, 1);
for (pi, col) in cols.iter().enumerate() {
if !col.need_logit {
continue;
}
let hh = &bp.heads[pi];
p.set_pipeline(&bp.block_pl);
p.set_bind_group(0, &hh.block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&bp.merge_pl);
p.set_bind_group(0, &hh.merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
}
}
if c.stage_last {
enc.copy_buffer_to_buffer(&bp.chosen_k, 0, &bp.chosen_staging, 0, (bp.k * 4) as u64);
} else {
enc.copy_buffer_to_buffer(
&bp.cur,
0,
&bp.hidden_staging,
0,
(cols.len() * c.hidden * 4) as u64,
);
}
enc.finish()
}
/// Génération mono-séquence ASYNC — le jumeau navigateur de `server::generate_once`
/// (k=1, préfill par colonnes puis décode étroit, greedy, arrêt sur `eos`). Les
/// soumissions passent par [`Self::stage_step_submit`] ; la relecture est synchrone en
/// natif (poll bloquant) et `await`ée sur wasm (WebGPU n'a pas de poll bloquant) — mêmes
/// octets, même décodage. La gate native `generate_async_gate` exige l'égalité BIT À BIT
/// avec `generate_once` : ce chemin n'a pas le droit d'être « presque » le même moteur.
///
/// `on_token(&généré)` est rappelé après chaque token accepté — le streaming de la page.
pub async fn generate_async(
&self,
ctx: &GpuCtx,
prompt: &[u32],
max_tokens: usize,
eos: &[u32],
mut on_token: impl FnMut(&[u32]),
) -> Result<Vec<u32>> {
anyhow::ensure!(!prompt.is_empty(), "empty prompt");
let cap_blocks = (prompt.len() + max_tokens).div_ceil(PAGE_BLK) as u32;
anyhow::ensure!(
(cap_blocks as usize) <= PAGE_ROW,
"prompt + max_tokens dépasse la capacité par séquence"
);
self.reset(ctx);
self.write_btab_row(ctx, 1, &(0..cap_blocks).collect::<Vec<_>>());
let bp = self.make_batch_plan(ctx, 1);
// Préfill : toutes les colonnes soumises d'affilée (la file GPU est ordonnée), un
// SEUL readback — la dernière colonne porte need_logit.
for (pos, &t) in prompt.iter().enumerate() {
let col = BatchCol::text(t, pos as u32, 1, pos + 1 == prompt.len());
self.stage_step_submit(ctx, &bp, &[col], None, false)?;
}
macro_rules! read_tok {
() => {{
#[cfg(not(target_arch = "wasm32"))]
let out = self.read_stage_out(ctx, &bp, 1)?;
#[cfg(target_arch = "wasm32")]
let out = self.read_stage_out_async(&bp, 1).await?;
match out {
StageBatchOut::Tokens(t) => t[0],
StageBatchOut::Hidden(_) => {
anyhow::bail!("generate_async: le modèle doit finir par la tête")
}
}
}};
}
let mut tok = read_tok!();
let mut out = Vec::with_capacity(max_tokens);
let mut pos = prompt.len() as u32;
while out.len() < max_tokens {
out.push(tok);
on_token(&out);
if eos.contains(&tok) {
break;
}
let col = BatchCol::text(tok, pos, 1, true);
self.stage_step_submit(ctx, &bp, &[col], None, false)?;
tok = read_tok!();
pos += 1;
}
Ok(out)
}
/// Prep + encode + submit a stage step, WITHOUT reading the result back. Split out of
/// `batch_stage_step_inner` so the wasm/browser worker can submit here and then read
/// ASYNCHRONOUSLY (WebGPU has no synchronous readback). On native this is exactly the
/// pre-readback half of the old monolith — the shard/serving suites gate the equivalence.
/// Test-only: run a k=1 batch plan step by step for ONE (token, pos), reading column 0's
/// residual `bp.cur[0..hidden]` after every dispatch. Returns `(tag, l2_norm, all_finite)`
/// per step — the first non-finite / collapsed entry names the offending kernel. Mirrors the
/// metadata prep of `stage_step_submit` for a single fresh column (no history).
#[doc(hidden)]
pub fn probe_col0_norms(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
token: u32,
pos: u32,
btrow: u32,
) -> Result<Vec<(&'static str, [f32; 4], bool)>> {
let c = &self.w.cfg;
// Global-theta tables serve the full-attention layers — the WIDE heads on an edge model.
let rope = RopeStager::new(c, c.head_dim_full());
let (cs, sn) = rope.col([0; 3], pos);
let mut fed = vec![0u32; bp.k];
fed[0] = token;
let mut cm = vec![0u32; bp.k * 2];
cm[0] = pos;
cm[1] = btrow;
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
ctx.queue
.write_buffer(&bp.fed, 0, bytemuck::cast_slice(&fed));
ctx.queue.write_buffer(
&bp.isrc,
0,
bytemuck::cast_slice(&vec![-1i32; (CHAIN_MAX + 1) * bp.k]),
);
ctx.queue
.write_buffer(&bp.cmeta_k, 0, bytemuck::cast_slice(&cm));
ctx.queue
.write_buffer(&bp.csk, 0, bytemuck::cast_slice(&cs));
ctx.queue
.write_buffer(&bp.snk, 0, bytemuck::cast_slice(&sn));
// Embed the token into cur col 0 (gather if available, else a row copy).
{
let mut enc = ctx.device.create_command_encoder(&Default::default());
if let Some((pl, bg)) = bp.gather.as_ref() {
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups((c.hidden as u32).div_ceil(64), 1, 1);
} else {
let (src, off) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(src, off, &bp.cur, 0, (c.hidden as u64) * 4);
}
ctx.queue.submit([enc.finish()]);
}
let mut out = Vec::new();
for step in bp.plan.iter() {
let Step::D {
pl,
bg,
gx,
gy,
tag,
} = step
else {
continue;
};
let mut enc = ctx.device.create_command_encoder(&Default::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
ctx.queue.submit([enc.finish()]);
let v = ctx.read(&bp.cur, c.hidden)?;
let finite = v.iter().all(|x| x.is_finite());
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
// The attention-chain intermediates (post-rope Q, pre-o-proj attention out): cur
// only changes at the residual add, so these are what BISECT inside the chain.
let l2 = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt();
let nq = l2(&ctx.read(&bp.qn, ((bp.qn.size() / 4) as usize).min(65536))?);
let na = l2(&ctx.read(&bp.attn_o, ((bp.attn_o.size() / 4) as usize).min(65536))?);
let nv = l2(&ctx.read(&bp.qkv, ((bp.qkv.size() / 4) as usize).min(65536))?);
out.push((*tag, [n, nq, na, nv], finite));
}
Ok(out)
}
/// [`Self::probe_col0_norms`] en ASYNC — la même sonde pas-à-pas, exécutable dans le
/// navigateur : elle NOMME le premier kernel dont la sortie s'effondre sur un backend
/// donné (BrowserWebGpu inclus), là où le natif est gaté vert.
#[doc(hidden)]
#[cfg(target_arch = "wasm32")]
pub async fn probe_col0_norms_async(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
token: u32,
pos: u32,
btrow: u32,
) -> Result<Vec<(&'static str, [f32; 4], bool, String)>> {
let c = &self.w.cfg;
let rope = RopeStager::new(c, c.head_dim_full());
let (cs, sn) = rope.col([0; 3], pos);
let mut fed = vec![0u32; bp.k];
fed[0] = token;
let mut cm = vec![0u32; bp.k * 2];
cm[0] = pos;
cm[1] = btrow;
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
ctx.queue
.write_buffer(&bp.fed, 0, bytemuck::cast_slice(&fed));
ctx.queue.write_buffer(
&bp.isrc,
0,
bytemuck::cast_slice(&vec![-1i32; (CHAIN_MAX + 1) * bp.k]),
);
ctx.queue
.write_buffer(&bp.cmeta_k, 0, bytemuck::cast_slice(&cm));
ctx.queue
.write_buffer(&bp.csk, 0, bytemuck::cast_slice(&cs));
ctx.queue
.write_buffer(&bp.snk, 0, bytemuck::cast_slice(&sn));
{
let mut enc = ctx.device.create_command_encoder(&Default::default());
if let Some((pl, bg)) = bp.gather.as_ref() {
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups((c.hidden as u32).div_ceil(64), 1, 1);
} else {
let (src, off) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(src, off, &bp.cur, 0, (c.hidden as u64) * 4);
}
ctx.queue.submit([enc.finish()]);
}
let mut out = Vec::new();
for step in bp.plan.iter() {
let Step::D {
pl,
bg,
gx,
gy,
tag,
} = step
else {
continue;
};
let scope = ctx.device.push_error_scope(wgpu::ErrorFilter::Validation);
let mut enc = ctx.device.create_command_encoder(&Default::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
ctx.queue.submit([enc.finish()]);
let verr = scope
.pop()
.await
.map(|e| e.to_string())
.unwrap_or_default();
let v = ctx.read_async(&bp.cur, c.hidden).await?;
let finite = v.iter().all(|x| x.is_finite());
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
let l2 = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt();
let nq = l2(&ctx
.read_async(&bp.qn, ((bp.qn.size() / 4) as usize).min(65536))
.await?);
let na = l2(&ctx
.read_async(&bp.attn_o, ((bp.attn_o.size() / 4) as usize).min(65536))
.await?);
let nv = l2(&ctx
.read_async(&bp.qkv, ((bp.qkv.size() / 4) as usize).min(65536))
.await?);
out.push((*tag, [n, nq, na, nv], finite, verr));
}
Ok(out)
}
pub fn stage_step_submit(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
hidden_in: Option<&[f32]>,
cur_ready: bool,
) -> Result<()> {
let t_prep0 = wclock::Instant::now();
assert!(cols.len() <= bp.k, "batch overflow");
let c = &self.w.cfg;
assert_eq!(
hidden_in.is_none(),
c.stage_first || cur_ready,
"stage 0 takes tokens; later stages take the residual stream"
);
let hd = c.head_dim;
// Global-theta tables serve the full-attention layers — the WIDE heads on an edge model.
let rope = RopeStager::new(c, c.head_dim_full());
let trash_row = (self.dn_slots - 1) as u32;
let mut fed = vec![0u32; bp.k];
// Every chained round, not just round 0: the gather indexes this by `cnt[0]*stride + col`
// exactly as it does `fed`. Rounds >= 1 carry GENERATED tokens, which are always text, so
// they stay -1 — but they must be WRITTEN -1, not left as whatever the buffer held.
let mut isrc = vec![-1i32; (CHAIN_MAX + 1) * bp.k];
let mut cm = vec![0u32; bp.k * 2];
let mut coss = Vec::with_capacity(bp.k * hd);
let mut sins = Vec::with_capacity(bp.k * hd);
let mut coss_l = Vec::new();
let mut sins_l = Vec::new();
let local = matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
);
for p in 0..bp.k {
let (tok, pos, row, mpos) = match cols.get(p) {
Some(col) => (col.token, col.pos, col.btrow, col.mpos),
None => (0, 0, trash_row, [0; 3]),
};
fed[p] = tok;
isrc[p] = cols.get(p).map_or(-1, |col| col.embed_row);
cm[p * 2] = pos;
cm[p * 2 + 1] = row;
let (cs, sn) = rope.col(mpos, pos);
coss.extend_from_slice(&cs);
sins.extend_from_slice(&sn);
if local {
let (cl, sl) = rope_cs(c.head_dim_sliding(), c.rope_local_theta, pos as usize);
coss_l.extend_from_slice(&cl);
sins_l.extend_from_slice(&sl);
}
}
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
// `OSFKB_CHOSEN_CANARY=1`: pre-fill the chosen slots with a sentinel — a readback slot
// still holding it means "argmax merge never wrote this column" (vs a genuine id-0 pick
// from dead/NaN logits). Diagnostic only.
if c.stage_last && std::env::var("OSFKB_CHOSEN_CANARY").is_ok() {
ctx.queue.write_buffer(
&bp.chosen_k,
0,
bytemuck::cast_slice(&vec![0xDEAD_BEEFu32; bp.k]),
);
}
ctx.queue
.write_buffer(&bp.fed, 0, bytemuck::cast_slice(&fed));
ctx.queue
.write_buffer(&bp.isrc, 0, bytemuck::cast_slice(&isrc));
ctx.queue
.write_buffer(&bp.cmeta_k, 0, bytemuck::cast_slice(&cm));
ctx.queue
.write_buffer(&bp.csk, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&bp.snk, 0, bytemuck::cast_slice(&sins));
if local {
ctx.queue
.write_buffer(&bp.csk_l, 0, bytemuck::cast_slice(&coss_l));
ctx.queue
.write_buffer(&bp.snk_l, 0, bytemuck::cast_slice(&sins_l));
}
if let Some(hidden) = hidden_in {
anyhow::ensure!(
hidden.len() == cols.len() * c.hidden,
"hidden hop size: got {}, want {}",
hidden.len(),
cols.len() * c.hidden
);
ctx.queue
.write_buffer(&bp.cur, 0, bytemuck::cast_slice(hidden));
}
let t_prep = t_prep0.elapsed().as_secs_f64();
let t_enc0 = wclock::Instant::now();
let use_gather = c.stage_first && bp.gather.is_some();
// `OSFKB_VISION_TRACE=1`: which embedding source each step used. The bug this found was a
// gather reading PAST the initialized region of `isrc` on chained decode rounds, so the
// per-step image-column count is exactly the thing worth being able to see.
if std::env::var("OSFKB_VISION_TRACE").is_ok() {
let img_cols = cols.iter().filter(|x| x.embed_row >= 0).count();
eprintln!(
"[vision] arena={} rows gather={} cols={} image_cols={img_cols}",
self.img_rows,
if use_gather { "yes" } else { "row-copies" },
cols.len(),
);
}
if use_gather {
let toks: Vec<u32> = cols.iter().map(|col| col.token).collect();
ctx.queue
.write_buffer(&bp.fed, 0, bytemuck::cast_slice(&toks));
// The image-placeholder columns' rows in the vision arena. All `-1` for text, which is
// what a text-only model's gather ignores entirely (it does not bind this).
let src: Vec<i32> = cols.iter().map(|col| col.embed_row).collect();
ctx.queue
.write_buffer(&bp.isrc, 0, bytemuck::cast_slice(&src));
}
// The need_logit pattern as a u64 bitmask — columns past bit 63 are dropped from the
// mask (a `1 << 64` shift would panic in debug and WRAP in release, silently aliasing
// keys), so reuse below is additionally capped at 64 columns.
let mask: u64 = cols
.iter()
.enumerate()
.filter(|(i, col)| *i < 64 && col.need_logit)
.fold(0, |m, (i, _)| m | (1 << i));
// Content-static stages: everything per-round arrives via `write_buffer` (hidden, fed,
// cmeta, rope tables), so the command buffer is reusable across identical
// (ncols, need_logit) rounds on EVERY stage except the stage-first copy fallback,
// which bakes token offsets into buffer-copy commands. Wider-than-64-column rounds are
// never cached: their need_logit pattern doesn't fit the key.
let reusable =
self.cmdbuf_reuse && (!c.stage_first || bp.gather.is_some()) && cols.len() <= 64;
let key = (bp.serial, cols.len() as u32, mask);
let cached = if reusable {
bp.cmdbuf
.borrow_mut()
.take()
.filter(|(k, _)| *k == key)
.map(|(_, cb)| cb)
} else {
None
};
let hit = cached.is_some();
let cb = match cached {
Some(cb) => cb,
None => self.encode_step_cmdbuf(ctx, bp, cols, use_gather),
};
ctx.queue.submit([cb]);
if hit {
self.cmdbuf_hits.set(self.cmdbuf_hits.get() + 1);
}
// Pre-encode the next identical round during this round's GPU window.
if reusable {
let next = self.encode_step_cmdbuf(ctx, bp, cols, use_gather);
*bp.cmdbuf.borrow_mut() = Some((key, next));
}
let t_enc = t_enc0.elapsed().as_secs_f64();
*self.step_breakdown.borrow_mut() = (t_prep, t_enc, wclock::Instant::now());
Ok(())
}
/// The submitted stage step's output staging buffer + its byte length (tokens on the last
/// stage, else the residual stream).
fn stage_out_staging<'a>(
&self,
bp: &'a BatchPlan,
cols_len: usize,
) -> (&'a wgpu::Buffer, usize) {
let c = &self.w.cfg;
if c.stage_last {
(&bp.chosen_staging, bp.k * 4)
} else {
(&bp.hidden_staging, cols_len * c.hidden * 4)
}
}
/// Decode a mapped staging slice into the stage output. Caller holds the map alive.
fn stage_out_from_mapped(&self, mapped: &[u8]) -> StageBatchOut {
if self.w.cfg.stage_last {
StageBatchOut::Tokens(bytemuck::cast_slice(mapped).to_vec())
} else {
StageBatchOut::Hidden(bytemuck::cast_slice(mapped).to_vec())
}
}
/// Read back a submitted stage step's output — the NATIVE blocking path (`map_async` +
/// `device.poll(Wait)`). The wasm sibling is [`Self::read_stage_out_async`].
pub fn read_stage_out(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols_len: usize,
) -> Result<StageBatchOut> {
let (staging, bytes) = self.stage_out_staging(bp, cols_len);
let slice = staging.slice(..bytes as u64);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.map_err(|e| anyhow::anyhow!("stage staging: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let out = self.stage_out_from_mapped(&slice.get_mapped_range().expect("mapped range"));
staging.unmap();
Ok(out)
}
/// Read back a submitted stage step's output — the wasm/browser path. WebGPU has no
/// blocking poll, so this AWAITS the `map_async` completion (the callback fires as the JS
/// event loop drives the WebGPU queue). Same bytes/decoding as the native path.
#[cfg(target_arch = "wasm32")]
pub async fn read_stage_out_async(
&self,
bp: &BatchPlan,
cols_len: usize,
) -> Result<StageBatchOut> {
let (staging, bytes) = self.stage_out_staging(bp, cols_len);
let slice = staging.slice(..bytes as u64);
let (tx, rx) = futures_channel::oneshot::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
rx.await
.map_err(|e| anyhow::anyhow!("stage map canceled: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let out = self.stage_out_from_mapped(&slice.get_mapped_range().expect("mapped range"));
staging.unmap();
Ok(out)
}
fn batch_stage_step_inner(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
hidden_in: Option<&[f32]>,
cur_ready: bool,
) -> Result<StageBatchOut> {
self.stage_step_submit(ctx, bp, cols, hidden_in, cur_ready)?;
self.read_stage_out(ctx, bp, cols.len())
}
/// Chained decode: run `steps` batch steps in ONE submit with on-GPU token feedback (argmax →
/// next gather via a buffer copy; per-step metadata pre-uploaded and indexed by the bumped
/// step counter). Every column must be a DECODE column (`cols[i]` at its current position);
/// returns `steps × k` chosen tokens (`out[s][col]`). Removes the per-step CPU sync that
/// dominates unchained batch stepping — the batch analogue of the solo loop's submit chain.
pub fn batch_step_chained(
&self,
ctx: &GpuCtx,
bp: &BatchPlan,
cols: &[BatchCol],
steps: usize,
) -> Result<Vec<Vec<u32>>> {
assert!(cols.len() <= bp.k, "batch overflow");
assert!((1..=CHAIN_MAX).contains(&steps), "chain length");
let c = &self.w.cfg;
let hd = c.head_dim;
let trash_row = (self.dn_slots - 1) as u32;
let k = bp.k;
let mut fed0 = vec![0u32; k];
let mut cm = vec![0u32; steps * k * 2];
let mut coss = Vec::with_capacity(steps * k * hd);
let mut sins = Vec::with_capacity(steps * k * hd);
let mut coss_l = Vec::new();
let mut sins_l = Vec::new();
let local = matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
);
// Global-theta tables serve the full-attention layers — the WIDE heads on an edge model.
let rope = RopeStager::new(c, c.head_dim_full());
for st in 0..steps {
for p in 0..k {
let (tok, pos, row, mpos) = match cols.get(p) {
Some(col) => (
col.token,
col.pos + st as u32,
col.btrow,
// Chained rounds advance the position by `st`. A chained column is always a
// DECODE column, whose 3-D position is uniform, so this walks all three axes
// together — which is what generation after an image does.
col.mpos.map(|v| v + st as u32),
),
None => (0, 0, trash_row, [0; 3]),
};
if st == 0 {
fed0[p] = tok;
}
cm[(st * k + p) * 2] = pos;
cm[(st * k + p) * 2 + 1] = row;
let (cs, sn) = rope.col(mpos, pos);
coss.extend_from_slice(&cs);
sins.extend_from_slice(&sn);
if local {
let (cl, sl) = rope_cs(c.head_dim_sliding(), c.rope_local_theta, pos as usize);
coss_l.extend_from_slice(&cl);
sins_l.extend_from_slice(&sl);
}
}
}
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
ctx.queue
.write_buffer(&bp.fed, 0, bytemuck::cast_slice(&fed0));
// Chained rounds are DECODE columns — generated text, never image placeholders. Written for
// EVERY round (see GATHER_K_MM): an undefined round reads 0, i.e. "image row 0".
ctx.queue.write_buffer(
&bp.isrc,
0,
bytemuck::cast_slice(&vec![-1i32; (CHAIN_MAX + 1) * bp.k]),
);
ctx.queue
.write_buffer(&bp.cmeta_k, 0, bytemuck::cast_slice(&cm));
ctx.queue
.write_buffer(&bp.csk, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&bp.snk, 0, bytemuck::cast_slice(&sins));
if local {
ctx.queue
.write_buffer(&bp.csk_l, 0, bytemuck::cast_slice(&coss_l));
ctx.queue
.write_buffer(&bp.snk_l, 0, bytemuck::cast_slice(&sins_l));
}
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for st in 0..steps {
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
let (gather_pl, gather_bg) = bp
.gather
.as_ref()
.expect("chained decode needs the on-GPU gather (embed exceeds binding cap)");
p.set_pipeline(gather_pl);
p.set_bind_group(0, gather_bg, &[]);
p.dispatch_workgroups((c.hidden as u32).div_ceil(64), k as u32, 1);
for step in &bp.plan {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
} else if let Step::Di {
pl,
bg,
args,
offset,
..
} = step
{
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
}
let (lm_bg, lm_gx, lm_gy) = bp.lm.as_ref().expect("chained decode has a head");
p.set_pipeline(&bp.lm_pl);
p.set_bind_group(0, lm_bg, &[]);
p.dispatch_workgroups(*lm_gx, *lm_gy, 1);
for hh in &bp.heads {
p.set_pipeline(&bp.block_pl);
p.set_bind_group(0, &hh.block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&bp.merge_pl);
p.set_bind_group(0, &hh.merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
p.set_pipeline(&bp.bump_pl);
p.set_bind_group(0, &bp.bump_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
// This step's chosen tokens feed the next step's gather.
if st + 1 < steps {
enc.copy_buffer_to_buffer(
&bp.chosen_k,
(st * k * 4) as u64,
&bp.fed,
((st + 1) * k * 4) as u64,
(k * 4) as u64,
);
}
}
enc.copy_buffer_to_buffer(
&bp.chosen_k,
0,
&bp.chosen_staging,
0,
(steps * k * 4) as u64,
);
ctx.queue.submit([enc.finish()]);
let slice = bp.chosen_staging.slice(..(steps * k * 4) as u64);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.map_err(|e| anyhow::anyhow!("chosen staging: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let flat: Vec<u32> =
bytemuck::cast_slice(&slice.get_mapped_range().expect("mapped range")).to_vec();
bp.chosen_staging.unmap();
Ok(flat.chunks(k).map(|ch| ch.to_vec()).collect())
}
/// Read one batch column's logits back (temperature sampling for that request) — copies only
/// that column's slice, not the whole `[k, vocab]` buffer.
pub fn read_batch_logits(&self, ctx: &GpuCtx, bp: &BatchPlan, col: usize) -> Result<Vec<f32>> {
let vocab = self.w.cfg.vocab;
let staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("logits_col_staging"),
size: (vocab * 4) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
enc.copy_buffer_to_buffer(
&bp.logits_k,
(col * vocab * 4) as u64,
&staging,
0,
(vocab * 4) as u64,
);
ctx.queue.submit([enc.finish()]);
let slice = staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.map_err(|e| anyhow::anyhow!("logits staging: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let out: Vec<f32> =
bytemuck::cast_slice(&slice.get_mapped_range().expect("mapped range")).to_vec();
staging.unmap();
Ok(out)
}
/// Overwrite one sequence slot's block-table row (`row` ∈ 1..=MAX_SLOTS).
pub fn write_btab_row(&self, ctx: &GpuCtx, row: u32, blocks: &[u32]) {
assert!(
(1..=MAX_SLOTS as u32).contains(&row),
"btab row out of range"
);
assert!(blocks.len() <= PAGE_ROW, "block table overflow");
ctx.queue.write_buffer(
&self.btab,
(row as u64) * (PAGE_ROW as u64) * 4,
bytemuck::cast_slice(blocks),
);
}
/// Number of KV pool blocks (the last one is the trash block — never allocate it).
pub fn pool_blocks(&self) -> usize {
self.pool_tokens / PAGE_BLK
}
/// Per-sequence (prompt + generated) token cap the scheduler enforces (from [`EngineOpts`]).
pub fn max_seq_tokens(&self) -> usize {
self.max_seq_tokens
}
/// Which GEMV family the M=1 plan carries: `"lcpp-split"` (guaranteed-32-subgroup Vulkan —
/// the batch plans' family, norm-split), `"v3-sg"` (Metal subgroup family), or `"tree"`
/// (portable workgroup-tree fallback). Solo, constrained, and per-token forward all replay
/// this plan, so the family is a property of the engine instance, not the call.
pub fn m1_family(&self) -> &'static str {
self.m1_family
}
/// Opt this engine's LM head into the fast small-width kernel, for SINGLE-STREAM decoding only.
///
/// Measured on the 35B-A3B (V100S, locked clocks, quiet box): the head is 30% of the verify and
/// the shipped `tiled` kernel runs it at 248 GB/s, while `lcpp` runs it at **803 GB/s** — 71% of
/// HBM peak. End-to-end that is **71.5 → 88.7 tok/s duo mono (+24%)**. The fast kernel is
/// compiled on every 32-subgroup adapter and simply never selected, because nothing in-tree sets
/// `OSFKB_HEAD_NBAR` (`bench/P7_the_lm_head_is_the_cost.md`).
///
/// **Single-stream ONLY.** `lcpp` is not bitwise-equal to `tiled` (different accumulation order;
/// `tests/head_bands.rs` gates it on tolerance for that reason). Choosing the head by `ncols`
/// would make a sequence's logits differ between solo and batched decoding — the exact bitwise
/// batch-invariance this engine guarantees. A single stream has no batch to be invariant with,
/// so it may take the fast kernel; batched serving may not, until `lcpp`'s reduction is made
/// width-independent and gated bitwise.
///
/// Call BEFORE `make_batch_plan*`: the plan clones the selected pipeline at build time.
pub fn set_single_stream_head(&mut self, on: bool) {
self.q4lm.set_single_stream(on);
}
/// Pre-encoded command-buffer cache hits across all batch plans (see `BatchPlan::cmdbuf`).
pub fn cmdbuf_hits(&self) -> u64 {
self.cmdbuf_hits.get()
}
/// Turn P2 command-buffer reuse off (or back on) for this engine; the `OSFKB_CMDBUF` kill
/// switch only sets the initial value.
///
/// Exposed because reuse has no *output* oracle otherwise: replaying a cached command buffer
/// and comparing it against ITSELF cannot catch a cmdbuf that bakes in a stale offset — every
/// replay would be identically wrong. Disabling reuse re-encodes from scratch each round,
/// which IS the oracle, and running both ways in one process is what lets the gate compare
/// them bitwise.
pub fn set_cmdbuf_reuse(&mut self, on: bool) {
self.cmdbuf_reuse = on;
}
/// Test-only bisection hook: upload the batch, then replay gather + the first `nsteps`
/// spec-plan steps; returns the k-column residual stream (`k·hidden`). Pairs with
/// [`Self::m1_steps_for_tests`] to locate the first diverging plan step.
#[doc(hidden)]
pub fn spec_steps_for_tests(
&self,
ctx: &GpuCtx,
sp: &SpecPlan,
fed: &[u32],
pos0: usize,
nsteps: usize,
) -> Result<Vec<f32>> {
self.upload_spec_batch(ctx, sp, fed, pos0);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(&sp.gather_pl);
p.set_bind_group(0, &sp.gather_bg, &[]);
p.dispatch_workgroups((self.w.cfg.hidden as u32).div_ceil(64), sp.k as u32, 1);
for step in sp.plan.iter().take(nsteps) {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
}
ctx.queue.submit([enc.finish()]);
ctx.read(&sp.spec_cur, sp.k * self.w.cfg.hidden)
}
/// Test-only: solo M=1 analogue of [`Self::probe_col0_norms`] — per-step (tag, l2, finite)
/// of `self.s.cur` for one (token, pos). Layer-aligned norms let a caller diff the batch
/// residual against the solo one and localize the first diverging layer.
#[doc(hidden)]
pub fn probe_m1_norms(
&self,
ctx: &GpuCtx,
token: u32,
pos: usize,
) -> Result<Vec<(&'static str, f32, bool)>> {
let c = &self.w.cfg;
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
{
let mut enc = ctx.device.create_command_encoder(&Default::default());
let (esrc, eoff) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(esrc, eoff, &self.s.cur, 0, (c.hidden as u64) * 4);
ctx.queue.submit([enc.finish()]);
}
let mut out = Vec::new();
for step in self.plan.iter() {
let Step::D {
pl,
bg,
gx,
gy,
tag,
} = step
else {
continue;
};
let mut enc = ctx.device.create_command_encoder(&Default::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
ctx.queue.submit([enc.finish()]);
let v = ctx.read(&self.s.cur, c.hidden)?;
let finite = v.iter().all(|x| x.is_finite());
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
out.push((*tag, n, finite));
}
Ok(out)
}
/// Test-only bisection hook: gather `(token, pos)` then replay the first `nsteps` M=1 plan
/// steps (skipping `Step::Embed`, which the explicit copy here replaces); returns `cur`.
#[doc(hidden)]
pub fn m1_steps_for_tests(
&self,
ctx: &GpuCtx,
token: u32,
pos: usize,
nsteps: usize,
) -> Result<Vec<f32>> {
let c = &self.w.cfg;
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
let (esrc, eoff) = self.w.embed_row_src(token);
enc.copy_buffer_to_buffer(esrc, eoff, &self.s.cur, 0, (c.hidden as u64) * 4);
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
for step in self
.plan
.iter()
.filter(|s| matches!(s, Step::D { .. }))
.take(nsteps)
{
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
}
ctx.queue.submit([enc.finish()]);
ctx.read(&self.s.cur, c.hidden)
}
/// Test-only: upload `fed` + rope tables for `pos0` and replay ONLY the batched forward
/// (gather + k-column plan) once, returning the k-column final norm (`hnorm`, `k·hidden`).
/// The bitwise gate compares each column against a sequential M=1 forward's `hnorm`.
#[doc(hidden)]
pub fn spec_forward_for_tests(
&self,
ctx: &GpuCtx,
sp: &SpecPlan,
fed: &[u32],
pos0: usize,
) -> Result<Vec<f32>> {
anyhow::ensure!(fed.len() == sp.k, "fed must have k tokens");
self.upload_spec_batch(ctx, sp, fed, pos0);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
self.encode_spec_forward(&mut p, sp);
}
ctx.queue.submit([enc.finish()]);
ctx.read(&sp.hnorm, sp.k * self.w.cfg.hidden)
}
/// Test-only: the M=1 final-norm buffer (after a [`Self::forward`]) for bitwise comparison.
#[doc(hidden)]
pub fn read_hnorm_for_tests(&self, ctx: &GpuCtx) -> Result<Vec<f32>> {
ctx.read(&self.s.hnorm, self.w.cfg.hidden)
}
/// Test-only: TurboQuant probe — the f32 K/V vectors for (layer, pos, kv-head) alongside the
/// quantized nibbles + scales the GPU wrote, so a test can compare against the CPU quantizer.
#[doc(hidden)]
#[allow(clippy::type_complexity)]
pub fn tq_probe_for_tests(
&self,
ctx: &GpuCtx,
li: usize,
pos: usize,
head: usize,
) -> Result<Option<(Vec<f32>, Vec<f32>, Vec<u32>, Vec<f32>, Vec<u32>)>> {
let Some(t) = &self.tq else { return Ok(None) };
let Some((kvq, kvs, ksig)) = &t.caches[li] else {
return Ok(None);
};
let (hd, nkv) = (self.w.cfg.head_dim, self.w.cfg.n_kv_heads);
let kf = ctx.read_f16(&self.k_cache[li], MAX_T * nkv * hd)?;
let vf = ctx.read_f16(&self.v_cache[li], MAX_T * nkv * hd)?;
let base = pos * nkv * hd + head * hd;
let qwords = ctx.read_u32(kvq, MAX_T * nkv * 2 * (hd / 8))?;
let scales = ctx.read(kvs, MAX_T * nkv * 3)?;
let qbase = ((pos * nkv + head) * 2) * (hd / 8);
Ok(Some((
kf[base..base + hd].to_vec(),
vf[base..base + hd].to_vec(),
qwords[qbase..qbase + 2 * (hd / 8)].to_vec(),
scales[(pos * nkv + head) * 3..(pos * nkv + head) * 3 + 3].to_vec(),
{
let sw = ctx.read_u32(ksig, MAX_T * nkv * (hd / 32))?;
sw[(pos * nkv + head) * (hd / 32)..(pos * nkv + head + 1) * (hd / 32)].to_vec()
},
)))
}
/// Test-only: read a named M=1 scratch buffer (debug bisection).
#[doc(hidden)]
pub fn read_scratch_for_tests(&self, ctx: &GpuCtx, name: &str) -> Result<Vec<f32>> {
let c = &self.w.cfg;
let (buf, len) = match name {
"gate" => (&self.s.gate, c.intermediate),
"mlpout" => (&self.s.mlpout, c.hidden),
"opout" => (&self.s.opout, c.hidden),
"qkv" => (&self.s.qkv, (c.n_heads + 2 * c.n_kv_heads) * c.head_dim),
"attn" => (&self.s.attn, c.n_heads * c.head_dim),
"normed" => (&self.s.normed, c.hidden),
"moe_sel" => (&self.s.moe_sel, c.top_k_experts.max(1)),
"moe_wsel" => (&self.s.moe_wsel, c.top_k_experts.max(1)),
"moe_gate" => (&self.s.moe_gate, c.moe_intermediate.max(1)),
"hnorm" => (&self.s.hnorm, c.hidden),
_ => anyhow::bail!("unknown scratch {name}"),
};
ctx.read(buf, len)
}
/// Upload the per-batch inputs: fed tokens, per-position RoPE tables, and the shared
/// position/attention uniform (`kp.w = pos0+1`; column p derives its own T/pos from it).
fn upload_spec_batch(&self, ctx: &GpuCtx, sp: &SpecPlan, fed: &[u32], pos0: usize) {
let c = &self.w.cfg;
let hd = c.head_dim;
ctx.queue
.write_buffer(&sp.fed, 0, bytemuck::cast_slice(fed));
let mut coss = Vec::with_capacity(sp.k * hd);
let mut sins = Vec::with_capacity(sp.k * hd);
for p in 0..sp.k {
let (cs, sn) = rope_cs_padded(hd, c.rotary_dim, c.rope_theta, pos0 + p);
coss.extend_from_slice(&cs);
sins.extend_from_slice(&sn);
}
ctx.queue
.write_buffer(&sp.csk, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&sp.snk, 0, bytemuck::cast_slice(&sins));
if matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
) {
let mut cl = Vec::with_capacity(sp.k * hd);
let mut sl = Vec::with_capacity(sp.k * hd);
for p in 0..sp.k {
let (cs, sn) = rope_cs(c.head_dim_sliding(), c.rope_local_theta, pos0 + p);
cl.extend_from_slice(&cs);
sl.extend_from_slice(&sn);
}
ctx.queue
.write_buffer(&sp.csk_l, 0, bytemuck::cast_slice(&cl));
ctx.queue
.write_buffer(&sp.snk_l, 0, bytemuck::cast_slice(&sl));
}
ctx.queue.write_buffer(
&sp.attn_tk,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
hd as u32,
(pos0 + 1) as u32,
)),
);
let cm: Vec<u32> = (0..sp.k).flat_map(|p| [(pos0 + p) as u32, 0u32]).collect();
ctx.queue
.write_buffer(&sp.cmeta_k, 0, bytemuck::cast_slice(&cm));
}
/// Record the batched forward (gather + k-column plan) into an open pass.
fn encode_spec_forward(&self, p: &mut wgpu::ComputePass<'_>, sp: &SpecPlan) {
p.set_pipeline(&sp.gather_pl);
p.set_bind_group(0, &sp.gather_bg, &[]);
p.dispatch_workgroups((self.w.cfg.hidden as u32).div_ceil(64), sp.k as u32, 1);
for step in &sp.plan {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
}
/// Speculative grammar-constrained decode: verify k drafted tokens per batched plan replay.
///
/// Per batch: (a) the pre-pass walks the on-GPU grammar FSM through the k draft tokens —
/// snapshotting the state per column and building each column's compacted allowed-id list; (b)
/// ONE k-column forward (weights read once per batch, not once per token); (c) per column, the
/// grammar-sparse head + lowest-id argmax → `chosen[p]`. One 4·k-byte readback then accepts the
/// longest draft prefix (`chosen[p] == draft[p]`) plus the first correction — every batch emits
/// 1..=k tokens that are BYTE-IDENTICAL to what the sequential constrained loop would emit
/// (per-column kernels are bitwise-equal and the pick logic is unchanged). On a partial accept,
/// the next batch first restores the FSM-state / conv-ring snapshots of the accepted point (the
/// KV cache needs no rollback — stale positions are positionally overwritten).
///
/// `draft(all_tokens) -> k drafts`; `sg` carries the grammar's pipelines/buffers (the FSM state
/// must be freshly seeded by the caller, and the prompt prefilled via [`Self::prefill_cached`]
/// when `prefilled`). Returns `(generated, decode_tok_s, prefill_ms, draft_accept_rate)`.
#[allow(clippy::too_many_arguments)]
pub fn decode_constrained_spec(
&self,
ctx: &GpuCtx,
sp: &SpecPlan,
prompt: &[u32],
ngen: usize,
sg: &SpecGrammar<'_>,
draft: &mut dyn FnMut(&[u32]) -> Vec<u32>,
prefilled: bool,
) -> Result<(Vec<u32>, f64, f64, f64)> {
anyhow::ensure!(!prompt.is_empty(), "spec decode needs a non-empty prompt");
anyhow::ensure!(
sg.snap_adv_bgs.len() == sp.k,
"snap_adv_bgs must have one entry per draft column"
);
if !prefilled {
self.reset(ctx);
// Build the prompt KV/conv state with M=1 forwards (positions 0..n-2; the last prompt
// token is fed by the first batch's column 0).
for (pos, &tok) in prompt[..prompt.len() - 1].iter().enumerate() {
self.forward_only(ctx, tok, pos);
}
}
let k = sp.k;
let dp = self.decode_pipelines(ctx);
let mut all: Vec<u32> = prompt.to_vec();
let mut out: Vec<u32> = Vec::new();
let mut pending_fix: Option<usize> = None;
let (mut drafted, mut accepted) = (0usize, 0usize);
let t = wclock::Instant::now();
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
let prefill_secs = t.elapsed().as_secs_f64();
while out.len() < ngen {
let pos0 = all.len() - 1;
anyhow::ensure!(
pos0 + k <= 512,
"spec batch would exceed the attention score cap (T ≤ 512)"
);
// An EMPTY draft list is the drafter's stop signal (it tracks the CPU grammar and knows
// the query is complete/dead) — the sequential loop instead burns its whole budget on
// post-Done sentinel tokens that the caller's re-walk discards.
let mut ds = draft(&all);
if ds.is_empty() {
break;
}
// Columns beyond the real drafts are fed PAD tokens whose continuations are garbage —
// verification and emission must never look past real+1 columns (chosen[real] is the
// last one fed a real token).
let real = ds.len().min(k);
ds.resize(k, 1);
let mut fed = Vec::with_capacity(k);
fed.push(all[pos0]);
fed.extend_from_slice(&ds[..k - 1]);
self.upload_spec_batch(ctx, sp, &fed, pos0);
ctx.queue
.write_buffer(&sp.drafts, 0, bytemuck::cast_slice(&ds));
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if let Some(j) = pending_fix.take() {
// Roll back to the accepted point: FSM state ← snap[j], conv rings ← snap slot j,
// then advance the FSM by the correction token (chosen[j] from the last batch).
enc.copy_buffer_to_buffer(
sg.snap_buf,
j as u64 * sg.state_bytes,
sg.state_buf,
0,
sg.state_bytes,
);
enc.copy_buffer_to_buffer(&sp.chosen_k, j as u64 * 4, sg.chosen_in, 0, 4);
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(sg.advance_pl);
p.set_bind_group(0, sg.advance_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
let stride = self.conv_bytes();
for (li, snap) in &sp.conv_snapk {
enc.copy_buffer_to_buffer(
snap,
j as u64 * stride,
&self.conv_state[*li],
0,
stride,
);
}
}
{
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
// Pre-pass: the k 1-thread snap_advance steps lay down the per-column FSM snapshot
// chain, ONE batched mask dispatch fills all k mask columns from those snapshots,
// then each column compacts its slice and freezes its indirect grid.
for bg in sg.snap_adv_bgs {
p.set_pipeline(sg.snap_adv_pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
p.set_pipeline(sg.mask_k_pl);
p.set_bind_group(0, sg.mask_k_bg, &[]);
p.dispatch_workgroups(sg.mask_wg, k as u32, 1);
for pick in &sp.picks {
p.set_pipeline(&dp.compact);
p.set_bind_group(0, &pick.compact_bg, &[]);
p.dispatch_workgroups((self.w.cfg.vocab as u32).div_ceil(256), 1, 1);
p.set_pipeline(&dp.sargs);
p.set_bind_group(0, &pick.sargs_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
self.encode_spec_forward(&mut p, sp);
for pick in &sp.picks {
pick.encode_pick_col(&mut p, dp);
}
}
// Readback staged INSIDE this submit — one submit + one map per batch, not two.
enc.copy_buffer_to_buffer(&sp.chosen_k, 0, &sp.chosen_staging, 0, (k * 4) as u64);
ctx.queue.submit([enc.finish()]);
let chosen: Vec<u32> = {
let slice = sp.chosen_staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.map_err(|e| anyhow::anyhow!("staging map: {e:?}"))?
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let data = slice.get_mapped_range().expect("mapped range");
let out: Vec<u32> = bytemuck::cast_slice(&data).to_vec();
drop(data);
sp.chosen_staging.unmap();
out
}; // the per-batch sync point
let mut j = 0usize;
while j < real && chosen[j] == ds[j] {
j += 1;
}
if std::env::var("OSFKB_SPEC_DEBUG").is_ok() {
eprintln!("[batch] pos0={pos0} real={real} ds={ds:?} chosen={chosen:?} j={j}");
}
drafted += real;
accepted += j;
// j == real (all real drafts accepted): chosen[real] is still a VALID continuation —
// column `real` was fed the last real draft — and the FSM fixup advances by it, so it
// MUST be emitted (dropping it desyncs the grammar state from the token stream).
let emitted = &chosen[..(j + 1).min(real + 1).min(k)];
all.extend_from_slice(emitted);
out.extend_from_slice(emitted);
if j < k {
pending_fix = Some(j);
}
}
out.truncate(ngen);
let secs = t.elapsed().as_secs_f64();
let decode_secs = (secs - prefill_secs).max(1e-9);
Ok((
// Rate over tokens actually EMITTED (the drafter's early stop means that is usually
// fewer than the budget — the sequential loop burns the full budget instead).
out.clone(),
out.len() as f64 / decode_secs,
prefill_secs * 1000.0,
accepted as f64 / drafted.max(1) as f64,
))
}
#[allow(clippy::too_many_arguments)]
pub fn decode_constrained(
&self,
ctx: &GpuCtx,
prompt: &[u32],
ngen: usize,
mask_buf: &wgpu::Buffer,
chosen_buf: &wgpu::Buffer,
mask_pl: &wgpu::ComputePipeline,
mask_bg: &wgpu::BindGroup,
mask_wg: u32,
advance_pl: &wgpu::ComputePipeline,
advance_bg: &wgpu::BindGroup,
prefilled: bool,
) -> Result<(Vec<u32>, f64, f64)> {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
// When `prefilled`, the caller already built the prompt's KV + conv state via a prefix-cached
// prefill (reusing the fixed instruction/example prefix shared across queries), so skip the
// reset — which would clobber it — and start the loop at the last prompt token below.
if !prefilled {
self.reset(ctx);
}
let total = prompt.len() + ngen;
// The grammar governs only the GENERATED suffix. For the prefill region the loop must build the
// prompt's KV (gather + plan) WITHOUT masking/argmax — otherwise it would overwrite the prompt
// tokens with the grammar's pick. `tokens[pos+1]` is generated iff `pos + 1 >= gen_start`.
let gen_start = prompt.len().max(1);
let _ts = wclock::Instant::now();
// Pipelines are engine-cached (compiling WGSL is ~150ms/shader on Metal; recompiling per call
// dominated the per-query wall-clock). The clone is an Arc refcount bump, NOT a recompile.
let dp = self.decode_pipelines(ctx);
let gather_pl = dp.gather.clone();
let tokens = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("tokens"),
size: (total * 4) as u64,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut init = vec![0u32; total];
init[..prompt.len()].copy_from_slice(prompt);
ctx.queue
.write_buffer(&tokens, 0, bytemuck::cast_slice(&init));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, prompt.len() as u32, c.vocab as u32, h as u32)),
);
let gather_bg = make_bg(
ctx,
&gather_pl,
&[&self.w.embed, &tokens, &self.s.cur],
&posu,
);
// The grammar-sparse pick: the mask (state-only, so it runs BEFORE the forward) is compacted
// into an allowed-id list, the tied lm_head reads ONLY those rows (indirect grid), and a
// lowest-id argmax writes `chosen` — the full-vocab head + full-vocab masked argmax are gone.
let sp = SparsePick::new(ctx, dp, &self.w, &self.s.hnorm, mask_buf, chosen_buf, &posu);
eprintln!(
"[dc] setup_ms(pipelines+buffers+bg)={}",
_ts.elapsed().as_millis()
);
let t = wclock::Instant::now();
// Split prefill vs decode wall-clock so the caller can report a *decode-only* rate — a short
// generation behind a long prompt is otherwise dominated by the one-time prompt prefill.
let mut prefill_secs = -1.0f64;
// With an external prefix-cached prefill, the prompt KV for positions `0..gen_start-1` is
// already built; start at the last prompt token (it produces the first output's logits).
let loop_start = if prefilled {
gen_start.saturating_sub(1)
} else {
0
};
for pos in loop_start..total - 1 {
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
ctx.queue.write_buffer(
&posu,
0,
bytemuck::cast_slice(&u32x4(
pos as u32,
prompt.len() as u32,
c.vocab as u32,
h as u32,
)),
);
// `tokens[pos+1]` is a generated token (grammar-masked) once we're past the prompt; before
// that we only build the prompt's KV and leave the prompt token in place.
let generating = pos + 1 >= gen_start;
if generating && prefill_secs < 0.0 {
// First generated token: flush the queued prefill GPU work and stamp its wall-clock so
// `secs - prefill_secs` below is the decode-loop time alone.
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
prefill_secs = t.elapsed().as_secs_f64();
}
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
// [if generating] grammar mask (from the FSM state advanced last iteration) →
// compact → indirect args, THEN the forward (gather + plan minus its full-vocab
// head, KV written by qkrc) → sparse head over the allowed rows → lowest-id argmax
// → chosen, all in one pass (Metal serializes the deps). Prefill positions skip the
// head too — their logits were never read.
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
if generating {
p.set_pipeline(mask_pl);
p.set_bind_group(0, mask_bg, &[]);
p.dispatch_workgroups(mask_wg, 1, 1);
sp.encode_compact(&mut p, dp, c.vocab as u32);
}
p.set_pipeline(&gather_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
// The plan's LAST step is the full-vocab Q4 lm_head (see the tail of `new`); the
// sparse pick replaces it here.
for step in &self.plan[..self.plan.len() - 1] {
if let Step::D { pl, bg, gx, gy, .. } = step {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
}
if generating {
sp.encode_pick(&mut p, dp);
}
}
if generating {
// chosen → tokens[pos+1] (next gather) — a transfer, so outside the pass.
enc.copy_buffer_to_buffer(chosen_buf, 0, &tokens, ((pos + 1) * 4) as u64, 4);
{
// advance the grammar FSM by the chosen token (reads chosen_buf).
let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
p.set_pipeline(advance_pl);
p.set_bind_group(0, advance_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
}
ctx.queue.submit([enc.finish()]);
}
let toks = ctx.read_u32(&tokens, total)?;
let secs = t.elapsed().as_secs_f64();
if prefill_secs < 0.0 {
prefill_secs = secs; // nothing generated (zero budget / immediate stop)
}
let decode_secs = (secs - prefill_secs).max(1e-9);
Ok((
toks[gen_start..].to_vec(),
ngen as f64 / decode_secs, // decode-only tok/s (prompt prefill excluded)
prefill_secs * 1000.0, // prompt-prefill wall-clock, in ms
))
}
/// GPU tracing: record one forward with per-compute-pass timestamps and print the GPU µs of each
/// pass (a "pass" = one run of dispatches between copies) + the total. Shows where the GPU time
/// actually goes (vs the ~3.5ms f32 bandwidth floor).
pub fn trace(&self, ctx: &GpuCtx) {
if !ctx.timestamps {
println!("(timestamp queries unavailable on this device)");
return;
}
// One pass per dispatch so each kernel gets its own timestamp (granular per-op GPU time).
let npass = self
.plan
.iter()
.filter(|s| matches!(s, Step::D { .. }))
.count() as u32;
let count = npass * 2;
let qs = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
label: Some("ts"),
ty: wgpu::QueryType::Timestamp,
count,
});
let resolve = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("ts-resolve"),
size: (count * 8) as u64,
usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
// warm
let mut e = ctx.device.create_command_encoder(&Default::default());
self.record(&mut e, 1, 10);
ctx.queue.submit([e.finish()]);
let _ = ctx.read(&self.s.logits, self.w.cfg.vocab);
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
let _ = h;
let (coss, sins) = rope_cs(hd, c.rope_theta, 10);
ctx.queue
.write_buffer(&self.s.cosb, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&self.s.sinb, 0, bytemuck::cast_slice(&sins));
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, 11)),
);
let mut enc = ctx.device.create_command_encoder(&Default::default());
let (mut i, mut p) = (0usize, 0u32);
while i < self.plan.len() {
match &self.plan[i] {
Step::Embed => {
let (esrc, eoff) = self.w.embed_row_src(1);
enc.copy_buffer_to_buffer(esrc, eoff, &self.s.cur, 0, (h * 4) as u64);
i += 1;
}
Step::D { pl, bg, gx, gy, .. } => {
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: None,
timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
query_set: &qs,
beginning_of_pass_write_index: Some(2 * p),
end_of_pass_write_index: Some(2 * p + 1),
}),
});
p += 1;
pass.set_pipeline(pl);
pass.set_bind_group(0, bg, &[]);
pass.dispatch_workgroups(*gx, *gy, 1);
i += 1;
}
Step::Di {
pl,
bg,
args,
offset,
..
} => {
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: None,
timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
query_set: &qs,
beginning_of_pass_write_index: Some(2 * p),
end_of_pass_write_index: Some(2 * p + 1),
}),
});
p += 1;
pass.set_pipeline(pl);
pass.set_bind_group(0, bg, &[]);
pass.dispatch_workgroups_indirect(args, *offset);
i += 1;
}
}
}
enc.resolve_query_set(&qs, 0..count, &resolve, 0);
let staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: None,
size: (count * 8) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
enc.copy_buffer_to_buffer(&resolve, 0, &staging, 0, (count * 8) as u64);
ctx.queue.submit([enc.finish()]);
let slice = staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv().unwrap().unwrap();
let ticks: Vec<u64> =
bytemuck::cast_slice::<u8, u64>(&slice.get_mapped_range().expect("mapped range"))
.to_vec();
let per = ctx.ts_period as f64 / 1000.0; // ns/tick → µs
let mut total = 0.0;
println!("--- GPU per-pass trace (pos=10) ---");
for pi in 0..npass as usize {
let us = ticks[2 * pi + 1].wrapping_sub(ticks[2 * pi]) as f64 * per;
total += us;
println!(" pass {pi:2}: {us:7.1} µs");
}
println!(
" TOTAL GPU: {total:.0} µs/forward ({:.0} tok/s pure GPU)",
1e6 / total
);
}
/// Diagnostic: time `n` plan replays in a SINGLE submit + one readback — isolates the pure GPU
/// per-forward cost from the per-token submit/sync overhead. Returns tok/s.
pub fn bench_batched(&self, ctx: &GpuCtx, n: usize) -> f64 {
let c = &self.w.cfg;
let (h, hd, nh, nkv) = (c.hidden, c.head_dim, c.n_heads, c.n_kv_heads);
let _ = h;
let (pos, token) = (10usize, 1u32);
self.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.attn_t,
0,
bytemuck::cast_slice(&u32x4(nh as u32, nkv as u32, hd as u32, (pos + 1) as u32)),
);
ctx.queue
.write_buffer(&self.cmeta1, 0, bytemuck::cast_slice(&[pos as u32, 0u32]));
let t0 = wclock::Instant::now();
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for _ in 0..n {
if self.w.cfg.edge.ple_dim > 0 {
// PLE gather input for this step (queue-ordered before the submit below).
ctx.queue
.write_buffer(&self.s.ple_tok, 0, bytemuck::cast_slice(&[token]));
}
self.record(&mut enc, token, pos);
}
ctx.queue.submit([enc.finish()]);
let _ = ctx.read(&self.s.logits, c.vocab);
n as f64 / t0.elapsed().as_secs_f64()
}
}
/// Cached per-`k` resources for the 1-SUBMIT mono draft chain (see [`MtpEngine::draft_chain`]).
/// The old per-step loop needed one submit per draft step because `queue.write_buffer` is
/// queue-ordered (a step's rope/attn/cmeta write would otherwise land before EVERY step's
/// dispatches); staging the per-step metas in round buffers and copying them into the plan's
/// buffers IN-ENCODER removes that constraint. Bind groups depend only on the step index (the
/// per-round position rides in the staged bytes), so they persist across rounds — and the whole
/// round body becomes round-invariant for a fixed `k`, letting [`MtpEngine::draft_chain`]
/// pre-encode the NEXT round's body during the current round's GPU window (`pending`).
struct DraftRound {
k: usize,
/// `[k, rope_words]` f32 — per-step cos|sin (+ the local-theta pair on Gemma) copied into
/// the plan's rope buffers between step records.
rope_stage: wgpu::Buffer,
rope_words: usize,
/// `[k, 8]` u32 — per-step `attn_t` vec4 + `cmeta1` pair (16 B + 8 B slices, 32 B stride).
meta_stage: wgpu::Buffer,
/// Persistent MAP_READ staging for the `tok` chain readback (one map per round; the copy
/// rides the round body, so no separate readback submit).
tok_staging: wgpu::Buffer,
/// Per-step bind groups: norm_h, gather, norm_e, fc_h, fc_e, am_block, am_merge.
steps: Vec<[wgpu::BindGroup; 7]>,
/// The next round's pre-encoded body (encoded while the previous round's GPU work ran).
pending: Option<wgpu::CommandBuffer>,
}
/// One serving draft-KV prefill entry (see [`MtpEngine::draft_pairs`]): the MAIN model's hidden
/// at `cur_col` merges with `token`'s embedding and the draft layer's KV lands at `pos + 1`.
/// `mpos` is the 3-D M-RoPE angle AT that KV position (`[pos + 1; 3]` for text; a multimodal
/// sequence passes the true vision-plan position so the draft ropes like the reference MTP).
#[derive(Clone, Copy, Debug)]
pub struct DraftPair {
pub btrow: u32,
pub cur_col: usize,
pub token: u32,
pub pos: usize,
pub mpos: [u32; 3],
}
/// One serving draft-chain request (see [`MtpEngine::draft_chain_batch_seeded`]): chain from
/// `first_token` emitted at KV position `pos0`; draft `i` occupies `pos0 + 1 + i`. `mpos0` is the
/// M-RoPE angle at `pos0` — drafts advance it uniformly (`+1+i` on every axis: generated tokens
/// are text, but on a multimodal sequence the angle base is `next_pos`-relative, NOT `pos0`).
#[derive(Clone, Copy, Debug)]
pub struct DraftReq {
pub btrow: u32,
pub first_token: u32,
pub pos0: usize,
pub mpos0: [u32; 3],
}
/// Self-speculative MTP draft engine: the checkpoint's `mtp.*` head run as a standalone
/// ONE-LAYER MODEL through the generic plan machinery (see [`Weights::mtp_engine_weights`] —
/// any architecture's draft head takes this path), plus the fc merge recorded in front of the
/// layer each draft step: `cur = fc_hidden·norm_h(prev) + fc_emb·norm_e(embed(token))`. The
/// chain state `prev` is the draft layer's own pre-norm residual stream (the MAIN model's for
/// the first draft), exactly the DeepSeek-V3 MTP recurrence the reference vLLM deployment runs.
pub struct MtpEngine {
/// The 1-layer draft model (own plan, own KV pool, shared embed/head buffers).
pub gpu: Lfm2Gpu,
fc_hidden: crate::weights::Q4,
fc_emb: crate::weights::Q4,
pre_norm_hidden: wgpu::Buffer,
pre_norm_emb: wgpu::Buffer,
embed: wgpu::Buffer,
/// Chain state: pre-norm residual stream feeding the next draft's fc merge —
/// SLOT-SLABBED `[1 + MAX_SLOTS, hidden]` (slab 0 = the single-stream/M=1 chain; serving
/// seeds slab `slot` from the accepted verify column).
prev_hidden: wgpu::Buffer,
/// Compact per-round staging `[k_batch, hidden]` (slot slabs gathered by column order).
prev_round: wgpu::Buffer,
toks_round: wgpu::Buffer,
/// Lazily-built batch plans for the 1-layer draft model: two widths (8, 32) — the shared
/// 248k-row head dominates a draft step and its column-tiles scale with the PLAN width,
/// so small rounds must not pay the 32-wide geometry.
bp_draft8: std::sync::OnceLock<BatchPlan>,
bp_draft32: std::sync::OnceLock<BatchPlan>,
h_normed: wgpu::Buffer,
e_raw: wgpu::Buffer,
e_normed: wgpu::Buffer,
tok: wgpu::Buffer,
gather_pl: wgpu::ComputePipeline,
gather_cols_pl: wgpu::ComputePipeline,
/// β-probe (`OSFKB_MTP_BETA_PROBE=1`): per-depth top-3 draft candidates of the LAST chain —
/// measures top-k coverage β_d(k), the input to the Sequoia tree DP. Probe-only readback.
last_top3: std::cell::RefCell<Vec<[u32; 3]>>,
argmax_block_pl: wgpu::ComputePipeline,
argmax_merge_pl: wgpu::ComputePipeline,
am_pv: wgpu::Buffer,
am_pi: wgpu::Buffer,
/// 1-submit chain resources, rebuilt when the requested depth changes (see [`DraftRound`]).
round: std::cell::RefCell<Option<DraftRound>>,
}
impl MtpEngine {
/// Build from a loaded main model; `None` when the checkpoint ships no `mtp.*` head.
/// `OSFKB_MTP_FR_VOCAB=<N>` enables frequency-ranked draft-head compression (see
/// [`Self::new_capped`]).
pub fn new(ctx: &GpuCtx, main: &Weights) -> Option<Self> {
let cap = std::env::var("OSFKB_MTP_FR_VOCAB")
.ok()
.and_then(|v| v.parse().ok());
Self::new_capped(ctx, main, cap)
}
/// [`Self::new`] with an explicit frequency-ranked (FR) draft-vocab cap: the DRAFT model's
/// vocab is clamped to the first `N` rows of the shared Q4 head ([`quantize_q4_0`]
/// (crate::weights::quantize_q4_0) is row-major, so a row prefix is a byte prefix — no
/// remap). BPE ids are merge-rank ≈ frequency ordered, so the prefix keeps the hot ids
/// while cutting the head read — the dominant traffic of a draft step — by `vocab/N`.
/// SOUND BY CONSTRUCTION: the verify model is untouched and every emitted token comes from
/// the verify argmax, so decode output stays bitwise plain-greedy; a draft id outside the
/// main argmax is simply rejected. Acceptance E is the only thing at stake — measure it.
pub fn new_capped(
ctx: &GpuCtx,
main: &Weights,
draft_vocab_cap: Option<usize>,
) -> Option<Self> {
Self::new_opts(ctx, main, draft_vocab_cap, EngineOpts::sized_from_env())
}
/// [`Self::new`] with the draft engine's KV pool matched to the MAIN engine's: the scheduler
/// mirrors its block tables into the draft KV, so the draft buffers must be addressable by
/// every block id the main allocator can hand out. Use this from serving whenever the main
/// engine was built with non-default [`EngineOpts`] (`--vram-gb`).
pub fn new_matching(ctx: &GpuCtx, main: &Lfm2Gpu) -> Option<Self> {
let cap = std::env::var("OSFKB_MTP_FR_VOCAB")
.ok()
.and_then(|v| v.parse().ok());
let opts = EngineOpts {
pool_tokens: main.pool_blocks() * PAGE_BLK,
max_seq_tokens: main.max_seq_tokens(),
};
Self::new_opts(ctx, &main.w, cap, opts)
}
fn new_opts(
ctx: &GpuCtx,
main: &Weights,
draft_vocab_cap: Option<usize>,
opts: EngineOpts,
) -> Option<Self> {
let mut w = main.mtp_engine_weights()?;
if let Some(n) = draft_vocab_cap {
w.cfg.vocab = w.cfg.vocab.min(n.max(1));
}
let m = main.mtp.as_ref().expect("mtp checked");
let h = main.cfg.hidden;
let (fc_hidden, fc_emb) = (m.fc_hidden.clone(), m.fc_emb.clone());
let (pre_norm_hidden, pre_norm_emb) = (m.pre_norm_hidden.clone(), m.pre_norm_emb.clone());
// The DRAFT model's gather table — the head's own embed when it ships one (GLM-OCR),
// else the shared main table (`mtp_engine_weights` resolved this already).
let embed = w.embed.clone();
let gpu = Lfm2Gpu::new_with_opts(ctx, w, opts);
Some(Self {
gpu,
fc_hidden,
fc_emb,
pre_norm_hidden,
pre_norm_emb,
embed,
prev_hidden: ctx.empty((1 + MAX_SLOTS) * h),
prev_round: ctx.empty(Self::K_DRAFT_BATCH * h),
toks_round: ctx.empty(Self::K_DRAFT_BATCH),
bp_draft8: std::sync::OnceLock::new(),
bp_draft32: std::sync::OnceLock::new(),
h_normed: ctx.empty(Self::K_DRAFT_BATCH * h),
e_raw: ctx.empty(Self::K_DRAFT_BATCH * h),
e_normed: ctx.empty(Self::K_DRAFT_BATCH * h),
// Chain buffer: tok[0] = seed, tok[i+1] = draft i (argmax writes, gather reads —
// the whole chain runs without an intermediate readback).
tok: ctx.empty(32),
gather_pl: pipeline(ctx, "mtp_gather", GATHER),
gather_cols_pl: pipeline(ctx, "mtp_gather_cols", GATHER_COLS),
last_top3: std::cell::RefCell::new(Vec::new()),
argmax_block_pl: pipeline(ctx, "mtp_argmax_block", ARGMAX_BLOCK),
argmax_merge_pl: pipeline(ctx, "mtp_argmax_merge", ARGMAX_MERGE),
am_pv: ctx.empty(256),
am_pi: ctx.empty(256),
round: std::cell::RefCell::new(None),
})
}
/// Chain `n` greedy drafts; the chain state must hold the MAIN model's pre-norm hidden at
/// the position that emitted `first_token` (seed via [`Self::write_prev_hidden`] or a
/// buffer copy into [`Self::prev_hidden_buf`]) (absolute position `pos0`); draft `i` occupies position
/// `pos0 + 1 + i` in the draft layer's own KV. Returns the drafted token ids. Rejected
/// positions need no cleanup — the next chain overwrites them positionally.
///
/// ONE submit + one map per round (see [`DraftRound`]): per-step metas ride staged buffers
/// copied in-encoder, and the round body is pre-encoded during the previous round's GPU
/// window. The β-probe (`OSFKB_MTP_BETA_PROBE=1`) needs per-step logits readbacks, so it
/// takes the old per-step-submit path ([`Self::draft_chain_probe`]) — measurement only.
pub fn draft_chain(
&self,
ctx: &GpuCtx,
first_token: u32,
pos0: usize,
n: usize,
) -> Result<Vec<u32>> {
if std::env::var("OSFKB_MTP_BETA_PROBE").is_ok() {
return self.draft_chain_probe(ctx, first_token, pos0, n);
}
self.draft_chain_fast(ctx, None, first_token, pos0, n)
}
/// [`Self::draft_chain`] with the chain seed folded into the SAME submit: column `col` of
/// `main_bp.cur` (the MAIN model's pre-norm residual at the position that emitted
/// `first_token`) is copied into the chain slab by a lead command buffer — this replaces
/// the separate seed submit the OP_MTP handler used ([`Lfm2Gpu::copy_cur_col_to`]).
pub fn draft_chain_seeded(
&self,
ctx: &GpuCtx,
main_bp: &BatchPlan,
col: usize,
first_token: u32,
pos0: usize,
n: usize,
) -> Result<Vec<u32>> {
let h = self.gpu.w.cfg.hidden as u64;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
enc.copy_buffer_to_buffer(
&main_bp.cur,
col as u64 * h * 4,
&self.prev_hidden,
0,
h * 4,
);
let seed = enc.finish();
if std::env::var("OSFKB_MTP_BETA_PROBE").is_ok() {
ctx.queue.submit([seed]);
return self.draft_chain_probe(ctx, first_token, pos0, n);
}
self.draft_chain_fast(ctx, Some(seed), first_token, pos0, n)
}
/// The 1-submit chain body: queued writes for the round's metas (rope tables, `attn_t` T,
/// `cmeta1` pos — they flush at submit, BEFORE the command buffers execute), then
/// `[seed?, body]` in one submit, then the single persistent-staging map for the drafted
/// ids. The next round's body is encoded while this round's GPU work runs.
fn draft_chain_fast(
&self,
ctx: &GpuCtx,
seed: Option<wgpu::CommandBuffer>,
first_token: u32,
pos0: usize,
n: usize,
) -> Result<Vec<u32>> {
let c = &self.gpu.w.cfg;
assert!((1..32).contains(&n), "draft chain capped by the tok buffer");
let mut slot = self.round.borrow_mut();
if slot.as_ref().map(|r| r.k) != Some(n) {
*slot = Some(self.build_round(ctx, n));
}
let r = slot.as_mut().expect("round built");
let local_rope = r.rope_words == 4 * c.head_dim;
let mut rope = Vec::with_capacity(n * r.rope_words);
let mut meta = Vec::with_capacity(n * 8);
for i in 0..n {
let pos = pos0 + 1 + i;
let (coss, sins) = rope_cs_padded(c.head_dim, c.rotary_dim, c.rope_theta, pos);
rope.extend_from_slice(&coss);
rope.extend_from_slice(&sins);
if local_rope {
let (cl, sl) = rope_cs(c.head_dim_sliding(), c.rope_local_theta, pos);
rope.extend_from_slice(&cl);
rope.extend_from_slice(&sl);
}
meta.extend_from_slice(&[
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
pos as u32,
0,
0,
0,
]);
}
ctx.queue
.write_buffer(&r.rope_stage, 0, bytemuck::cast_slice(&rope));
ctx.queue
.write_buffer(&r.meta_stage, 0, bytemuck::cast_slice(&meta));
ctx.queue
.write_buffer(&self.tok, 0, bytemuck::cast_slice(&[first_token]));
let body = r
.pending
.take()
.unwrap_or_else(|| self.encode_round_body(ctx, r));
ctx.queue.submit(seed.into_iter().chain([body]));
r.pending = Some(self.encode_round_body(ctx, r));
let slice = r.tok_staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send(res);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.unwrap()
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let data = slice.get_mapped_range().expect("mapped range");
let toks: Vec<u32> = bytemuck::cast_slice::<u8, u32>(&data)[1..=n].to_vec();
drop(data);
r.tok_staging.unmap();
Ok(toks)
}
/// Build the per-`k` 1-submit resources: staging buffers plus the step-index-keyed bind
/// groups (positions ride the staged bytes, so nothing here depends on the round).
fn build_round(&self, ctx: &GpuCtx, k: usize) -> DraftRound {
let c = &self.gpu.w.cfg;
let (h, hd) = (c.hidden, c.head_dim);
let local_rope = matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
);
let rope_words = if local_rope { 4 * hd } else { 2 * hd };
let stage = |label: &str, bytes: u64| {
ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size: bytes,
usage: wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
})
};
let rope_stage = stage("mtp_rope_stage", (k * rope_words * 4) as u64);
let meta_stage = stage("mtp_meta_stage", (k * 8 * 4) as u64);
let tok_staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("mtp_tok_staging"),
size: self.tok.size(),
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mfk = uni(
ctx,
bytemuck::cast_slice(&[h as f32, 1.0f32, c.eps, 0.0f32]),
);
let fc_h_u = uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, 1)));
let fc_e_u = uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, 1)));
let gemv = &self.gpu.spec_pls.gemv_nbar;
let steps = (0..k)
.map(|i| {
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(i as u32, 0, c.vocab as u32, h as u32)),
);
[
make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.prev_hidden, &self.pre_norm_hidden, &self.h_normed],
&mfk,
),
make_bg(
ctx,
&self.gather_pl,
&[&self.embed, &self.tok, &self.e_raw],
&posu,
),
make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk,
),
make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&self.gpu.s.cur,
],
&fc_h_u,
),
make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&self.gpu.s.cur,
],
&fc_e_u,
),
make_bg(
ctx,
&self.argmax_block_pl,
&[&self.gpu.s.logits, &self.am_pv, &self.am_pi],
&posu,
),
make_bg(
ctx,
&self.argmax_merge_pl,
&[&self.am_pv, &self.am_pi, &self.tok],
&posu,
),
]
})
.collect();
DraftRound {
k,
rope_stage,
rope_words,
meta_stage,
tok_staging,
steps,
pending: None,
}
}
/// Encode one full `k`-step chain body (round-invariant for a fixed `k`): per step, the
/// stage→plan-buffer meta copies, the fc merge, the layer replay, the chain-state copy and
/// the on-GPU argmax; the tail copies `tok` into the persistent readback staging. The same
/// pipelines/bind groups/dispatch shapes as the per-step path — byte-identical inputs land
/// in the same buffers before the same dispatches, so drafts are bitwise-unchanged.
fn encode_round_body(&self, ctx: &GpuCtx, r: &DraftRound) -> wgpu::CommandBuffer {
let c = &self.gpu.w.cfg;
let (h, hd) = (c.hidden, c.head_dim);
let local_rope = r.rope_words == 4 * hd;
let gemv = &self.gpu.spec_pls.gemv_nbar;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (i, bgs) in r.steps.iter().enumerate() {
let ro = (i * r.rope_words * 4) as u64;
let hb = (hd * 4) as u64;
enc.copy_buffer_to_buffer(&r.rope_stage, ro, &self.gpu.s.cosb, 0, hb);
enc.copy_buffer_to_buffer(&r.rope_stage, ro + hb, &self.gpu.s.sinb, 0, hb);
if local_rope {
enc.copy_buffer_to_buffer(&r.rope_stage, ro + 2 * hb, &self.gpu.s.cosb_l, 0, hb);
enc.copy_buffer_to_buffer(&r.rope_stage, ro + 3 * hb, &self.gpu.s.sinb_l, 0, hb);
}
let mo = (i * 32) as u64;
enc.copy_buffer_to_buffer(&r.meta_stage, mo, &self.gpu.attn_t, 0, 16);
enc.copy_buffer_to_buffer(&r.meta_stage, mo + 16, &self.gpu.cmeta1, 0, 8);
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &bgs[0], &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(&self.gather_pl);
p.set_bind_group(0, &bgs[1], &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &bgs[2], &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &bgs[3], &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), 1, 1);
p.set_bind_group(0, &bgs[4], &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), 1, 1);
}
self.gpu.record(&mut enc, 0, 0);
enc.copy_buffer_to_buffer(&self.gpu.s.cur, 0, &self.prev_hidden, 0, (h * 4) as u64);
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.argmax_block_pl);
p.set_bind_group(0, &bgs[5], &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&self.argmax_merge_pl);
p.set_bind_group(0, &bgs[6], &[]);
p.dispatch_workgroups(1, 1, 1);
}
}
enc.copy_buffer_to_buffer(&self.tok, 0, &r.tok_staging, 0, self.tok.size());
enc.finish()
}
/// The β-probe chain: one submit + one full-vocab logits readback PER STEP (top-3 capture).
/// Kept as the measurement path; the production path is [`Self::draft_chain_fast`].
fn draft_chain_probe(
&self,
ctx: &GpuCtx,
first_token: u32,
pos0: usize,
n: usize,
) -> Result<Vec<u32>> {
let c = &self.gpu.w.cfg;
let h = c.hidden;
assert!(n < 32, "draft chain capped by the tok buffer");
// Zero-readback chain: tok[0] seeds, each step's argmax writes tok[i+1] on-GPU and the
// next step's gather reads it — one poll + one 4·n-byte read for the whole chain.
ctx.queue
.write_buffer(&self.tok, 0, bytemuck::cast_slice(&[first_token]));
for i in 0..n {
let pos = pos0 + 1 + i;
self.gpu.upload_rope(ctx, pos);
ctx.queue.write_buffer(
&self.gpu.attn_t,
0,
bytemuck::cast_slice(&u32x4(
c.n_heads as u32,
c.n_kv_heads as u32,
c.head_dim as u32,
(pos + 1) as u32,
)),
);
ctx.queue.write_buffer(
&self.gpu.cmeta1,
0,
bytemuck::cast_slice(&[pos as u32, 0u32]),
);
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, 1.0f32, c.eps, 0.0f32]));
// One uniform serves gather (reads tok[i]), argmax block (vocab), merge (writes i+1).
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(i as u32, 0, c.vocab as u32, h as u32)),
);
let norm_h_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.prev_hidden, &self.pre_norm_hidden, &self.h_normed],
&mfk(h as f32),
);
let gather_bg = make_bg(
ctx,
&self.gather_pl,
&[&self.embed, &self.tok, &self.e_raw],
&posu,
);
let norm_e_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk(h as f32),
);
let gemv = &self.gpu.spec_pls.gemv_nbar;
let fc_h_bg = make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&self.gpu.s.cur,
],
&uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, 1))),
);
let fc_e_bg = make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&self.gpu.s.cur,
],
&uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, 1))),
);
let am_block_bg = make_bg(
ctx,
&self.argmax_block_pl,
&[&self.gpu.s.logits, &self.am_pv, &self.am_pi],
&posu,
);
let am_merge_bg = make_bg(
ctx,
&self.argmax_merge_pl,
&[&self.am_pv, &self.am_pi, &self.tok],
&posu,
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_h_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(&self.gather_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_e_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &fc_h_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), 1, 1);
p.set_bind_group(0, &fc_e_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), 1, 1);
}
self.gpu.record(&mut enc, 0, pos);
enc.copy_buffer_to_buffer(&self.gpu.s.cur, 0, &self.prev_hidden, 0, (h * 4) as u64);
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.argmax_block_pl);
p.set_bind_group(0, &am_block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&self.argmax_merge_pl);
p.set_bind_group(0, &am_merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
ctx.queue.submit([enc.finish()]);
if std::env::var("OSFKB_MTP_BETA_PROBE").is_ok() {
// Probe: top-3 of this step's logits (slow readback — measurement only).
let lg = ctx.read(&self.gpu.s.logits, c.vocab)?;
let mut best = [(f32::NEG_INFINITY, 0u32); 3];
for (ti, v) in lg.iter().enumerate() {
if *v > best[2].0 {
best[2] = (*v, ti as u32);
if best[2].0 > best[1].0 {
best.swap(1, 2);
}
if best[1].0 > best[0].0 {
best.swap(0, 1);
}
}
}
if i == 0 {
self.last_top3.borrow_mut().clear();
}
self.last_top3
.borrow_mut()
.push([best[0].1, best[1].1, best[2].1]);
}
}
let toks = ctx.read_u32(&self.tok, n + 1)?;
Ok(toks[1..=n].to_vec())
}
/// The β-probe's per-depth top-3 candidates from the last [`Self::draft_chain`] call.
pub fn last_top3(&self) -> Vec<[u32; 3]> {
self.last_top3.borrow().clone()
}
/// Serving draft batch width (µbatch slots per round ≤ this).
pub const K_DRAFT_BATCH: usize = 32;
/// The narrowest draft plan that fits `ncols` (see the field docs).
fn draft_plan(&self, ctx: &GpuCtx, ncols: usize) -> &BatchPlan {
if ncols <= 8 {
self.bp_draft8
.get_or_init(|| self.gpu.make_batch_plan(ctx, 8))
} else {
self.bp_draft32
.get_or_init(|| self.gpu.make_batch_plan(ctx, Self::K_DRAFT_BATCH))
}
}
/// Seed slot `slot`'s chain slab from column `col` of the MAIN model's last verify µbatch
/// (`bp.cur` PRE-final-norm residual) — on-GPU copy, no host round-trip.
pub fn seed_slot_from_col(&self, ctx: &GpuCtx, main_bp: &BatchPlan, col: usize, slot: usize) {
let h = self.gpu.w.cfg.hidden as u64;
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
enc.copy_buffer_to_buffer(
&main_bp.cur,
col as u64 * h * 4,
&self.prev_hidden,
slot as u64 * h * 4,
h * 4,
);
ctx.queue.submit([enc.finish()]);
}
/// Draft-KV PREFILL pairs: run the draft layer once per `(slot, cur_col, next_token, pos)`
/// pair — the hidden comes straight from `main_bp.cur\[cur_col\]` (NOT the chain slabs, which
/// stay untouched), the KV lands at `pos + 1`, outputs are discarded. This is the serving
/// counterpart of the single-stream MTP prompt prefill: without it the draft attends over
/// zeros for the whole prompt and acceptance collapses.
pub fn draft_pairs(
&self,
ctx: &GpuCtx,
main_bp: &BatchPlan,
pairs: &[DraftPair],
) -> Result<()> {
let c = &self.gpu.w.cfg;
let h = c.hidden;
let ncols = pairs.len();
if ncols == 0 {
return Ok(());
}
assert!(ncols <= Self::K_DRAFT_BATCH);
let bp = self.draft_plan(ctx, ncols);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (p, pair) in pairs.iter().enumerate() {
enc.copy_buffer_to_buffer(
&main_bp.cur,
pair.cur_col as u64 * h as u64 * 4,
&self.prev_round,
(p * h * 4) as u64,
(h * 4) as u64,
);
}
ctx.queue.submit([enc.finish()]);
let toks: Vec<u32> = pairs.iter().map(|p| p.token).collect();
ctx.queue
.write_buffer(&self.toks_round, 0, bytemuck::cast_slice(&toks));
self.fc_merge_round(ctx, bp, ncols);
let cols: Vec<BatchCol> = pairs
.iter()
.map(|p| BatchCol {
token: 0,
pos: (p.pos + 1) as u32,
btrow: p.btrow,
need_logit: false,
mpos: p.mpos,
embed_row: -1,
})
.collect();
// Layer replay writes the pairs' KV; outputs are never consumed, so SUBMIT ONLY — no
// readback sync. Later draft submissions observe the KV by queue order.
self.gpu.stage_step_submit(ctx, bp, &cols, None, true)?;
Ok(())
}
/// The fc-merge dispatch block shared by the batched chain and the prefill pairs: reads
/// `prev_round` + `toks_round`, leaves the merged residual in `bp.cur\[..ncols\]`.
fn fc_merge_round(&self, ctx: &GpuCtx, bp: &BatchPlan, ncols: usize) {
let c = &self.gpu.w.cfg;
let h = c.hidden;
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, ncols as f32, c.eps, 0.0f32]));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, c.vocab as u32, h as u32)),
);
let norm_h_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.prev_round, &self.pre_norm_hidden, &self.h_normed],
&mfk(h as f32),
);
let gather_bg = make_bg(
ctx,
&self.gather_cols_pl,
&[&self.embed, &self.toks_round, &self.e_raw],
&posu,
);
let norm_e_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk(h as f32),
);
let gemv = &self.gpu.spec_pls.gemv_nbar;
let fc_h_bg = make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, ncols as u32)),
),
);
let fc_e_bg = make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, ncols as u32)),
),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_h_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(&self.gather_cols_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), ncols as u32, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_e_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &fc_h_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
p.set_bind_group(0, &fc_e_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
}
ctx.queue.submit([enc.finish()]);
}
/// Batched draft chains for serving: one entry per ACTIVE slot `(slot, first_token, pos0)`,
/// all chained TOGETHER for `n` steps through the 1-layer draft model's batch plan (per-slot
/// KV via the mirrored block tables). Returns per-request drafts. CPU-chained v1: one
/// `chosen` readback per step, amortized over every slot in the round.
pub fn draft_chain_batch(
&self,
ctx: &GpuCtx,
reqs: &[(u32, u32, usize)],
n: usize,
) -> Result<Vec<Vec<u32>>> {
let reqs: Vec<DraftReq> = reqs
.iter()
.map(|&(btrow, first_token, pos0)| DraftReq {
btrow,
first_token,
pos0,
mpos0: [pos0 as u32; 3],
})
.collect();
self.draft_chain_batch_seeded(ctx, None, &[], &reqs, n)
}
/// [`Self::draft_chain_batch`] with the per-span SEEDS folded into the first step's copy
/// encoder: `seeds[i] = (slot, main col)` copies `main_bp.cur[col] → prev_hidden[slot]`
/// ahead of the slab gather in the SAME submission — replacing one tiny submit per span
/// (`seed_slot_from_col`) with zero extra submissions. Byte-identical by construction:
/// the same copies in the same queue order.
pub fn draft_chain_batch_seeded(
&self,
ctx: &GpuCtx,
main_bp: Option<&BatchPlan>,
seeds: &[(u32, usize)],
reqs: &[DraftReq],
n: usize,
) -> Result<Vec<Vec<u32>>> {
let c = &self.gpu.w.cfg;
let h = c.hidden;
let ncols = reqs.len();
assert!(ncols > 0 && ncols <= Self::K_DRAFT_BATCH);
assert!(
seeds.is_empty() || main_bp.is_some(),
"seed copies read the main plan's residual stream"
);
let bp = self.draft_plan(ctx, ncols);
let mut toks: Vec<u32> = reqs.iter().map(|r| r.first_token).collect();
let mut out = vec![Vec::with_capacity(n); ncols];
for step in 0..n {
// Gather slot slabs → compact round buffer (column order = reqs order).
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
if step == 0
&& let Some(mp) = main_bp
{
for (slot, col) in seeds {
enc.copy_buffer_to_buffer(
&mp.cur,
*col as u64 * h as u64 * 4,
&self.prev_hidden,
*slot as u64 * h as u64 * 4,
(h * 4) as u64,
);
}
}
for (p, r) in reqs.iter().enumerate() {
enc.copy_buffer_to_buffer(
&self.prev_hidden,
u64::from(r.btrow) * h as u64 * 4,
&self.prev_round,
(p * h * 4) as u64,
(h * 4) as u64,
);
}
ctx.queue.submit([enc.finish()]);
ctx.queue
.write_buffer(&self.toks_round, 0, bytemuck::cast_slice(&toks));
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, ncols as f32, c.eps, 0.0f32]));
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, c.vocab as u32, h as u32)),
);
let norm_h_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.prev_round, &self.pre_norm_hidden, &self.h_normed],
&mfk(h as f32),
);
let gather_bg = make_bg(
ctx,
&self.gather_cols_pl,
&[&self.embed, &self.toks_round, &self.e_raw],
&posu,
);
let norm_e_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk(h as f32),
);
let gemv = &self.gpu.spec_pls.gemv_nbar;
let fc_h_bg = make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, ncols as u32)),
),
);
let fc_e_bg = make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, ncols as u32)),
),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_h_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(&self.gather_cols_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), ncols as u32, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_e_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &fc_h_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
p.set_bind_group(0, &fc_e_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
}
ctx.queue.submit([enc.finish()]);
// The layer replay + head/argmax (per-column pos/btab rope handled inside).
let cols: Vec<BatchCol> = reqs
.iter()
.map(|r| BatchCol {
token: 0,
pos: (r.pos0 + 1 + step) as u32,
btrow: r.btrow,
need_logit: true,
mpos: r.mpos0.map(|m| m + 1 + step as u32),
embed_row: -1,
})
.collect();
let drafted = match self.gpu.batch_stage_step_cur_ready(ctx, bp, &cols)? {
StageBatchOut::Tokens(t) => t,
StageBatchOut::Hidden(_) => anyhow::bail!("draft model must be stage_last"),
};
// Persist the post-layer residual back to the slot slabs for the next step.
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for (p, r) in reqs.iter().enumerate() {
enc.copy_buffer_to_buffer(
&bp.cur,
(p * h * 4) as u64,
&self.prev_hidden,
u64::from(r.btrow) * h as u64 * 4,
(h * 4) as u64,
);
}
ctx.queue.submit([enc.finish()]);
for (sidx, d) in drafted.iter().take(ncols).enumerate() {
out[sidx].push(*d);
toks[sidx] = *d;
}
}
Ok(out)
}
/// The serving draft chain as ONE submission — the batched twin of the mono [`DraftRound`]
/// 1-submit design, and the fix for the measured serving-MTP loss (the v1 batched chain
/// paid one submit+readback PER STEP; on a GPU shared with the vision tower those ~4 syncs
/// per round buried the acceptance win: 65 vs 227 tok/s).
///
/// `seeds[i]` = the MAIN plan column whose PRE-final-norm residual (`main_bp.cur`) seeds
/// request `i`'s chain — copied in-encoder, so the chain state lives entirely in the draft
/// plan's `cur` across steps (no slot slabs, no host round-trips). Per-step rope/meta ride
/// staged buffers copied in-encoder; each step's argmax feeds the next step's token gather
/// on-GPU; ONE staging map at the end returns all `n × reqs` drafts.
///
/// Bitwise-identical drafts to the per-step path: the same dispatches consume byte-identical
/// inputs from the same buffers in the same order — only the submission granularity changes.
pub fn draft_chain_round(
&self,
ctx: &GpuCtx,
main_bp: &BatchPlan,
seeds: &[usize],
reqs: &[DraftReq],
n: usize,
) -> Result<Vec<Vec<u32>>> {
let c = &self.gpu.w.cfg;
let (h, hd) = (c.hidden, c.head_dim);
let ncols = reqs.len();
assert!(ncols > 0 && ncols <= Self::K_DRAFT_BATCH && seeds.len() == ncols);
assert!((1..32).contains(&n));
debug_assert!(
!matches!(
c.arch,
crate::weights::Arch::Gemma3 | crate::weights::Arch::Gemma4
),
"dual-theta local rope is not staged here (no Gemma draft heads exist)"
);
let bp = self.draft_plan(ctx, ncols);
let k = bp.k;
let trash_row = (self.gpu.dn_slots - 1) as u32;
// Stage the per-step tables: rope cos/sin `[n][k][hd]` and cmeta `[n][k][2]` — idle plan
// columns park on the trash block-table row exactly as `stage_step_submit` parks them.
// Global-theta tables serve the full-attention layers — the WIDE heads on an edge model.
let rope = RopeStager::new(c, c.head_dim_full());
let mut coss = Vec::with_capacity(n * k * hd);
let mut sins = Vec::with_capacity(n * k * hd);
let mut meta = Vec::with_capacity(n * k * 2);
for s in 0..n {
for p in 0..k {
match reqs.get(p) {
Some(r) => {
let pos = (r.pos0 + 1 + s) as u32;
let mpos = r.mpos0.map(|m| m + 1 + s as u32);
let (cs, sn) = rope.col(mpos, pos);
coss.extend_from_slice(&cs);
sins.extend_from_slice(&sn);
meta.extend_from_slice(&[pos, r.btrow]);
}
None => {
coss.extend(std::iter::repeat_n(0f32, hd));
sins.extend(std::iter::repeat_n(0f32, hd));
meta.extend_from_slice(&[0, trash_row]);
}
}
}
}
let stage = |label: &str, bytes: u64| {
ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size: bytes,
usage: wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
})
};
let cs_stage = stage("mtp_round_cs", (n * k * hd * 4) as u64);
let sn_stage = stage("mtp_round_sn", (n * k * hd * 4) as u64);
let meta_stage = stage("mtp_round_meta", (n * k * 2 * 4) as u64);
let drafts_staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("mtp_round_drafts"),
size: (n * k * 4) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
ctx.queue
.write_buffer(&cs_stage, 0, bytemuck::cast_slice(&coss));
ctx.queue
.write_buffer(&sn_stage, 0, bytemuck::cast_slice(&sins));
ctx.queue
.write_buffer(&meta_stage, 0, bytemuck::cast_slice(&meta));
ctx.queue
.write_buffer(&bp.cnt, 0, bytemuck::cast_slice(&[0u32]));
let first: Vec<u32> = reqs.iter().map(|r| r.first_token).collect();
ctx.queue
.write_buffer(&self.toks_round, 0, bytemuck::cast_slice(&first));
// Round-invariant bind groups (per-step variation rides the staged copies above).
let mfk = uni(
ctx,
bytemuck::cast_slice(&[h as f32, ncols as f32, c.eps, 0.0f32]),
);
let posu = uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, c.vocab as u32, h as u32)),
);
let norm_h_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&bp.cur, &self.pre_norm_hidden, &self.h_normed],
&mfk,
);
let gather_bg = make_bg(
ctx,
&self.gather_cols_pl,
&[&self.embed, &self.toks_round, &self.e_raw],
&posu,
);
let norm_e_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk,
);
let gemv = &self.gpu.spec_pls.gemv_nbar;
let fc_h_bg = make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, ncols as u32)),
),
);
let fc_e_bg = make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&bp.cur,
],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, ncols as u32)),
),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
// Chain seeds: the main plan's pre-final-norm residual at each accepted column.
for (i, &col) in seeds.iter().enumerate() {
enc.copy_buffer_to_buffer(
&main_bp.cur,
col as u64 * h as u64 * 4,
&bp.cur,
(i * h * 4) as u64,
(h * 4) as u64,
);
}
for s in 0..n {
let ro = (s * k * hd * 4) as u64;
enc.copy_buffer_to_buffer(&cs_stage, ro, &bp.csk, 0, (k * hd * 4) as u64);
enc.copy_buffer_to_buffer(&sn_stage, ro, &bp.snk, 0, (k * hd * 4) as u64);
enc.copy_buffer_to_buffer(
&meta_stage,
(s * k * 2 * 4) as u64,
&bp.cmeta_k,
0,
(k * 2 * 4) as u64,
);
{
// fc merge: cur ← fc_h·norm(cur) + fc_e·norm(embed[tok]) — the DeepSeek MTP
// recurrence; the previous step's post-layer residual is still in `cur`.
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_h_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(&self.gather_cols_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), ncols as u32, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_e_bg, &[]);
p.dispatch_workgroups(ncols as u32, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &fc_h_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
p.set_bind_group(0, &fc_e_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(16), (ncols as u32).div_ceil(4), 1);
}
{
// Layer replay + head + per-column argmax — the draft plan's own step body.
let mut p = enc.begin_compute_pass(&Default::default());
for step in &bp.plan {
match step {
Step::D { pl, bg, gx, gy, .. } => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups(*gx, *gy, 1);
}
Step::Di {
pl,
bg,
args,
offset,
..
} => {
p.set_pipeline(pl);
p.set_bind_group(0, bg, &[]);
p.dispatch_workgroups_indirect(args, *offset);
}
Step::Embed => {} // stage_first only; the draft model has no embed step
}
}
let (lm_bg, lm_gx, lm_gy) = bp.lm.as_ref().expect("draft model is stage_last");
p.set_pipeline(&bp.lm_pl);
p.set_bind_group(0, lm_bg, &[]);
p.dispatch_workgroups(*lm_gx, *lm_gy, 1);
for hh in bp.heads.iter().take(ncols) {
p.set_pipeline(&bp.block_pl);
p.set_bind_group(0, &hh.block_bg, &[]);
p.dispatch_workgroups(256, 1, 1);
p.set_pipeline(&bp.merge_pl);
p.set_bind_group(0, &hh.merge_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
}
}
// The step's argmax feeds the next step's gather AND lands in the round staging.
enc.copy_buffer_to_buffer(&bp.chosen_k, 0, &self.toks_round, 0, (ncols * 4) as u64);
enc.copy_buffer_to_buffer(
&bp.chosen_k,
0,
&drafts_staging,
(s * k * 4) as u64,
(k * 4) as u64,
);
}
ctx.queue.submit([enc.finish()]);
let slice = drafts_staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send(res);
});
let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
rx.recv()
.unwrap()
.map_err(|e| anyhow::anyhow!("map_async: {e:?}"))?;
let data = slice.get_mapped_range().expect("mapped range");
let flat: &[u32] = bytemuck::cast_slice(&data);
let out = (0..ncols)
.map(|i| (0..n).map(|s| flat[s * k + i]).collect())
.collect();
drop(data);
drafts_staging.unmap();
Ok(out)
}
/// TEST PROBE: run only the fc-merge pre-steps for `tok` and return `s.cur` (the layer
/// input) — isolates the merge from the layer when bisecting a draft-chain divergence.
pub fn fc_merge_for_tests(&self, ctx: &GpuCtx, tok: u32) -> Result<Vec<f32>> {
let c = &self.gpu.w.cfg;
let h = c.hidden;
ctx.queue
.write_buffer(&self.tok, 0, bytemuck::cast_slice(&[tok]));
let mfk = |a: f32| uni(ctx, bytemuck::cast_slice(&[a, 1.0f32, c.eps, 0.0f32]));
let norm_h_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.prev_hidden, &self.pre_norm_hidden, &self.h_normed],
&mfk(h as f32),
);
let gather_bg = make_bg(
ctx,
&self.gather_pl,
&[&self.embed, &self.tok, &self.e_raw],
&uni(
ctx,
bytemuck::cast_slice(&u32x4(0, 0, c.vocab as u32, h as u32)),
),
);
let norm_e_bg = make_bg(
ctx,
&self.gpu.spec_pls.rmsnorm,
&[&self.e_raw, &self.pre_norm_emb, &self.e_normed],
&mfk(h as f32),
);
let gemv = &self.gpu.spec_pls.gemv_k;
let fc_h_bg = make_bg(
ctx,
gemv,
&[
&self.fc_hidden.scales,
&self.fc_hidden.quants,
&self.h_normed,
&self.gpu.s.cur,
],
&uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 0, 1))),
);
let fc_e_bg = make_bg(
ctx,
gemv,
&[
&self.fc_emb.scales,
&self.fc_emb.quants,
&self.e_normed,
&self.gpu.s.cur,
],
&uni(ctx, bytemuck::cast_slice(&u32x4(h as u32, h as u32, 1, 1))),
);
let mut enc = ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
{
let mut p = enc.begin_compute_pass(&Default::default());
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_h_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(&self.gather_pl);
p.set_bind_group(0, &gather_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(64), 1, 1);
p.set_pipeline(&self.gpu.spec_pls.rmsnorm);
p.set_bind_group(0, &norm_e_bg, &[]);
p.dispatch_workgroups(1, 1, 1);
p.set_pipeline(gemv);
p.set_bind_group(0, &fc_h_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(self.gpu.gemv_rows_per_wg), 1, 1);
p.set_bind_group(0, &fc_e_bg, &[]);
p.dispatch_workgroups((h as u32).div_ceil(self.gpu.gemv_rows_per_wg), 1, 1);
}
ctx.queue.submit([enc.finish()]);
ctx.read(&self.gpu.s.cur, h)
}
/// Seed the chain state from a CPU-side hidden (test/bring-up path).
pub fn write_prev_hidden(&self, ctx: &GpuCtx, hidden: &[f32]) {
ctx.queue
.write_buffer(&self.prev_hidden, 0, bytemuck::cast_slice(hidden));
}
/// The chain-state buffer (the serving integration copies the main model's pre-norm hidden
/// column straight into it, no host round-trip).
pub fn prev_hidden_buf(&self) -> &wgpu::Buffer {
&self.prev_hidden
}
}
#[cfg(test)]
mod paged_kv_geometry {
use super::*;
/// Every paged-KV shader must take its block-table stride from [`PAGE_ROW`], never a literal.
///
/// The stride used to be `256u`, written out in three separate WGSL sources next to a Rust
/// constant, with a comment saying they MUST match. Nothing enforced it. A shader left on the
/// old stride would not fail — it would index the block table with the wrong row pitch and
/// silently read a different sequence's KV, which is the worst possible failure: fluent,
/// confident, wrong output.
///
/// So the sources now carry a `{{PROW}}` placeholder that is INVALID WGSL, and [`page_geom`]
/// fills it in. A compile site that forgets the call dies at pipeline creation. This test
/// guards the other direction: that nobody re-hardcodes the literal back into the source.
#[test]
fn paged_shaders_carry_the_placeholder_not_a_literal() {
// The two raw-string shaders carry the placeholder and are patched by `page_geom`.
for (name, src) in [("ATTN_K", ATTN_K), ("ATTN_TILED", ATTN_TILED)] {
assert!(
src.contains("const PROW: u32 = {{PROW}}u;"),
"{name} declares PROW without the placeholder — the stride is hardcoded again"
);
let patched = page_geom(src);
assert!(
patched.contains(&format!("const PROW: u32 = {PAGE_ROW}u;")),
"{name}: page_geom did not substitute PAGE_ROW"
);
assert!(
!patched.contains("{{PROW}}"),
"{name}: a placeholder survived patching — the shader would not compile"
);
}
// `qk_norm_rope_cache_k_src` is `format!`-generated, and `format!` UNESCAPES `{{…}}` — so a
// `{{PROW}}` placeholder cannot survive it (this test caught exactly that). It interpolates
// PAGE_ROW itself instead, and must therefore already be correct without patching.
let generated = qk_norm_rope_cache_k_src(1e-6, true);
assert!(
generated.contains(&format!("const PROW: u32 = {PAGE_ROW}u;")),
"qk_norm_rope_cache_k did not interpolate PAGE_ROW"
);
}
/// The addressable capacity and the cap the scheduler enforces are the same number.
#[test]
fn the_scheduler_cannot_admit_more_than_the_shaders_can_address() {
assert_eq!(PAGE_ROW * PAGE_BLK, MAX_T);
}
}
#[cfg(test)]
mod mrope {
use super::{mrope_axis_map, mrope_cs_padded, rope_cs_padded};
/// NuExtract-3 / Qwen3.5-VL geometry: head_dim 256, partial 0.25 ⇒ rotary 64 ⇒ half 32,
/// and `mrope_section` [11, 11, 10] sums to exactly that 32.
const HD: usize = 256;
const RD: usize = 64;
const SECTION: [usize; 3] = [11, 11, 10];
const THETA: f32 = 10_000_000.0;
/// The safety argument for switching a production text model onto the M-RoPE path: on text, the
/// three axes carry the SAME position, so the axis map picks between three equal numbers and the
/// tables come out bit-for-bit identical. Not "close" — identical. If this ever fails, every
/// text-only workload on this engine has silently changed.
#[test]
fn mrope_is_bitwise_identical_to_1d_rope_on_text() {
let axes = mrope_axis_map(SECTION, RD / 2, false);
for pos in [0u32, 1, 7, 63, 1024, 4095, 16_383] {
let (c1, s1) = rope_cs_padded(HD, RD, THETA, pos as usize);
let (c3, s3) = mrope_cs_padded(HD, RD, THETA, &axes, [pos; 3]);
assert_eq!(c1, c3, "cos differs at text position {pos}");
assert_eq!(s1, s3, "sin differs at text position {pos}");
}
}
/// The section reduces to a clean T H W T H W … interleave — worth pinning, because the
/// implementation derives it from the section rather than hardcoding `i % 3`, and a future
/// checkpoint with a different section must not silently reuse this layout.
#[test]
fn the_11_11_10_section_interleaves_t_h_w() {
let axes = mrope_axis_map(SECTION, RD / 2, false);
assert_eq!(axes.len(), 32);
for (i, &a) in axes.iter().enumerate() {
assert_eq!(
a as usize,
i % 3,
"frequency {i} should rotate by axis {}",
i % 3
);
}
// …and each axis gets exactly the width the section advertises.
for (axis, want) in SECTION.iter().enumerate() {
let got = axes.iter().filter(|&&a| a as usize == axis).count();
assert_eq!(
got, *want,
"axis {axis} claims {want} frequencies, got {got}"
);
}
}
/// A non-uniform position MUST change the tables — otherwise the previous test is vacuous and we
/// would be shipping a no-op that "passes" while images attend with text geometry.
#[test]
fn a_3d_position_actually_rotates_differently() {
let axes = mrope_axis_map(SECTION, RD / 2, false);
let (c_text, _) = mrope_cs_padded(HD, RD, THETA, &axes, [40; 3]);
let (c_img, _) = mrope_cs_padded(HD, RD, THETA, &axes, [40, 41, 42]);
assert_ne!(
c_text, c_img,
"an image's (t,h,w) must not rotate like text"
);
// The T axis carries the same value (40) in both, so EVERY T frequency must be untouched.
// That is the exact invariant; the converse is not, because at θ=1e7 the high-index
// frequencies have an inv_freq so small that positions 40 and 42 produce the same f32 cosine.
// So H/W frequencies are only *required* to move where the angle is actually resolvable.
let half = RD / 2;
let mut moved = [0usize; 3];
for i in 0..half {
if axes[i] == 0 {
assert_eq!(
c_text[i], c_img[i],
"frequency {i} rotates by T, whose position is the same in both — it must not move"
);
} else if c_text[i] != c_img[i] {
moved[axes[i] as usize] += 1;
}
}
assert!(
moved[1] > 0,
"no H frequency moved — the h axis is not wired up"
);
assert!(
moved[2] > 0,
"no W frequency moved — the w axis is not wired up"
);
}
}