diskformer 0.0.1 → 0.3.0

package/engines/qwen3_moe_ssd.js

@@ -0,0 +1,1041 @@
+/* qwen3_moe_ssd.js — Qwen3-MoE (Qwen3-30B-A3B) in a browser tab, larger than RAM: the dense
+ * weights, routers and KV cache live on the GPU; the 6,144 routed experts (≈31 GB at Q8_0)
+ * live in OPFS and stream into a GPU slot pool as the router asks for them.
+ *
+ * Rung 2a of the SSD-streaming ladder (plan/pending.md). From-scratch WebGPU engine; the
+ * kernels are ports of browser-big-fast-lab's custom-kernels engine (an Apache-2.0 fork of
+ * tylerstraub/gemma4-webgpu): in-shader Q8_0 GEMV (int8 + f16 scale per 32-block, f32
+ * accumulate), GPU softmax + top-k router, slot-indexed batched expert GEMVs, MoE combine.
+ *
+ *   const m = await Qwen3MoeSsd.load('Qwen/Qwen3-30B-A3B-GGUF', { fetch, onProgress, poolBytes });
+ *   for await (const { text } of m.generate(messages, { maxNewTokens })) …
+ *
+ * Load = (first time only) ingest the GGUF into OPFS in the engine's own layout, then upload
+ * the dense part to the GPU. OPFS layout under localmind-ssd/<key>/:
+ *   header.bin     the GGUF header bytes (metadata + tokenizer), re-parsed at every load
+ *   dense.bin      every non-expert tensor; Q8_0 split losslessly into int8 + f16-scale planes
+ *   experts.bin    one 5,013,504-byte record per (layer, expert), gate‖up‖down contiguous
+ *   manifest.json  source (repo, revision, sha256, size) + every offset above
+ * Q8_0 is never re-quantized: each block's f16 scale and 32 int8 values are copied bit-exact,
+ * so the engine computes the same dot products llama.cpp does, up to summation order.
+ */
+
+import { OpfsReaderPool, readOpfsText } from '../src/opfs-reader.js';
+import { ExpertStreamer } from '../src/expert-stream.js';
+// The store layer is diskformer.js's (github.com/NakliTechie/diskformer.js), byte-identical copies: gguf.js, ingest.js.
+import { GGML, Q8_BLOCK, Q4_BLOCK, Q6K_BLOCK, parseGguf, tensorBytes } from '../src/gguf.js';
+import { ingestGguf as ingestStore, ingestProgress, fileFetch, storeKey, removeStore } from '../src/ingest.js';
+export { GGML, Q4_BLOCK, Q6K_BLOCK, parseGguf, tensorBytes, splitQ4 } from '../src/gguf.js';
+export { ingestProgress } from '../src/ingest.js';
+
+export const QWEN3_30B_A3B = {
+  repo: 'Qwen/Qwen3-30B-A3B-GGUF',
+  file: 'Qwen3-30B-A3B-Q8_0.gguf',
+  revision: 'e4d4bafdfb96a411a163846265362aceb0b9c63a',
+  sha256: '4ad960d180b16f56024f5b704697e5dd5b0837167c2e515ef0569abfc599743c',
+  size: 32483931648,
+};
+const FORMAT = 'localmind-qwen3moe-ssd/1';
+const OPFS_ROOT = 'localmind-ssd';
+
+// ── GGUF ────────────────────────────────────────────────────────────────────
+// qwen3moe (rung 2a) and qwen35moe (Qwen3.5/3.6 MoE, rung 2b: qwen35_moe_ssd.js) share the
+// expert layout; qwen35moe adds the Gated DeltaNet (ssm.*), partial RoPE and a shared expert.
+export function configFromGguf(kv) {
+  const a = kv['general.architecture'];
+  if (a !== 'qwen3moe' && a !== 'qwen35moe') throw new Error(`expected a qwen3moe or qwen35moe GGUF, got ${a}`);
+  const g = (k) => kv[`${a}.${k}`];
+  const extra = a !== 'qwen35moe' ? {} : {
+    ropeDims: g('rope.dimension_count'), shexpFf: g('expert_shared_feed_forward_length'),
+    attnInterval: g('full_attention_interval') || 4,
+    ssm: { dConv: g('ssm.conv_kernel'), dInner: g('ssm.inner_size'), dState: g('ssm.state_size'), vHeads: g('ssm.time_step_rank'), kHeads: g('ssm.group_count') },
+  };
+  return {
+    ...extra,
+    arch: a,
+    layers: g('block_count'), hidden: g('embedding_length'),
+    heads: g('attention.head_count'), kvHeads: g('attention.head_count_kv'),
+    headDim: g('attention.key_length') || g('embedding_length') / g('attention.head_count'),
+    experts: g('expert_count'), topK: g('expert_used_count'), expertFf: g('expert_feed_forward_length'),
+    ropeTheta: g('rope.freq_base'), eps: g('attention.layer_norm_rms_epsilon'),
+    contextLength: g('context_length'),
+    vocab: kv['tokenizer.ggml.tokens'].length,
+    bos: kv['tokenizer.ggml.bos_token_id'], eos: kv['tokenizer.ggml.eos_token_id'],
+  };
+}
+
+// The engine's OPFS layout, derived from the header alone (ingest and load agree on it).
+export function planLayout(gguf) {
+  const cfg = configFromGguf(gguf.kv);
+  const H = cfg.hidden, F = cfg.expertFf;
+  const qBytes = F * H;                  // one projection's int8 plane
+  const sBytes = (F * H / 32) * 2;       // its f16 scale plane
+  const parts = {
+    guQ: { off: 0, bytes: 2 * qBytes },               // gate rows 0..F-1, up rows F..2F-1
+    guS: { off: 2 * qBytes, bytes: 2 * sBytes },
+    dQ: { off: 2 * qBytes + 2 * sBytes, bytes: qBytes },
+    dS: { off: 3 * qBytes + 2 * sBytes, bytes: sBytes },
+  };
+  const record = 3 * qBytes + 3 * sBytes;
+  const dense = {};
+  let off = 0;
+  const take = (n) => { const o = off; off += Math.ceil(n / 256) * 256; return o; };
+  for (const t of gguf.tensors) {
+    if (/_exps\.weight$/.test(t.name)) continue;
+    const n = t.dims.reduce((a, b) => a * b, 1);
+    if (t.type === GGML.Q8_0) {
+      dense[t.name] = { type: 'q8', dims: t.dims, q: { off: take(n), bytes: n }, s: { off: take(n / 16), bytes: n / 16 } };
+    } else if (t.type === GGML.F32) {
+      dense[t.name] = { type: 'f32', dims: t.dims, raw: { off: take(n * 4), bytes: n * 4 } };
+    } else throw new Error(`dense tensor ${t.name}: unsupported type ${t.type}`);
+  }
+  return {
+    config: cfg,
+    experts: { file: 'experts.bin', record, parts, layers: cfg.layers, perLayer: cfg.experts, bytes: record * cfg.layers * cfg.experts },
+    dense: { file: 'dense.bin', bytes: off, tensors: dense },
+  };
+}
+
+// ── Ingest: GGUF (over HTTP Range) → OPFS, in the engine layout ─────────────
+// Every byte range of the GGUF the engine keeps, as units small enough to transform in
+// memory: one unit per expert slice, row bands for dense Q8_0, whole F32 tensors.
+function planUnits(gguf, layout) {
+  const units = [];
+  const { record, parts, perLayer } = layout.experts;
+  const qBytes = parts.dQ.bytes, sBytes = parts.dS.bytes;
+  for (const t of gguf.tensors) {
+    const abs = gguf.dataStart + t.offset;
+    const m = /^blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight$/.exec(t.name);
+    if (m) {
+      const layer = Number(m[1]), which = m[2];
+      if (t.type !== GGML.Q8_0) throw new Error(`${t.name}: experts must be Q8_0`);
+      const slice = tensorBytes(t) / perLayer;
+      const qOff = which === 'gate' ? parts.guQ.off : which === 'up' ? parts.guQ.off + qBytes : parts.dQ.off;
+      const sOff = which === 'gate' ? parts.guS.off : which === 'up' ? parts.guS.off + sBytes : parts.dS.off;
+      for (let e = 0; e < perLayer; e++) {
+        const rec = (layer * perLayer + e) * record;
+        units.push({ src: abs + e * slice, len: slice, file: 'experts', q: rec + qOff, s: rec + sOff });
+      }
+      continue;
+    }
+    const d = layout.dense.tensors[t.name];
+    if (d.type === 'f32') { units.push({ src: abs, len: d.raw.bytes, file: 'dense', raw: d.raw.off }); continue; }
+    const cols = t.dims[0], rows = d.q.bytes / cols, rowBytes = (cols / 32) * Q8_BLOCK;
+    const band = Math.max(1, Math.floor((8 << 20) / rowBytes));
+    for (let r = 0; r < rows; r += band) {
+      const nr = Math.min(band, rows - r);
+      units.push({ src: abs + r * rowBytes, len: nr * rowBytes, file: 'dense', q: d.q.off + r * cols, s: d.s.off + r * cols / 16 });
+    }
+  }
+  units.sort((a, b) => a.src - b.src);
+  return units;
+}
+
+// The ingest is diskformer's ingestGguf with this engine's store folder and layout tag (kept, so stores written
+// before the split stay valid); `plan` lets another architecture supply its own layout (Gemma 4: gemma4_moe_ssd.js).
+export function ingestGguf(opts) {
+  return ingestStore({ ...opts, root: opts.root || OPFS_ROOT, format: FORMAT, plan: opts.plan || { layout: planLayout, units: planUnits } });
+}
+
+// ── Tokenizer: byte-level BPE from the GGUF vocab (tokenizer.ggml.model = gpt2, pre = qwen2) ──
+// Same pre-tokenizer regex llama.cpp uses for LLAMA_VOCAB_PRE_TYPE_QWEN2.
+const QWEN2_PRE = /(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
+// LLAMA_VOCAB_PRE_TYPE_QWEN35 (Qwen3.5/3.6): combining marks (\p{M}) count as part of a word.
+const QWEN35_PRE = /(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
+
+function byteToUnicode() {
+  const bs = [];
+  for (let i = 33; i <= 126; i++) bs.push(i);
+  for (let i = 161; i <= 172; i++) bs.push(i);
+  for (let i = 174; i <= 255; i++) bs.push(i);
+  const cs = bs.slice();
+  let n = 0;
+  for (let b = 0; b < 256; b++) if (!bs.includes(b)) { bs.push(b); cs.push(256 + n); n++; }
+  const enc = new Array(256), dec = new Map();
+  bs.forEach((b, i) => { enc[b] = String.fromCodePoint(cs[i]); dec.set(String.fromCodePoint(cs[i]), b); });
+  return { enc, dec };
+}
+
+export class BpeTokenizer {
+  constructor(kv) {
+    this.tokens = kv['tokenizer.ggml.tokens'];
+    this.preRe = kv['tokenizer.ggml.pre'] === 'qwen35' ? QWEN35_PRE : QWEN2_PRE;
+    const types = kv['tokenizer.ggml.token_type'] || [];
+    this.ids = new Map(this.tokens.map((t, i) => [t, i]));
+    this.ranks = new Map((kv['tokenizer.ggml.merges'] || []).map((m, i) => [m, i]));
+    this.special = [];
+    this.isSpecial = new Uint8Array(this.tokens.length);
+    for (let i = 0; i < this.tokens.length; i++) if (types[i] === 3 || types[i] === 4) { this.special.push(this.tokens[i]); this.isSpecial[i] = 1; }
+    this.special.sort((a, b) => b.length - a.length);
+    const esc = (s) => s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+    this.specialRe = this.special.length ? new RegExp(`(${this.special.map(esc).join('|')})`) : null;
+    const { enc, dec } = byteToUnicode();
+    this.byteEnc = enc; this.byteDec = dec;
+    this.cache = new Map();
+    this.utf8 = new TextEncoder();
+  }
+  bpe(word) {
+    const hit = this.cache.get(word);
+    if (hit) return hit;
+    let parts = Array.from(word);
+    while (parts.length > 1) {
+      let best = -1, bestRank = Infinity;
+      for (let i = 0; i < parts.length - 1; i++) {
+        const r = this.ranks.get(parts[i] + ' ' + parts[i + 1]);
+        if (r !== undefined && r < bestRank) { bestRank = r; best = i; }
+      }
+      if (best < 0) break;
+      parts = [...parts.slice(0, best), parts[best] + parts[best + 1], ...parts.slice(best + 2)];
+    }
+    const ids = parts.map((p) => {
+      const id = this.ids.get(p);
+      if (id === undefined) throw new Error(`BPE: no token for ${JSON.stringify(p)}`);
+      return id;
+    });
+    if (this.cache.size < 50000) this.cache.set(word, ids);
+    return ids;
+  }
+  // parseSpecial: special-token text (e.g. <|im_start|>) maps to its single id.
+  encode(text, { parseSpecial = true } = {}) {
+    const out = [];
+    const pieces = parseSpecial && this.specialRe ? text.split(this.specialRe) : [text];
+    for (let i = 0; i < pieces.length; i++) {
+      const piece = pieces[i];
+      if (!piece) continue;
+      if (parseSpecial && i % 2 === 1) { out.push(this.ids.get(piece)); continue; }
+      for (const m of piece.matchAll(this.preRe)) {
+        let s = '';
+        for (const b of this.utf8.encode(m[0])) s += this.byteEnc[b];
+        out.push(...this.bpe(s));
+      }
+    }
+    return out;
+  }
+  tokenBytes(id) {
+    const t = this.tokens[id];
+    if (this.isSpecial[id]) return this.utf8.encode(t);
+    const bytes = [];
+    for (const ch of t) { const b = this.byteDec.get(ch); if (b !== undefined) bytes.push(b); }
+    return Uint8Array.from(bytes);
+  }
+  decode(ids) {
+    const chunks = ids.map((id) => this.tokenBytes(id));
+    const n = chunks.reduce((a, c) => a + c.length, 0), all = new Uint8Array(n);
+    let o = 0; for (const c of chunks) { all.set(c, o); o += c.length; }
+    return new TextDecoder().decode(all);
+  }
+}
+
+export function chatPrompt(messages, { enableThinking = true } = {}) {
+  let s = '';
+  for (const m of messages) s += `<|im_start|>${m.role}\n${m.content}<|im_end|>\n`;
+  s += '<|im_start|>assistant\n';
+  if (!enableThinking) s += '<think>\n\n</think>\n\n';
+  return s;
+}
+
+// ── WGSL ────────────────────────────────────────────────────────────────────
+// Activations f32 throughout; Q8_0 weights as an int8 plane (u32-packed) + an f16 scale per
+// 32 values, dequantized exactly in f32 inside the GEMV; KV cache f16 (llama.cpp's default).
+export const TOK = 'struct Tok { token: u32, pos: u32, seqLen: u32, pad: u32 }';
+export const I8 = `fn i8at(w: u32, lane: u32) -> f32 { let b = (w >> (lane * 8u)) & 0xffu; return f32(select(i32(b), i32(b) - 256, b >= 128u)); }`;
+const WGSL = {
+  embedQ8: `enable f16;
+${TOK}
+struct P { n: u32 }
+@group(0) @binding(0) var<storage, read> q: array<u32>;
+@group(0) @binding(1) var<storage, read> s: array<f16>;
+@group(0) @binding(2) var<storage, read_write> y: array<f32>;
+@group(0) @binding(3) var<uniform> tok: Tok;
+@group(0) @binding(4) var<uniform> p: P;
+${I8}
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let i = g.x; if (i >= p.n) { return; }
+  let e = tok.token * p.n + i;
+  y[i] = i8at(q[e >> 2u], e & 3u) * f32(s[e >> 5u]);
+}`,
+  rmsnorm: `
+struct P { n: u32, eps: f32 }
+@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> p: P;
+var<workgroup> red: array<f32, 256>;
+@compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
+  let t = l.x; var acc = 0.0;
+  for (var i = t; i < p.n; i += 256u) { let v = x[i]; acc += v * v; }
+  red[t] = acc; workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
+  let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
+  for (var i = t; i < p.n; i += 256u) { y[i] = (x[i] * scale) * w[i]; }
+}`,
+  // 4 output rows per workgroup; input read once per k. Dispatch ceil(M/4) workgroups (2-D ok).
+  matmulQ8: `enable f16;
+struct P { M: u32, N: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> qw: array<u32>;
+@group(0) @binding(2) var<storage, read> qs: array<f16>;
+@group(0) @binding(3) var<storage, read_write> y: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+${I8}
+var<workgroup> part: array<f32, 1024>;
+@compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
+  let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let N = p.N; let M = p.M; let nb = N / 32u;
+  if (m0 >= M) { return; }
+  let r0 = min(m0, M - 1u); let r1 = min(m0 + 1u, M - 1u); let r2 = min(m0 + 2u, M - 1u); let r3 = min(m0 + 3u, M - 1u);
+  let q0 = (r0 * N) >> 2u; let q1 = (r1 * N) >> 2u; let q2 = (r2 * N) >> 2u; let q3 = (r3 * N) >> 2u;
+  var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
+  for (var k = t; k < N; k += 256u) {
+    let xv = x[k]; let wi = k >> 2u; let ln = k & 3u; let b = k >> 5u;
+    a0 += i8at(qw[q0 + wi], ln) * f32(qs[r0 * nb + b]) * xv;
+    a1 += i8at(qw[q1 + wi], ln) * f32(qs[r1 * nb + b]) * xv;
+    a2 += i8at(qw[q2 + wi], ln) * f32(qs[r2 * nb + b]) * xv;
+    a3 += i8at(qw[q3 + wi], ln) * f32(qs[r3 * nb + b]) * xv;
+  }
+  part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
+  workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) {
+    if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
+    workgroupBarrier();
+  }
+  if (t < 4u && m0 + t < M) { y[m0 + t] = part[t * 256u]; }
+}`,
+  matmulF32: `
+struct P { M: u32, N: u32 }
+@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> p: P;
+var<workgroup> red: array<f32, 256>;
+@compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
+  let m = wg.x; let t = l.x; var acc = 0.0;
+  for (var k = t; k < p.N; k += 256u) { acc += w[m * p.N + k] * x[k]; }
+  red[t] = acc; workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
+  if (t == 0u) { y[m] = red[0]; }
+}`,
+  // Per-head RMSNorm (QK-norm) then NeoX RoPE (pairs i, i+hd/2), in place. One workgroup per head.
+  // rope[0..hd/2) = cos, rope[hd/2..hd) = sin for this token's position (computed on the CPU).
+  qkNormRope: `
+struct P { hd: u32, eps: f32 }
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+@group(0) @binding(1) var<storage, read> w: array<f32>;
+@group(0) @binding(2) var<storage, read> rope: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+var<workgroup> red: array<f32, 64>;
+@compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
+  let base = wg.x * p.hd; let t = l.x; let half = p.hd / 2u;
+  var acc = 0.0;
+  for (var i = t; i < p.hd; i += 64u) { let v = x[base + i]; acc += v * v; }
+  red[t] = acc; workgroupBarrier();
+  for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
+  let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
+  for (var i = t; i < half; i += 64u) {
+    let a = (x[base + i] * scale) * w[i];
+    let b = (x[base + i + half] * scale) * w[i + half];
+    let c = rope[i]; let sn = rope[half + i];
+    x[base + i] = a * c - b * sn;
+    x[base + i + half] = a * sn + b * c;
+  }
+}`,
+  kvStore: `enable f16;
+${TOK}
+struct P { n: u32 }
+@group(0) @binding(0) var<storage, read> k: array<f32>;
+@group(0) @binding(1) var<storage, read> v: array<f32>;
+@group(0) @binding(2) var<storage, read_write> kc: array<f16>;
+@group(0) @binding(3) var<storage, read_write> vc: array<f16>;
+@group(0) @binding(4) var<uniform> tok: Tok;
+@group(0) @binding(5) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let i = g.x; if (i >= p.n) { return; }
+  kc[tok.pos * p.n + i] = f16(k[i]);
+  vc[tok.pos * p.n + i] = f16(v[i]);
+}`,
+  attnScore: `enable f16;
+${TOK}
+struct P { heads: u32, kvHeads: u32, hd: u32, scale: f32 }
+@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_write> sc: array<f32>;
+@group(0) @binding(3) var<uniform> tok: Tok;
+@group(0) @binding(4) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let n = tok.seqLen; let idx = g.x; if (idx >= p.heads * n) { return; }
+  let h = idx / n; let pos = idx % n; let kvh = h / (p.heads / p.kvHeads);
+  let qo = h * p.hd; let ko = pos * p.kvHeads * p.hd + kvh * p.hd;
+  var acc = 0.0;
+  for (var d = 0u; d < p.hd; d++) { acc += q[qo + d] * f32(kc[ko + d]); }
+  sc[h * n + pos] = acc * p.scale;
+}`,
+  softmax: `
+${TOK}
+@group(0) @binding(0) var<storage, read_write> sc: array<f32>;
+@group(0) @binding(1) var<uniform> tok: Tok;
+var<workgroup> red: array<f32, 256>;
+@compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
+  let n = tok.seqLen; let base = wg.x * n; let t = l.x;
+  var mx = -3.4e38;
+  for (var i = t; i < n; i += 256u) { mx = max(mx, sc[base + i]); }
+  red[t] = mx; workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
+  let m = red[0]; workgroupBarrier();
+  var sum = 0.0;
+  for (var i = t; i < n; i += 256u) { let e = exp(sc[base + i] - m); sc[base + i] = e; sum += e; }
+  red[t] = sum; workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
+  let inv = 1.0 / red[0];
+  for (var i = t; i < n; i += 256u) { sc[base + i] = sc[base + i] * inv; }
+}`,
+  attnOut: `enable f16;
+${TOK}
+struct P { heads: u32, kvHeads: u32, hd: u32, pad: u32 }
+@group(0) @binding(0) var<storage, read> pr: array<f32>;
+@group(0) @binding(1) var<storage, read> vc: array<f16>;
+@group(0) @binding(2) var<storage, read_write> y: array<f32>;
+@group(0) @binding(3) var<uniform> tok: Tok;
+@group(0) @binding(4) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let idx = g.x; if (idx >= p.heads * p.hd) { return; }
+  let n = tok.seqLen; let h = idx / p.hd; let d = idx % p.hd; let kvh = h / (p.heads / p.kvHeads);
+  let stride = p.kvHeads * p.hd;
+  var acc = 0.0;
+  for (var pos = 0u; pos < n; pos++) { acc += pr[h * n + pos] * f32(vc[pos * stride + kvh * p.hd + d]); }
+  y[idx] = acc;
+}`,
+  addInPlace: `
+struct P { n: u32 }
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+@group(0) @binding(1) var<storage, read> y: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let i = g.x; if (i < p.n) { x[i] = x[i] + y[i]; }
+}`,
+  // Router: softmax over n logits (n ≤ 256), greedy top-k (ties → lowest index), weights
+  // renormalized to sum 1 (Qwen3-MoE norm_topk_prob). sel = k ids then k f32 weights (bitcast).
+  // One workgroup; each of the k picks is a parallel argmax over the per-thread probabilities.
+  // Every probability is the same f32 expression the serial version evaluated, so the picks and
+  // weights are the same; the serial version cost ~0.1 ms per call on an M4 Pro.
+  topk: `
+struct P { n: u32, k: u32 }
+@group(0) @binding(0) var<storage, read> lg: array<f32>;
+@group(0) @binding(1) var<storage, read_write> sel: array<u32>;
+@group(0) @binding(2) var<uniform> p: P;
+var<workgroup> red: array<f32, 256>;
+var<workgroup> bi: array<u32, 256>;
+var<workgroup> gmax: f32;
+var<workgroup> gsum: f32;
+var<workgroup> picks: array<u32, 32>;
+var<workgroup> probs: array<f32, 32>;
+@compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
+  let t = l.x; let n = p.n;
+  red[t] = select(-3.4e38, lg[min(t, n - 1u)], t < n); workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
+  if (t == 0u) { gmax = red[0]; } workgroupBarrier();
+  red[t] = select(0.0, exp(lg[min(t, n - 1u)] - gmax), t < n); workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
+  if (t == 0u) { gsum = red[0]; } workgroupBarrier();
+  var mine = select(-1.0, exp(lg[min(t, n - 1u)] - gmax) / gsum, t < n);
+  for (var j = 0u; j < p.k; j++) {
+    red[t] = mine; bi[t] = t; workgroupBarrier();
+    for (var s = 128u; s > 0u; s >>= 1u) {
+      if (t < s) { let o = red[t + s]; if (o > red[t] || (o == red[t] && bi[t + s] < bi[t])) { red[t] = o; bi[t] = bi[t + s]; } }
+      workgroupBarrier();
+    }
+    if (t == 0u) { picks[j] = bi[0]; probs[j] = red[0]; }
+    workgroupBarrier();
+    if (t == picks[j]) { mine = -1.0; }
+    workgroupBarrier();
+  }
+  if (t == 0u) {
+    var wsum = 0.0;
+    for (var j = 0u; j < p.k; j++) { wsum += probs[j]; }
+    for (var j = 0u; j < p.k; j++) { sel[j] = picks[j]; sel[p.k + j] = bitcast<u32>(probs[j] / wsum); }
+  }
+}`,
+  // Routed experts, all k at once: workgroup z = top-k position, slot = slots[z] picks the
+  // expert's rows in the pool buffer (rows_per_expert per slot). Input is shared (gate/up) or
+  // per-position (down). Output region z × M.
+  expertQ8: `enable f16;
+struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
+@group(0) @binding(0) var<storage, read> x: array<f32>;
+@group(0) @binding(1) var<storage, read> qw: array<u32>;
+@group(0) @binding(2) var<storage, read> qs: array<f16>;
+@group(0) @binding(3) var<storage, read_write> y: array<f32>;
+@group(0) @binding(4) var<uniform> p: P;
+@group(0) @binding(5) var<storage, read> slots: array<u32>;
+${I8}
+var<workgroup> part: array<f32, 1024>;
+@compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
+  let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let N = p.N; let M = p.M; let nb = N / 32u;
+  if (m0 >= M) { return; }
+  let rb = slots[z] * p.rowsPerExpert;
+  let xb = select(0u, z * N, p.inputPerSlot == 1u);
+  let r0 = rb + min(m0, M - 1u); let r1 = rb + min(m0 + 1u, M - 1u); let r2 = rb + min(m0 + 2u, M - 1u); let r3 = rb + min(m0 + 3u, M - 1u);
+  let q0 = (r0 * N) >> 2u; let q1 = (r1 * N) >> 2u; let q2 = (r2 * N) >> 2u; let q3 = (r3 * N) >> 2u;
+  var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
+  for (var k = t; k < N; k += 256u) {
+    let xv = x[xb + k]; let wi = k >> 2u; let ln = k & 3u; let b = k >> 5u;
+    a0 += i8at(qw[q0 + wi], ln) * f32(qs[r0 * nb + b]) * xv;
+    a1 += i8at(qw[q1 + wi], ln) * f32(qs[r1 * nb + b]) * xv;
+    a2 += i8at(qw[q2 + wi], ln) * f32(qs[r2 * nb + b]) * xv;
+    a3 += i8at(qw[q3 + wi], ln) * f32(qs[r3 * nb + b]) * xv;
+  }
+  part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
+  workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) {
+    if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
+    workgroupBarrier();
+  }
+  if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
+}`,
+  // gu = k × [gate F | up F] → act = k × F of SiLU(gate)·up.
+  siluMulMoe: `
+struct P { F: u32, k: u32 }
+@group(0) @binding(0) var<storage, read> gu: array<f32>;
+@group(0) @binding(1) var<storage, read_write> act: array<f32>;
+@group(0) @binding(2) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let i = g.x; if (i >= p.k * p.F) { return; }
+  let z = i / p.F; let j = i % p.F; let g0 = gu[z * 2u * p.F + j];
+  act[i] = (g0 / (1.0 + exp(-g0))) * gu[z * 2u * p.F + p.F + j];
+}`,
+  // x[h] += Σ_j w_j · out[j·H + h]   (the expert sum, then the residual add — llama.cpp's order)
+  moeAccum: `
+struct P { H: u32, k: u32 }
+@group(0) @binding(0) var<storage, read> out: array<f32>;
+@group(0) @binding(1) var<storage, read> sel: array<u32>;
+@group(0) @binding(2) var<storage, read_write> x: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
+  let h = g.x; if (h >= p.H) { return; }
+  var acc = 0.0;
+  for (var j = 0u; j < p.k; j++) { acc += bitcast<f32>(sel[p.k + j]) * out[j * p.H + h]; }
+  x[h] = x[h] + acc;
+}`,
+  argmax: `
+struct P { n: u32 }
+@group(0) @binding(0) var<storage, read> lg: array<f32>;
+@group(0) @binding(1) var<storage, read_write> res: array<u32>;
+@group(0) @binding(2) var<uniform> p: P;
+var<workgroup> bv: array<f32, 256>;
+var<workgroup> bi: array<u32, 256>;
+@compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
+  let t = l.x; var v = -3.4e38; var ix = 0u;
+  for (var i = t; i < p.n; i += 256u) { let x = lg[i]; if (x > v) { v = x; ix = i; } }
+  bv[t] = v; bi[t] = ix; workgroupBarrier();
+  for (var s = 128u; s > 0u; s >>= 1u) {
+    if (t < s) { let o = bv[t + s]; if (o > bv[t] || (o == bv[t] && bi[t + s] < bi[t])) { bv[t] = o; bi[t] = bi[t + s]; } }
+    workgroupBarrier();
+  }
+  if (t == 0u) { res[0] = bi[0]; }
+}`,
+};
+
+export const STORAGE = 0x80, COPY_SRC = 0x04, COPY_DST = 0x08, UNIFORM = 0x40, MAP_READ = 0x01;
+
+// ── Engine ──────────────────────────────────────────────────────────────────
+export class Qwen3MoeSsd {
+  // A subclass for another architecture overrides these (gemma4_moe_ssd.js).
+  static configFrom(kv) { return configFromGguf(kv); }
+  static get ingestPlan() { return undefined; }
+  static tokenizerFrom(kv) { return new BpeTokenizer(kv); }
+
+  static async load(modelId = QWEN3_30B_A3B.repo, opts = {}) {
+    const { onProgress = () => {}, signal, localFile } = opts;
+    const src = opts.source || (modelId === QWEN3_30B_A3B.repo || !modelId ? QWEN3_30B_A3B : { repo: modelId, file: opts.file, revision: opts.revision || 'main' });
+    // localFile: the same GGUF already on disk (a File the user picked), ingested instead of downloaded. A size
+    // other than the pinned file's is refused before any write; the GGUF header check in the ingest plan follows.
+    if (localFile && src.size && localFile.size !== src.size) {
+      throw new Error(`${localFile.name} is ${localFile.size.toLocaleString()} bytes, not ${src.file} (${src.size.toLocaleString()} bytes). Pick that file.`);
+    }
+    const fetchFn = localFile ? fileFetch(localFile) : (opts.fetch || ((u, i) => fetch(u, i)));
+    const url = opts.url || `https://huggingface.co/${src.repo}/resolve/${src.revision}/${src.file}`;
+    const key = opts.key || storeKey(opts.file || src.file || url.split('/').pop());
+    const root = opts.root || OPFS_ROOT;              // the OPFS folder for this engine's stores
+    const dir = `${root}/${key}`;
+
+    onProgress({ status: 'init' });
+    if (!navigator.gpu) throw new Error('WebGPU is not available');
+    const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
+    if (!adapter) throw new Error('no WebGPU adapter');
+    if (!adapter.features.has('shader-f16')) throw new Error('this GPU lacks shader-f16');
+    const L = adapter.limits;
+    // subgroups: optional; an engine may use subgroup kernels when the adapter offers them.
+    const features = ['shader-f16', ...(adapter.features.has('subgroups') ? ['subgroups'] : [])];
+    const device = await adapter.requestDevice({
+      requiredFeatures: features,
+      requiredLimits: { maxBufferSize: L.maxBufferSize, maxStorageBufferBindingSize: L.maxStorageBufferBindingSize, maxStorageBuffersPerShaderStage: L.maxStorageBuffersPerShaderStage, maxComputeWorkgroupsPerDimension: L.maxComputeWorkgroupsPerDimension },
+    });
+    device.lost.then((info) => console.warn('qwen3-moe-ssd: GPU device lost:', info.message));
+
+    let manifest = null;
+    try { manifest = JSON.parse(await readOpfsText(`${dir}/manifest.json`) || 'null'); } catch (_) { manifest = null; }
+    if (opts.reingest || !manifest || !manifest.complete || manifest.format !== FORMAT) {
+      manifest = await ingestGguf({
+        url, key, root, fetch: fetchFn, signal, source: { repo: src.repo, file: src.file, revision: src.revision, sha256: src.sha256, size: src.size, ...(localFile ? { localFile: localFile.name } : {}) },
+        onProgress: (e) => onProgress(ingestProgress(e)),
+        plan: this.ingestPlan,
+      });
+    }
+    const headerFile = await (await (await navigator.storage.getDirectory()).getDirectoryHandle(root)).getDirectoryHandle(key);
+    const headerBytes = new Uint8Array(await (await (await headerFile.getFileHandle('header.bin')).getFile()).arrayBuffer());
+    const gguf = parseGguf(headerBytes);
+    const engine = new this(device, manifest, gguf, { ...opts, dir, key, url, adapterInfo: adapter.info || {} });
+    await engine.init(onProgress);
+    return engine;
+  }
+
+  constructor(device, manifest, gguf, opts) {
+    this.device = device;
+    this.manifest = manifest;
+    this.cfg = this.constructor.configFrom(gguf.kv);
+    this.tokenizer = this.constructor.tokenizerFrom(gguf.kv);
+    this.opts = opts;
+    this.maxCtx = Math.min(opts.maxCtx || 4096, this.cfg.contextLength);
+    this.gpuBytes = { dense: 0, pool: 0, kv: 0, act: 0 };
+    this.pipelines = {};
+    this.position = 0;
+    this.cached = [];          // token ids whose K/V are in the cache, in order
+    // Prefetch: routers of the next `lookahead` layers guess their experts from this layer's
+    // post-attention state. Measured 2026-10-04 (full model, 4 GB pool, M4 Pro): off 4.41,
+    // lookahead 1 5.02, 2 5.48, 4 4.56, 6 3.90 tok/s — guesses beyond 2 layers waste SSD bandwidth.
+    this.prefetch = opts.prefetch ?? true;
+    this.lookahead = Math.max(1, Math.floor(opts.lookahead ?? 2));
+    this.resetCounters();
+  }
+
+  buffer(size, usage, cat = 'act') {
+    const b = this.device.createBuffer({ size: Math.ceil(size / 4) * 4, usage });
+    this.gpuBytes[cat] = (this.gpuBytes[cat] || 0) + b.size;
+    return b;
+  }
+  uniform(words) {
+    const b = this.buffer(Math.max(16, words.length * 4), UNIFORM | COPY_DST, 'act');
+    const ab = new ArrayBuffer(b.size), u = new Uint32Array(ab), f = new Float32Array(ab);
+    words.forEach((w, i) => { if (typeof w === 'object') f[i] = w.f; else u[i] = w; });
+    this.device.queue.writeBuffer(b, 0, ab);
+    return b;
+  }
+  // The kernel sources; a subclass adds its own.
+  get kernels() { return WGSL; }
+  pipeline(name) {
+    if (!this.pipelines[name]) {
+      const code = this.kernels[name];
+      if (!code) throw new Error(`no kernel ${name}`);
+      const module = this.device.createShaderModule({ code, label: name });
+      this.pipelines[name] = this.device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' }, label: name });
+    }
+    return this.pipelines[name];
+  }
+  bind(name, buffers) {
+    return this.device.createBindGroup({
+      layout: this.pipeline(name).getBindGroupLayout(0),
+      entries: buffers.map((b, i) => ({ binding: i, resource: { buffer: b } })),
+    });
+  }
+
+  async init(onProgress) {
+    const dev = this.device;
+    // Fail early on a shader that does not compile.
+    dev.pushErrorScope('validation');
+    for (const k of Object.keys(this.kernels)) this.pipeline(k);
+    const err = await dev.popErrorScope();
+    if (err) throw new Error(`WGSL: ${err.message}`);
+    await this.uploadDense(onProgress);
+    await this.initPool();
+    this.initBuffers();
+    await dev.queue.onSubmittedWorkDone();
+  }
+
+  // Dense weights: OPFS → GPU (this.w[name] = { q, s, rows, cols } or { raw }).
+  async uploadDense(onProgress) {
+    const dev = this.device, m = this.manifest;
+    const t0 = performance.now();
+    const reader = await OpfsReaderPool.open(`${this.opts.dir}/${m.dense.file}`, { workers: 4 });
+    const names = Object.keys(m.dense.tensors);
+    this.w = {};
+    let doneT = 0;
+    const CH = 64 << 20;
+    const upload = async (gbuf, part) => {
+      for (let o = 0; o < part.bytes; o += CH) {
+        const n = Math.min(CH, part.bytes - o);
+        const r = await reader.read(part.off + o, n);
+        if (r.got !== n) throw new Error(`dense read short at ${part.off + o}`);
+        dev.queue.writeBuffer(gbuf, o, r.buf, 0, n);
+      }
+    };
+    const queue = [];
+    for (const name of names) {
+      const d = m.dense.tensors[name];
+      if (d.type === 'q8' || d.type === 'q4') {   // q4: Gemma 4 (two values per byte)
+        const q = this.buffer(d.q.bytes, STORAGE | COPY_DST, 'dense'), s = this.buffer(d.s.bytes, STORAGE | COPY_DST, 'dense');
+        this.w[name] = { q, s, rows: (d.type === 'q4' ? 2 * d.q.bytes : d.q.bytes) / d.dims[0], cols: d.dims[0] };
+        queue.push(() => upload(q, d.q), () => upload(s, d.s));
+      } else {
+        const raw = this.buffer(d.raw.bytes, STORAGE | COPY_DST, 'dense');
+        this.w[name] = { raw };
+        queue.push(() => upload(raw, d.raw));
+      }
+    }
+    let qi = 0;
+    await Promise.all([0, 1, 2, 3].map(async () => {
+      while (qi < queue.length) { await queue[qi++](); doneT++; onProgress({ status: 'weights', kind: 'tensors', loaded: doneT, total: queue.length }); }
+    }));
+    await reader.close();
+    await dev.queue.onSubmittedWorkDone();
+    this.denseUploadSecs = (performance.now() - t0) / 1000;
+  }
+
+  // Bytes for the expert cache. With opts.gpuBudgetBytes, the cache gets what the budget leaves after the weights
+  // that stay on the GPU, the KV cache and working buffers; a budget too small for those is refused. Otherwise
+  // opts.poolBytes, default 4 GiB.
+  poolBudget() {
+    const o = this.opts;
+    if (!o.gpuBudgetBytes) return o.poolBytes || 4 * 2 ** 30;
+    const gb = (n) => (n / 1e9).toFixed(2) + ' GB';
+    const fixed = this.gpuBytes.dense + this.kvBytes() + (o.workBytes ?? 512 * 2 ** 20);
+    const rest = o.gpuBudgetBytes - fixed;
+    const min = this.cfg.topK * 2 * this.manifest.experts.record;
+    if (rest < min) throw new Error(`GPU budget ${gb(o.gpuBudgetBytes)} is too small: weights that stay on the GPU, the KV cache and working buffers need ${gb(fixed)}, and the expert cache at least ${gb(min)} more`);
+    return rest;
+  }
+  // KV cache bytes at this.maxCtx (f16 K and V for every layer); engines with other layouts override it.
+  kvBytes() { const c = this.cfg; return 2 * 2 * c.layers * this.maxCtx * c.kvHeads * c.headDim; }
+
+  // Expert slot pool + the streamer that feeds it.
+  async initPool() {
+    const c = this.cfg, dev = this.device, xp = this.manifest.experts;
+    const maxSlotsByBinding = Math.floor(dev.limits.maxStorageBufferBindingSize / xp.parts.guQ.bytes);
+    this.poolSlots = Math.max(c.topK * 2, Math.min(maxSlotsByBinding, Math.floor(this.poolBudget() / xp.record), c.layers * c.experts));
+    this.pool = {};
+    for (const [k, part] of Object.entries(xp.parts)) this.pool[k] = this.buffer(this.poolSlots * part.bytes, STORAGE | COPY_DST, 'pool');
+    this.expertReader = await OpfsReaderPool.open(`${this.opts.dir}/${xp.file}`, { workers: this.opts.readers || 4 });
+    this.xs = new ExpertStreamer({
+      device: dev, reader: this.expertReader, recordBytes: xp.record,
+      recordOffset: (layer, e) => (layer * xp.perLayer + e) * xp.record,
+      parts: ['guQ', 'guS', 'dQ', 'dS'].map((k) => ({ buffer: this.pool[k], srcOffset: xp.parts[k].off, bytes: xp.parts[k].bytes })),
+      slots: this.poolSlots, numLayers: c.layers, numExperts: c.experts,
+      ...(this.opts.uploadRing !== undefined ? { uploadRing: this.opts.uploadRing } : {}),
+      ...(this.opts.evict ? { evict: this.opts.evict, hotHalfLife: this.opts.hotHalfLife } : {}),
+    });
+    this.gpuBytes.staging = this.xs.ringBytes();
+  }
+
+  // KV cache (f16), activations, uniforms and every layer's bind groups.
+  initBuffers() {
+    const c = this.cfg;
+    const kvn = c.kvHeads * c.headDim;
+    this.kc = []; this.vc = [];
+    for (let l = 0; l < c.layers; l++) {
+      this.kc.push(this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv'));
+      this.vc.push(this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv'));
+    }
+    const H = c.hidden, QN = c.heads * c.headDim, F = c.expertFf, K = c.topK;
+    if (2 * K * 4 > 256) throw new Error(`top-k ${K} > 32: selections would overlap their 256-byte readback regions`);
+    if (c.experts > 256) throw new Error(`${c.experts} experts: the router top-k kernel handles at most 256`);
+    const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST);
+    this.a = {
+      x: A(H), xn: A(H), q: A(QN), k: A(kvn), v: A(kvn), att: A(QN), o: A(H), sc: A(c.heads * this.maxCtx),
+      rl: A(c.experts), sel: A(2 * K), pxn: A(H), prl: A(c.experts), psel: A(64 * this.lookahead), slots: A(K),
+      gu: A(K * 2 * F), act: A(K * F), dn: A(K * H), logits: A(c.vocab), am: A(4),
+      rope: A(c.headDim),
+    };
+    this.tok = this.buffer(16, UNIFORM | COPY_DST);
+    // Readback: this layer's selection at 0, guess d (layers ahead) at 256 × d bytes.
+    this.rbSel = this.buffer(256 * (1 + this.lookahead), MAP_READ | COPY_DST);
+    this.rbLogits = this.buffer(c.vocab * 4, MAP_READ | COPY_DST);
+    this.rbArg = this.buffer(16, MAP_READ | COPY_DST);
+
+    const u = {
+      nH: this.uniform([H]), rmsH: this.uniform([H, { f: c.eps }]), mmQ: this.uniform([QN, H]), mmKV: this.uniform([kvn, H]),
+      mmO: this.uniform([H, QN]), qk: this.uniform([c.headDim, { f: c.eps }]), kvn: this.uniform([kvn]),
+      att: this.uniform([c.heads, c.kvHeads, c.headDim, { f: 1 / Math.sqrt(c.headDim) }]), attOut: this.uniform([c.heads, c.kvHeads, c.headDim, 0]),
+      router: this.uniform([c.experts, H]), topk: this.uniform([c.experts, K]),
+      gu: this.uniform([2 * F, H, 2 * F, 0]), dn: this.uniform([H, F, H, 1]), silu: this.uniform([F, K]), acc: this.uniform([H, K]),
+      lm: this.uniform([c.vocab, H]), am: this.uniform([c.vocab]),
+    };
+    const a = this.a, W = this.w;
+    this.g = { embed: this.bind('embedQ8', [W['token_embd.weight'].q, W['token_embd.weight'].s, a.x, this.tok, u.nH]) };
+    this.layers = [];
+    for (let l = 0; l < c.layers; l++) {
+      const p = (n) => W[`blk.${l}.${n}.weight`];
+      this.layers.push({
+        rmsA: this.bind('rmsnorm', [a.x, p('attn_norm').raw, a.xn, u.rmsH]),
+        q: this.bind('matmulQ8', [a.xn, p('attn_q').q, p('attn_q').s, a.q, u.mmQ]),
+        k: this.bind('matmulQ8', [a.xn, p('attn_k').q, p('attn_k').s, a.k, u.mmKV]),
+        v: this.bind('matmulQ8', [a.xn, p('attn_v').q, p('attn_v').s, a.v, u.mmKV]),
+        ropeQ: this.bind('qkNormRope', [a.q, p('attn_q_norm').raw, a.rope, u.qk]),
+        ropeK: this.bind('qkNormRope', [a.k, p('attn_k_norm').raw, a.rope, u.qk]),
+        kv: this.bind('kvStore', [a.k, a.v, this.kc[l], this.vc[l], this.tok, u.kvn]),
+        score: this.bind('attnScore', [a.q, this.kc[l], a.sc, this.tok, u.att]),
+        soft: this.bind('softmax', [a.sc, this.tok]),
+        attOut: this.bind('attnOut', [a.sc, this.vc[l], a.att, this.tok, u.attOut]),
+        o: this.bind('matmulQ8', [a.att, p('attn_output').q, p('attn_output').s, a.o, u.mmO]),
+        addO: this.bind('addInPlace', [a.x, a.o, u.nH]),
+        rmsF: this.bind('rmsnorm', [a.x, p('ffn_norm').raw, a.xn, u.rmsH]),
+        router: this.bind('matmulF32', [a.xn, p('ffn_gate_inp').raw, a.rl, u.router]),
+        topk: this.bind('topk', [a.rl, a.sel, u.topk]),
+        // Routers of layers l+1..l+lookahead on this layer's post-attention state: the prefetch
+        // guesses, each written to its own 256-byte region of psel.
+        pf: Array.from({ length: Math.min(this.lookahead, c.layers - 1 - l) }, (_, i) => {
+          const at = (n) => W[`blk.${l + 1 + i}.${n}.weight`];
+          return {
+            rms: this.bind('rmsnorm', [a.x, at('ffn_norm').raw, a.pxn, u.rmsH]),
+            router: this.bind('matmulF32', [a.pxn, at('ffn_gate_inp').raw, a.prl, u.router]),
+            topk: this.device.createBindGroup({ layout: this.pipeline('topk').getBindGroupLayout(0), entries: [
+              { binding: 0, resource: { buffer: a.prl } },
+              { binding: 1, resource: { buffer: a.psel, offset: 256 * i, size: 8 * K } },
+              { binding: 2, resource: { buffer: u.topk } },
+            ] }),
+          };
+        }),
+        gu: this.bind('expertQ8', [a.xn, this.pool.guQ, this.pool.guS, a.gu, u.gu, a.slots]),
+        silu: this.bind('siluMulMoe', [a.gu, a.act, u.silu]),
+        dn: this.bind('expertQ8', [a.act, this.pool.dQ, this.pool.dS, a.dn, u.dn, a.slots]),
+        acc: this.bind('moeAccum', [a.dn, a.sel, a.x, u.acc]),
+      });
+    }
+    this.g.rmsOut = this.bind('rmsnorm', [a.x, W['output_norm.weight'].raw, a.xn, u.rmsH]);
+    this.g.lm = this.bind('matmulQ8', [a.xn, W['output.weight'].q, W['output.weight'].s, a.logits, u.lm]);
+    this.g.am = this.bind('argmax', [a.logits, a.am, u.am]);
+  }
+
+  // ── forward ───────────────────────────────────────────────────────────────
+  dispatch(pass, name, group, x, y = 1, z = 1) {
+    pass.setPipeline(this.pipeline(name));
+    pass.setBindGroup(0, group);
+    if (x > 65535) { y = Math.ceil(x / 32768); x = 32768; }
+    pass.dispatchWorkgroups(x, y, z);
+  }
+  encodeAttention(pass, l, seqLen) {
+    const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
+    const QN = c.heads * c.headDim, kvn = c.kvHeads * c.headDim;
+    d('rmsnorm', g.rmsA, 1);
+    d('matmulQ8', g.q, Math.ceil(QN / 4)); d('matmulQ8', g.k, Math.ceil(kvn / 4)); d('matmulQ8', g.v, Math.ceil(kvn / 4));
+    d('qkNormRope', g.ropeQ, c.heads); d('qkNormRope', g.ropeK, c.kvHeads);
+    d('kvStore', g.kv, Math.ceil(kvn / 256));
+    d('attnScore', g.score, Math.ceil(c.heads * seqLen / 256));
+    d('softmax', g.soft, c.heads);
+    d('attnOut', g.attOut, Math.ceil(QN / 256));
+    d('matmulQ8', g.o, Math.ceil(c.hidden / 4));
+    d('addInPlace', g.addO, Math.ceil(c.hidden / 256));
+    if (this.prefetch) for (const p of g.pf) { d('rmsnorm', p.rms, 1); d('matmulF32', p.router, c.experts); d('topk', p.topk, 1); }
+    d('rmsnorm', g.rmsF, 1);
+    d('matmulF32', g.router, c.experts);
+    d('topk', g.topk, 1);
+  }
+  // The token's embedding into a.x, and the head (final norm, logits, argmax when wanted): a subclass
+  // with another embedding or head format overrides these two.
+  encodeEmbed(pass) { this.dispatch(pass, 'embedQ8', this.g.embed, Math.ceil(this.cfg.hidden / 256)); }
+  encodeHead(pass, want) {
+    this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);
+    this.dispatch(pass, 'matmulQ8', this.g.lm, Math.ceil(this.cfg.vocab / 4));
+    if (want === 'argmax') this.dispatch(pass, 'argmax', this.g.am, 1);
+  }
+  encodeExperts(pass, l) {
+    const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
+    d('expertQ8', g.gu, Math.ceil(2 * c.expertFf / 4), 1, c.topK);
+    d('siluMulMoe', g.silu, Math.ceil(c.topK * c.expertFf / 256));
+    d('expertQ8', g.dn, Math.ceil(c.hidden / 4), 1, c.topK);
+    d('moeAccum', g.acc, Math.ceil(c.hidden / 256));
+  }
+
+  writeTokenUniforms(token, pos) {
+    const c = this.cfg;
+    this.device.queue.writeBuffer(this.tok, 0, new Uint32Array([token, pos, pos + 1, 0]));
+    // NeoX RoPE angles the way llama.cpp's Metal kernel forms them: f32 pos × base^(-2i/hd).
+    const half = c.headDim / 2, r = new Float32Array(c.headDim), inv = Math.fround(-1 / c.headDim);
+    for (let i = 0; i < half; i++) {
+      const th = Math.fround(pos * Math.fround(Math.pow(c.ropeTheta, Math.fround(inv * 2 * i))));
+      r[i] = Math.cos(th); r[half + i] = Math.sin(th);
+    }
+    this.device.queue.writeBuffer(this.a.rope, 0, r);
+  }
+
+  resetCounters() {
+    this.counters = { tokens: 0, gpuWaitMs: 0, ensureMs: 0, encodeMs: 0, wallMs: 0, predHits: 0, predTotal: 0 };
+    if (this.xs) this.xs.resetStats();
+  }
+
+  // → { ids, preds }: this layer's top-k ids and `nPred` guesses for the layers after it.
+  async readSel(nPred) {
+    const K = this.cfg.topK, bytes = 256 * (1 + nPred);
+    const t0 = performance.now();
+    await this.rbSel.mapAsync(1, 0, bytes);
+    this.counters.gpuWaitMs += performance.now() - t0;
+    const u = new Uint32Array(this.rbSel.getMappedRange(0, bytes).slice(0));
+    this.rbSel.unmap();
+    const preds = [];
+    for (let i = 1; i <= nPred; i++) preds.push(u.subarray(64 * i, 64 * i + K));
+    return { ids: u.subarray(0, K), preds };
+  }
+  // Encodes the readback of layer l's selection and its guesses.
+  copySel(enc, l) {
+    enc.copyBufferToBuffer(this.a.sel, 0, this.rbSel, 0, 8 * this.cfg.topK);
+    const n = this.prefetch ? this.layers[l].pf.length : 0;
+    if (n) enc.copyBufferToBuffer(this.a.psel, 0, this.rbSel, 256, 256 * n);
+    return n;
+  }
+
+  // One token through all layers at position this.position. want: 'none' | 'argmax' | 'logits'.
+  async step(token, want = 'argmax') {
+    const c = this.cfg, dev = this.device, K = c.topK;
+    if (this.position >= this.maxCtx) throw new Error(`context full (${this.maxCtx} tokens)`);
+    const tStart = performance.now();
+    const pos = this.position, seqLen = pos + 1;
+    this.writeTokenUniforms(token, pos);
+    let te = performance.now();
+    let enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
+    this.encodeEmbed(pass);
+    this.encodeAttention(pass, 0, seqLen);
+    pass.end();
+    let nPred = this.copySel(enc, 0);
+    dev.queue.submit([enc.finish()]);
+    this.counters.encodeMs += performance.now() - te;
+    let predIds = null;
+    for (let l = 0; l < c.layers; l++) {
+      const { ids, preds } = await this.readSel(nPred);
+      if (predIds) { // accuracy of the one-layer-ahead guess made one layer earlier
+        let hit = 0; for (const e of ids) if (predIds.includes(e)) hit++;
+        this.counters.predHits += hit; this.counters.predTotal += K;
+      }
+      const tq = performance.now();
+      const ensuring = this.xs.ensure(l, ids);
+      preds.forEach((p, i) => this.xs.prefetch(l + 1 + i, p));   // nearest layer first
+      predIds = preds.length ? Array.from(preds[0]) : null;
+      const slots = await ensuring;
+      this.counters.ensureMs += performance.now() - tq;
+      dev.queue.writeBuffer(this.a.slots, 0, slots);
+      te = performance.now();
+      enc = dev.createCommandEncoder(); pass = enc.beginComputePass();
+      this.encodeExperts(pass, l);
+      if (l + 1 < c.layers) {
+        this.encodeAttention(pass, l + 1, seqLen);
+        pass.end();
+        nPred = this.copySel(enc, l + 1);
+      } else {
+        if (want !== 'none') this.encodeHead(pass, want);
+        pass.end();
+        if (want === 'argmax') enc.copyBufferToBuffer(this.a.am, 0, this.rbArg, 0, 4);
+        if (want === 'logits') enc.copyBufferToBuffer(this.a.logits, 0, this.rbLogits, 0, c.vocab * 4);
+      }
+      dev.queue.submit([enc.finish()]);
+      this.counters.encodeMs += performance.now() - te;
+      this.xs.release(slots);
+    }
+    this.position++;
+    this.cached.push(token);
+    let result = null;
+    const tw = performance.now();
+    if (want === 'argmax') {
+      await this.rbArg.mapAsync(1, 0, 4);
+      result = new Uint32Array(this.rbArg.getMappedRange(0, 4))[0];
+      this.rbArg.unmap();
+    } else if (want === 'logits') {
+      await this.rbLogits.mapAsync(1);
+      result = new Float32Array(this.rbLogits.getMappedRange().slice(0));
+      this.rbLogits.unmap();
+    } else {
+      await dev.queue.onSubmittedWorkDone();
+    }
+    this.counters.gpuWaitMs += performance.now() - tw;
+    this.counters.tokens++;
+    this.counters.wallMs += performance.now() - tStart;
+    return result;
+  }
+
+  reset() { this.position = 0; this.cached = []; }
+
+  // Feeds `ids` after the cached prefix they share (re-prefilling only what changed); returns
+  // the logits/argmax of the last token.
+  async prefill(ids, want = 'argmax') {
+    let common = 0;
+    while (common < ids.length - 1 && common < this.cached.length && this.cached[common] === ids[common]) common++;
+    if (common < this.cached.length) { this.position = common; this.cached.length = common; }
+    let r = null;
+    for (let i = common; i < ids.length; i++) r = await this.step(ids[i], i === ids.length - 1 ? want : 'none');
+    return r;
+  }
+
+  // Diagnostics for the llama.cpp comparison: greedy continuation of raw token ids, with the
+  // top-n logits at every generated position.
+  async greedy(ids, n, { top = 5 } = {}) {
+    this.reset();
+    const out = [], tops = [];
+    let logits = await this.prefill(ids, 'logits');
+    for (let i = 0; i < n; i++) {
+      const order = topN(logits, top);
+      let mx = -Infinity, sum = 0;
+      for (let j = 0; j < logits.length; j++) if (logits[j] > mx) mx = logits[j];
+      for (let j = 0; j < logits.length; j++) sum += Math.exp(logits[j] - mx);
+      const lse = mx + Math.log(sum);
+      for (const o of order) o.logprob = o.logit - lse;
+      tops.push(order);
+      const next = order[0].id;
+      out.push(next);
+      if (i + 1 < n) logits = await this.step(next, 'logits');
+    }
+    return { ids: out, tops };
+  }
+
+  async *generate(messages, { maxNewTokens = 512, signal, enableThinking = true } = {}) {
+    const ids = this.tokenizer.encode(this.chatPrompt(messages, { enableThinking }));
+    const stops = new Set(this.stopTokenIds);
+    let next = await this.prefill(ids, 'argmax');
+    const outIds = [];
+    for (let i = 0; i < maxNewTokens; i++) {
+      if (signal && signal.aborted) break;
+      if (stops.has(next)) break;
+      outIds.push(next);
+      yield { text: this.tokenizer.decode(outIds), token: next, tokens: outIds.length };
+      if (this.position >= this.maxCtx) break;
+      next = await this.step(next, 'argmax');
+    }
+  }
+
+  chatPrompt(messages, opts) { return chatPrompt(messages, opts); }
+  get stopTokenIds() { return [this.cfg.eos, this.tokenizer.ids.get('<|im_end|>'), this.tokenizer.ids.get('<|endoftext|>')].filter((t) => t !== undefined); }
+
+  // Raw ids of the reasoning delimiters, for hosts that re-mark the thought block.
+  get thinkOpenTokenId() { return this.tokenizer.ids.get('<think>') ?? null; }
+  get thinkCloseTokenId() { return this.tokenizer.ids.get('</think>') ?? null; }
+
+  stats() {
+    const s = this.xs.stats, c = this.counters, n = Math.max(1, c.tokens);
+    return {
+      tokens: c.tokens, tokPerSec: c.tokens / (c.wallMs / 1000),
+      msPerToken: c.wallMs / n, gpuWaitMsPerToken: c.gpuWaitMs / n, ensureMsPerToken: c.ensureMs / n, encodeMsPerToken: c.encodeMs / n,
+      readMsPerToken: s.readMs / n, uploadMsPerToken: s.uploadMs / n,
+      hitRate: (s.hits + s.lateHits) / Math.max(1, s.hits + s.lateHits + s.misses),
+      hits: s.hits, lateHits: s.lateHits, misses: s.misses, evictions: s.evictions,
+      bytesReadPerToken: s.bytesRead / n, prefetchIssued: s.prefetchIssued, prefetchUsed: s.prefetchUsed,
+      prefetchAccuracy: c.predTotal ? c.predHits / c.predTotal : null,
+      hitRateByLayer: Array.from(s.hitsByLayer, (h, l) => +(h / Math.max(1, h + s.missesByLayer[l])).toFixed(3)),
+      poolSlots: this.poolSlots, resident: this.xs.resident(), gpuBytes: { ...this.gpuBytes },
+    };
+  }
+
+  async warmup() {}
+
+  async dispose() {
+    try { await this.xs.drain(); } catch (_) {}
+    try { await this.expertReader.close(); } catch (_) {}
+    try { this.device.destroy(); } catch (_) {}
+  }
+}
+
+export function topN(logits, n) {
+  const best = [];
+  for (let i = 0; i < logits.length; i++) {
+    const v = logits[i];
+    if (best.length < n || v > best[best.length - 1].logit) {
+      best.push({ id: i, logit: v });
+      best.sort((a, b) => b.logit - a.logit || a.id - b.id);
+      if (best.length > n) best.pop();
+    }
+  }
+  return best;
+}
+
+export async function removeIngest(key) { await removeStore(key, { root: OPFS_ROOT }); }