diskformer 0.5.2 → 0.5.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/CHANGELOG.md CHANGED
@@ -1,5 +1,30 @@
1
1
  # Changelog
2
2
 
3
+ ## 0.5.4 (2026-10-22)
4
+ - `diskformer/llama`: the GPU budget plan picks the physical batch size (llama.cpp's `n_ubatch`) as the largest of
5
+ 512, 256, 128 and 64 that leaves room for one token's experts. The compute buffers scale with it (Gemma 4 26B-A4B:
6
+ 542 MiB at 512, 271 at 256, 135 at 128), so a tight budget now runs instead of failing: Gemma 4 runs at 3 GB
7
+ (8 slots, batch 256; 8/9 replies identical to llama.cpp in Chrome and natively), where it needed about 3.4 GB. A
8
+ smaller batch slows long prompts (649 tokens: 95 tok/s at 512, 46 at 128), so ample budgets keep 512.
9
+ `DF_UBATCH` fixes the size. `scripts/chrome-gate.mjs --budget GB`.
10
+ - Flint 40B-A2.5B (RecursiveSI, architecture `flint`, Q4_1, 26.0 GB) runs on the llama.cpp path and is in the demo
11
+ list. Paging is exact (9/9 replies identical at 8 and 32 slots natively); in Chrome at a 4 GB budget (26 slots),
12
+ 6/9 replies are identical to llama.cpp on the CPU, and the other 3 part at word choices inside the reasoning,
13
+ 1.8-3.4 tok/s. 17 architectures now pass both checks.
14
+ - Browser build (`ggml/dist`) rebuilt on the fork after merging 178 upstream commits (`ecaa47b29`, adds Flint),
15
+ Emscripten 6.0.13.
16
+ - `scripts/sync-llama-fork.sh` puts Homebrew on PATH, so the weekly launchd run finds cmake (the 2026-10-18 run failed
17
+ without it).
18
+ - Engines: `gemma4_moe_ssd.js` synced from LocalMind (embedding-kernel hooks); chat gate 9/9.
19
+
20
+ ## 0.5.3 (2026-10-15)
21
+ - Engines synced from LocalMind: the Q4_0 and Q6_K kernels are shared from `engines/qwen3_moe_ssd.js`;
22
+ `engines/qwen35_moe_ssd.js` reads Q4_0 files and carries MiniCPM-V 4.7's image encoder
23
+ (`engines/minicpmv_vision.js`, new). Gemma 4 26B-A4B chat gate: 9/9 replies identical to llama.cpp at a 2.5 GB
24
+ budget (2.02 GB on the GPU).
25
+ - docs/coverage.md: Velvet-Eclipse 4x12B (Mixtral-style `llama` MoE) runs with paged experts, too slowly to be
26
+ practical (0.74 tok/s, 1.4 GB read per token).
27
+
3
28
  ## 0.5.2 (2026-10-14)
4
29
  - `diskformer/llama` works when imported from a CDN: a `Llama` on another origin starts its worker from a same-origin
5
30
  stub that imports `worker.mjs`, because browsers refuse a cross-origin worker URL. Checked from a localhost page
package/README.md CHANGED
@@ -2,7 +2,7 @@
2
2
 
3
3
  **Run a mixture-of-experts model larger than your GPU memory in a browser tab, from one import.** llama.cpp runs on
4
4
  WebGPU in a worker; the GGUF stays on disk in the browser's private storage (OPFS), and only the experts each token
5
- needs are paged into GPU memory. 16 architectures are checked against llama.cpp, from Qwen3.6 35B to gpt-oss 20B.
5
+ needs are paged into GPU memory. 17 architectures are checked against llama.cpp, from Qwen3.6 35B to gpt-oss 20B.
6
6
 
7
7
  Plain ES modules, no dependencies, no build step. Chrome or Edge with WebGPU, on any origin (CDN or self-hosted).
8
8
  Weights stay in the user's own browser storage. The disk tier underneath (`src/`) also runs two hand-written engines.
