Spaces:
Running
Running
pk_sft_v1 F16 (97.5M parameters, 221 MB) with ember.cpp 9d993f48de: greedy F16 = F32, WebAssembly = native, WebGPU = WebAssembly 40/40 (experiments/release_pk_sft_v1_q8_2026-10-28)
8e107b9 verified | // The llama forward pass on WebGPU, for ember's browser chat (prototype; LOCAL_LLM_STACK.md | |
| // section 7, Research item 1). ember's WebAssembly keeps tokenising, the chat template, sampling | |
| // and the stop rules; this module turns token ids into next-token logits on the GPU. | |
| // | |
| // Weights stay resident on the GPU (f32; Q8_0/f16 GGUFs are dequantised at load). Each token is | |
| // one command buffer of about 13 dispatches per layer, built from bind groups made once at load; | |
| // only a 16-byte step uniform (position, token id) changes per token, and only the last prompt | |
| // token's logits are read back. The maths follows ember's CPU path (model.cpp, utils.cpp): RMSNorm | |
| // (1/sqrt(mean(x^2) + eps)), RoPE on adjacent pairs with theta^(-i/head_dim), causal softmax | |
| // attention scaled by 1/sqrt(head_dim) with grouped KV heads, SwiGLU. Sums are f32 in a different | |
| // order, so logits match the CPU to rounding, not bit for bit. | |
| import { GGML_Q8_0, f16ToF32, readGguf, tensorF32 } from './gguf.js'; | |
| const WG = 64; | |
| const MAX_CTX = 2048; // attention's workgroup score array | |
| const SHADERS = { | |
| embed: ` | |
| struct S { pos: u32, token: u32, n: u32, _p: u32 }; | |
| @group(0) @binding(0) var<uniform> st: S; | |
| @group(0) @binding(1) var<storage, read> E: array<f32>; | |
| @group(0) @binding(2) var<storage, read_write> x: array<f32>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(global_invocation_id) g: vec3u) { | |
| if (g.x < st.n) { x[g.x] = E[st.token * st.n + g.x]; } | |
| }`, | |
| rmsnorm: ` | |
| struct P { n: u32, eps: f32, _a: u32, _b: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> x: array<f32>; | |
| @group(0) @binding(2) var<storage, read> w: array<f32>; | |
| @group(0) @binding(3) var<storage, read_write> y: array<f32>; | |
| var<workgroup> part: array<f32, 256>; | |
| @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3u) { | |
| var s = 0.0; | |
| for (var i = l.x; i < p.n; i += 256u) { s += x[i] * x[i]; } | |
| part[l.x] = s; | |
| workgroupBarrier(); | |
| for (var k = 128u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| let inv = 1.0 / sqrt(part[0] / f32(p.n) + p.eps); | |
| for (var i = l.x; i < p.n; i += 256u) { y[i] = x[i] * inv * w[i]; } | |
| }`, | |
| matvec: ` | |
| struct P { rows: u32, cols: u32, accumulate: u32, _p: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> W: array<f32>; | |
| @group(0) @binding(2) var<storage, read> x: array<f32>; | |
| @group(0) @binding(3) var<storage, read_write> y: array<f32>; | |
| var<workgroup> part: array<f32, ${WG}>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) { | |
| let row = wg.x; | |
| let base = row * p.cols; | |
| var s = 0.0; | |
| for (var c = l.x; c < p.cols; c += ${WG}u) { s += W[base + c] * x[c]; } | |
| part[l.x] = s; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| if (l.x == 0u) { if (p.accumulate != 0u) { y[row] = y[row] + part[0]; } else { y[row] = part[0]; } } | |
| }`, | |
| // Q8_0 kept quantised: per 32-weight block an f32 scale and 8 words of 4 int8. Each weight is | |
| // scale * q in f32, the value ember's dequantise-at-load produces, then the same dot as matvec. | |
| matvec_q8: ` | |
| struct P { rows: u32, cols: u32, accumulate: u32, _p: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> Q: array<u32>; // [rows][cols / 4] packed int8 | |
| @group(0) @binding(2) var<storage, read> D: array<f32>; // [rows][cols / 32] scales | |
| @group(0) @binding(3) var<storage, read> x: array<f32>; | |
| @group(0) @binding(4) var<storage, read_write> y: array<f32>; | |
| var<workgroup> part: array<f32, ${WG}>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) { | |
| let row = wg.x; | |
| let words = p.cols / 4u; | |
| var s = 0.0; | |
| for (var w = l.x; w < words; w += ${WG}u) { | |
| let packed = Q[row * words + w]; | |
| let sc = D[row * (p.cols / 32u) + w / 8u]; | |
| let c = 4u * w; | |
| for (var b = 0u; b < 4u; b++) { | |
| let q = f32((i32(packed << (24u - 8u * b))) >> 24u); | |
| s += (sc * q) * x[c + b]; | |
| } | |
| } | |
| part[l.x] = s; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| if (l.x == 0u) { if (p.accumulate != 0u) { y[row] = y[row] + part[0]; } else { y[row] = part[0]; } } | |
| }`, | |
| rope_kv: ` | |
| struct S { pos: u32, token: u32, n: u32, _p: u32 }; | |
| struct P { n_heads: u32, n_kv: u32, head_dim: u32, _p: u32 }; | |
| @group(0) @binding(0) var<uniform> st: S; | |
| @group(0) @binding(1) var<uniform> p: P; | |
| @group(0) @binding(2) var<storage, read> cs: array<f32>; // [ctx][head_dim/2][cos, sin] | |
| @group(0) @binding(3) var<storage, read_write> q: array<f32>; | |
| @group(0) @binding(4) var<storage, read> k: array<f32>; | |
| @group(0) @binding(5) var<storage, read> v: array<f32>; | |
| @group(0) @binding(6) var<storage, read_write> kc: array<f32>; // [ctx][n_kv * head_dim] | |
| @group(0) @binding(7) var<storage, read_write> vc: array<f32>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(global_invocation_id) g: vec3u) { | |
| let half = p.head_dim / 2u; | |
| let kvd = p.n_kv * p.head_dim; | |
| let i = g.x; | |
| if (i < p.n_heads * half) { // rotate q in place | |
| let h = i / half; let j = i % half; | |
| let a = h * p.head_dim + 2u * j; | |
| let c = cs[2u * (st.pos * half + j)]; let s = cs[2u * (st.pos * half + j) + 1u]; | |
| let x0 = q[a]; let x1 = q[a + 1u]; | |
| q[a] = x0 * c - x1 * s; q[a + 1u] = x0 * s + x1 * c; | |
| } | |
| if (i < p.n_kv * half) { // rotate k into the cache | |
| let h = i / half; let j = i % half; | |
| let a = h * p.head_dim + 2u * j; | |
| let c = cs[2u * (st.pos * half + j)]; let s = cs[2u * (st.pos * half + j) + 1u]; | |
| let x0 = k[a]; let x1 = k[a + 1u]; | |
| kc[st.pos * kvd + a] = x0 * c - x1 * s; kc[st.pos * kvd + a + 1u] = x0 * s + x1 * c; | |
| } | |
| if (i < kvd) { vc[st.pos * kvd + i] = v[i]; } | |
| }`, | |
| attention: ` | |
| struct S { pos: u32, token: u32, n: u32, _p: u32 }; | |
| struct P { n_heads: u32, n_kv: u32, head_dim: u32, _p: u32 }; | |
| @group(0) @binding(0) var<uniform> st: S; | |
| @group(0) @binding(1) var<uniform> p: P; | |
| @group(0) @binding(2) var<storage, read> q: array<f32>; | |
| @group(0) @binding(3) var<storage, read> kc: array<f32>; | |
| @group(0) @binding(4) var<storage, read> vc: array<f32>; | |
| @group(0) @binding(5) var<storage, read_write> o: array<f32>; | |
| var<workgroup> sc: array<f32, ${MAX_CTX}>; | |
| var<workgroup> part: array<f32, ${WG}>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) { | |
| let h = wg.x; | |
| let kvh = h / (p.n_heads / p.n_kv); | |
| let kvd = p.n_kv * p.head_dim; | |
| let qo = h * p.head_dim; let ko = kvh * p.head_dim; | |
| let n = st.pos + 1u; | |
| let scale = 1.0 / sqrt(f32(p.head_dim)); | |
| var m = -3.4e38; | |
| for (var t = l.x; t < n; t += ${WG}u) { | |
| var s = 0.0; | |
| for (var d = 0u; d < p.head_dim; d++) { s += q[qo + d] * kc[t * kvd + ko + d]; } | |
| s = s * scale; sc[t] = s; m = max(m, s); | |
| } | |
| part[l.x] = m; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] = max(part[l.x], part[l.x + k]); } workgroupBarrier(); } | |
| let mx = part[0]; | |
| workgroupBarrier(); | |
| var z = 0.0; | |
| for (var t = l.x; t < n; t += ${WG}u) { let e = exp(sc[t] - mx); sc[t] = e; z += e; } | |
| part[l.x] = z; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| let zz = part[0]; | |
| for (var d = l.x; d < p.head_dim; d += ${WG}u) { | |
| var acc = 0.0; | |
| for (var t = 0u; t < n; t++) { acc += sc[t] * vc[t * kvd + ko + d]; } | |
| o[qo + d] = acc / zz; | |
| } | |
| }`, | |
| // Mixture of experts, as ember's moe_ffn_f32_cpu (utils.cpp): top-k router logits (ties to the | |
| // lower index), gates = a softmax over the k selected logits (renorm), experts in ascending index. | |
| moe_select: ` | |
| struct P { E: u32, k: u32, _a: u32, _b: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> r: array<f32>; | |
| @group(0) @binding(2) var<storage, read_write> sel: array<u32>; | |
| @group(0) @binding(3) var<storage, read_write> gate: array<f32>; | |
| @compute @workgroup_size(1) fn main() { | |
| var taken: array<bool, 256>; | |
| var pick: array<u32, 16>; | |
| for (var j = 0u; j < p.k; j++) { | |
| var best = 0xffffffffu; | |
| for (var e = 0u; e < p.E; e++) { | |
| if (!taken[e] && (best == 0xffffffffu || r[e] > r[best])) { best = e; } | |
| } | |
| taken[best] = true; pick[j] = best; | |
| } | |
| for (var a = 1u; a < p.k; a++) { // ascending expert index | |
| let v = pick[a]; var b = a; | |
| while (b > 0u && pick[b - 1u] > v) { pick[b] = pick[b - 1u]; b--; } | |
| pick[b] = v; | |
| } | |
| var mx = -3.4e38; | |
| for (var j = 0u; j < p.k; j++) { mx = max(mx, r[pick[j]]); } | |
| var z = 0.0; | |
| for (var j = 0u; j < p.k; j++) { z += exp(r[pick[j]] - mx); } | |
| for (var j = 0u; j < p.k; j++) { sel[j] = pick[j]; gate[j] = exp(r[pick[j]] - mx) / z; } | |
| }`, | |
| matvec_exp: ` | |
| struct P { rows: u32, cols: u32, _a: u32, _b: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> sel: array<u32>; | |
| @group(0) @binding(2) var<storage, read> W: array<f32>; // [E][rows][cols] | |
| @group(0) @binding(3) var<storage, read> x: array<f32>; | |
| @group(0) @binding(4) var<storage, read_write> y: array<f32>; // [k][rows] | |
| var<workgroup> part: array<f32, ${WG}>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) { | |
| let row = wg.x; let j = wg.y; | |
| let base = (sel[j] * p.rows + row) * p.cols; | |
| var s = 0.0; | |
| for (var c = l.x; c < p.cols; c += ${WG}u) { s += W[base + c] * x[c]; } | |
| part[l.x] = s; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| if (l.x == 0u) { y[j * p.rows + row] = part[0]; } | |
| }`, | |
| moe_down: ` | |
| struct P { rows: u32, cols: u32, k: u32, _a: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read> sel: array<u32>; | |
| @group(0) @binding(2) var<storage, read> gate: array<f32>; | |
| @group(0) @binding(3) var<storage, read> W: array<f32>; // [E][rows = d][cols = d_ff] | |
| @group(0) @binding(4) var<storage, read> h: array<f32>; // [k][d_ff] | |
| @group(0) @binding(5) var<storage, read_write> x: array<f32>; // residual stream, += sum_j gate_j * down_j h_j | |
| var<workgroup> part: array<f32, ${WG}>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) l: vec3u) { | |
| let row = wg.x; | |
| var acc = 0.0; | |
| for (var j = 0u; j < p.k; j++) { | |
| let base = (sel[j] * p.rows + row) * p.cols; | |
| var s = 0.0; | |
| for (var c = l.x; c < p.cols; c += ${WG}u) { s += W[base + c] * h[j * p.cols + c]; } | |
| part[l.x] = s; | |
| workgroupBarrier(); | |
| for (var k = ${WG / 2}u; k > 0u; k >>= 1u) { if (l.x < k) { part[l.x] += part[l.x + k]; } workgroupBarrier(); } | |
| acc += gate[j] * part[0]; | |
| workgroupBarrier(); | |
| } | |
| if (l.x == 0u) { x[row] = x[row] + acc; } | |
| }`, | |
| swiglu: ` | |
| struct P { n: u32, _a: u32, _b: u32, _c: u32 }; | |
| @group(0) @binding(0) var<uniform> p: P; | |
| @group(0) @binding(1) var<storage, read_write> g: array<f32>; | |
| @group(0) @binding(2) var<storage, read> u: array<f32>; | |
| @compute @workgroup_size(${WG}) fn main(@builtin(global_invocation_id) i: vec3u) { | |
| if (i.x < p.n) { let a = g[i.x]; g[i.x] = a / (1.0 + exp(-a)) * u[i.x]; } | |
| }`, | |
| }; | |
| // Test hook: one weight matrix through matvec (dequantised f32) and matvec_q8 (kept quantised) | |
| // with the same input x; returns both outputs and a JS f64 reference. | |
| export async function testMatvecQ8(device, ggufBuffer, name, x) { | |
| const g = readGguf(ggufBuffer), t = g.tensors[name]; | |
| const cols = t.dims[0], rows = t.dims[1], S = GPUBufferUsage; | |
| const buf = (arr, extra = 0) => { | |
| const b = device.createBuffer({ size: Math.max(16, arr.byteLength), usage: S.STORAGE | S.COPY_DST | S.COPY_SRC | extra }); | |
| device.queue.writeBuffer(b, 0, arr.buffer, arr.byteOffset, arr.byteLength); | |
| return b; | |
| }; | |
| const f = tensorF32(g, name); | |
| const nb = t.n / 32, raw = new Uint8Array(g.buf, t.start, 34 * nb), dv = new DataView(g.buf, t.start); | |
| const q = new Uint8Array(32 * nb), sc = new Float32Array(nb); | |
| for (let b = 0; b < nb; b++) { sc[b] = f16ToF32(dv.getUint16(34 * b, true)); q.set(raw.subarray(34 * b + 2, 34 * b + 34), 32 * b); } | |
| const P = device.createBuffer({ size: 16, usage: S.UNIFORM | S.COPY_DST }); | |
| device.queue.writeBuffer(P, 0, new Uint32Array([rows, cols, 0, 0])); | |
| const X = buf(x), Y1 = buf(new Float32Array(rows)), Y2 = buf(new Float32Array(rows)); | |
| const run = (code, bufs, out) => { | |
| const pl = device.createComputePipeline({ layout: 'auto', compute: { module: device.createShaderModule({ code }), entryPoint: 'main' } }); | |
| const bg = device.createBindGroup({ layout: pl.getBindGroupLayout(0), entries: bufs.map((b, i) => ({ binding: i, resource: { buffer: b } })) }); | |
| const e = device.createCommandEncoder(), c = e.beginComputePass(); | |
| c.setPipeline(pl); c.setBindGroup(0, bg); c.dispatchWorkgroups(rows); c.end(); | |
| device.queue.submit([e.finish()]); | |
| }; | |
| run(SHADERS.matvec, [P, buf(f), X, Y1]); | |
| run(SHADERS.matvec_q8, [P, buf(q), buf(sc), X, Y2]); | |
| const read = async (b) => { | |
| const r = device.createBuffer({ size: 4 * rows, usage: S.MAP_READ | S.COPY_DST }); | |
| const e = device.createCommandEncoder(); e.copyBufferToBuffer(b, 0, r, 0, 4 * rows); device.queue.submit([e.finish()]); | |
| await r.mapAsync(GPUMapMode.READ); const v = new Float32Array(r.getMappedRange().slice(0)); r.unmap(); return v; | |
| }; | |
| const ref = new Float64Array(rows); | |
| for (let i = 0; i < rows; i++) { let s = 0; for (let c = 0; c < cols; c++) s += f[i * cols + c] * x[c]; ref[i] = s; } | |
| return { rows, cols, f32: await read(Y1), q8: await read(Y2), ref }; | |
| } | |
| export async function createLlama(device, ggufBuffer) { | |
| const g = readGguf(ggufBuffer), M = g.meta; | |
| if (M['general.architecture'] !== 'llama') throw new Error('WebGPU path: llama GGUFs only'); | |
| const d = M['llama.embedding_length'], L = M['llama.block_count'], ff = M['llama.feed_forward_length']; | |
| const nh = M['llama.attention.head_count'], nkv = M['llama.attention.head_count_kv'] || nh; | |
| const eps = M['llama.attention.layer_norm_rms_epsilon'], theta = M['llama.rope.freq_base'] || 10000; | |
| const ctx = Math.min(M['llama.context_length'], MAX_CTX), hd = d / nh, kvd = nkv * hd; | |
| const moe = !!g.tensors['blk.0.ffn_gate_inp.weight']; | |
| const E = moe ? (M['llama.expert_count'] || g.tensors['blk.0.ffn_gate_inp.weight'].n / d) : 0; | |
| const topk = moe ? (M['llama.expert_used_count'] || 2) : 0; | |
| const ffE = moe ? (M['llama.expert_feed_forward_length'] || g.tensors['blk.0.ffn_gate_exps.weight'].n / (E * d)) : 0; | |
| if (moe && (E > 256 || topk > 16)) throw new Error('WebGPU path: at most 256 experts, top-16'); | |
| // Shapes come from the metadata: GGUFs written before the dims fix record 2-D tensors as | |
| // [out, in] (bytes are row-major [out][in] either way), so dims order cannot be trusted. | |
| const vocab = M['llama.vocab_size'] || (M['tokenizer.ggml.tokens'] || []).length || | |
| Math.max(...g.tensors['token_embd.weight'].dims); | |
| const S = GPUBufferUsage; | |
| const storage = (arr) => { | |
| const b = device.createBuffer({ size: Math.max(16, arr.byteLength), usage: S.STORAGE | S.COPY_DST }); | |
| device.queue.writeBuffer(b, 0, arr.buffer, arr.byteOffset, arr.byteLength); | |
| return b; | |
| }; | |
| const scratch = (n) => device.createBuffer({ size: Math.max(16, 4 * n), usage: S.STORAGE | S.COPY_SRC | S.COPY_DST }); | |
| const uniform = (words) => { | |
| const b = device.createBuffer({ size: 16, usage: S.UNIFORM | S.COPY_DST }); | |
| device.queue.writeBuffer(b, 0, words); | |
| return b; | |
| }; | |
| const u4 = (a, b = 0, c = 0, e = 0) => new Uint32Array([a, b, c, e]); | |
| const weight = (name) => storage(tensorF32(g, name)); | |
| // A matrix for matvec: Q8_0 stays quantised (repacked to scales + packed int8), others go to f32. | |
| const mat = (name) => { | |
| const t = g.tensors[name]; | |
| if (t.type !== GGML_Q8_0) return { q8: false, buf: weight(name) }; | |
| const nb = t.n / 32, raw = new Uint8Array(g.buf, t.start, 34 * nb), dv = new DataView(g.buf, t.start); | |
| const q = new Uint8Array(32 * nb), sc = new Float32Array(nb); | |
| for (let b = 0; b < nb; b++) { | |
| sc[b] = f16ToF32(dv.getUint16(34 * b, true)); | |
| q.set(raw.subarray(34 * b + 2, 34 * b + 34), 32 * b); | |
| } | |
| return { q8: true, Q: storage(q), D: storage(sc) }; | |
| }; | |
| const mv = (P, m, x, y, n) => m.q8 ? ['matvec_q8', bind('matvec_q8', [P, m.Q, m.D, x, y]), n] | |
| : ['matvec', bind('matvec', [P, m.buf, x, y]), n]; | |
| const outName = g.tensors['output.weight'] ? 'output.weight' : 'token_embd.weight'; // tied embeddings | |
| // RoPE table exactly as ember builds it: f32 freq = 1 / theta^(i / head_dim), angle pos * freq. | |
| const cs = new Float32Array(ctx * hd); | |
| for (let pos = 0; pos < ctx; pos++) | |
| for (let i = 0; i < hd; i += 2) { | |
| const freq = Math.fround(1 / Math.fround(Math.pow(theta, i / hd))); | |
| const val = Math.fround(pos * freq); | |
| cs[pos * hd + i] = Math.cos(val); cs[pos * hd + i + 1] = Math.sin(val); | |
| } | |
| const W = { | |
| embd: weight('token_embd.weight'), outNorm: weight('output_norm.weight'), out: mat(outName), rope: storage(cs), | |
| layers: Array.from({ length: L }, (_, l) => ({ | |
| attnNorm: weight(`blk.${l}.attn_norm.weight`), q: mat(`blk.${l}.attn_q.weight`), k: mat(`blk.${l}.attn_k.weight`), | |
| v: mat(`blk.${l}.attn_v.weight`), o: mat(`blk.${l}.attn_output.weight`), ffnNorm: weight(`blk.${l}.ffn_norm.weight`), | |
| ...(moe ? { router: mat(`blk.${l}.ffn_gate_inp.weight`), gateE: weight(`blk.${l}.ffn_gate_exps.weight`), | |
| upE: weight(`blk.${l}.ffn_up_exps.weight`), downE: weight(`blk.${l}.ffn_down_exps.weight`) } | |
| : { gate: mat(`blk.${l}.ffn_gate.weight`), up: mat(`blk.${l}.ffn_up.weight`), | |
| down: mat(`blk.${l}.ffn_down.weight`) }), | |
| kc: scratch(ctx * kvd), vc: scratch(ctx * kvd), | |
| })), | |
| }; | |
| const A = { x: scratch(d), xn: scratch(d), q: scratch(d), k: scratch(kvd), v: scratch(kvd), att: scratch(d), | |
| g: scratch(moe ? topk * ffE : ff), u: scratch(moe ? topk * ffE : ff), logits: scratch(vocab), | |
| r: scratch(Math.max(E, 4)), sel: scratch(16), gate: scratch(16) }; | |
| const readback = device.createBuffer({ size: 4 * vocab, usage: S.MAP_READ | S.COPY_DST }); | |
| const step = device.createBuffer({ size: 16, usage: S.UNIFORM | S.COPY_DST }); | |
| const pipe = {}; | |
| for (const [k, code] of Object.entries(SHADERS)) | |
| pipe[k] = device.createComputePipeline({ layout: 'auto', compute: { module: device.createShaderModule({ code }), entryPoint: 'main' } }); | |
| const bind = (k, bufs) => device.createBindGroup({ layout: pipe[k].getBindGroupLayout(0), | |
| entries: bufs.map((b, i) => ({ binding: i, resource: { buffer: b } })) }); | |
| const f32u = (n, e) => { const a = new ArrayBuffer(16); new Uint32Array(a)[0] = n; new Float32Array(a)[1] = e; return a; }; | |
| const pNorm = uniform(f32u(d, eps)), pHeads = uniform(u4(nh, nkv, hd)), pFF = uniform(u4(ff)); | |
| const pMat = (rows, cols, acc) => uniform(u4(rows, cols, acc ? 1 : 0)); | |
| const pQ = pMat(d, d, false), pKV = pMat(kvd, d, false), pO = pMat(d, d, true), pUp = pMat(ff, d, false), | |
| pDown = pMat(d, ff, true), pOut = pMat(vocab, d, false), pEmb = step; | |
| const pRouter = moe ? pMat(E, d, false) : null, pSel = moe ? uniform(u4(E, topk)) : null, | |
| pExp = moe ? uniform(u4(ffE, d)) : null, pMoeDown = moe ? uniform(u4(d, ffE, topk)) : null, | |
| pFFk = moe ? uniform(u4(topk * ffE)) : null; | |
| const ffnSteps = (lw) => moe ? [ | |
| mv(pRouter, lw.router, A.xn, A.r, E), | |
| ['moe_select', bind('moe_select', [pSel, A.r, A.sel, A.gate]), 1], | |
| ['matvec_exp', bind('matvec_exp', [pExp, A.sel, lw.gateE, A.xn, A.g]), [ffE, topk]], | |
| ['matvec_exp', bind('matvec_exp', [pExp, A.sel, lw.upE, A.xn, A.u]), [ffE, topk]], | |
| ['swiglu', bind('swiglu', [pFFk, A.g, A.u]), Math.ceil(topk * ffE / WG)], | |
| ['moe_down', bind('moe_down', [pMoeDown, A.sel, A.gate, lw.downE, A.g, A.x]), d], | |
| ] : [ | |
| mv(pUp, lw.gate, A.xn, A.g, ff), | |
| mv(pUp, lw.up, A.xn, A.u, ff), | |
| ['swiglu', bind('swiglu', [pFF, A.g, A.u]), Math.ceil(ff / WG)], | |
| mv(pDown, lw.down, A.g, A.x, d), | |
| ]; | |
| // The per-token program: [pipeline, bind group, workgroups] in order. | |
| const prog = [['embed', bind('embed', [pEmb, W.embd, A.x]), Math.ceil(d / WG)]]; | |
| for (const lw of W.layers) { | |
| prog.push(['rmsnorm', bind('rmsnorm', [pNorm, A.x, lw.attnNorm, A.xn]), 1], | |
| mv(pQ, lw.q, A.xn, A.q, d), | |
| mv(pKV, lw.k, A.xn, A.k, kvd), | |
| mv(pKV, lw.v, A.xn, A.v, kvd), | |
| ['rope_kv', bind('rope_kv', [step, pHeads, W.rope, A.q, A.k, A.v, lw.kc, lw.vc]), Math.ceil(Math.max(d / 2, kvd) / WG)], | |
| ['attention', bind('attention', [step, pHeads, A.q, lw.kc, lw.vc, A.att]), nh], | |
| mv(pO, lw.o, A.att, A.x, d), | |
| ['rmsnorm', bind('rmsnorm', [pNorm, A.x, lw.ffnNorm, A.xn]), 1], | |
| ...ffnSteps(lw)); | |
| } | |
| prog.push(['rmsnorm', bind('rmsnorm', [pNorm, A.x, W.outNorm, A.xn]), 1], | |
| mv(pOut, W.out, A.xn, A.logits, vocab)); | |
| function submitToken(token, pos, wantLogits) { | |
| device.queue.writeBuffer(step, 0, u4(pos, token, d)); | |
| const enc = device.createCommandEncoder(), pass = enc.beginComputePass(); | |
| for (const [k, bg, n] of prog) { | |
| pass.setPipeline(pipe[k]); pass.setBindGroup(0, bg); | |
| if (Array.isArray(n)) pass.dispatchWorkgroups(n[0], n[1]); else pass.dispatchWorkgroups(n); | |
| } | |
| pass.end(); | |
| if (wantLogits) enc.copyBufferToBuffer(A.logits, 0, readback, 0, 4 * vocab); | |
| device.queue.submit([enc.finish()]); | |
| } | |
| return { | |
| config: { d, L, ff: moe ? ffE : ff, nh, nkv, hd, ctx, vocab, eps, theta, experts: E, topk, | |
| dispatchesPerToken: prog.length }, | |
| // Run tokens at positions startPos.. and return the next-token logits after the last one. | |
| async forward(tokens, startPos) { | |
| if (startPos + tokens.length > ctx) throw new Error('past the context'); | |
| tokens.forEach((t, i) => submitToken(t, startPos + i, i === tokens.length - 1)); | |
| await readback.mapAsync(GPUMapMode.READ); | |
| const out = new Float32Array(readback.getMappedRange().slice(0)); | |
| readback.unmap(); | |
| return out; | |
| }, | |
| }; | |
| } | |