mlx: tighten up gemma4 moe loading code (#16964)

This change allows .experts.gate_proj / .up_proj / .down_proj tensor names to each
be used for both quantized (i.e. nvfp4 and mxfp8) and non-quantized (bf16) models.
Previous to this only non-quantized models used that tensor naming scheme.

x/models/gemma4/gemma4.go

@@ -243,6 +243,46 @@ func collectExpertProjection(tensors map[string]*mlx.Array, cfg *TextConfig, pre
 	return out
 }
 
+// loadFusedExperts populates moe with a stacked-3D fused gate_up + down expert
+// pair, choosing the quantized GatherQMM path when scale sidecars are present
+// and the dense GatherMM path otherwise. gateUpW/downW are the stacked tensors
+// and gateUpKey/downKey are their map keys, used to locate the _scale/_qbias
+// sidecars and per-tensor quant metadata. Shared by the .experts./.moe.
+// pre-stacked names and the .moe.switch_mlp. names so all fused layouts load
+// identically.
+func loadFusedExperts(tensors map[string]*mlx.Array, cfg *TextConfig, moe *MoEBlock, gateUpW *mlx.Array, gateUpKey string, downW *mlx.Array, downKey string) {
+	moe.UseFusedGateUp = true
+
+	gateUpScales := firstNonNil(tensors, gateUpKey+"_scale", gateUpKey+".scale")
+	downScales := firstNonNil(tensors, downKey+"_scale", downKey+".scale")
+	if gateUpScales == nil || downScales == nil {
+		// Unquantized: keep fused and transpose for GatherMM ([experts, in, out]).
+		moe.GateUpWeight = transposeForGatherMM(gateUpW)
+		moe.DownWeight = transposeForGatherMM(downW)
+		return
+	}
+
+	// Quantized (affine/nvfp4/mxfp8): keep fused as a single tensor for one
+	// GatherQMM call instead of two. No transpose. nvfp4/mxfp8 carry no qbias,
+	// and this path does not apply nvfp4's optional double-scale global_scale.
+	moe.UseQuantized = true
+	moe.GateUpWeightQ = gateUpW
+	moe.GateUpScales = gateUpScales
+	moe.GateUpBiases = firstNonNil(tensors, gateUpKey+"_qbias", gateUpKey+".bias")
+	moe.DownWeightQ = downW
+	moe.DownScales = downScales
+	moe.DownBiases = firstNonNil(tensors, downKey+"_qbias", downKey+".bias")
+
+	moe.GateUpGroupSize, moe.GateUpBits, moe.QuantMode = model.ResolveLinearQuantParams(
+		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode,
+		cfg.TensorQuant, gateUpKey, gateUpW, gateUpScales,
+	)
+	moe.DownGroupSize, moe.DownBits, moe.DownQuantMode = model.ResolveLinearQuantParams(
+		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode,
+		cfg.TensorQuant, downKey, downW, downScales,
+	)
+}
+
 // Router implements Gemma 4's expert routing mechanism.
 type Router struct {
 	Proj  nn.LinearLayer // [hidden_size -> num_experts]
@@ -761,32 +801,32 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 
 			moe := &MoEBlock{PerExpertScale: perExpertScale}
 
-			// Check for pre-stacked tensors (unquantized HF format).
+			// Check for pre-stacked tensors (BF16 or quantized HF format).
-			// Try .experts. first (new weight drop), fall back to .moe. (old format).
+			// Try .experts. first (new weight drop), fall back to .moe. (old
-			gateUpW := tensors[layerPrefix+".experts.gate_up_proj"]
+			// format). Track the matched key so loadFusedExperts can find the
+			// quant sidecars.
+			gateUpKey := layerPrefix + ".experts.gate_up_proj"
+			gateUpW := tensors[gateUpKey]
 			if gateUpW == nil {
-				gateUpW = tensors[layerPrefix+".moe.gate_up_proj"]
+				gateUpKey = layerPrefix + ".moe.gate_up_proj"
+				gateUpW = tensors[gateUpKey]
 			}
 			gateW := tensors[layerPrefix+".experts.gate_proj"]
 			if gateW == nil {
 				gateW = tensors[layerPrefix+".moe.gate_proj"]
 			}
 			if gateUpW != nil {
-				// Fused gate+up: split along dim 1, transpose for GatherMM.
+				// Fused gate+up (dense or quantized): split happens at matmul time.
-				dims := gateUpW.Dims()
+				downKey := layerPrefix + ".experts.down_proj"
-				half := int32(dims[1] / 2)
+				downW := tensors[downKey]
-				gateSlice := sliceAxis1(gateUpW, 0, half)
-				upSlice := sliceAxis1(gateUpW, half, int32(dims[1]))
-				moe.GateWeight = transposeForGatherMM(gateSlice)
-				moe.UpWeight = transposeForGatherMM(upSlice)
-				downW := tensors[layerPrefix+".experts.down_proj"]
 				if downW == nil {
-					downW = tensors[layerPrefix+".moe.down_proj"]
+					downKey = layerPrefix + ".moe.down_proj"
+					downW = tensors[downKey]
 				}
 				if downW == nil {
 					return fmt.Errorf("layer %d: missing MoE down_proj with fused gate_up_proj", i)
 				}
-				moe.DownWeight = transposeForGatherMM(downW)
+				loadFusedExperts(tensors, m.TextConfig, moe, gateUpW, gateUpKey, downW, downKey)
 			} else if gateW != nil {
 				// Separate gate_proj and up_proj (older format). Transpose for GatherMM.
 				moe.GateWeight = transposeForGatherMM(gateW)
@@ -814,8 +854,8 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 					return fmt.Errorf("layer %d: missing switch_mlp down_proj", i)
 				}
 
-				// Check for quantized weights (scales present).
+				// Resolve base keys: the scale/bias suffix attaches to whichever
-				// The scale key depends on whether the tensor has .weight suffix.
+				// of .weight / bare matched the weight tensor.
 				gateUpKey := layerPrefix + ".moe.switch_mlp.gate_up_proj.weight"
 				if tensors[gateUpKey] == nil {
 					gateUpKey = layerPrefix + ".moe.switch_mlp.gate_up_proj"
@@ -824,47 +864,7 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 				if tensors[downKey] == nil {
 					downKey = layerPrefix + ".moe.switch_mlp.down_proj"
 				}
-				gateUpScales := firstNonNil(tensors, gateUpKey+"_scale", gateUpKey+".scale")
+				loadFusedExperts(tensors, m.TextConfig, moe, switchGateUp, gateUpKey, switchDown, downKey)
-				downScales := firstNonNil(tensors, downKey+"_scale", downKey+".scale")
-
-				if gateUpScales != nil && downScales != nil {
-					// Quantized: keep fused gate_up as single tensor for GatherQMM.
-					// One fused call instead of two separate gate+up calls.
-					gateUpBiases := firstNonNil(tensors, gateUpKey+"_qbias", gateUpKey+".bias")
-					downBiases := firstNonNil(tensors, downKey+"_qbias", downKey+".bias")
-
-					moe.GateUpWeightQ = switchGateUp
-					moe.GateUpScales = gateUpScales
-					moe.GateUpBiases = gateUpBiases
-					moe.DownWeightQ = switchDown
-					moe.DownScales = downScales
-					if downBiases != nil {
-						moe.DownBiases = downBiases
-					}
-
-					groupSize, bits, mode := model.ResolveLinearQuantParams(
-						m.QuantGroupSize, m.QuantBits, m.QuantMode,
-						m.TensorQuant, gateUpKey, switchGateUp, gateUpScales,
-					)
-					moe.UseQuantized = true
-					moe.UseFusedGateUp = true
-					moe.GateUpGroupSize = groupSize
-					moe.GateUpBits = bits
-					moe.QuantMode = mode
-
-					dGroupSize, dBits, dMode := model.ResolveLinearQuantParams(
-						m.QuantGroupSize, m.QuantBits, m.QuantMode,
-						m.TensorQuant, downKey, switchDown, downScales,
-					)
-					moe.DownGroupSize = dGroupSize
-					moe.DownBits = dBits
-					moe.DownQuantMode = dMode
-				} else {
-					// Unquantized switch_mlp: keep fused and transpose for GatherMM.
-					moe.GateUpWeight = transposeForGatherMM(switchGateUp)
-					moe.UseFusedGateUp = true
-					moe.DownWeight = transposeForGatherMM(switchDown)
-				}
 			} else {
 				// Per-expert tensors (from create path).
 				// Try separate gate_proj/up_proj first, then fused gate_up_proj.