@@ -53,8 +53,8 @@ once. Measured in Chrome on a MacBook Pro M4 Pro, 24 GB, 2026-10-13 (the recordi
53
53
  | OLMoE 1B-7B, Q4_K_M | 4.21 GB | 3.92 GB | 22.4 tok/s | 9/9 identical |
54
54
  | LFM2 8B-A1B, Q4_K_M | 5.04 GB | 3.89 GB | 17.1 tok/s | 9/9 identical |
55
55
 
56
- 16 architectures pass both checks (paging exact natively, then Chrome against llama.cpp): Gemma 4, Qwen3, Qwen3-Next, Qwen3.6,
57
- Qwen3-VL, GLM-4.7-Flash, Nemotron 3, ERNIE 4.5, gpt-oss, North, LFM2/2.5, Mellum2, Ling 2.0/3.0, Granite 4.0 H, OLMoE. [docs/coverage.md](docs/coverage.md).
56
+ 17 architectures pass both checks (paging exact natively, then Chrome against llama.cpp): Gemma 4, Qwen3, Qwen3-Next, Qwen3.6,
57
+ Qwen3-VL, GLM-4.7-Flash, Nemotron 3, ERNIE 4.5, gpt-oss, North, LFM2/2.5, Mellum2, Ling 2.0/3.0, Granite 4.0 H, OLMoE, Flint. [docs/coverage.md](docs/coverage.md).
58
58
 
59
59
  ## Why
60
60
 
@@ -17,7 +17,7 @@
17
17
  * and `ffn_down_exps` [F, H, E], both Q4_0, plus `ffn_down_exps.scale` [E].
18
18
  */
19
19
 
20
- import { Qwen3MoeSsd, GGML, Q4_BLOCK, Q6K_BLOCK, tensorBytes, TOK, STORAGE, COPY_SRC, COPY_DST, MAP_READ } from './qwen3_moe_ssd.js';
20
+ import { Qwen3MoeSsd, GGML, Q4_BLOCK, Q6K_BLOCK, tensorBytes, TOK, STORAGE, COPY_SRC, COPY_DST, MAP_READ, Q4_KERNELS, Q4_SG_KERNELS, Q4_BATCH_KERNELS } from './qwen3_moe_ssd.js';
21
21
  import { OpfsReaderPool } from '../src/opfs-reader.js';
22
22
  import { QWEN35_BATCH_KERNELS } from './qwen35_moe_ssd.js';
23
23
 
@@ -226,37 +226,6 @@ export class GemmaTokenizer {
226
226
  }
227
227
  }
228
228
 
