@tryhamster/gerbil 1.1.3 → 1.2.1

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