From 710292ff4f191d8da9f6a4230804fbc693338d4a Mon Sep 17 00:00:00 2001 From: Patrick Devine Date: Mon, 29 Jun 2026 21:15:08 -0700 Subject: [PATCH] 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 | 114 +++++++++++++++++++------------------- 1 file changed, 57 insertions(+), 57 deletions(-) diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index 9a17193b..ff77864d 100644 --- a/x/models/gemma4/gemma4.go +++ b/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). - // Try .experts. first (new weight drop), fall back to .moe. (old format). - gateUpW := tensors[layerPrefix+".experts.gate_up_proj"] + // Check for pre-stacked tensors (BF16 or quantized HF format). + // Try .experts. first (new weight drop), fall back to .moe. (old + // 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. - dims := gateUpW.Dims() - half := int32(dims[1] / 2) - 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"] + // Fused gate+up (dense or quantized): split happens at matmul time. + downKey := layerPrefix + ".experts.down_proj" + downW := tensors[downKey] 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). - // The scale key depends on whether the tensor has .weight suffix. + // Resolve base keys: the scale/bias suffix attaches to whichever + // 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") - 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) - } + loadFusedExperts(tensors, m.TextConfig, moe, switchGateUp, gateUpKey, switchDown, downKey) } else { // Per-expert tensors (from create path). // Try separate gate_proj/up_proj first, then fused gate_up_proj.