229
- // ── WGSL ────────────────────────────────────────────────────────────────────
230
- // Q4_0 as two planes (ingest's splitQ4): nibbles, 16 bytes per 32 values (byte j holds values j and j + 16),
231
- // and one f16 scale per 32. Value = (nibble − 8) · scale, as ggml dequantizes it. Each thread reads one u32 of
232
- // nibbles: values 4m..4m+3 (low nibbles) and 16+4m..16+4m+3 (high nibbles) of block b.
233
- const Q4DOT = `
234
- fn q4dot(wd: u32, d: f32, xb: u32) -> f32 {
235
- let lo = (f32(wd & 15u) - 8.0) * xin[xb] + (f32((wd >> 8u) & 15u) - 8.0) * xin[xb + 1u]
236
- + (f32((wd >> 16u) & 15u) - 8.0) * xin[xb + 2u] + (f32((wd >> 24u) & 15u) - 8.0) * xin[xb + 3u];
237
- let hi = (f32((wd >> 4u) & 15u) - 8.0) * xin[xb + 16u] + (f32((wd >> 12u) & 15u) - 8.0) * xin[xb + 17u]
238
- + (f32((wd >> 20u) & 15u) - 8.0) * xin[xb + 18u] + (f32((wd >> 28u) & 15u) - 8.0) * xin[xb + 19u];
239
- return d * (lo + hi);
240
- }`;
241
- // Q6_K (ggml block_q6_K: ql[128], qh[64], int8 scales[16], f16 d per 256 values). The GGUF's 210-byte blocks
242
- // sit at 2-byte alignment, so at load q6kPlanes splits the table on the GPU into aligned planes: ql (32 words
243
- // per block), qh (16), scales (4) and d as f32. q6v reads value t of block bi exactly as dequantize_row_q6_K
244
- // forms it: (d · scale) · (q − 32).
245
- const Q6P = `
246
- @group(0) @binding(0) var<storage, read> ql: array<u32>;
247
- @group(0) @binding(1) var<storage, read> qh: array<u32>;
248
- @group(0) @binding(2) var<storage, read> sc: array<u32>;
249
- @group(0) @binding(3) var<storage, read> dd: array<f32>;
250
- fn sc8(bi: u32, si: u32) -> f32 { let s8 = (sc[bi * 4u + (si >> 2u)] >> ((si & 3u) * 8u)) & 0xffu; return f32(select(i32(s8), i32(s8) - 256, s8 >= 128u)); }
251
- fn q6v(bi: u32, t: u32) -> f32 {
252
- let n = t >> 7u; let r = t & 127u; let qd = r >> 5u; let l = r & 31u;
253
- let qb = n * 64u + l + (qd & 1u) * 32u;
254
- let lo = (ql[bi * 32u + (qb >> 2u)] >> ((qb & 3u) * 8u)) & 0xffu;
255
- let q4 = select(lo & 15u, lo >> 4u, qd >= 2u);
256
- let hb = n * 32u + l;
257
- let h = (qh[bi * 16u + (hb >> 2u)] >> ((hb & 3u) * 8u + qd * 2u)) & 3u;
258
- return dd[bi] * sc8(bi, n * 8u + (l >> 4u) + qd * 2u) * f32(i32(q4 | (h << 4u)) - 32);
259
- }`;
260
229
  // tanh with its argument clamped: tanh(±10) is ±1 in f32, and a naive exp-based tanh would give NaN past ~44.
261
230
  const TANH = 'fn tanhc(a: f32) -> f32 { return tanh(clamp(a, -10.0, 10.0)); }';
