@tryhamster/gerbil 1.6.0 → 1.6.1

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

@@ -3066,34 +3066,39 @@ function generateLfm2Graph(rawConfig, dtype, groupSizeOverride, kvDtype) {
 		});
 	}
 	let prevOutput = "embed_out";
+	let pendingNorm1Out = null;
 	for (let i = 0; i < num_layers; i++) {
 		const prefix = `layer${i}`;
-		const operatorNormW = CANONICAL_KEYS.lfm2OperatorNorm(i);
-		addTensor({
-			name: operatorNormW,
-			shape: [hidden_size],
-			dtype: "f32",
-			storage: "constant",
-			safetensorsKey: operatorNormW
-		});
-		const norm1Out = `${prefix}_norm1_out`;
-		addTensor({
-			name: norm1Out,
-			shape: ["T", hidden_size],
-			dtype: "f32",
-			storage: "activation"
-		});
-		addNode({
-			id: `${prefix}_norm1`,
-			opType: "RMSNorm",
-			inputs: [prevOutput, operatorNormW],
-			outputs: [norm1Out],
-			attributes: {
-				hidden_size,
-				eps: rms_norm_eps,
-				seq_len_tensor: prevOutput
-			}
-		});
+		let norm1Out;
+		if (pendingNorm1Out) norm1Out = pendingNorm1Out;
+		else {
+			const operatorNormW = CANONICAL_KEYS.lfm2OperatorNorm(i);
+			addTensor({
+				name: operatorNormW,
+				shape: [hidden_size],
+				dtype: "f32",
+				storage: "constant",
+				safetensorsKey: operatorNormW
+			});
+			norm1Out = `${prefix}_norm1_out`;
+			addTensor({
+				name: norm1Out,
+				shape: ["T", hidden_size],
+				dtype: "f32",
+				storage: "activation"
+			});
+			addNode({
+				id: `${prefix}_norm1`,
+				opType: "RMSNorm",
+				inputs: [prevOutput, operatorNormW],
+				outputs: [norm1Out],
+				attributes: {
+					hidden_size,
+					eps: rms_norm_eps,
+					seq_len_tensor: prevOutput
+				}
+			});
+		}
 		const mixerOut = isAttentionLayer(i) ? buildAttentionBlock(i, prefix, norm1Out) : buildShortConvBlock(i, prefix, norm1Out);
 		const resid1Out = `${prefix}_resid1`;
 		const ffnNormW = CANONICAL_KEYS.lfm2FfnNorm(i);
@@ -3180,16 +3185,51 @@ function generateLfm2Graph(rawConfig, dtype, groupSizeOverride, kvDtype) {
 			dtype: "f32",
 			storage: "activation"
 		});
-		addNode({
-			id: `${prefix}_resid2`,
-			opType: "Add",
-			inputs: [resid1Out, mlpOut],
-			outputs: [resid2Out],
-			attributes: {
-				count_tensor: resid1Out,
-				hidden_size
-			}
-		});
+		if (i + 1 < num_layers) {
+			const nextNormW = CANONICAL_KEYS.lfm2OperatorNorm(i + 1);
+			addTensor({
+				name: nextNormW,
+				shape: [hidden_size],
+				dtype: "f32",
+				storage: "constant",
+				safetensorsKey: nextNormW
+			});
+			const nextNorm1Out = `layer${i + 1}_norm1_out`;
+			addTensor({
+				name: nextNorm1Out,
+				shape: ["T", hidden_size],
+				dtype: "f32",
+				storage: "activation"
+			});
+			addNode({
+				id: `${prefix}_resid2_norm1`,
+				opType: "ResidualRMSNorm",
+				inputs: [
+					resid1Out,
+					mlpOut,
+					nextNormW
+				],
+				outputs: [resid2Out, nextNorm1Out],
+				attributes: {
+					hidden_size,
+					eps: rms_norm_eps,
+					seq_len_tensor: resid1Out
+				}
+			});
+			pendingNorm1Out = nextNorm1Out;
+		} else {
+			addNode({
+				id: `${prefix}_resid2`,
+				opType: "Add",
+				inputs: [resid1Out, mlpOut],
+				outputs: [resid2Out],
+				attributes: {
+					count_tensor: resid1Out,
+					hidden_size
+				}
+			});
+			pendingNorm1Out = null;
+		}
 		prevOutput = resid2Out;
 	}
 	function buildShortConvBlock(i, prefix, norm1Out) {
@@ -3203,43 +3243,6 @@ function generateLfm2Graph(rawConfig, dtype, groupSizeOverride, kvDtype) {
 			storage: "activation"
 		});
 		addLinearOp(`${prefix}_conv_in_proj`, norm1Out, inProjW, [bcxDim, hidden_size], bcxOut, hidden_size, bcxDim);
-		const cOut = `${prefix}_conv_C`;
-		addTensor({
-			name: cOut,
-			shape: ["T", conv_dim],
-			dtype: "f32",
-			storage: "activation"
-		});
-		addNode({
-			id: `${prefix}_conv_split_C`,
-			opType: "SliceCols",
-			inputs: [bcxOut],
-			outputs: [cOut],
-			attributes: {
-				in_width: bcxDim,
-				out_width: conv_dim,
-				col_offset: conv_dim
-			}
-		});
-		const bxOut = `${prefix}_conv_Bx`;
-		addTensor({
-			name: bxOut,
-			shape: ["T", conv_dim],
-			dtype: "f32",
-			storage: "activation"
-		});
-		addNode({
-			id: `${prefix}_conv_gate_in`,
-			opType: "MulCols",
-			inputs: [bcxOut],
-			outputs: [bxOut],
-			attributes: {
-				in_width: bcxDim,
-				out_width: conv_dim,
-				off_a: 0,
-				off_b: 2 * conv_dim
-			}
-		});
 		const convW = CANONICAL_KEYS.lfm2ConvWeight(i);
 		addTensor({
 			name: convW,
@@ -3264,29 +3267,30 @@ function generateLfm2Graph(rawConfig, dtype, groupSizeOverride, kvDtype) {
 		});
 		addNode({
 			id: `${prefix}_conv1d`,
-			opType: "CausalConv1dGated",
+			opType: "CausalConv1dGatedBcx",
 			inputs: [
-				bxOut,
+				bcxOut,
 				convW,
-				convState,
-				cOut
+				convState
 			],
 			outputs: [yOut],
 			attributes: {
 				channels: conv_dim,
 				kernel_size: conv_kernel,
-				seq_len_tensor: bxOut
+				in_width: bcxDim,
+				seq_len_tensor: bcxOut
 			}
 		});
 		addNode({
 			id: `${prefix}_conv_state_update`,
-			opType: "ConvStateUpdate",
-			inputs: [bxOut],
+			opType: "ConvStateUpdateBcx",
+			inputs: [bcxOut],
 			outputs: [convState],
 			attributes: {
 				channels: conv_dim,
 				state_size: conv_kernel - 1,
-				T_tensor: bxOut
+				in_width: bcxDim,
+				T_tensor: bcxOut
 			}
 		});
 		const outProjW = CANONICAL_KEYS.lfm2ConvOutProj(i);
@@ -6120,4 +6124,4 @@ function isArchitectureSupported(name) {
 
 //#endregion
 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
+//# sourceMappingURL=architectures-CpmH4Gha.mjs.map