Something went wrong. Try again.
Anonymize your writing style. Zig WASM engine detects authorship markers, fine-tuned LLM rewrites to remove them. Runs entirely in-browser. fantasma.qstorage.quilibrium.com
wasm privacy qwen zig
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879// WebGPU inference engine for Qwen3.5-0.8B (QwenTheBard)// GPU handles large matrix-vector multiplies; JS handles norms, activations, attention, DeltaNet recurrence.
// ============================================================// Model constants// ============================================================
const VOCAB_SIZE = 248320;const HIDDEN_SIZE = 1024;const NUM_LAYERS = 24;const INTERMEDIATE_SIZE = 3584;const RMS_NORM_EPS = 1e-6;const EOS_TOKEN = 248046;
// Full attention (layers 3,7,11,15,19,23)const FULL_ATTN_LAYERS = new Set([3, 7, 11, 15, 19, 23]);const NUM_HEADS = 8;const NUM_KV_HEADS = 2;const HEAD_DIM = 256;const ROTARY_DIM = 64;const ROPE_THETA = 10_000_000;const KV_GROUPS = NUM_HEADS / NUM_KV_HEADS; // 4const ATTN_OUTPUT_DIM = NUM_HEADS * HEAD_DIM; // 2048
// DeltaNetconst DN_NUM_HEADS = 16;const DN_KEY_DIM = 128;const DN_VALUE_DIM = 128;const DN_CONV_KERNEL = 4;const DN_CONV_DIM = DN_NUM_HEADS * (DN_KEY_DIM * 2 + DN_VALUE_DIM); // 6144const DN_TOTAL_VALUE = DN_NUM_HEADS * DN_VALUE_DIM; // 2048const DN_TOTAL_KEY = DN_NUM_HEADS * DN_KEY_DIM; // 2048
// Buffer sizingconst MAX_OUTPUT_SIZE = VOCAB_SIZE; // 248320 floats for logitsconst MAX_INPUT_SIZE = DN_CONV_DIM; // 6144 floats (largest projection input after MLP gate*up)const MAX_SEQ_LEN = 2048;
// ============================================================// WGSL compute shader for Q4 matrix-vector multiply// ============================================================
const Q4_MATVEC_SHADER = /* wgsl */ `struct Params { M: u32, // output rows N: u32, // input cols groups_per_row: u32,}
@group(0) @binding(0) var<storage, read> weights: array<u32>;@group(0) @binding(1) var<storage, read> input: array<f32>;@group(0) @binding(2) var<storage, read_write> output: array<f32>;@group(0) @binding(3) var<uniform> params: Params;
// Q4 group: 68 bytes = 17 u32s// [0]: scale(f16) | zero(f16) packed as u32// [1..16]: 64 bytes of packed nibbles (128 elements, 2 per byte)
@compute @workgroup_size(256)fn main(@builtin(global_invocation_id) gid: vec3u) { let row = gid.x; if (row >= params.M) { return; }
let groups_per_row = params.groups_per_row; let group_u32s = 17u; // 68 bytes / 4 let row_offset = row * groups_per_row * group_u32s;
var acc: f32 = 0.0; var col: u32 = 0u;
for (var g: u32 = 0u; g < groups_per_row; g++) { let base = row_offset + g * group_u32s;
// Unpack scale and zero from f16 pair let sz = weights[base]; let scale_bits = sz & 0xFFFFu; let zero_bits = (sz >> 16u) & 0xFFFFu; let scale = unpack2x16float(sz).x; let zero = unpack2x16float(sz).y;
// 64 bytes of nibbles = 16 u32s, each u32 has 8 nibbles = 8 values for (var w: u32 = 0u; w < 16u; w++) { let packed = weights[base + 1u + w]; // Each u32 holds 8 nibbles (4 bytes, 2 nibbles per byte) for (var n: u32 = 0u; n < 8u; n++) { let nibble = (packed >> (n * 4u)) & 0xFu; let val = f32(nibble) * scale + zero; acc += val * input[col]; col++; } } }
output[row] = acc;}`;
// ============================================================// Q4MF binary format parser// ============================================================
interface TensorMeta { shape: number[]; dtype: number; // 0=Q4, 1=F32, 2=BF16 data: Uint8Array;}
function parseQ4MF(buffer: ArrayBuffer): Map<string, TensorMeta> { const view = new DataView(buffer); const bytes = new Uint8Array(buffer); let offset = 0;
// Header: magic(u32), version(u32), num_tensors(u32), group_size(u32) const magic = view.getUint32(offset, true); offset += 4; if (magic !== 0x46344D51) throw new Error(`Invalid Q4MF magic: 0x${magic.toString(16)}`); const version = view.getUint32(offset, true); offset += 4; if (version !== 1) throw new Error(`Unsupported Q4MF version: ${version}`); const numTensors = view.getUint32(offset, true); offset += 4; /* const groupSize = */ view.getUint32(offset, true); offset += 4;
const tensors = new Map<string, TensorMeta>(); const decoder = new TextDecoder();
for (let t = 0; t < numTensors; t++) { const nameLen = view.getUint16(offset, true); offset += 2; const name = decoder.decode(bytes.subarray(offset, offset + nameLen)); offset += nameLen; const ndim = bytes[offset]; offset += 1; const shape: number[] = []; for (let d = 0; d < ndim; d++) { shape.push(view.getUint32(offset, true)); offset += 4; } const dtype = bytes[offset]; offset += 1; const dataLen = view.getUint32(offset, true); offset += 4; const data = bytes.subarray(offset, offset + dataLen); offset += dataLen; tensors.set(name, { shape, dtype, data }); }
return tensors;}
// ============================================================// CPU helpers// ============================================================
function bf16ToF32(bits: number): number { const buf = new ArrayBuffer(4); const u32 = new Uint32Array(buf); u32[0] = bits << 16; return new Float32Array(buf)[0];}
function f16ToF32(bits: number): number { const sign = (bits >> 15) & 1; const exp = (bits >> 10) & 0x1F; const frac = bits & 0x3FF; if (exp === 0) { // Subnormal or zero return (sign ? -1 : 1) * Math.pow(2, -14) * (frac / 1024); } if (exp === 31) { return frac ? NaN : (sign ? -Infinity : Infinity); } return (sign ? -1 : 1) * Math.pow(2, exp - 15) * (1 + frac / 1024);}
/** Convert a tensor's raw data to Float32Array based on dtype. */function toF32Array(meta: TensorMeta): Float32Array { if (meta.dtype === 1) { // F32 — copy to aligned buffer since subarray offset may not be 4-byte aligned const out = new Float32Array(meta.data.byteLength / 4); out.set(new Float32Array(meta.data.slice().buffer)); return out; } if (meta.dtype === 2) { // BF16 const count = meta.data.byteLength / 2; const out = new Float32Array(count); const view = new DataView(meta.data.buffer, meta.data.byteOffset, meta.data.byteLength); for (let i = 0; i < count; i++) { out[i] = bf16ToF32(view.getUint16(i * 2, true)); } return out; } throw new Error(`Cannot directly convert Q4 tensor to F32`);}
/** Dequantize one row of a Q4 tensor on CPU. */function dequantQ4Row(data: Uint8Array, row: number, numCols: number, groupSize: number): Float32Array { const groupsPerRow = Math.ceil(numCols / groupSize); const bytesPerGroup = 68; // 2(scale) + 2(zero) + 64(nibbles) const rowOffset = row * groupsPerRow * bytesPerGroup; const out = new Float32Array(numCols); const view = new DataView(data.buffer, data.byteOffset); let col = 0;
for (let g = 0; g < groupsPerRow; g++) { const base = rowOffset + g * bytesPerGroup; const scale = f16ToF32(view.getUint16(base, true)); const zero = f16ToF32(view.getUint16(base + 2, true));
// 64 bytes of packed nibbles for (let b = 0; b < 64 && col < numCols; b++) { const byte = data[base + 4 + b]; const lo = byte & 0xF; const hi = (byte >> 4) & 0xF; out[col++] = lo * scale + zero; if (col < numCols) { out[col++] = hi * scale + zero; } } }
return out;}
function sigmoid(x: number): number { return 1.0 / (1.0 + Math.exp(-x));}
function silu(x: number): number { return x * sigmoid(x);}
function softplus(x: number): number { return Math.log(1.0 + Math.exp(x));}
/** RMSNorm: out[i] = x[i] * rsqrt(mean(x^2) + eps) * (1 + weight[i]) */function rmsNorm(out: Float32Array, x: Float32Array, weight: Float32Array, eps: number, offset = 0, len?: number): void { const n = len ?? x.length; let sumSq = 0; for (let i = 0; i < n; i++) sumSq += x[offset + i] * x[offset + i]; const rms = 1.0 / Math.sqrt(sumSq / n + eps); for (let i = 0; i < n; i++) { out[offset + i] = x[offset + i] * rms * (1.0 + weight[i]); }}
/** Gated RMSNorm: out = rmsNorm(x, weight_direct) * silu(gate). weight used directly (NOT 1+w). */function rmsNormGated(out: Float32Array, outOff: number, x: Float32Array, xOff: number, gate: Float32Array, gateOff: number, weight: Float32Array, eps: number, len: number): void { let sumSq = 0; for (let i = 0; i < len; i++) sumSq += x[xOff + i] * x[xOff + i]; const rms = 1.0 / Math.sqrt(sumSq / len + eps); for (let i = 0; i < len; i++) { out[outOff + i] = x[xOff + i] * rms * weight[i] * silu(gate[gateOff + i]); }}
/** In-place L2 normalization. */function l2Norm(x: Float32Array, offset: number, len: number): void { let sumSq = 0; for (let i = 0; i < len; i++) sumSq += x[offset + i] * x[offset + i]; const inv = 1.0 / (Math.sqrt(sumSq) + 1e-12); for (let i = 0; i < len; i++) x[offset + i] *= inv;}
/** Apply rotary position embedding (rotate_half style) to first rotaryDim elements. */function applyRope(x: Float32Array, xOff: number, cos: Float32Array, sin: Float32Array, cosOff: number, rotaryDim: number): void { const half = rotaryDim / 2; for (let i = 0; i < half; i++) { const x0 = x[xOff + i]; const x1 = x[xOff + i + half]; x[xOff + i] = x0 * cos[cosOff + i] - x1 * sin[cosOff + i]; x[xOff + i + half] = x1 * cos[cosOff + i + half] + x0 * sin[cosOff + i + half]; }}
// ============================================================// GPUEngine// ============================================================
const GROUP_SIZE = 128;
export class GPUEngine { // WebGPU objects private device!: GPUDevice; private q4Pipeline!: GPUComputePipeline; private bindGroupLayout!: GPUBindGroupLayout;
// Weight buffers on GPU (Q4 data) private weightBuffers = new Map<string, GPUBuffer>();
// Parsed weight metadata private weightMeta = new Map<string, { shape: number[]; dtype: number }>();
// Small unquantized weights on CPU private cpuWeights = new Map<string, Float32Array>();
// Embedding Q4 data kept on CPU for row-level dequant private embedData!: Uint8Array; private embedCols!: number;
// Working GPU buffers private inputBuf!: GPUBuffer; private outputBufs: GPUBuffer[] = []; private paramsBufs: GPUBuffer[] = []; private stagingBuf!: GPUBuffer;
// CPU-side state vectors private hidden = new Float32Array(HIDDEN_SIZE); private residual = new Float32Array(HIDDEN_SIZE); private normed = new Float32Array(HIDDEN_SIZE);
// KV caches for full attention layers (6 layers) // Each stores [maxSeq * kvHeads * headDim] private kvKeys: Float32Array[] = []; private kvValues: Float32Array[] = []; private kvLen: number[] = [];
// DeltaNet states (18 layers) // Recurrent state: [heads * keyDim * valueDim] per layer private deltaStates: Float32Array[] = []; // Conv state: [convDim * (kernel-1)] per layer private convStates: Float32Array[] = [];
// RoPE tables private ropeCos!: Float32Array; private ropeSin!: Float32Array;
private constructor() {}
/** Initialize WebGPU and create pipeline. Returns null if WebGPU unavailable. */ static async create(): Promise<GPUEngine | null> { if (!navigator.gpu) { console.warn("[gpu] navigator.gpu not available (need secure context: localhost or HTTPS)"); return null; } const adapter = await navigator.gpu.requestAdapter(); if (!adapter) { console.warn("[gpu] No GPU adapter found — check chrome://gpu and chrome://flags/#enable-unsafe-webgpu"); return null; } const device = await adapter.requestDevice({ requiredLimits: { maxStorageBufferBindingSize: 256 * 1024 * 1024, // 256MB for large Q4 weight tensors maxBufferSize: 256 * 1024 * 1024, // 256MB }, });
const engine = new GPUEngine(); engine.device = device;
// Create bind group layout engine.bindGroupLayout = device.createBindGroupLayout({ entries: [ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }, { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }, { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }, { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } }, ], });
// Create compute pipeline const shaderModule = device.createShaderModule({ code: Q4_MATVEC_SHADER }); engine.q4Pipeline = device.createComputePipeline({ layout: device.createPipelineLayout({ bindGroupLayouts: [engine.bindGroupLayout] }), compute: { module: shaderModule, entryPoint: 'main' }, });
// Pre-allocate working buffers engine.inputBuf = device.createBuffer({ size: MAX_INPUT_SIZE * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); // Pool of output + param buffers for batched matvec (up to 4 concurrent ops) for (let i = 0; i < 4; i++) { engine.outputBufs.push(device.createBuffer({ size: MAX_OUTPUT_SIZE * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC, })); engine.paramsBufs.push(device.createBuffer({ size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, })); } engine.stagingBuf = device.createBuffer({ size: MAX_OUTPUT_SIZE * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, });
// Precompute RoPE tables engine.ropeCos = new Float32Array(MAX_SEQ_LEN * ROTARY_DIM); engine.ropeSin = new Float32Array(MAX_SEQ_LEN * ROTARY_DIM); const halfDim = ROTARY_DIM / 2; for (let pos = 0; pos < MAX_SEQ_LEN; pos++) { for (let i = 0; i < halfDim; i++) { const freq = 1.0 / Math.pow(ROPE_THETA, (2 * i) / ROTARY_DIM); const angle = pos * freq; const c = Math.cos(angle); const s = Math.sin(angle); engine.ropeCos[pos * ROTARY_DIM + i] = c; engine.ropeCos[pos * ROTARY_DIM + halfDim + i] = c; engine.ropeSin[pos * ROTARY_DIM + i] = s; engine.ropeSin[pos * ROTARY_DIM + halfDim + i] = s; } }
// Initialize KV caches and DeltaNet states for (let layer = 0; layer < NUM_LAYERS; layer++) { if (FULL_ATTN_LAYERS.has(layer)) { engine.kvKeys.push(new Float32Array(MAX_SEQ_LEN * NUM_KV_HEADS * HEAD_DIM)); engine.kvValues.push(new Float32Array(MAX_SEQ_LEN * NUM_KV_HEADS * HEAD_DIM)); engine.kvLen.push(0); } else { engine.deltaStates.push(new Float32Array(DN_NUM_HEADS * DN_KEY_DIM * DN_VALUE_DIM)); engine.convStates.push(new Float32Array(DN_CONV_DIM * (DN_CONV_KERNEL - 1))); } }
return engine; }
/** Parse Q4MF binary and upload Q4 weights to GPU, keep small weights on CPU. */ async loadWeights(q4mfData: ArrayBuffer): Promise<void> { console.log(`[gpu] Parsing Q4MF (${(q4mfData.byteLength / 1024 / 1024).toFixed(1)} MB)...`); const tensors = parseQ4MF(q4mfData); console.log(`[gpu] ${tensors.size} tensors parsed, uploading to GPU...`); const prefix = 'model.language_model.'; let gpuBytes = 0; let cpuCount = 0;
// Small Q4 tensors that are accessed element-wise on CPU (not via GPU matvec) const cpuDequantSuffixes = ['conv1d.weight', 'A_log', 'dt_bias'];
for (const [name, meta] of tensors) { this.weightMeta.set(name, { shape: meta.shape, dtype: meta.dtype });
if (meta.dtype === 0) { // Check if this Q4 tensor should be dequantized for CPU access const needsCpu = cpuDequantSuffixes.some(s => name.endsWith(s)) || name === prefix + 'embed_tokens.weight';
if (needsCpu && name !== prefix + 'embed_tokens.weight') { // Small Q4 tensor used element-wise — dequant all rows to CPU const rows = meta.shape[0]; const cols = meta.shape.length > 1 ? meta.shape.reduce((a, b) => a * b) / rows : 1; const totalElements = meta.shape.reduce((a, b) => a * b); const flat = new Float32Array(totalElements); for (let r = 0; r < rows; r++) { const row = dequantQ4Row(meta.data, r, cols, GROUP_SIZE); flat.set(row, r * cols); } this.cpuWeights.set(name, flat); cpuCount++; }
// Upload to GPU for matvec const bufSize = Math.ceil(meta.data.byteLength / 4) * 4; const gpuBuf = this.device.createBuffer({ size: bufSize, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, mappedAtCreation: true, }); new Uint8Array(gpuBuf.getMappedRange()).set(meta.data); gpuBuf.unmap(); this.weightBuffers.set(name, gpuBuf); gpuBytes += bufSize;
// Keep embed_tokens data on CPU for row dequant if (name === prefix + 'embed_tokens.weight') { this.embedData = meta.data; this.embedCols = meta.shape[1]; } } else { // F32 or BF16 — keep on CPU this.cpuWeights.set(name, toF32Array(meta)); cpuCount++; } } console.log(`[gpu] Weights loaded: ${this.weightBuffers.size} on GPU (${(gpuBytes / 1024 / 1024).toFixed(1)} MB), ${cpuCount} on CPU`); }
/** Dequantize one embedding row on CPU. */ private getEmbedding(tokenId: number): Float32Array { return dequantQ4Row(this.embedData, tokenId, this.embedCols, GROUP_SIZE); }
/** Batched GPU Q4 matrix-vector multiply: runs multiple matvecs in a single GPU submission + readback. */ async matvecBatch( input: Float32Array, ops: Array<{ weightName: string; outputSize: number }>, ): Promise<Float32Array[]> { // Upload input once this.device.queue.writeBuffer(this.inputBuf, 0, input.buffer, input.byteOffset, input.byteLength);
// Upload per-op params for (let i = 0; i < ops.length; i++) { const meta = this.weightMeta.get(ops[i].weightName)!; const M = ops[i].outputSize; const N = meta.shape[meta.shape.length - 1]; this.device.queue.writeBuffer(this.paramsBufs[i], 0, new Uint32Array([M, N, Math.ceil(N / GROUP_SIZE)])); }
const commandEncoder = this.device.createCommandEncoder(); let totalFloats = 0; const offsets: number[] = [];
for (let i = 0; i < ops.length; i++) { const wBuf = this.weightBuffers.get(ops[i].weightName); if (!wBuf) throw new Error(`Weight not found on GPU: ${ops[i].weightName}`);
const bindGroup = this.device.createBindGroup({ layout: this.bindGroupLayout, entries: [ { binding: 0, resource: { buffer: wBuf } }, { binding: 1, resource: { buffer: this.inputBuf } }, { binding: 2, resource: { buffer: this.outputBufs[i] } }, { binding: 3, resource: { buffer: this.paramsBufs[i] } }, ], });
const pass = commandEncoder.beginComputePass(); pass.setPipeline(this.q4Pipeline); pass.setBindGroup(0, bindGroup); pass.dispatchWorkgroups(Math.ceil(ops[i].outputSize / 256)); pass.end();
offsets.push(totalFloats); totalFloats += ops[i].outputSize; }
// Copy all outputs to staging buffer at different offsets for (let i = 0; i < ops.length; i++) { commandEncoder.copyBufferToBuffer( this.outputBufs[i], 0, this.stagingBuf, offsets[i] * 4, ops[i].outputSize * 4, ); }
this.device.queue.submit([commandEncoder.finish()]);
// Single readback for all results const totalBytes = totalFloats * 4; await this.stagingBuf.mapAsync(GPUMapMode.READ, 0, totalBytes); const mapped = this.stagingBuf.getMappedRange(0, totalBytes); const allData = new Float32Array(totalFloats); allData.set(new Float32Array(mapped)); this.stagingBuf.unmap();
const results: Float32Array[] = []; for (let i = 0; i < ops.length; i++) { results.push(allData.slice(offsets[i], offsets[i] + ops[i].outputSize)); } return results; }
/** GPU Q4 matrix-vector multiply: output = weights[M, N] @ input[N]. */ async matvec(weightName: string, input: Float32Array, outputSize: number): Promise<Float32Array> { const [result] = await this.matvecBatch(input, [{ weightName, outputSize }]); return result; }
/** Helper to get a CPU weight by name, with prefix. */ private cpu(name: string): Float32Array { const w = this.cpuWeights.get(name); if (!w) throw new Error(`CPU weight not found: ${name}`); return w; }
/** Helper to get the full weight name for a layer tensor. */ private layerName(layer: number, suffix: string): string { return `model.language_model.layers.${layer}.${suffix}`; }
/** Run forward pass for a single token. Returns logits [VOCAB_SIZE]. */ async forward(tokenId: number, pos: number): Promise<Float32Array> { const prefix = 'model.language_model.';
// 1. Embedding lookup (CPU dequant of one row) const embed = this.getEmbedding(tokenId); this.hidden.set(embed);
// Track which full-attn and delta-net layer index we're on let fullAttnIdx = 0; let deltaIdx = 0;
// 2. Layer loop for (let layer = 0; layer < NUM_LAYERS; layer++) { const isFullAttn = FULL_ATTN_LAYERS.has(layer);
// Save residual this.residual.set(this.hidden);
// Pre-attention RMSNorm const lnW = this.cpu(this.layerName(layer, 'input_layernorm.weight')); rmsNorm(this.normed, this.hidden, lnW, RMS_NORM_EPS);
if (isFullAttn) { await this.fullAttentionForward(layer, fullAttnIdx, pos); fullAttnIdx++; } else { await this.deltaNetForward(layer, deltaIdx); deltaIdx++; }
// Residual add for (let i = 0; i < HIDDEN_SIZE; i++) this.hidden[i] += this.residual[i];
// Save residual for MLP this.residual.set(this.hidden);
// Post-attention RMSNorm const postLnW = this.cpu(this.layerName(layer, 'post_attention_layernorm.weight')); rmsNorm(this.normed, this.hidden, postLnW, RMS_NORM_EPS);
// MLP (SwiGLU) await this.mlpForward(layer);
// Residual add for (let i = 0; i < HIDDEN_SIZE; i++) this.hidden[i] += this.residual[i]; }
// 3. Final RMSNorm const finalNormW = this.cpu(prefix + 'norm.weight'); rmsNorm(this.normed, this.hidden, finalNormW, RMS_NORM_EPS);
// 4. LM head (tied weights — reuse embed_tokens) const logits = await this.matvec(prefix + 'embed_tokens.weight', this.normed, VOCAB_SIZE); return logits; }
/** Full attention forward for one layer. Writes result into this.hidden. */ private async fullAttentionForward(layer: number, kvIdx: number, pos: number): Promise<void> { const ln = (s: string) => this.layerName(layer, s);
// Batch Q/K/V projections (all from same normed input → single GPU submission) const [qBuf, kBuf, vBuf] = await this.matvecBatch(this.normed, [ { weightName: ln('self_attn.q_proj.weight'), outputSize: NUM_HEADS * HEAD_DIM * 2 }, { weightName: ln('self_attn.k_proj.weight'), outputSize: NUM_KV_HEADS * HEAD_DIM }, { weightName: ln('self_attn.v_proj.weight'), outputSize: NUM_KV_HEADS * HEAD_DIM }, ]);
// Split Q into query and gate const query = new Float32Array(ATTN_OUTPUT_DIM); const gate = new Float32Array(ATTN_OUTPUT_DIM); for (let h = 0; h < NUM_HEADS; h++) { const srcOff = h * HEAD_DIM * 2; const dstOff = h * HEAD_DIM; for (let d = 0; d < HEAD_DIM; d++) { query[dstOff + d] = qBuf[srcOff + d]; gate[dstOff + d] = qBuf[srcOff + HEAD_DIM + d]; } }
// QK-norm per head const qNormW = this.cpu(ln('self_attn.q_norm.weight')); const kNormW = this.cpu(ln('self_attn.k_norm.weight')); for (let h = 0; h < NUM_HEADS; h++) { rmsNorm(query, query, qNormW, RMS_NORM_EPS, h * HEAD_DIM, HEAD_DIM); } for (let h = 0; h < NUM_KV_HEADS; h++) { rmsNorm(kBuf, kBuf, kNormW, RMS_NORM_EPS, h * HEAD_DIM, HEAD_DIM); }
// Apply partial RoPE const cosOff = pos * ROTARY_DIM; for (let h = 0; h < NUM_HEADS; h++) { applyRope(query, h * HEAD_DIM, this.ropeCos, this.ropeSin, cosOff, ROTARY_DIM); } for (let h = 0; h < NUM_KV_HEADS; h++) { applyRope(kBuf, h * HEAD_DIM, this.ropeCos, this.ropeSin, cosOff, ROTARY_DIM); }
// Append K, V to cache const kvKeyDim = NUM_KV_HEADS * HEAD_DIM; const seqPos = this.kvLen[kvIdx]; this.kvKeys[kvIdx].set(kBuf, seqPos * kvKeyDim); this.kvValues[kvIdx].set(vBuf, seqPos * kvKeyDim); this.kvLen[kvIdx] = seqPos + 1; const seqLen = this.kvLen[kvIdx];
// Compute attention per query head (GQA) const attnOut = new Float32Array(ATTN_OUTPUT_DIM); const scale = 1.0 / Math.sqrt(HEAD_DIM); const scores = new Float32Array(seqLen);
for (let h = 0; h < NUM_HEADS; h++) { const kvH = Math.floor(h / KV_GROUPS); const qOff = h * HEAD_DIM;
// Attention scores let maxScore = -Infinity; for (let t = 0; t < seqLen; t++) { const kOff = t * kvKeyDim + kvH * HEAD_DIM; let dot = 0; for (let d = 0; d < HEAD_DIM; d++) { dot += query[qOff + d] * this.kvKeys[kvIdx][kOff + d]; } dot *= scale; scores[t] = dot; if (dot > maxScore) maxScore = dot; }
// Softmax let sumExp = 0; for (let t = 0; t < seqLen; t++) { scores[t] = Math.exp(scores[t] - maxScore); sumExp += scores[t]; } const invSum = sumExp > 0 ? 1.0 / sumExp : 0; for (let t = 0; t < seqLen; t++) scores[t] *= invSum;
// Weighted sum of values for (let d = 0; d < HEAD_DIM; d++) attnOut[qOff + d] = 0; for (let t = 0; t < seqLen; t++) { const vOff = t * kvKeyDim + kvH * HEAD_DIM; const w = scores[t]; for (let d = 0; d < HEAD_DIM; d++) { attnOut[qOff + d] += w * this.kvValues[kvIdx][vOff + d]; } } }
// Apply output gate: attnOut *= sigmoid(gate) for (let i = 0; i < ATTN_OUTPUT_DIM; i++) { attnOut[i] *= sigmoid(gate[i]); }
// Output projection this.hidden.set(await this.matvec(ln('self_attn.o_proj.weight'), attnOut, HIDDEN_SIZE)); }
/** DeltaNet forward for one layer. Writes result into this.hidden. */ private async deltaNetForward(layer: number, dnIdx: number): Promise<void> { const ln = (s: string) => this.layerName(layer, s);
// Batch all DeltaNet projections (all from same normed input → single GPU submission) const [qkv, z, aRaw, bRaw] = await this.matvecBatch(this.normed, [ { weightName: ln('linear_attn.in_proj_qkv.weight'), outputSize: DN_CONV_DIM }, { weightName: ln('linear_attn.in_proj_z.weight'), outputSize: DN_TOTAL_VALUE }, { weightName: ln('linear_attn.in_proj_a.weight'), outputSize: DN_NUM_HEADS }, { weightName: ln('linear_attn.in_proj_b.weight'), outputSize: DN_NUM_HEADS }, ]);
// Causal conv1d step const convW = this.cpu(ln('linear_attn.conv1d.weight')); const convState = this.convStates[dnIdx]; const stateCols = DN_CONV_KERNEL - 1;
for (let ch = 0; ch < DN_CONV_DIM; ch++) { const csBase = ch * stateCols;
// Compute conv output before updating state let convOut = 0; for (let k = 0; k < stateCols; k++) { convOut += convState[csBase + k] * convW[ch * DN_CONV_KERNEL + k]; } convOut += qkv[ch] * convW[ch * DN_CONV_KERNEL + stateCols];
// Shift state left, append new input for (let j = 0; j + 1 < stateCols; j++) { convState[csBase + j] = convState[csBase + j + 1]; } convState[csBase + stateCols - 1] = qkv[ch];
// SiLU activation qkv[ch] = silu(convOut); }
// Split QKV const qStart = 0; const kStart = DN_TOTAL_KEY; const vStart = DN_TOTAL_KEY * 2;
// Compute decay and beta const aLog = this.cpu(ln('linear_attn.A_log')); const dtBias = this.cpu(ln('linear_attn.dt_bias')); const gVals = new Float32Array(DN_NUM_HEADS); const betaVals = new Float32Array(DN_NUM_HEADS); for (let h = 0; h < DN_NUM_HEADS; h++) { gVals[h] = -Math.exp(aLog[h]) * softplus(aRaw[h] + dtBias[h]); betaVals[h] = sigmoid(bRaw[h]); }
// DeltaNet recurrence per head const stateMatrix = this.deltaStates[dnIdx]; const dnOut = new Float32Array(DN_TOTAL_VALUE); const scaleQ = 1.0 / Math.sqrt(DN_KEY_DIM); const qHead = new Float32Array(DN_KEY_DIM); const kHead = new Float32Array(DN_KEY_DIM);
for (let h = 0; h < DN_NUM_HEADS; h++) { // Extract Q and K, L2-normalize for (let d = 0; d < DN_KEY_DIM; d++) { qHead[d] = qkv[qStart + h * DN_KEY_DIM + d]; kHead[d] = qkv[kStart + h * DN_KEY_DIM + d]; } l2Norm(qHead, 0, DN_KEY_DIM); l2Norm(kHead, 0, DN_KEY_DIM);
// Scale query for (let d = 0; d < DN_KEY_DIM; d++) qHead[d] *= scaleQ;
const smBase = h * DN_KEY_DIM * DN_VALUE_DIM;
// Decay state const decay = Math.exp(gVals[h]); for (let i = 0; i < DN_KEY_DIM * DN_VALUE_DIM; i++) { stateMatrix[smBase + i] *= decay; }
// kv_mem[v] = sum_d(S[d,v] * k[d]) const kvMem = new Float32Array(DN_VALUE_DIM); for (let v = 0; v < DN_VALUE_DIM; v++) { let sum = 0; for (let d = 0; d < DN_KEY_DIM; d++) { sum += stateMatrix[smBase + d * DN_VALUE_DIM + v] * kHead[d]; } kvMem[v] = sum; }
// delta = (v - kv_mem) * beta const beta = betaVals[h]; for (let v = 0; v < DN_VALUE_DIM; v++) { const delta = (qkv[vStart + h * DN_VALUE_DIM + v] - kvMem[v]) * beta; // Update state: S[d,v] += k[d] * delta for (let d = 0; d < DN_KEY_DIM; d++) { stateMatrix[smBase + d * DN_VALUE_DIM + v] += kHead[d] * delta; } }
// Output: out[v] = sum_d(S[d,v] * q[d]) const outOff = h * DN_VALUE_DIM; for (let v = 0; v < DN_VALUE_DIM; v++) { let sum = 0; for (let d = 0; d < DN_KEY_DIM; d++) { sum += stateMatrix[smBase + d * DN_VALUE_DIM + v] * qHead[d]; } dnOut[outOff + v] = sum; } }
// Gated RMS norm per head const normW = this.cpu(ln('linear_attn.norm.weight')); for (let h = 0; h < DN_NUM_HEADS; h++) { const off = h * DN_VALUE_DIM; rmsNormGated(dnOut, off, dnOut, off, z, off, normW, RMS_NORM_EPS, DN_VALUE_DIM); }
// Output projection this.hidden.set(await this.matvec(ln('linear_attn.out_proj.weight'), dnOut, HIDDEN_SIZE)); }
/** SwiGLU MLP forward. Writes result into this.hidden. */ private async mlpForward(layer: number): Promise<void> { const ln = (s: string) => this.layerName(layer, s);
// Batch gate + up projections (both from same normed input → single GPU submission) const [gateOut, upOut] = await this.matvecBatch(this.normed, [ { weightName: ln('mlp.gate_proj.weight'), outputSize: INTERMEDIATE_SIZE }, { weightName: ln('mlp.up_proj.weight'), outputSize: INTERMEDIATE_SIZE }, ]);
// SwiGLU: silu(gate) * up for (let i = 0; i < INTERMEDIATE_SIZE; i++) { gateOut[i] = silu(gateOut[i]) * upOut[i]; }
this.hidden.set(await this.matvec(ln('mlp.down_proj.weight'), gateOut, HIDDEN_SIZE)); }
/** Reset all caches for a new generation. */ reset(): void { for (let i = 0; i < this.kvKeys.length; i++) { this.kvKeys[i].fill(0); this.kvValues[i].fill(0); this.kvLen[i] = 0; } for (let i = 0; i < this.deltaStates.length; i++) { this.deltaStates[i].fill(0); this.convStates[i].fill(0); } }}