262
231
  const GELU = `${TANH}
@@ -265,124 +234,6 @@ fn gelu(x: f32) -> f32 { return 0.5 * x * (1.0 + tanhc(0.7978845608 * x * (1.0 +
265
234
  const WIN = 'fn winStart(n: u32, win: u32) -> u32 { return select(0u, n - win, win > 0u && n > win); }';
266
235
 
267
236
  const G4 = {
268
- // The GGUF's Q6_K bytes (raw) → aligned planes; one thread per block.
269
- q6kPlanes: `
270
- struct P { blocks: u32 }
271
- @group(0) @binding(0) var<storage, read> raw: array<u32>;
272
- @group(0) @binding(1) var<storage, read_write> ql: array<u32>;
273
- @group(0) @binding(2) var<storage, read_write> qh: array<u32>;
274
- @group(0) @binding(3) var<storage, read_write> sc: array<u32>;
275
- @group(0) @binding(4) var<storage, read_write> dd: array<f32>;
276
- @group(0) @binding(5) var<uniform> p: P;
277
- fn byteAt(i: u32) -> u32 { return (raw[i >> 2u] >> ((i & 3u) * 8u)) & 0xffu; }
278
- fn wordAt(i: u32) -> u32 { return byteAt(i) | (byteAt(i + 1u) << 8u) | (byteAt(i + 2u) << 16u) | (byteAt(i + 3u) << 24u); }
279
- @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>) {
280
- let b = g.y * ng.x * 256u + g.x; if (b >= p.blocks) { return; }
281
- let o = b * 210u;
282
- for (var i = 0u; i < 32u; i++) { ql[b * 32u + i] = wordAt(o + 4u * i); }
283
- for (var i = 0u; i < 16u; i++) { qh[b * 16u + i] = wordAt(o + 128u + 4u * i); }
284
- for (var i = 0u; i < 4u; i++) { sc[b * 4u + i] = wordAt(o + 192u + 4u * i); }
285
- dd[b] = unpack2x16float(byteAt(o + 208u) | (byteAt(o + 209u) << 8u)).x;
286
- }`,
287
- embedQ6K: `
288
- ${TOK}
289
- struct P { n: u32, nb: u32, scale: f32 }
290
- ${Q6P}
291
- @group(0) @binding(4) var<storage, read_write> y: array<f32>;
292
- @group(0) @binding(5) var<uniform> tok: Tok;
293
- @group(0) @binding(6) var<uniform> p: P;
294
- @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
295
- let i = g.x; if (i >= p.n) { return; }
296
- y[i] = q6v(tok.token * p.nb + (i >> 8u), i & 255u) * p.scale;
297
- }`,
298
- // Logits over the Q6_K planes: 4 rows per 256-thread workgroup, 64 threads per row; thread j takes quad j
299
- // (4 consecutive values: one ql word, one qh word, one scale) of every block of the row.
300
- matmulQ6K: `
301
- struct P { M: u32, N: u32 }
302
- ${Q6P}
303
- @group(0) @binding(4) var<storage, read> x: array<vec4<f32>>;
304
- @group(0) @binding(5) var<storage, read_write> y: array<f32>;
305
- @group(0) @binding(6) var<uniform> p: P;
306
- var<workgroup> part: array<f32, 256>;
307
- @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>) {
308
- let j = l.x & 63u; let m = (wg.y * ng.x + wg.x) * 4u + (l.x >> 6u); let mr = min(m, p.M - 1u);
309
- let nb = p.N / 256u; let n = j >> 5u; let w = j & 31u; let qd = w >> 3u; let k = w & 7u;
310
- var acc = 0.0;
311
- for (var bb = 0u; bb < nb; bb++) {
312
- let bi = mr * nb + bb;
313
- let lo = unpack4xU8(ql[bi * 32u + n * 16u + (qd & 1u) * 8u + k]);
314
- let q4 = select(lo & vec4<u32>(15u), lo >> vec4<u32>(4u), qd >= 2u);
315
- let h = (unpack4xU8(qh[bi * 16u + n * 8u + k]) >> vec4<u32>(qd * 2u)) & vec4<u32>(3u);
316
- let q = vec4<f32>(vec4<i32>(q4 | (h << vec4<u32>(4u))) - vec4<i32>(32));
317
- acc += (dd[bi] * sc8(bi, n * 8u + (k >> 2u) + qd * 2u)) * dot(q, x[bb * 64u + j]);
318
- }
319
- part[l.x] = acc; workgroupBarrier();
320
- for (var s = 32u; s > 0u; s >>= 1u) { if (j < s) { part[l.x] += part[l.x + s]; } workgroupBarrier(); }
321
- if (j == 0u && m < p.M) { y[m] = part[l.x]; }
322
- }`,
323
- // Dense Q4_0 GEMV, 4 output rows per workgroup (2-D dispatch ok).
324
- matmulQ4: `enable f16;
325
- struct P { M: u32, N: u32 }
326
- @group(0) @binding(0) var<storage, read> xin: array<f32>;
327
- @group(0) @binding(1) var<storage, read> qw: array<u32>;
328
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
329
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
330
- @group(0) @binding(4) var<uniform> p: P;
331
- ${Q4DOT}
332
- var<workgroup> part: array<f32, 1024>;
333
- @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>) {
334
- let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
335
- if (m0 >= M) { return; }
336
- 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);
337
- var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
338
- for (var w = t; w < nw; w += 256u) {
339
- let b = w >> 2u; let xb = b * 32u + (w & 3u) * 4u;
340
- a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
341
- a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
342
- a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
343
- a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
344
- }
345
- part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
346
- workgroupBarrier();
347
- for (var s = 128u; s > 0u; s >>= 1u) {
348
- 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]; }
349
- workgroupBarrier();
350
- }
351
- if (t < 4u && m0 + t < M) { y[m0 + t] = part[t * 256u]; }
352
- }`,
353
- // Routed experts, all k at once (workgroup z = top-k position); slots[z] picks the expert's rows in the pool.
354
- expertQ4: `enable f16;
355
- struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
356
- @group(0) @binding(0) var<storage, read> xin: array<f32>;
357
- @group(0) @binding(1) var<storage, read> qw: array<u32>;
358
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
359
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
360
- @group(0) @binding(4) var<uniform> p: P;
361
- @group(0) @binding(5) var<storage, read> slots: array<u32>;
362
- ${Q4DOT}
363
- var<workgroup> part: array<f32, 1024>;
364
- @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>) {
365
- let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
366
- if (m0 >= M) { return; }
367
- let rb = slots[z] * p.rowsPerExpert;
368
- let xo = select(0u, z * p.N, p.inputPerSlot == 1u);
369
- 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);
370
- var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
371
- for (var w = t; w < nw; w += 256u) {
372
- let b = w >> 2u; let xb = xo + b * 32u + (w & 3u) * 4u;
373
- a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
374
- a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
375
- a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
376
- a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
377
- }
378
- part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
379
- workgroupBarrier();
380
- for (var s = 128u; s > 0u; s >>= 1u) {
381
- 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]; }
382
- workgroupBarrier();
383
- }
384
- if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
385
- }`,
386
237
  // gu = k × [gate F | up F] → act = k × F of GELU(gate)·up (ggml_geglu_split).
387
238
  geglu: `
