@tryhamster/gerbil 1.4.1 → 1.6.0

package/dist/{moonshine-stt-BGmctcdO.mjs → moonshine-stt-bExSnd2o.mjs}

@@ -27,7 +27,7 @@ async function initGPU(options) {
 	if (!adapter) throw new Error("No WebGPU adapter found. Your GPU may not support WebGPU, or another tab may be using it. Try closing other tabs.");
 	const hasF16 = adapter.features.has("shader-f16");
 	const hasSubgroups = adapter.features.has("subgroups");
-	const hasTimestamp = typeof process !== "undefined" && process.env?.GERBIL_PROFILE != null && adapter.features.has("timestamp-query");
+	const hasTimestamp = typeof process !== "undefined" && (process.env?.GERBIL_PROFILE != null || process.env?.GERBIL_TS_TOTAL != null) && adapter.features.has("timestamp-query");
 	const requiredFeatures = [];
 	if (hasF16) requiredFeatures.push("shader-f16");
 	if (hasSubgroups) requiredFeatures.push("subgroups");
@@ -3864,8 +3864,10 @@ fn main(
 `;
 const WGSL_SWIGLU_MATVEC_INT4 = `\
 // CONSTRAINT: N_TILE must equal workgroup_size / K_THREADS
-const N_TILE: u32 = 8u;
-const K_THREADS: u32 = 32u;
+// 16/16 matches the plain-matvec sweep winner (8/32 measured slower at every
+// production shape; this fused kernel had never been re-swept).
+const N_TILE: u32 = 16u;
+const K_THREADS: u32 = 16u;
 
 struct Params {
   K: u32,
@@ -3969,8 +3971,8 @@ fn main(
 }
 `;
 const WGSL_DUAL_MATVEC_INT4 = `\
-const N_TILE: u32 = 8u;
-const K_THREADS: u32 = 32u;
+const N_TILE: u32 = 16u;
+const K_THREADS: u32 = 16u;
 
 struct Params {
   K: u32,
@@ -4470,20 +4472,45 @@ struct Params {
 @group(0) @binding(6) var<storage, read_write> output: array<f32>;
 @group(0) @binding(7) var<storage, read> params: Params;
 
-// Shared memory for Q and K vectors (L2-normalized) and reduction
+// OCCUPANCY FAN-OUT: one workgroup per (head, d-chunk) instead of per head.
+// Each head's state column d is touched ONLY by the workgroup owning d's chunk
+// (kv_dot for lane d reads state[:, d]; the rank-1 update writes state[:, d]),
+// so chunks are fully independent across timesteps — no cross-workgroup sync.
+// 16 heads → 128 workgroups (vs 16), each 128 threads = D_LANES(16) × K_SPLITS(8).
+// Requires val_dim % (D_CHUNKS*... ) handled via bounds guards; key_dim <= 128.
+
+const D_CHUNKS: u32 = 8u;   // workgroups per head (must match JS dispatch)
+const D_LANES: u32 = 16u;   // d lanes per workgroup
+const K_SPLITS: u32 = 8u;   // threads cooperating on each d lane's k loop
+
+// Shared memory for Q and K vectors (L2-normalized), reduction, partials
 var<workgroup> shared_q: array<f32, 128>;
 var<workgroup> shared_k: array<f32, 128>;
 var<workgroup> shared_norm: array<f32, 128>;
+var<workgroup> shared_partial: array<f32, 128>;  // [D_LANES][K_SPLITS]
+var<workgroup> shared_kvmem: array<f32, D_LANES>;
 
 @compute @workgroup_size(128)
 fn main(
   @builtin(local_invocation_id) lid: vec3u,
   @builtin(workgroup_id) wid: vec3u,
 ) {
-  let h = wid.x;  // uniform within workgroup
+  let h = wid.x / D_CHUNKS;
+  let chunk = wid.x % D_CHUNKS;
   if (h >= params.num_heads) { return; }
 
-  let d = lid.x;  // 0..127, all threads participate in barriers
+  let tid = lid.x;             // 0..127
+  let d_local = tid % D_LANES; // lane within chunk
+  let k_split = tid / D_LANES; // 0..K_SPLITS-1
+  let d_per_chunk = (params.val_dim + D_CHUNKS - 1u) / D_CHUNKS;
+  // This lane's output dim (may exceed val_dim for ragged sizes — guarded below).
+  // d_per_chunk > D_LANES is unsupported (val_dim > 128 with D_CHUNKS=8).
+  let d = chunk * d_per_chunk + d_local;
+  let d_ok = d < params.val_dim && d_local < d_per_chunk;
+
+  let k_slice = (params.key_dim + K_SPLITS - 1u) / K_SPLITS;
+  let k_begin = k_split * k_slice;
+  let k_end = min(k_begin + k_slice, params.key_dim);
 
   let A_val = -exp(A_log[h]);
   let state_base = h * params.key_dim * params.val_dim;
@@ -4497,37 +4524,48 @@ fn main(
   for (var t: u32 = 0u; t < params.T; t = t + 1u) {
     let qkv_base = t * params.qkv_dim;
 
-    // ── Load Q and K into shared memory ──
-    shared_q[d] = qkv_conv[qkv_base + q_off + d];
-    shared_k[d] = qkv_conv[qkv_base + k_off + d];
+    // ── Load Q and K into shared memory (lane tid covers key_dim <= 128) ──
+    if (tid < params.key_dim) {
+      shared_q[tid] = qkv_conv[qkv_base + q_off + tid];
+      shared_k[tid] = qkv_conv[qkv_base + k_off + tid];
+      shared_norm[tid] = shared_q[tid] * shared_q[tid];
+    } else {
+      shared_norm[tid] = 0.0;
+    }
     workgroupBarrier();
 
     // ── L2-normalize Q (parallel reduction) ──
-    shared_norm[d] = shared_q[d] * shared_q[d];
-    workgroupBarrier();
     var stride: u32 = 64u;
     while (stride > 0u) {
-      if (d < stride) { shared_norm[d] += shared_norm[d + stride]; }
+      if (tid < stride) { shared_norm[tid] += shared_norm[tid + stride]; }
       workgroupBarrier();
       stride /= 2u;
     }
     let q_inv_norm = 1.0 / sqrt(shared_norm[0] + 1e-12);
-    shared_q[d] *= q_inv_norm;
+    workgroupBarrier();
 
     // ── L2-normalize K (parallel reduction) ──
-    shared_norm[d] = shared_k[d] * shared_k[d];
+    if (tid < params.key_dim) {
+      shared_q[tid] *= q_inv_norm;
+      shared_norm[tid] = shared_k[tid] * shared_k[tid];
+    } else {
+      shared_norm[tid] = 0.0;
+    }
     workgroupBarrier();
     stride = 64u;
     while (stride > 0u) {
-      if (d < stride) { shared_norm[d] += shared_norm[d + stride]; }
+      if (tid < stride) { shared_norm[tid] += shared_norm[tid + stride]; }
       workgroupBarrier();
       stride /= 2u;
     }
     let k_inv_norm = 1.0 / sqrt(shared_norm[0] + 1e-12);
-    shared_k[d] *= k_inv_norm;
+    workgroupBarrier();
+    if (tid < params.key_dim) {
+      shared_k[tid] *= k_inv_norm;
+    }
     workgroupBarrier();
 
-    // ── Compute decay and beta ──
+    // ── Compute decay and beta (redundantly per thread — trivial scalars) ──
     let a_val = a_input[t * params.num_heads + h];
     let b_val = b_input[t * params.num_heads + h];
     let dt_in = a_val + dt_bias[h];
@@ -4536,32 +4574,52 @@ fn main(
     let decay = exp(max(g, -80.0));
     let beta = 1.0 / (1.0 + exp(max(-b_val, -80.0)));
 
-    // ── 1+2. Retrieve from memory (decay is a per-step scalar, so it factors
-    // out of the dot product): kv_mem[d] = decay * Σ_k state[k,d] * k_norm[k] ──
-    let v_val = qkv_conv[qkv_base + v_off + d];
+    // ── 1+2. kv_mem[d] = decay * Σ_k state[k,d] * k_norm[k], k split 8 ways ──
     var kv_dot: f32 = 0.0;
-    for (var k: u32 = 0u; k < params.key_dim; k = k + 1u) {
-      kv_dot += f32(ssm_state[state_base + k * params.val_dim + d]) * shared_k[k];
+    if (d_ok) {
+      for (var k: u32 = k_begin; k < k_end; k = k + 1u) {
+        kv_dot += f32(ssm_state[state_base + k * params.val_dim + d]) * shared_k[k];
+      }
     }
-    let kv_mem = kv_dot * decay;
-
-    // ── 3. Delta rule: delta[d] = (v[d] - kv_mem[d]) * beta ──
-    let delta_val = (v_val - kv_mem) * beta;
+    shared_partial[d_local * K_SPLITS + k_split] = kv_dot;
+    workgroupBarrier();
+    if (k_split == 0u) {
+      var acc: f32 = 0.0;
+      for (var j: u32 = 0u; j < K_SPLITS; j = j + 1u) {
+        acc += shared_partial[d_local * K_SPLITS + j];
+      }
+      shared_kvmem[d_local] = acc;
+    }
+    workgroupBarrier();
 
-    // ── 4+5. Fused decay + rank-1 update + output in ONE pass over state
-    // (one read + one write instead of three reads + two writes):
-    //   state'[k,d] = state[k,d] * decay + k_norm[k] * delta[d]
-    //   out[d]      = Σ_k state'[k,d] * q_norm[k] * scale ──
-    var out_val: f32 = 0.0;
-    for (var k: u32 = 0u; k < params.key_dim; k = k + 1u) {
-      let idx = state_base + k * params.val_dim + d;
-      let s = f32(ssm_state[idx]) * decay + shared_k[k] * delta_val;
-      ssm_state[idx] = f16(s);
-      out_val += s * shared_q[k];
+    // ── 3. Delta rule (each of the d's 8 threads recomputes — cheap) ──
+    var delta_val: f32 = 0.0;
+    if (d_ok) {
+      let v_val = qkv_conv[qkv_base + v_off + d];
+      let kv_mem = shared_kvmem[d_local] * decay;
+      delta_val = (v_val - kv_mem) * beta;
     }
-    output[t * out_stride + h * params.val_dim + d] = out_val * scale;
 
-    workgroupBarrier();  // sync before next timestep loads shared_q/k
+    // ── 4+5. Fused decay + rank-1 update + output, k split 8 ways ──
+    var out_dot: f32 = 0.0;
+    if (d_ok) {
+      for (var k: u32 = k_begin; k < k_end; k = k + 1u) {
+        let idx = state_base + k * params.val_dim + d;
+        let s = f32(ssm_state[idx]) * decay + shared_k[k] * delta_val;
+        ssm_state[idx] = f16(s);
+        out_dot += s * shared_q[k];
+      }
+    }
+    shared_partial[d_local * K_SPLITS + k_split] = out_dot;
+    workgroupBarrier();
+    if (k_split == 0u && d_ok) {
+      var acc: f32 = 0.0;
+      for (var j: u32 = 0u; j < K_SPLITS; j = j + 1u) {
+        acc += shared_partial[d_local * K_SPLITS + j];
+      }
+      output[t * out_stride + h * params.val_dim + d] = acc * scale;
+    }
+    workgroupBarrier();  // sync before next timestep reuses shared arrays
   }
 }
 `;
@@ -6165,7 +6223,7 @@ const SWIGLU_MATVEC_INT4_SPEC = {
 	getDispatchSize(op) {
 		const N = op.attributes.N;
 		return [
-			cdiv(N, 8),
+			cdiv(N, 16),
 			1,
 			1
 		];
@@ -6200,7 +6258,7 @@ const DUAL_MATVEC_INT4_SPEC = {
 	getDispatchSize(op) {
 		const N = op.attributes.N;
 		return [
-			cdiv(N, 8),
+			cdiv(N, 16),
 			1,
 			1
 		];
@@ -7008,7 +7066,14 @@ const mambaSSMSpec = {
 };
 const MAMBA_SSM_F16_SPEC = {
 	...mambaSSMSpec,
-	shaderCode: WGSL_MAMBA_SSM_F16
+	shaderCode: WGSL_MAMBA_SSM_F16,
+	getDispatchSize(op) {
+		return [
+			op.attributes.num_heads * 8,
+			1,
+			1
+		];
+	}
 };
 const kvCacheAppendSpec = {
 	shaderCode: WGSL_KV_CACHE_APPEND,
@@ -8908,7 +8973,39 @@ var Executor = class Executor {
 		}
 		const encoder = this.ctx.device.createCommandEncoder({ label: "decode_step" });
 		if (tokenId === null) encoder.copyBufferToBuffer(this.argmaxResultBuffer, 0, this.inputIdsBuffer, 0, 4);
-		const pass = encoder.beginComputePass({ label: "decode_pass" });
+		const g = globalThis;
+		if (process.env?.GERBIL_TS_TOTAL != null) {
+			g.__tsTotal ??= {
+				ns: 0,
+				steps: 0,
+				seq: 0,
+				inFlight: false
+			};
+			g.__tsTotal.seq++;
+		}
+		const tsTotal = this.ctx.hasTimestamp && typeof process !== "undefined" && process.env?.GERBIL_TS_TOTAL != null && g.__tsTotal.seq % 8 === 0 && !g.__tsTotal.inFlight;
+		if (tsTotal && !this.querySet) {
+			this.querySet = this.ctx.device.createQuerySet({
+				type: "timestamp",
+				count: 2
+			});
+			this.queryResolveBuf = this.ctx.device.createBuffer({
+				size: 16,
+				usage: 516
+			});
+			this.queryReadbackBuf = this.ctx.device.createBuffer({
+				size: 16,
+				usage: 9
+			});
+		}
+		const pass = encoder.beginComputePass(tsTotal ? {
+			label: "decode_pass",
+			timestampWrites: {
+				querySet: this.querySet,
+				beginningOfPassWriteIndex: 0,
+				endOfPassWriteIndex: 1
+			}
+		} : { label: "decode_pass" });
 		for (let i = 0; i < this.decodeEntries.length; i++) {
 			const entry = this.decodeEntries[i];
 			pass.setPipeline(entry.pipeline);
@@ -8922,7 +9019,27 @@ var Executor = class Executor {
 		}
 		pass.end();
 		encoder.copyBufferToBuffer(this.argmaxResultBuffer, 0, this.decodeReadbacks[slot], 0, 4);
+		if (tsTotal) {
+			encoder.resolveQuerySet(this.querySet, 0, 2, this.queryResolveBuf, 0);
+			encoder.copyBufferToBuffer(this.queryResolveBuf, 0, this.queryReadbackBuf, 0, 16);
+		}
 		this.ctx.device.queue.submit([encoder.finish()]);
+		if (tsTotal) {
+			const rb = this.queryReadbackBuf;
+			const acc = g.__tsTotal;
+			acc.inFlight = true;
+			(async () => {
+				try {
+					await rb.mapAsync(MAP_MODE_READ$1, 0, 16);
+					const ts = new BigUint64Array(rb.getMappedRange(0, 16).slice(0));
+					rb.unmap();
+					acc.ns += Number(ts[1] - ts[0]);
+					acc.steps += 1;
+				} catch {} finally {
+					acc.inFlight = false;
+				}
+			})();
+		}
 		this.seqPos += 1;
 	}
 	/** Read back the token produced by the pipelined step that used `slot`. */
@@ -10013,6 +10130,47 @@ function repackMLX(mlx) {
 		zeros
 	};
 }
+/**
+* Dequantize an MLX affine-quantized tensor of arbitrary bit width to f32.
+*
+* MLX packs the unsigned integer codes as a CONTIGUOUS little-endian bitstream
+* (verified against ml-explore/mlx `extract_bits`): code `i` occupies bits
+* `[i*bits, i*bits+bits)` of the per-row byte stream, straddling byte/word
+* boundaries when `bits` does not divide 8 evenly. Power-of-two bit widths
+* (4/8) divide evenly so codes never straddle; 6-bit packs 4 codes per 3 bytes.
+*
+* Dequant (affine mode): `w = scale * q + bias`, with one (scale, bias) pair per
+* group of `groupSize` codes along K (the last axis), groups laid out row-major.
+*
+* The weight buffer is provided as a Uint32Array (how the store holds it); we
+* read it as a flat little-endian byte view so the bitstream is contiguous
+* across the whole row regardless of the 4-byte word boundaries.
+*
+* Returns a flat Float32Array of shape [N, K] in row-major order — directly
+* consumable by `quantizeInt4` for repacking into Gerbil's 4-bit kernels.
+*/
+function dequantizeMLX(mlx) {
+	const { weight, scales: mlxScales, biases: mlxBiases, N, K, groupSize, bits } = mlx;
+	const groupsPerRow = Math.ceil(K / groupSize);
+	const mask = (1 << bits) - 1;
+	const bytes = new Uint8Array(weight.buffer, weight.byteOffset, weight.byteLength);
+	const bytesPerRow = Math.ceil(K * bits / 8);
+	const out = new Float32Array(N * K);
+	for (let row = 0; row < N; row++) {
+		const rowByteBase = row * bytesPerRow;
+		const rowGroupBase = row * groupsPerRow;
+		const rowOutBase = row * K;
+		for (let col = 0; col < K; col++) {
+			const bitPos = col * bits;
+			const byteIdx = rowByteBase + (bitPos >> 3);
+			const bitOff = bitPos & 7;
+			const q = (bytes[byteIdx] | bytes[byteIdx + 1] << 8 | bytes[byteIdx + 2] << 16) >>> bitOff & mask;
+			const g = rowGroupBase + Math.floor(col / groupSize);
+			out[rowOutBase + col] = mlxScales[g] * q + mlxBiases[g];
+		}
+	}
+	return out;
+}
 
 //#endregion
 //#region src/gpu/opfs-cache.ts
@@ -10479,6 +10637,19 @@ const UNICODE_TO_BYTE = /* @__PURE__ */ new Map();
 }
 /** SentencePiece metaspace marker (U+2581, "▁") used by T5/Unigram tokenizers. */
 const METASPACE = "▁";
+/**
+* Matches the Qwen3.5-style empty-think prefill that a chat template emits after
+* `add_generation_prompt` (e.g. `{{- '<think>\n\n</think>\n\n' }}`). Used to
+* decide whether to inject the empty think block — LFM2.5 templates reference
+* `<think>` only to re-render prior assistant `.thinking` content and must NOT
+* be matched, so we look for the empty `<think></think>` pair specifically.
+*
+* The gap between the tags must tolerate BOTH real whitespace and the LITERAL
+* backslash-escape text (`\n`, `\r`, `\t`) that survives in the raw Jinja string
+* source — Qwen3.5 ships the prefill as the two-character `\n` sequences, not
+* actual newlines, so a plain `\s*` would miss it and drop Qwen's think block.
+*/
+const EMPTY_THINK_PREFILL = /<think>(?:\s|\\[nrt])*<\/think>/;
 var Tokenizer = class Tokenizer {
 	vocab;
 	vocabReverse;
@@ -10742,7 +10913,7 @@ var Tokenizer = class Tokenizer {
 		for (const msg of messages) parts.push(`<|im_start|>${msg.role}\n${msg.content}<|im_end|>\n`);
 		if (addGenPrompt) {
 			parts.push(`<|im_start|>assistant\n`);
-			if (template ? template.includes("<think>") : this.addedTokens.has("<think>")) parts.push(`<think>\n\n</think>\n\n`);
+			if (template ? EMPTY_THINK_PREFILL.test(template) : this.addedTokens.has("<think>")) parts.push(`<think>\n\n</think>\n\n`);
 		}
 		return parts.join("");
 	}
@@ -12322,6 +12493,12 @@ async function loadModel(options) {
 	if (isMLX) {
 		onProgress?.(96, 100, "Repacking MLX weights...");
 		const mlxGS = quantConfig?.group_size ?? 64;
+		const mlxDefaultBits = quantConfig?.bits ?? 4;
+		const bitsForBase = (baseName) => {
+			if (!quantConfig) return mlxDefaultBits;
+			for (const [key, val] of Object.entries(quantConfig)) if (val && typeof val === "object" && typeof val.bits === "number" && (key === baseName || key.endsWith(`.${baseName}`))) return val.bits;
+			return mlxDefaultBits;
+		};
 		for (const node of graph.nodes) {
 			if (node.opType !== "MatMulInt4" && node.opType !== "EmbeddingInt4") continue;
 			let N, K;
@@ -12341,26 +12518,54 @@ async function loadModel(options) {
 			const mlxScales = await store.get(`${baseName}.scales`);
 			const mlxBiases = await store.get(`${baseName}.biases`);
 			if (!mlxWeight || !mlxScales || !mlxBiases) continue;
-			const repacked = repackMLX({
-				weight: mlxWeight.data instanceof Uint32Array ? mlxWeight.data : new Uint32Array(mlxWeight.data.buffer, mlxWeight.data.byteOffset, mlxWeight.data.byteLength / 4),
-				scales: mlxScales.data instanceof Float32Array ? mlxScales.data : new Float32Array(mlxScales.data.buffer, mlxScales.data.byteOffset, mlxScales.data.byteLength / 4),
-				biases: mlxBiases.data instanceof Float32Array ? mlxBiases.data : new Float32Array(mlxBiases.data.buffer, mlxBiases.data.byteOffset, mlxBiases.data.byteLength / 4),
-				N,
-				K,
-				groupSize: mlxGS
-			});
-			await store.set(qTensorName, {
-				data: repacked.packed,
-				shape: [repacked.packed.length]
-			});
-			await store.set(scalesTensorName, {
-				data: repacked.scales,
-				shape: [repacked.scales.length]
-			});
-			await store.set(zerosTensorName, {
-				data: repacked.zeros,
-				shape: [repacked.zeros.length]
-			});
+			const weightData = mlxWeight.data instanceof Uint32Array ? mlxWeight.data : new Uint32Array(mlxWeight.data.buffer, mlxWeight.data.byteOffset, mlxWeight.data.byteLength / 4);
+			const scalesData = mlxScales.data instanceof Float32Array ? mlxScales.data : new Float32Array(mlxScales.data.buffer, mlxScales.data.byteOffset, mlxScales.data.byteLength / 4);
+			const biasesData = mlxBiases.data instanceof Float32Array ? mlxBiases.data : new Float32Array(mlxBiases.data.buffer, mlxBiases.data.byteOffset, mlxBiases.data.byteLength / 4);
+			const srcBits = bitsForBase(baseName);
+			if (srcBits === 4) {
+				const repacked = repackMLX({
+					weight: weightData,
+					scales: scalesData,
+					biases: biasesData,
+					N,
+					K,
+					groupSize: mlxGS
+				});
+				await store.set(qTensorName, {
+					data: repacked.packed,
+					shape: [repacked.packed.length]
+				});
+				await store.set(scalesTensorName, {
+					data: repacked.scales,
+					shape: [repacked.scales.length]
+				});
+				await store.set(zerosTensorName, {
+					data: repacked.zeros,
+					shape: [repacked.zeros.length]
+				});
+			} else {
+				const { packed, scales, zeros } = quantizeInt4(dequantizeMLX({
+					weight: weightData,
+					scales: scalesData,
+					biases: biasesData,
+					N,
+					K,
+					groupSize: mlxGS,
+					bits: srcBits
+				}), mlxGS);
+				await store.set(qTensorName, {
+					data: packed,
+					shape: [packed.length]
+				});
+				await store.set(scalesTensorName, {
+					data: scales,
+					shape: [scales.length]
+				});
+				await store.set(zerosTensorName, {
+					data: zeros,
+					shape: [zeros.length]
+				});
+			}
 			await store.delete(sourceName);
 			await store.delete(`${baseName}.scales`);
 			await store.delete(`${baseName}.biases`);
@@ -13220,4 +13425,4 @@ function selectGraphWeights(graph, weights) {
 
 //#endregion
 export { verifyGPU as S, createStorageBuffer as _, loadMoonshine as a, getOrCreatePipeline as b, quantizeBackboneInt4 as c, Tokenizer as d, Executor as f, createBindGroup as g, clearPipelineCache as h, loadModel as i, quantizeKaniBackbone as l, MATMUL_BIAS_F16C_SPEC as m, MoonshineEncoderExecutor as n, loadOuteTTS as o, KERNEL_REGISTRY as p, loadKaniTTS as r, loadParlerTTS as s, MoonshineSTT as t, remapPrunedToken as u, createUniformBuffer as v, initGPU as x, destroyBuffers as y };
-//# sourceMappingURL=moonshine-stt-BGmctcdO.mjs.map
+//# sourceMappingURL=moonshine-stt-bExSnd2o.mjs.map