Spaces:
Running
Running
| // WebGPU runtime of the VQ query encoder (docs/superpowers/specs/2026-09-08-vq-webgpu-runtime-design.md, section 6). | |
| // | |
| // One forward pass of the Qwen3-shaped encoder for one query (T <= max_len tokens), batch 1, no KV cache, no output head: | |
| // x0 = Q2_K rows (embed.js, CPU) -> per block: rmsnorm -> fwht_rotate(1024) -> q, k, v (vq_matmul) -> q/k head-norm + RoPE | |
| // -> attention -> fwht_rotate(2048) -> o (vq_matmul, += residual) -> rmsnorm -> fwht_rotate(1024) -> gate, up -> swiglu | |
| // -> fwht_rotate(3072) -> down (+= residual); final rmsnorm -> last-token pooling + L2 -> 1024-d readback. | |
| // 17 dispatches per block + 2, one command buffer per query. Activation buffers are allocated once for T_max tokens; the | |
| // bind groups are built once per block at load, so a query only writes the uniform (T), x0 and the RoPE table. | |
| // Dtypes: residual stream RESID_T (f32 default; 'f16' selectable -- scripts/vqw_layer_test.py measures both), norm / | |
| // attention / gate-up / swiglu activations f32, matmul inputs (rotation outputs) f16, matmul accumulation f32. | |
| // Kernels: kernels/*.wgsl with {{NAME}} placeholders filled here (shapes are compile-time constants; one pipeline per shape). | |
| // vq_matmul (M4): each workgroup decodes its 64 rows x 256-column block once (packed words staged coalesced, codebook in | |
| // workgroup memory) and reuses the decoded tile for a 32-token tile, so a query of T <= 32 tokens decodes every weight exactly | |
| // once; kernels/vq_matmul_v1.wgsl is the M2 kernel (decode per 16-token tile and lane), kept for A/B runs (tile.kernel = 'v1'). | |
| // debug: true in embed() ends the pass after every dispatch and copies the named intermediate to a staging buffer; the | |
| // names are those of scripts/vq_reference.py --dump-activations (T3), so the layer test can compare them one by one. | |
| import { f16ArrayToF32 } from './embed.js'; | |
| const KERNEL_FILES = ['rmsnorm', 'fwht_rotate', 'vq_matmul', 'vq_matmul_v1', 'qk_norm_rope', 'attention', 'swiglu', 'pool_normalize']; | |
| const U = 'uniform', RO = 'read-only-storage', RW = 'storage'; | |
| /** Adapter + device with shader-f16 (required). powerPreference 'low-power' prefers the integrated GPU. */ | |
| export async function requestDevice({ powerPreference = 'low-power', label = 'vqweb' } = {}) { | |
| if (!navigator.gpu) throw new Error('WebGPU is not available (navigator.gpu missing)'); | |
| const adapter = await navigator.gpu.requestAdapter({ powerPreference }); | |
| if (!adapter) throw new Error('no WebGPU adapter'); | |
| if (!adapter.features.has('shader-f16')) throw new Error('adapter has no shader-f16'); | |
| const info = { vendor: adapter.info?.vendor, architecture: adapter.info?.architecture, device: adapter.info?.device, description: adapter.info?.description, | |
| shader_f16: true, subgroups: adapter.features.has('subgroups'), timestamp_query: adapter.features.has('timestamp-query'), | |
| maxComputeWorkgroupStorageSize: adapter.limits.maxComputeWorkgroupStorageSize, maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize, | |
| maxBufferSize: adapter.limits.maxBufferSize, maxComputeInvocationsPerWorkgroup: adapter.limits.maxComputeInvocationsPerWorkgroup }; | |
| const requiredFeatures = ['shader-f16']; | |
| if (adapter.features.has('timestamp-query')) requiredFeatures.push('timestamp-query'); // per-kernel profiling (embed({profile: true})) | |
| const requiredLimits = {}; | |
| if (adapter.limits.maxComputeWorkgroupStorageSize > 16384) requiredLimits.maxComputeWorkgroupStorageSize = adapter.limits.maxComputeWorkgroupStorageSize; // larger vq_matmul token tiles | |
| const device = await adapter.requestDevice({ label, requiredFeatures, requiredLimits }); | |
| return { adapter, device, info }; | |
| } | |
| /** Load the kernel sources next to this module (browser). */ | |
| export async function loadKernelSources(baseUrl = new URL('./kernels/', import.meta.url)) { | |
| const out = {}; | |
| await Promise.all(KERNEL_FILES.map(async (n) => { | |
| const r = await fetch(new URL(`${n}.wgsl`, baseUrl)); | |
| if (!r.ok) throw new Error(`kernel ${n}.wgsl: ${r.status}`); | |
| out[n] = await r.text(); | |
| })); | |
| return out; | |
| } | |
| function fill(src, vars) { | |
| return src.replace(/\{\{(\w+)\}\}/g, (m, k) => { if (!(k in vars)) throw new Error(`template variable ${k} missing`); return String(vars[k]); }); | |
| } | |
| const f32lit = (x) => { const s = String(x); return /[.eE]/.test(s) ? s : s + '.0'; }; | |
| export class VqwRuntime { | |
| /** | |
| * @param {import('./container.js').VqwContainer} container | |
| * @param {GPUDevice} device | |
| * @param {object} opts resid: 'f32' | 'f16'; tmax: max tokens (default model.max_len); tile: vq_matmul tile | |
| * {rwg, tt, cb, mr, mt} (M4 kernel, default 64,32,64,4,4) or {kernel: 'v1', rwg, lanes, tpt} (M2 kernel); | |
| * sources: kernel sources (default: fetched next to this module) | |
| */ | |
| constructor(container, device, opts = {}) { | |
| this.c = container; this.device = device; this.m = container.model; | |
| this.resid = opts.resid ?? 'f32'; | |
| this.tmax = opts.tmax ?? this.m.maxLen; | |
| // M6: stored prompt K / V (container.prefix); opts.usePrefix = false runs the prompt through the quantised model instead | |
| this.P = (opts.usePrefix ?? true) && container.prefix ? container.prefix.tokens : 0; | |
| this.prefixIds = this.P ? container.prefix.ids : null; | |
| // vq_matmul tile (M4 kernel): rwg rows x tt tokens per workgroup, cb columns per staging step, mr x mt outputs per thread; | |
| // kernel: 'v1' selects the M2 kernel (kernels/vq_matmul_v1.wgsl, tile {rwg, lanes, tpt}) for back-to-back comparisons. | |
| const t = opts.tile || {}; | |
| this.tile = t.kernel === 'v1' ? Object.assign({ kernel: 'v1', rwg: 32, lanes: 8, tpt: 2 }, t) : Object.assign({ kernel: 'v2', rwg: 64, tt: 32, cb: 64, mr: 4, mt: 4 }, t); | |
| if (this.tile.kernel === 'v2') { | |
| const { rwg, tt, cb, mr, mt } = this.tile; | |
| const wgs = (rwg / mr) * (tt / mt), gs = cb / 4; | |
| if (256 % cb || cb % 4 || rwg % mr || tt % mt || wgs > 1024 || wgs % rwg || gs % (wgs / rwg)) throw new Error(`bad vq_matmul tile ${JSON.stringify(this.tile)}`); | |
| } | |
| this.sources = opts.sources || null; | |
| this.pipelines = new Map(); | |
| this.layouts = new Map(); | |
| this.dispatchesPerBlock = 17; | |
| if (this.m.headDim % 2 || this.m.heads % this.m.kvHeads) throw new Error('unsupported head configuration'); | |
| } | |
| static async create(container, device, opts = {}) { | |
| const rt = new VqwRuntime(container, device, opts); | |
| if (!rt.sources) rt.sources = await loadKernelSources(opts.kernelBase); | |
| rt._buildPipelines(); | |
| rt._allocate(); | |
| rt._bindAll(); | |
| await device.queue.onSubmittedWorkDone(); | |
| return rt; | |
| } | |
| // ---------------------------------------------------------------------------------------------- pipelines | |
| _layout(key, kinds) { | |
| let l = this.layouts.get(key); | |
| if (!l) { | |
| l = this.device.createBindGroupLayout({ label: key, entries: kinds.map((type, binding) => ({ binding, visibility: GPUShaderStage.COMPUTE, buffer: { type } })) }); | |
| this.layouts.set(key, l); | |
| } | |
| return l; | |
| } | |
| _pipeline(key, file, vars, kinds) { | |
| let p = this.pipelines.get(key); | |
| if (p) return p; | |
| const code = fill(this.sources[file], vars); | |
| const module = this.device.createShaderModule({ label: key, code }); | |
| const layout = this._layout(kinds.join(','), kinds); | |
| p = { pipeline: this.device.createComputePipeline({ label: key, layout: this.device.createPipelineLayout({ bindGroupLayouts: [layout] }), compute: { module, entryPoint: 'main' } }), layout, key, module }; | |
| this.pipelines.set(key, p); | |
| return p; | |
| } | |
| _buildPipelines() { | |
| const m = this.m, R = this.resid; | |
| this.pRms = this._pipeline(`rmsnorm/${R}`, 'rmsnorm', { N: m.hidden, EPS: f32lit(m.rmsEps), IN_T: R }, [U, RO, RO, RW]); | |
| this.pRot = {}; | |
| for (const n of this.c.rotation.sizes) { | |
| const B = this.c.rotation.block; | |
| this.pRot[n] = this._pipeline(`fwht/${n}`, 'fwht_rotate', { N: n, B, R: this.c.rotation.rounds, INV_SQRT_B: f32lit(1 / Math.sqrt(B)) }, [U, RO, RO, RO, RW]); | |
| } | |
| this.pQk = this._qkPipeline(true, true); | |
| this.pAttn = this._attnPipeline(false); | |
| this.pSwiglu = this._pipeline('swiglu', 'swiglu', { I: m.intermediate }, [U, RO, RO, RW]); | |
| this.pPool = this._pipeline('pool', 'pool_normalize', { N: m.hidden }, [U, RO, RW]); | |
| } | |
| _matmulPipeline(lin, add) { | |
| const { rows, cols, k, bits, words } = lin; | |
| const dim = lin.dim || 4; | |
| const t = this.tile; | |
| const OUT_T = add ? this.resid : 'f32'; | |
| // dim 4: one index per 4-column quad; dim 2: two consecutive indices per quad (the second one straddles at most one more word: | |
| // for W <= 12 the pair fits in 2*W <= 24 bits, so the shifted 32-bit window v (plus the next word) holds both) | |
| const CBT = dim === 4 ? 'vec4<f16>' : 'vec2<f16>'; | |
| const DECODE = dim === 4 ? 'ws[(g * MR + di) * NTR + dtr] = vec4<f32>(cbs[v & MASK]);' | |
| : 'let hi = select(0u, iw[min(base + wi + 1u, RWG * WPB - 1u)], sh > 0u) << ((32u - sh) & 31u); let v2 = (iw[base + wi] >> sh) | hi; ' | |
| + 'let i0 = v2 & MASK; let i1 = (v2 >> W) & MASK; ws[(g * MR + di) * NTR + dtr] = vec4<f32>(vec4<f16>(cbs[i0], cbs[i1]));'; | |
| const common = { ROWS: rows, COLS: cols, K: k, W: bits, MASK: (1 << bits) - 1, WORDS: words, OUT_T, DIM: dim, CBT, DECODE }; | |
| if (t.kernel === 'v1') { | |
| if (dim !== 4) throw new Error('vq_matmul v1 kernel supports dim 4 only'); | |
| const key = `vq_matmul/${rows}x${cols}/k${k}/${add ? 'add-' + this.resid : 'store'}/v1-${t.rwg}x${t.lanes}x${t.tpt}`; | |
| const OUT_STMT = add ? `y[t * ROWS + r] = ${OUT_T}(f32(y[t * ROWS + r]) + acc[k]);` : 'y[t * ROWS + r] = acc[k];'; | |
| return this._pipeline(key, 'vq_matmul_v1', { ...common, RWG: t.rwg, LANES: t.lanes, TPT: t.tpt, WGS: t.rwg * t.lanes, OUT_STMT }, [U, RO, RO, RO, RO, RW]); | |
| } | |
| if (rows % t.rwg) throw new Error(`vq_matmul: ${rows} rows not a multiple of rwg=${t.rwg}`); | |
| const key = `vq_matmul/${rows}x${cols}/k${k}/d${dim}/${add ? 'add-' + this.resid : 'store'}/${t.rwg}x${t.tt}x${t.cb}x${t.mr}x${t.mt}`; | |
| const OUT_STMT = add ? `y[t * ROWS + r] = ${OUT_T}(f32(y[t * ROWS + r]) + acc[i][j]);` : 'y[t * ROWS + r] = acc[i][j];'; | |
| return this._pipeline(key, 'vq_matmul', { ...common, RWG: t.rwg, TT: t.tt, CB: t.cb, MR: t.mr, MT: t.mt, WGS: (t.rwg / t.mr) * (t.tt / t.mt), OUT_STMT }, [U, RO, RO, RO, RO, RW]); | |
| } | |
| _qkPipeline(rope, inplace) { | |
| const m = this.m; | |
| const vars = { NH: m.heads, NKV: m.kvHeads, HD: m.headDim, HALF: m.headDim / 2, EPS: f32lit(m.rmsEps), DO_ROPE: rope ? 'true' : 'false' }; | |
| const kinds = [U, RW, RW, RO, RO, RO, RO]; | |
| if (inplace) { | |
| vars.OUT_BINDINGS = ''; | |
| vars.OUT_STMT = 'if (isq) { q[base + i] = an; q[base + i + HALF] = bn; } else { k[base + i] = an; k[base + i + HALF] = bn; }'; | |
| } else { | |
| vars.OUT_BINDINGS = '@group(0) @binding(7) var<storage, read_write> qo: array<f32>;\n@group(0) @binding(8) var<storage, read_write> ko: array<f32>;'; | |
| vars.OUT_STMT = 'if (isq) { qo[base + i] = an; qo[base + i + HALF] = bn; } else { ko[base + i] = an; ko[base + i + HALF] = bn; }'; | |
| kinds.push(RW, RW); | |
| } | |
| return this._pipeline(`qk/${rope ? 'rope' : 'norm'}/${inplace ? 'inplace' : 'out'}`, 'qk_norm_rope', vars, kinds); | |
| } | |
| _attnPipeline(probs) { | |
| const m = this.m; | |
| // scores array = stored prefix + query tokens; bindings 5 / 6 = prefix K / V (dummy 16-byte buffers when the container has none) | |
| const vars = { NH: m.heads, NKV: m.kvHeads, HD: m.headDim, PMAX: Math.max(this.P, 1), TMAX: this.P + this.tmax, SCALE: f32lit(Math.pow(m.headDim, -0.5)), WRITE_PROBS: probs ? 'true' : 'false', | |
| PROBS_BINDING: probs ? '@group(0) @binding(7) var<storage, read_write> probs: array<f32>;' : '', | |
| PROBS_STMT: probs ? 'probs[(h * T + t) * W + j] = select(0.0, sc[j] * inv, j <= last);' : '' }; | |
| const kinds = [U, RO, RO, RO, RW, RO, RO]; if (probs) kinds.push(RW); | |
| return this._pipeline(`attention/${probs ? 'probs' : 'plain'}/P${this.P}`, 'attention', vars, kinds); | |
| } | |
| // ---------------------------------------------------------------------------------------------- buffers | |
| _buf(label, bytes, extra = 0) { | |
| const b = this.device.createBuffer({ label, size: Math.max(16, (bytes + 3) & ~3), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | extra }); | |
| this.actBytes += b.size; | |
| return b; | |
| } | |
| _allocate() { | |
| const m = this.m, T = this.tmax, rb = this.resid === 'f16' ? 2 : 4; | |
| this.actBytes = 0; | |
| const H = m.hidden, QD = m.heads * m.headDim, KD = m.kvHeads * m.headDim, I = m.intermediate; | |
| this.b = { | |
| uni: this.device.createBuffer({ label: 'params', size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST }), | |
| x: this._buf('x (residual)', T * H * rb), h: this._buf('h (norm out)', T * H * 4), hr: this._buf('hr (rot, f16)', T * H * 2), | |
| q: this._buf('q', T * QD * 4), k: this._buf('k', T * KD * 4), v: this._buf('v', T * KD * 4), | |
| attn: this._buf('attn_out', T * QD * 4), orot: this._buf('o_rot (f16)', T * QD * 2), | |
| g: this._buf('gate', T * I * 4), u: this._buf('up', T * I * 4), a: this._buf('swiglu', T * I * 4), ar: this._buf('down_rot (f16)', T * I * 2), | |
| emb: this._buf('embedding', H * 4), cos: this._buf('rope cos', T * (m.headDim / 2) * 4), sin: this._buf('rope sin', T * (m.headDim / 2) * 4), | |
| stage: this.device.createBuffer({ label: 'readback', size: H * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }), | |
| }; | |
| if (!this.P) { this.b.kp0 = this._buf('prefix k (none)', 16); this.b.vp0 = this._buf('prefix v (none)', 16); } // bound but never read (PP = 0) | |
| this.ropeRows = 0; this.ropeOffset = -1; // rows / position offset of the cos/sin table already uploaded | |
| } | |
| _bg(p, buffers, label) { | |
| return this.device.createBindGroup({ label, layout: p.layout, entries: buffers.map((buffer, binding) => ({ binding, resource: { buffer } })) }); | |
| } | |
| /** [prefix K, prefix V] buffers of block li for the attention bind group (dummies when the runtime uses no prefix). */ | |
| _prefixPair(li) { | |
| if (!this.P) return [this.b.kp0, this.b.vp0]; | |
| const pk = this.c.prefixKV(li); | |
| return [pk.k, pk.v]; | |
| } | |
| _bindAll() { | |
| const c = this.c, b = this.b, m = this.m; | |
| const rotBG = (n, src, dst) => { const r = c.rot(n); return this._bg(this.pRot[n], [b.uni, src, r.signs, r.perms, dst], `rot${n}`); }; | |
| this.blocks = []; | |
| for (let li = 0; li < m.layers; li++) { | |
| const lin = (nm) => c.linear(`blk.${li}.${nm}`); | |
| const mm = (nm, src, dst, add) => { const l = lin(nm); const p = this._matmulPipeline(l, add); return { p, bg: this._bg(p, [b.uni, src, l.codebook, l.idx, l.scale, dst], `blk.${li}.${nm}`), rows: l.rows }; }; | |
| const blk = { | |
| attnNorm: this._bg(this.pRms, [b.uni, b.x, c.norm(`blk.${li}.attn_norm`), b.h], `blk.${li}.attn_norm`), | |
| rot1024: rotBG(m.hidden, b.h, b.hr), | |
| q: mm('q', b.hr, b.q, false), k: mm('k', b.hr, b.k, false), v: mm('v', b.hr, b.v, false), | |
| qk: this._bg(this.pQk, [b.uni, b.q, b.k, c.norm(`blk.${li}.attn_q_norm`), c.norm(`blk.${li}.attn_k_norm`), b.cos, b.sin], `blk.${li}.qk`), | |
| attn: this._bg(this.pAttn, [b.uni, b.q, b.k, b.v, b.attn, ...this._prefixPair(li)], `blk.${li}.attn`), | |
| rot2048: rotBG(m.heads * m.headDim, b.attn, b.orot), | |
| o: mm('o', b.orot, b.x, true), | |
| ffnNorm: this._bg(this.pRms, [b.uni, b.x, c.norm(`blk.${li}.ffn_norm`), b.h], `blk.${li}.ffn_norm`), | |
| gate: mm('gate', b.hr, b.g, false), up: mm('up', b.hr, b.u, false), | |
| swiglu: this._bg(this.pSwiglu, [b.uni, b.g, b.u, b.a], `blk.${li}.swiglu`), | |
| rot3072: rotBG(m.intermediate, b.a, b.ar), | |
| down: mm('down', b.ar, b.x, true), | |
| li, | |
| }; | |
| this.blocks.push(blk); | |
| } | |
| this.outNorm = this._bg(this.pRms, [b.uni, b.x, c.norm('output_norm'), b.h], 'output_norm'); | |
| this.pool = this._bg(this.pPool, [b.uni, b.h, b.emb], 'pool'); | |
| } | |
| _debugSetup() { | |
| if (this.dbg) return; | |
| const m = this.m, T = this.tmax, b = this.b; | |
| const qn = this._buf('q_norm (debug)', T * m.heads * m.headDim * 4), kn = this._buf('k_norm (debug)', T * m.kvHeads * m.headDim * 4); | |
| const probs = this._buf('attn_probs (debug)', m.heads * T * (this.P + T) * 4); | |
| const pQkNorm = this._qkPipeline(false, false), pAttnP = this._attnPipeline(true); | |
| this.dbg = { qn, kn, probs, pQkNorm, pAttnP, blocks: this.blocks.map((blk) => ({ | |
| qkNorm: this._bg(pQkNorm, [b.uni, b.q, b.k, this.c.norm(`blk.${blk.li}.attn_q_norm`), this.c.norm(`blk.${blk.li}.attn_k_norm`), b.cos, b.sin, qn, kn], `dbg qk ${blk.li}`), | |
| attn: this._bg(pAttnP, [b.uni, b.q, b.k, b.v, b.attn, ...this._prefixPair(blk.li), probs], `dbg attn ${blk.li}`), | |
| })) }; | |
| } | |
| // ---------------------------------------------------------------------------------------------- per query | |
| /** RoPE cos/sin rows 0..T-1 as f32, mimicking the reference's fp32 evaluation order (inv_freq, t*inv_freq, cos) with Math.fround. */ | |
| _ropeTable(T, offset = 0) { | |
| // rows t = 0..T-1 hold the positions offset + t (offset = stored prefix length when the prefix K / V are used) | |
| if (T <= this.ropeRows && offset === this.ropeOffset) return; | |
| const m = this.m, half = m.headDim / 2, fr = Math.fround; | |
| const cos = new Float32Array(T * half), sin = new Float32Array(T * half); | |
| const inv = new Float32Array(half); | |
| for (let j = 0; j < half; j++) inv[j] = fr(1 / fr(Math.pow(m.ropeTheta, fr((2 * j) / m.headDim)))); | |
| for (let t = 0; t < T; t++) for (let j = 0; j < half; j++) { const a = fr((offset + t) * inv[j]); cos[t * half + j] = Math.cos(a); sin[t * half + j] = Math.sin(a); } | |
| this.device.queue.writeBuffer(this.b.cos, 0, cos); this.device.queue.writeBuffer(this.b.sin, 0, sin); | |
| this.ropeRows = T; this.ropeOffset = offset; | |
| } | |
| /** M6: ids that start with the stored prefix are cut; returns {ids: query tokens, PP: prefix length used (0 = plain path)}. */ | |
| _splitPrefix(ids) { | |
| const pre = this.prefixIds; | |
| if (!pre || ids.length <= pre.length) return { ids, PP: 0 }; | |
| for (let i = 0; i < pre.length; i++) if (ids[i] !== pre[i]) return { ids, PP: 0 }; | |
| return { ids: ids.slice(pre.length), PP: pre.length }; | |
| } | |
| /** | |
| * Embed one token-id sequence (already re-tokenised onto the trimmed table). | |
| * @returns {Promise<{embedding: Float32Array, ms: number, T: number, captures?: Map<string, {data: Float32Array, shape: number[]}>}>} | |
| */ | |
| async embed(idsIn, { debug = false, profile = false } = {}) { | |
| const t0 = performance.now(); | |
| const { ids, PP } = this._splitPrefix(idsIn); // M6: the stored prompt K / V replace the prompt tokens (PP = 0: plain path) | |
| const T = ids.length, m = this.m, b = this.b, dev = this.device; | |
| if (T < 1 || T > this.tmax) throw new Error(`T=${T} outside 1..${this.tmax}`); | |
| profile = profile && dev.features.has('timestamp-query'); | |
| const nDispatch = m.layers * (this.dispatchesPerBlock + (debug ? 2 : 0)) + 2; | |
| const qset = profile ? dev.createQuerySet({ type: 'timestamp', count: 2 * nDispatch }) : null; | |
| const kinds = []; | |
| const x0 = this.resid === 'f16' ? this.c.tokens.lookupF16(ids) : this.c.tokens.lookup(ids); | |
| dev.queue.writeBuffer(b.x, 0, x0); | |
| this._ropeTable(T, PP); | |
| dev.queue.writeBuffer(b.uni, 0, new Uint32Array([T, debug ? 1 : 0, PP, 0])); | |
| if (debug) this._debugSetup(); | |
| const H = m.hidden, QD = m.heads * m.headDim, KD = m.kvHeads * m.headDim, I = m.intermediate; | |
| const enc = dev.createCommandEncoder({ label: 'query' }); | |
| const caps = []; | |
| let pass = null; | |
| // debug: copy the intermediate NOW (the encoder executes in order; h / hr / x are overwritten by later dispatches) | |
| const capture = (name, buffer, dtype, shape) => { | |
| if (!debug) return; | |
| const n = shape.reduce((p, q) => p * q, 1), bytes = n * (dtype === 'f16' ? 2 : 4); | |
| const stage = dev.createBuffer({ label: `capture ${name}`, size: Math.max(16, (bytes + 3) & ~3), usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); | |
| enc.copyBufferToBuffer(buffer, 0, stage, 0, (bytes + 3) & ~3); | |
| caps.push({ name, dtype, shape, bytes, stage }); | |
| }; | |
| const run = (p, bg, gx, gy = 1) => { | |
| if (!pass) { | |
| const i = kinds.length; | |
| pass = enc.beginComputePass(profile ? { timestampWrites: { querySet: qset, beginningOfPassWriteIndex: 2 * i, endOfPassWriteIndex: 2 * i + 1 } } : {}); | |
| } | |
| kinds.push(p.key.split('/')[0]); | |
| pass.setPipeline(p.pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(gx, gy); | |
| if (debug || profile) { pass.end(); pass = null; } | |
| }; | |
| const tokTile = this.tile.kernel === 'v1' ? this.tile.lanes * this.tile.tpt : this.tile.tt; | |
| const mm = (op) => run(op.p, op.bg, op.rows / this.tile.rwg, Math.ceil(T / tokTile)); | |
| const rT = this.resid === 'f16' ? 'f16' : 'f32'; | |
| capture('embed', b.x, rT, [T, H]); | |
| for (const blk of this.blocks) { | |
| const pre = `blk.${blk.li}.`; | |
| run(this.pRms, blk.attnNorm, T); capture(pre + 'attn_norm', b.h, 'f32', [T, H]); | |
| run(this.pRot[H], blk.rot1024, T); capture(pre + 'attn_rot', b.hr, 'f16', [T, H]); | |
| mm(blk.q); capture(pre + 'q', b.q, 'f32', [T, QD]); | |
| mm(blk.k); capture(pre + 'k', b.k, 'f32', [T, KD]); | |
| mm(blk.v); capture(pre + 'v', b.v, 'f32', [T, KD]); | |
| if (debug) { | |
| run(this.dbg.pQkNorm, this.dbg.blocks[blk.li].qkNorm, T, m.heads + m.kvHeads); | |
| capture(pre + 'q_norm', this.dbg.qn, 'f32', [T, m.heads, m.headDim]); capture(pre + 'k_norm', this.dbg.kn, 'f32', [T, m.kvHeads, m.headDim]); | |
| } | |
| run(this.pQk, blk.qk, T, m.heads + m.kvHeads); capture(pre + 'q_rope', b.q, 'f32', [T, m.heads, m.headDim]); capture(pre + 'k_rope', b.k, 'f32', [T, m.kvHeads, m.headDim]); | |
| if (debug) { run(this.dbg.pAttnP, this.dbg.blocks[blk.li].attn, T, m.heads); capture(pre + 'attn_probs', this.dbg.probs, 'f32', [m.heads, T, PP + T]); } | |
| else run(this.pAttn, blk.attn, T, m.heads); | |
| capture(pre + 'attn_out', b.attn, 'f32', [T, QD]); | |
| run(this.pRot[QD], blk.rot2048, T); capture(pre + 'o_rot', b.orot, 'f16', [T, QD]); | |
| mm(blk.o); capture(pre + 'resid_attn', b.x, rT, [T, H]); | |
| run(this.pRms, blk.ffnNorm, T); capture(pre + 'ffn_norm', b.h, 'f32', [T, H]); | |
| run(this.pRot[H], blk.rot1024, T); capture(pre + 'ffn_rot', b.hr, 'f16', [T, H]); | |
| mm(blk.gate); capture(pre + 'gate', b.g, 'f32', [T, I]); | |
| mm(blk.up); capture(pre + 'up', b.u, 'f32', [T, I]); | |
| run(this.pSwiglu, blk.swiglu, Math.ceil((T * I) / 256)); capture(pre + 'swiglu', b.a, 'f32', [T, I]); | |
| run(this.pRot[I], blk.rot3072, T); capture(pre + 'down_rot', b.ar, 'f16', [T, I]); | |
| mm(blk.down); capture(pre + 'resid_ffn', b.x, rT, [T, H]); | |
| } | |
| run(this.pRms, this.outNorm, T); capture('output_norm', b.h, 'f32', [T, H]); | |
| run(this.pPool, this.pool, 1); capture('embedding', b.emb, 'f32', [H]); | |
| if (pass) { pass.end(); pass = null; } | |
| enc.copyBufferToBuffer(b.emb, 0, b.stage, 0, H * 4); | |
| let qbuf = null, qstage = null; | |
| if (profile) { | |
| qbuf = dev.createBuffer({ label: 'timestamps', size: 16 * nDispatch, usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC }); | |
| qstage = dev.createBuffer({ label: 'timestamps readback', size: 16 * nDispatch, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); | |
| enc.resolveQuerySet(qset, 0, 2 * kinds.length, qbuf, 0); | |
| enc.copyBufferToBuffer(qbuf, 0, qstage, 0, 16 * kinds.length); | |
| } | |
| dev.queue.submit([enc.finish()]); | |
| await b.stage.mapAsync(GPUMapMode.READ); | |
| const embedding = new Float32Array(b.stage.getMappedRange().slice(0)); | |
| b.stage.unmap(); | |
| const out = { embedding, T, PP, ms: performance.now() - t0, dispatches: kinds.length }; | |
| if (profile) { | |
| await qstage.mapAsync(GPUMapMode.READ); | |
| const ts = new BigUint64Array(qstage.getMappedRange().slice(0)); | |
| qstage.unmap(); qstage.destroy(); qbuf.destroy(); qset.destroy(); | |
| const byKind = {}; let total = 0; | |
| kinds.forEach((k, i) => { const ms = Number(ts[2 * i + 1] - ts[2 * i]) / 1e6; byKind[k] = (byKind[k] || 0) + ms; total += ms; }); | |
| out.profile = { by_kind_ms: byKind, gpu_total_ms: total, span_ms: Number(ts[2 * kinds.length - 1] - ts[0]) / 1e6 }; | |
| } | |
| if (debug) { | |
| await Promise.all(caps.map((cp) => cp.stage.mapAsync(GPUMapMode.READ))); | |
| out.captures = new Map(); | |
| for (const cp of caps) { | |
| const all = cp.stage.getMappedRange(); | |
| const data = cp.dtype === 'f16' ? f16ArrayToF32(new Uint16Array(all, 0, cp.bytes / 2)) : new Float32Array(all.slice(0, cp.bytes)); | |
| out.captures.set(cp.name, { data, shape: cp.shape, dtype: cp.dtype }); | |
| cp.stage.unmap(); cp.stage.destroy(); | |
| } | |
| } | |
| return out; | |
| } | |
| destroy() { for (const v of Object.values(this.b)) v.destroy(); } | |
| } | |