388
239
  struct P { F: u32, k: u32 }
@@ -556,146 +407,12 @@ struct P { n: u32, off: u32 }
556
407
  }`,
557
408
  };
558
409
 
559
- // ── Subgroup Q4_0 GEMVs (decode) ─────────────────────────────────────────────────
560
- // One 32-lane subgroup per output row, 4 rows per 128-thread workgroup (the scalar kernels' dispatch
561
- // geometry). A lane takes whole blocks: one vec4<u32> of nibbles, unpack4xU8 of the low and high nibble
562
- // masks, vec4 dot products against the activations, the block scale once; subgroupAdd reduces the row.
563
- // Used only when the adapter's subgroups are exactly 32 wide.
564
- const Q4BLK = `
565
- fn q4blk(w: vec4<u32>, xb: u32) -> f32 {
566
- let m = 0x0f0f0f0fu;
567
- return dot(vec4<f32>(unpack4xU8(w.x & m)) - 8.0, x[xb]) + dot(vec4<f32>(unpack4xU8((w.x >> 4u) & m)) - 8.0, x[xb + 4u])
568
- + dot(vec4<f32>(unpack4xU8(w.y & m)) - 8.0, x[xb + 1u]) + dot(vec4<f32>(unpack4xU8((w.y >> 4u) & m)) - 8.0, x[xb + 5u])
569
- + dot(vec4<f32>(unpack4xU8(w.z & m)) - 8.0, x[xb + 2u]) + dot(vec4<f32>(unpack4xU8((w.z >> 4u) & m)) - 8.0, x[xb + 6u])
570
- + dot(vec4<f32>(unpack4xU8(w.w & m)) - 8.0, x[xb + 3u]) + dot(vec4<f32>(unpack4xU8((w.w >> 4u) & m)) - 8.0, x[xb + 7u]);
571
- }`;
572
- const G4SG = {
573
- matmulQ4: `enable f16;
574
- enable subgroups;
575
- struct P { M: u32, N: u32 }
576
- @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
577
- @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
578
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
579
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
580
- @group(0) @binding(4) var<uniform> p: P;
581
- ${Q4BLK}
582
- @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
583
- let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
584
- let nb = p.N / 32u; let rb = min(m, p.M - 1u) * nb;
585
- var acc = 0.0;
586
- for (var b = lane; b < nb; b += 32u) { acc += q4blk(qw[rb + b], b * 8u) * f32(qs[rb + b]); }
587
- let total = subgroupAdd(acc);
588
- if (lane == 0u && m < p.M) { y[m] = total; }
589
- }`,
590
- expertQ4: `enable f16;
591
- enable subgroups;
592
- struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
593
- @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
594
- @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
595
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
596
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
597
- @group(0) @binding(4) var<uniform> p: P;
598
- @group(0) @binding(5) var<storage, read> slots: array<u32>;
599
- ${Q4BLK}
600
- @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
601
- let z = wg.z; let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
602
- let nb = p.N / 32u; let rb = (slots[z] * p.rowsPerExpert + min(m, p.M - 1u)) * nb;
603
- let xo = select(0u, z * p.N / 4u, p.inputPerSlot == 1u);
604
- var acc = 0.0;
605
- for (var b = lane; b < nb; b += 32u) { acc += q4blk(qw[rb + b], xo + b * 8u) * f32(qs[rb + b]); }
606
- let total = subgroupAdd(acc);
607
- if (lane == 0u && m < p.M) { y[z * p.M + m] = total; }
608
- }`,
609
- };
610
410
 
611
411
  // ── Chunked prefill: T tokens per layer pass (rung 2b's pattern) ────────────────
612
412
  // Rows are tokens, token-major. Generic kernels come from rung 2b (rmsnormB, matmulF32T, topkB, kvStoreB,
613
413
  // softmaxB, attnOutB); these are Gemma's.
614
414
  const QB = 'struct Q { T: u32, pos0: u32, S: u32, pad: u32 }';
615
415
  const G4B = {
616
- embedQ6KB: `
617
- struct P { n: u32, nb: u32, scale: f32 }
618
- ${QB}
619
- ${Q6P}
620
- @group(0) @binding(4) var<storage, read> ids: array<u32>;
621
- @group(0) @binding(5) var<storage, read_write> y: array<f32>;
622
- @group(0) @binding(6) var<uniform> p: P;
623
- @group(0) @binding(7) var<uniform> qd: Q;
624
- @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
625
- let i = g.x; let t = g.y; if (i >= p.n || t >= qd.T) { return; }
626
- y[t * p.n + i] = q6v(ids[t] * p.nb + (i >> 8u), i & 255u) * p.scale;
627
- }`,
628
- // Tiled Q4_0 GEMM: y[t·M + m] = Σ_k x[t·N + k]·W[m][k]. 16 rows × 16 tokens per workgroup; each 32-wide
629
- // block of W (decoded from its nibbles) and of x goes through shared memory, so a weight serves 16 tokens.
630
- // Dispatch (ceil(M/16), ceil(T/16)).
631
- matmulQ4T: `enable f16;
632
- struct P { M: u32, N: u32 }
633
- ${QB}
634
- @group(0) @binding(0) var<storage, read> x: array<f32>;
635
- @group(0) @binding(1) var<storage, read> qw: array<u32>;
636
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
637
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
638
- @group(0) @binding(4) var<uniform> p: P;
639
- @group(0) @binding(5) var<uniform> qd: Q;
640
- var<workgroup> ws: array<f32, 512>;
641
- var<workgroup> xs: array<f32, 512>;
642
- @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
643
- let m0 = wg.x * 16u; let t0 = wg.y * 16u; let r = l.x / 16u; let c = l.x % 16u;
644
- let nb = p.N / 32u; let nw = p.N / 8u; var acc = 0.0;
645
- for (var b = 0u; b < nb; b++) {
646
- for (var i = l.x; i < 512u; i += 256u) {
647
- let rr = i / 32u; let kk = i % 32u;
648
- let m = m0 + rr; var wv = 0.0;
649
- if (m < p.M) {
650
- let wd = qw[m * nw + b * 4u + ((kk & 15u) >> 2u)];
651
- wv = (f32((wd >> ((kk & 3u) * 8u + (kk >> 4u) * 4u)) & 15u) - 8.0) * f32(qs[m * nb + b]);
652
- }
653
- ws[i] = wv;
654
- let t = t0 + rr; var xv = 0.0;
655
- if (t < qd.T) { xv = x[t * p.N + b * 32u + kk]; }
656
- xs[i] = xv;
657
- }
658
- workgroupBarrier();
659
- for (var kk = 0u; kk < 32u; kk++) { acc += ws[r * 32u + kk] * xs[c * 32u + kk]; }
660
- workgroupBarrier();
661
- }
662
- let m = m0 + r; let t = t0 + c;
663
- if (m < p.M && t < qd.T) { y[t * p.M + m] = acc; }
664
- }`,
665
- // Routed experts for every (token, pick) pair: z = t·K + j, slots[z] its pool slot. Input: the token's row
666
- // (gate/up, inputMode 0) or the pair's own row (down, inputMode 1).
667
- expertQ4B: `enable f16;
668
- struct P { M: u32, N: u32, rowsPerExpert: u32, inputMode: u32, K: u32 }
669
- @group(0) @binding(0) var<storage, read> xin: array<f32>;
670
- @group(0) @binding(1) var<storage, read> qw: array<u32>;
671
- @group(0) @binding(2) var<storage, read> qs: array<f16>;
672
- @group(0) @binding(3) var<storage, read_write> y: array<f32>;
673
- @group(0) @binding(4) var<uniform> p: P;
674
- @group(0) @binding(5) var<storage, read> slots: array<u32>;
675
- ${Q4DOT}
676
- var<workgroup> part: array<f32, 1024>;
677
- @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>) {
678
- let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
679
- if (m0 >= M) { return; }
680
- let rb = slots[z] * p.rowsPerExpert;
681
- let xo = select((z / p.K) * p.N, z * p.N, p.inputMode == 1u);
682
- 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);
683
- var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
684
- for (var w = t; w < nw; w += 256u) {
685
- let b = w >> 2u; let xb = xo + b * 32u + (w & 3u) * 4u;
686
- a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
687
- a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
688
- a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
689
- a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
690
- }
691
- part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
692
- workgroupBarrier();
693
- for (var s = 128u; s > 0u; s >>= 1u) {
694
- 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]; }
695
- workgroupBarrier();
696
- }
697
- if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
698
- }`,
699
416
  // Per-head RMS norm and NeoX RoPE in place for (head wg.x, token wg.y); rope holds hd values per token.
700
417
  qkNormRopeB: `
