@tryhamster/gerbil 1.1.2 → 1.2.0

package/README.md

@@ -123,12 +123,12 @@ For a copy-pasteable version see [`examples/benchmark.ts`](./examples/benchmark.
 
 ## React Quickstart
 
-`useEngine` (from `@tryhamster/gerbil/gpu/hooks`) owns the full engine lifecycle —
+`useEngine` (from `@tryhamster/gerbil/hooks`) owns the full engine lifecycle —
 load, unload, hot-swap on config change, and reference-counted sharing so multiple
 components never upload the same weights to the GPU twice.
 
 ```tsx
-import { useEngine } from "@tryhamster/gerbil/gpu/hooks";
+import { useEngine } from "@tryhamster/gerbil/hooks";
 
 function Chat() {
   const { complete, completion, isLoading, isGenerating, tps } = useEngine({
@@ -156,7 +156,7 @@ and `dispose`. Pass `enableVision: true` or `embedding: true` to load those moda
 until the engine is ready — and keep it warm across navigation (no GPU re-upload):
 
 ```tsx
-import { GerbilGate } from "@tryhamster/gerbil/gpu/hooks";
+import { GerbilGate } from "@tryhamster/gerbil/hooks";
 
 <GerbilGate
   model="mlx-community/Qwen3.5-0.8B-4bit"
@@ -199,10 +199,10 @@ const engine = await WebGPUEngine.create({ repo: "mlx-community/Qwen3.5-0.8B-4bi
 await engine.generateObject("…", { schema: { required: ["title"] } });
 ```
 
-In React, use `useObject` (from `@tryhamster/gerbil/gpu/hooks`):
+In React, use `useObject` (from `@tryhamster/gerbil/hooks`):
 
 ```tsx
-import { useObject } from "@tryhamster/gerbil/gpu/hooks";
+import { useObject } from "@tryhamster/gerbil/hooks";
 
 const { generate, object, isGenerating } = useObject<{ city: string }>();
 await generate("Extract the city from: I live in Paris", {
@@ -322,17 +322,19 @@ const skill = useSkill("my-skill");
 
 ## Tools & Agents
 
-Gerbil supports tool calling with Qwen3 models for agentic workflows:
+Gerbil supports tool calling with Qwen3 models for agentic workflows. A tool is a
+plain object (`AgentTool`) — give it a `name`, `description`, optional `parameters`,
+and an `execute` function:
 
 ```typescript
-import { defineTool, executeToolCall } from "@tryhamster/gerbil";
+import type { AgentTool } from "@tryhamster/gerbil/gpu";
 
-const weatherTool = defineTool({
+const weatherTool: AgentTool = {
   name: "get_weather",
   description: "Get weather for a city",
-  parameters: z.object({ city: z.string() }),
-  execute: async ({ city }) => `Weather in ${city}: 72°F, sunny`,
-});
+  parameters: { city: "string" },
+  execute: ({ city }) => `Weather in ${city}: 72°F, sunny`,
+};
 ```
 
 **Agentic loop, on-device.** `engine.generateWithTools` (and the `useAgent` React hook)
@@ -340,7 +342,7 @@ run the whole loop — generate → call a tool → feed the result back → rep
 step trace for UIs:
 
 ```tsx
-import { useAgent } from "@tryhamster/gerbil/gpu/hooks";
+import { useAgent } from "@tryhamster/gerbil/hooks";
 
 const { run, steps, answer, isRunning } = useAgent({
   model: "mlx-community/Qwen3.5-0.8B-4bit",
@@ -377,7 +379,7 @@ Inline autocomplete — `engine.autocomplete(prefix)` and the debounced `useAuto
 hook return a brief single-line continuation (low-latency defaults + cleanup):
 
 ```tsx
-import { useAutocomplete } from "@tryhamster/gerbil/gpu/hooks";
+import { useAutocomplete } from "@tryhamster/gerbil/hooks";
 
 const { suggestion, onInput, accept, dismiss } = useAutocomplete({
   model: "mlx-community/Qwen3.5-0.8B-4bit",
@@ -418,10 +420,10 @@ gerbil update                                 # Update to latest version
 ## Browser Usage
 
 Run LLMs directly in the browser with WebGPU — no server required. The React hooks
-live at `@tryhamster/gerbil/gpu/hooks` and run pure WebGPU compute:
+live at `@tryhamster/gerbil/hooks` and run pure WebGPU compute:
 
 ```tsx
-import { useChat } from "@tryhamster/gerbil/gpu/hooks";
+import { useChat } from "@tryhamster/gerbil/hooks";
 
 function Chat() {
   const { messages, send, isLoading, isGenerating } = useChat();
@@ -454,7 +456,7 @@ function Chat() {
 | **LangChain** | `@tryhamster/gerbil/langchain` | [📖 LangChain](https://gerbilsdk.com/docs/frameworks/langchain) |
 | **MCP Server** | `npx @tryhamster/gerbil serve --mcp` | [📖 MCP](https://gerbilsdk.com/docs/mcp) |
 
-**Native engine:** `import { WebGPUEngine } from "@tryhamster/gerbil/gpu"` (or `useEngine` from `@tryhamster/gerbil/gpu/hooks` for React) is the primary surface for text, vision, embeddings, and speech.
+**Native engine:** `import { WebGPUEngine } from "@tryhamster/gerbil/gpu"` (or `useEngine` from `@tryhamster/gerbil/hooks` for React) is the primary surface for text, vision, embeddings, and speech.
 
 ## Supported Models
 

package/dist/{architectures-wan0swXu.mjs → architectures-Ol7yUmKf.mjs}

@@ -148,6 +148,66 @@ function createDefaultHFKeyMapper() {
 	};
 }
 
+//#endregion
+//#region src/gpu/quantize.ts
+/**
+* On-the-fly INT4 quantization for WebGPU inference.
+*
+* Converts F32 weight matrices to 4-bit unsigned integers with per-group
+* scale/zero-point parameters. Matches the dequantization formula in the
+* MatMulInt4 WGSL kernel:
+*
+*   dequant_val = (raw_nibble - zero) * scale
+*
+* Packing: 8 × 4-bit values per u32, little-nibble-first order.
+* Grouping: contiguous elements in row-major (K × N) layout.
+*/
+const DEFAULT_GROUP_SIZE = 128;
+/**
+* Quantize a flat F32 weight array to INT4.
+*
+* @param data   - Source F32 weight data (row-major, K × N elements)
+* @param groupSize - Number of elements per quantization group (default 128)
+* @returns Packed INT4 data with per-group scales and zeros
+*/
+function quantizeInt4(data, groupSize = DEFAULT_GROUP_SIZE) {
+	const totalElements = data.length;
+	const numGroups = Math.ceil(totalElements / groupSize);
+	const packedSize = Math.ceil(totalElements / 8);
+	const packed = new Uint32Array(packedSize);
+	const scales = new Float32Array(numGroups);
+	const zeros = new Float32Array(numGroups);
+	for (let g = 0; g < numGroups; g++) {
+		const start = g * groupSize;
+		const end = Math.min(start + groupSize, totalElements);
+		let min = data[start];
+		let max = data[start];
+		for (let i = start + 1; i < end; i++) {
+			const v = data[i];
+			if (v < min) min = v;
+			if (v > max) max = v;
+		}
+		const range = max - min;
+		const scale = range > 0 ? range / 15 : 1;
+		const zero = range > 0 ? -min / scale : 0;
+		scales[g] = scale;
+		zeros[g] = zero;
+		const invRange = range > 0 ? 15 / range : 0;
+		for (let i = start; i < end; i++) {
+			const q = Math.round((data[i] - min) * invRange);
+			const clamped = q < 0 ? 0 : q > 15 ? 15 : q;
+			const packedIdx = i >>> 3;
+			const nibblePos = i & 7;
+			packed[packedIdx] |= clamped << nibblePos * 4;
+		}
+	}
+	return {
+		packed,
+		scales,
+		zeros
+	};
+}
+
 //#endregion
 //#region src/gpu/architectures/kani_tts.ts
 const KANI_TOKENISER_LENGTH = 64400;
@@ -195,19 +255,23 @@ const NANOCODEC_BASE_CHANNELS = 864;
 const NANOCODEC_INPUT_DIM = 16;
 function parseKaniConfig(rawConfig) {
 	const textVocabSize = rawConfig.text_vocab_size ?? KANI_TOKENISER_LENGTH;
+	const useLearnableRope = rawConfig.use_learnable_rope ?? false;
+	const archList = rawConfig.architectures ?? [];
+	const isKaniTTS2 = useLearnableRope || rawConfig.tokens_per_frame != null || rawConfig.audio_step != null || archList.includes("KaniTTS2ForCausalLM");
 	return {
 		textVocabSize,
 		vocabSize: rawConfig.vocab_size,
 		tokensPerFrame: rawConfig.tokens_per_frame ?? KANI_TOKENS_PER_FRAME,
 		audioStep: rawConfig.audio_step ?? 1,
-		useLearnableRope: rawConfig.use_learnable_rope ?? false,
+		useLearnableRope,
 		alphaMin: rawConfig.alpha_min ?? .1,
 		alphaMax: rawConfig.alpha_max ?? 2,
 		speakerEmbDim: rawConfig.speaker_emb_dim ?? 128,
 		audioTokensStart: textVocabSize + 10,
 		startOfSpeech: textVocabSize + 1,
 		endOfSpeech: textVocabSize + 2,
-		codebookSize: KANI_CODEBOOK_SIZE
+		codebookSize: KANI_CODEBOOK_SIZE,
+		frameLevelPositions: isKaniTTS2
 	};
 }
 /**
@@ -249,6 +313,10 @@ function kaniAttentionLayerIndices(rawConfig) {
 */
 function computeKaniPositions(tokenIds, cfg = parseKaniConfig({})) {
 	const pos = new Float32Array(tokenIds.length);
+	if (!cfg.frameLevelPositions) {
+		for (let t = 0; t < tokenIds.length; t++) pos[t] = t;
+		return pos;
+	}
 	let textCount = 0;
 	let audioTokenCount = 0;
 	for (let t = 0; t < tokenIds.length; t++) {
@@ -652,15 +720,25 @@ function computeKaniFfnDim(rawConfig) {
 	return ff;
 }
 /**
-* Generate the KaniTTS2 codec-LM backbone graph (LFM2-350M body, full-vocab logits,
+* Generate the KaniTTS codec-LM backbone graph (LFM2-350M body, full-vocab logits,
 * per-layer learnable MRoPE). Mirrors generateLfm2Graph block-for-block; the only
 * difference is that each attention layer rotates Q/K with the MRoPE op fed a
 * per-layer host cos/sin table instead of the position-counter RoPE op.
 *
-* f32 only (the checkpoint is bf16→f32; q4 would need the codec-LM head re-validated).
+* dtype:
+*   • undefined / "f32" — full precision (the checkpoint is bf16/f32 → f32).
+*   • "q4" — weight-only group-wise INT4 (W4A16) for the LFM2 linear projections
+*     (q/k/v/o, gate/up/down, short-conv in/out_proj) via MatMulInt4. The tied
+*     embedding + LM head stay at f32 on purpose: decoder-LM heads are the most
+*     quantization-sensitive part of a codec-LM and the embed table is the single
+*     buffer shared by both, so quantizing it would degrade audio quality for no
+*     load-size win that the int4 projections do not already capture. Backbone
+*     drops from ~1.5 GB (f32) to ~0.37 GB. The NanoCodec decoder is a separate
+*     graph and always runs f32.
 */
-function generateKaniTtsGraph(rawConfig, dtype, _groupSize, kvDtype) {
-	if (dtype === "q4") throw new Error("Kani-TTS-2 codec-LM backbone supports f32 only (load without dtype: 'q4').");
+function generateKaniTtsGraph(rawConfig, dtype, groupSize, kvDtype) {
+	const useQ4 = dtype === "q4";
+	const groupSz = groupSize ?? DEFAULT_GROUP_SIZE;
 	const hidden_size = rawConfig.hidden_size;
 	const num_layers = rawConfig.num_hidden_layers;
 	const num_heads = rawConfig.num_attention_heads;
@@ -711,6 +789,48 @@ function generateKaniTtsGraph(rawConfig, dtype, _groupSize, kvDtype) {
 		executionOrder.push(node.id);
 	};
 	const addLinear = (id, activationInput, weightName, weightShape, outputName, K, N) => {
+		if (useQ4) {
+			const wQ = `${weightName}.q`;
+			const wScales = `${weightName}.scales`;
+			const wZeros = `${weightName}.zeros`;
+			const totalElements = K * N;
+			addTensor({
+				name: wQ,
+				shape: [Math.ceil(totalElements / 8)],
+				dtype: "u32",
+				storage: "constant"
+			});
+			addTensor({
+				name: wScales,
+				shape: [Math.ceil(totalElements / groupSz)],
+				dtype: "f32",
+				storage: "constant"
+			});
+			addTensor({
+				name: wZeros,
+				shape: [Math.ceil(totalElements / groupSz)],
+				dtype: "f32",
+				storage: "constant"
+			});
+			addNode({
+				id,
+				opType: "MatMulInt4",
+				inputs: [
+					activationInput,
+					wQ,
+					wScales,
+					wZeros
+				],
+				outputs: [outputName],
+				attributes: {
+					M_tensor: activationInput,
+					K,
+					N,
+					group_size: groupSz
+				}
+			});
+			return;
+		}
 		addTensor({
 			name: weightName,
 			shape: weightShape,
@@ -1239,66 +1359,6 @@ function generateKaniTtsGraph(rawConfig, dtype, _groupSize, kvDtype) {
 	};
 }
 
-//#endregion
-//#region src/gpu/quantize.ts
-/**
-* On-the-fly INT4 quantization for WebGPU inference.
-*
-* Converts F32 weight matrices to 4-bit unsigned integers with per-group
-* scale/zero-point parameters. Matches the dequantization formula in the
-* MatMulInt4 WGSL kernel:
-*
-*   dequant_val = (raw_nibble - zero) * scale
-*
-* Packing: 8 × 4-bit values per u32, little-nibble-first order.
-* Grouping: contiguous elements in row-major (K × N) layout.
-*/
-const DEFAULT_GROUP_SIZE = 128;
-/**
-* Quantize a flat F32 weight array to INT4.
-*
-* @param data   - Source F32 weight data (row-major, K × N elements)
-* @param groupSize - Number of elements per quantization group (default 128)
-* @returns Packed INT4 data with per-group scales and zeros
-*/
-function quantizeInt4(data, groupSize = DEFAULT_GROUP_SIZE) {
-	const totalElements = data.length;
-	const numGroups = Math.ceil(totalElements / groupSize);
-	const packedSize = Math.ceil(totalElements / 8);
-	const packed = new Uint32Array(packedSize);
-	const scales = new Float32Array(numGroups);
-	const zeros = new Float32Array(numGroups);
-	for (let g = 0; g < numGroups; g++) {
-		const start = g * groupSize;
-		const end = Math.min(start + groupSize, totalElements);
-		let min = data[start];
-		let max = data[start];
-		for (let i = start + 1; i < end; i++) {
-			const v = data[i];
-			if (v < min) min = v;
-			if (v > max) max = v;
-		}
-		const range = max - min;
-		const scale = range > 0 ? range / 15 : 1;
-		const zero = range > 0 ? -min / scale : 0;
-		scales[g] = scale;
-		zeros[g] = zero;
-		const invRange = range > 0 ? 15 / range : 0;
-		for (let i = start; i < end; i++) {
-			const q = Math.round((data[i] - min) * invRange);
-			const clamped = q < 0 ? 0 : q > 15 ? 15 : q;
-			const packedIdx = i >>> 3;
-			const nibblePos = i & 7;
-			packed[packedIdx] |= clamped << nibblePos * 4;
-		}
-	}
-	return {
-		packed,
-		scales,
-		zeros
-	};
-}
-
 //#endregion
 //#region src/gpu/architectures/gemma3_encoder.ts
 function generateGemma3EncoderGraph(rawConfig, dtype, groupSizeOverride, _kvDtype, _embedding) {
@@ -6059,5 +6119,5 @@ function isArchitectureSupported(name) {
 }
 
 //#endregion
-export { parseKaniConfig as C, createDefaultHFKeyMapper as D, GEMMA4_VIS_KEYS as E, nanoCodecWeightMap as S, DTYPE_BYTES as T, generateNanoCodecDecoderGraph as _, generateMoonshineEncoderGraph as a, kaniLayerAlpha as b, DEFAULT_GROUP_SIZE as c, KANI_START_OF_HUMAN as d, audioTokensToCodes as f, generateKaniTtsGraph as g, foldNanoCodecWeightNorm as h, generateMoonshineDecoderGraph as i, quantizeInt4 as l, computeKaniPositions as m, isArchitectureSupported as n, moonshineEncoderFrames as o, buildKaniLayerCosSin as p, MOONSHINE_REMAINING_WORK as r, parseMoonshineConfig as s, generateGraph as t, KANI_END_OF_HUMAN as u, kaniAttentionLayerIndices as v, CANONICAL_KEYS as w, kaniSinTensor as x, kaniCosTensor as y };
-//# sourceMappingURL=architectures-wan0swXu.mjs.map
+export { quantizeInt4 as C, createDefaultHFKeyMapper as D, GEMMA4_VIS_KEYS as E, DEFAULT_GROUP_SIZE as S, DTYPE_BYTES as T, kaniCosTensor as _, generateMoonshineEncoderGraph as a, nanoCodecWeightMap as b, KANI_END_OF_HUMAN as c, buildKaniLayerCosSin as d, computeKaniPositions as f, kaniAttentionLayerIndices as g, generateNanoCodecDecoderGraph as h, generateMoonshineDecoderGraph as i, KANI_START_OF_HUMAN as l, generateKaniTtsGraph as m, isArchitectureSupported as n, moonshineEncoderFrames as o, foldNanoCodecWeightNorm as p, MOONSHINE_REMAINING_WORK as r, parseMoonshineConfig as s, generateGraph as t, audioTokensToCodes as u, kaniLayerAlpha as v, CANONICAL_KEYS as w, parseKaniConfig as x, kaniSinTensor as y };
+//# sourceMappingURL=architectures-Ol7yUmKf.mjs.map