701
418
  struct P { hd: u32, row: u32, eps: f32 }
@@ -853,8 +570,8 @@ export class Gemma4MoeSsd extends Qwen3MoeSsd {
853
570
 
854
571
  get kernels() {
855
572
  return (this._kernels ||= {
856
- ...super.kernels, ...G4, ...(this.subgroupKernels ? G4SG : {}),
857
- ...(this.batchPrefill ? { ...Object.fromEntries(REUSED_B.map((k) => [k, QWEN35_BATCH_KERNELS[k]])), ...G4B } : {}),
573
+ ...super.kernels, ...Q4_KERNELS, ...G4, ...(this.subgroupKernels ? Q4_SG_KERNELS : {}),
574
+ ...(this.batchPrefill ? { ...Object.fromEntries(REUSED_B.map((k) => [k, QWEN35_BATCH_KERNELS[k]])), ...Q4_BATCH_KERNELS, ...G4B } : {}),
858
575
  });
859
576
  }
860
577
  // Subgroup GEMVs replace the scalar Q4_0 ones when subgroups are exactly 32 wide (Apple, most others).
@@ -1133,7 +850,7 @@ export class Gemma4MoeSsd extends Qwen3MoeSsd {
1133
850
  tables(ls, null, this.b.ropeS); tables(lf, this.ropeFreqs, this.b.ropeF);
1134
851
  const selBytes = T * 2 * K * 4;
1135
852
  let enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
1136
- this.dispatch(pass, 'embedQ6KB', this.gB.embed, Math.ceil(H / 256), T);
853
+ this.dispatch(pass, this.embedKernelB || 'embedQ6KB', this.gB.embed, Math.ceil(H / 256), T); // embedKernelB: an engine with another table type (hy_v3 Q8_0)
1137
854
  this.encodeTrunkB(pass, 0, T, S);
1138
855
  pass.end();
1139
856
  enc.copyBufferToBuffer(this.b.sel, 0, this.rbSelB, 0, selBytes);
@@ -1206,7 +923,7 @@ export class Gemma4MoeSsd extends Qwen3MoeSsd {
1206
923
  return r;
1207
924
  }
1208
925
 
1209
- encodeEmbed(pass) { this.dispatch(pass, 'embedQ6K', this.g.embed, Math.ceil(this.cfg.hidden / 256)); }
926
+ encodeEmbed(pass) { this.dispatch(pass, this.embedKernel || 'embedQ6K', this.g.embed, Math.ceil(this.cfg.hidden / 256)); }
1210
927
  encodeHead(pass, want) {
1211
928
  const c = this.cfg;
1212
929
  this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);