mlx: apply global scales in Dequantize

The C-level dequantize accepts a global_scale argument but rejects it
on the Metal backend, so dequantize-fallback sites hand-rolled the
same post-multiply. Take the scale in the Go wrapper and apply it on
top of the op, cast back to the output dtype. The quantized embedding
passes its scale; laguna's expert paths keep their own multiplies,
which shape per-expert scales and differ on result dtype.

x/mlxrunner/mlx/ops_extra.go

@@ -45,7 +45,7 @@ func ToFP8(x *Array) *Array {
 	return out
 }
 
-func Dequantize(w, scales, biases *Array, groupSize, bits int, mode string) *Array {
+func Dequantize(w, scales, biases *Array, groupSize, bits int, mode string, globalScale *Array) *Array {
 	cMode := C.CString(mode)
 	defer C.free(unsafe.Pointer(cMode))
 	optGroupSize := C.mlx_optional_int{value: C.int(groupSize), has_value: true}
@@ -58,8 +58,18 @@ func Dequantize(w, scales, biases *Array, groupSize, bits int, mode string) *Arr
 	}
 
 	out := New("DEQUANTIZE")
-	var globalScale C.mlx_array
-	C.mlx_dequantize(&out.ctx, w.ctx, scales.ctx, b, optGroupSize, optBits, cMode, globalScale, optDtype, DefaultStream().ctx)
+	var noGlobalScale C.mlx_array
+	C.mlx_dequantize(&out.ctx, w.ctx, scales.ctx, b, optGroupSize, optBits, cMode, noGlobalScale, optDtype, DefaultStream().ctx)
+	if globalScale != nil {
+		// The C-level global_scale argument is rejected on Metal; apply it on top.
+		gs := globalScale
+		if gs.Size() > 1 {
+			// A vector scale is per-row; bind it to the weight's leading axis.
+			gs = Reshape(gs, int32(gs.Size()), 1)
+		}
+		outType := out.DType()
+		out = Mul(out, gs).AsType(outType)
+	}
 	return out
 }
 

x/mlxrunner/mlx/ops_extra_test.go

@@ -0,0 +1,80 @@
+package mlx
+
+import (
+	"math"
+	"testing"
+)
+
+// fp4Values decodes an fp4 (E2M1) code to its value.
+var fp4Values = [16]float32{0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6}
+
+func TestDequantizeGlobalScale(t *testing.T) {
+	skipIfNoMLX(t)
+	withMLXThread(t, func() {
+		testDequantizeGlobalScale(t)
+	})
+}
+
+// The quantized payload is built directly, the way an nvfp4 checkpoint ships
+// it: packed fp4 codes, e4m3 group-scale bytes, and a separate global scale.
+// Only the dequantize consumer path runs, so expectations are exact.
+func testDequantizeGlobalScale(t *testing.T) {
+	const rows, cols, group = 4, 64, 16
+	// Every group cycles through all 16 codes; group g of row r has scale
+	// 2^((r+g)%4-1), a power of two so every expected product is exact.
+	scaleOf := func(r, g int) float32 {
+		return float32(math.Ldexp(1, (r+g)%4-1))
+	}
+
+	packed := make([]uint32, rows*cols/8)
+	for i := range packed {
+		for j := range 8 {
+			packed[i] |= uint32((i*8+j)%16) << (4 * j)
+		}
+	}
+	scaleBits := make([]uint8, rows*(cols/group))
+	for i := range scaleBits {
+		exp := (i/(cols/group)+i%(cols/group))%4 - 1
+		scaleBits[i] = uint8((exp + 7) << 3)
+	}
+	wq := FromValues(packed, rows, cols/8)
+	scales := FromValues(scaleBits, rows, cols/group)
+
+	check := func(name string, got *Array, gs func(r int) float32) {
+		t.Helper()
+		g32 := got.AsType(DTypeFloat32)
+		Eval(g32)
+		values := g32.Floats()
+		if len(values) != rows*cols {
+			t.Errorf("%s: length = %d, want %d", name, len(values), rows*cols)
+			return
+		}
+		for i, v := range values {
+			r, c := i/cols, i%cols
+			if want := fp4Values[c%16] * scaleOf(r, c/group) * gs(r); v != want {
+				t.Errorf("%s[%d] = %v, want %v", name, i, v, want)
+				return
+			}
+		}
+	}
+
+	base := Dequantize(wq, scales, nil, group, 4, "nvfp4", nil)
+	check("no scale", base, func(int) float32 { return 1 })
+
+	perRow := []float32{0.5, 1, 2, 4}
+	cases := []struct {
+		name  string
+		scale *Array
+		gs    func(r int) float32
+	}{
+		{"scalar", FromValues([]float32{2}, 1), func(int) float32 { return 2 }},
+		{"perRow", FromValues(perRow, rows), func(r int) float32 { return perRow[r] }},
+	}
+	for _, tc := range cases {
+		got := Dequantize(wq, scales, nil, group, 4, "nvfp4", tc.scale)
+		if got.DType() != base.DType() {
+			t.Errorf("%s: dtype = %v, want %v", tc.name, got.DType(), base.DType())
+		}
+		check(tc.name, got, tc.gs)
+	}
+}

x/models/cohere2_moe/cohere2_moe.go

@@ -439,7 +439,7 @@ func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuanti
 	}
 
 	return &stackedExpertWeights{
-		Weight:    mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode),
+		Weight:    mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode, nil),
 		Bits:      bits,
 		GroupSize: groupSize,
 		Mode:      mode,

x/models/glm4_moe_lite/glm4_moe_lite.go

@@ -403,7 +403,7 @@ func loadExpertWeight(tensors map[string]*mlx.Array, path string, useQuantized b
 			return &ExpertWeight{Weight: w, Scales: scales, Biases: qbiases, Bits: bits, GroupSize: groupSize}
 		}
 
-		return &ExpertWeight{Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode)}
+		return &ExpertWeight{Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode, nil)}
 	}
 
 	return &ExpertWeight{Weight: w}
@@ -441,7 +441,7 @@ func loadStackedProjection(tensors map[string]*mlx.Array, base string, useQuanti
 		return &StackedExpertWeights{Weight: w, Scales: scales, Biases: qbiases, Bits: bits, GroupSize: groupSize}
 	}
 
-	return &StackedExpertWeights{Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode)}
+	return &StackedExpertWeights{Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode, nil)}
 }
 
 // loadStackedExperts loads a stacked expert projection by its .experts name,
@@ -532,7 +532,7 @@ func sanitizeMLAWeights(tensors map[string]*mlx.Array, prefix string, cfg *Confi
 			w,
 			scales,
 		)
-		w = mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode)
+		w = mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode, nil)
 	}
 
 	headDim := cfg.QKNopeHeadDim + cfg.VHeadDim

x/models/laguna/laguna.go

@@ -613,7 +613,7 @@ func denseExpertWeight(w *stackedExpertWeights) *mlx.Array {
 	}
 	weight := w.Weight
 	if w.Scales != nil {
-		weight = mlx.Dequantize(w.Weight, w.Scales, w.Biases, w.GroupSize, w.Bits, w.Mode)
+		weight = mlx.Dequantize(w.Weight, w.Scales, w.Biases, w.GroupSize, w.Bits, w.Mode, nil)
 		if w.GlobalScales != nil {
 			scale := w.GlobalScales
 			if scale.DType() != weight.DType() {
@@ -907,7 +907,7 @@ func collectPerExpertProjection(tensors map[string]*mlx.Array, cfg *Config, useQ
 				biases = append(biases, qb)
 			}
 		} else {
-			deq := mlx.Dequantize(w, s, qb, gs, b, m)
+			deq := mlx.Dequantize(w, s, qb, gs, b, m, nil)
 			if globalScale != nil {
 				deq = mlx.Mul(deq, globalScale)
 				globalScales = append(globalScales, globalScale)
@@ -968,7 +968,7 @@ func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuanti
 			freeTensorKeys(tensors, consumedKeys...)
 			return &stackedExpertWeights{Weight: w, Scales: s, Biases: qb, GlobalScales: globalScale, Bits: b, GroupSize: gs, Mode: m}
 		}
-		deq := mlx.Dequantize(w, s, qb, gs, b, m)
+		deq := mlx.Dequantize(w, s, qb, gs, b, m, nil)
 		if globalScale != nil {
 			deq = mlx.Mul(deq, globalScale)
 		}

x/models/nn/nn.go

@@ -190,12 +190,7 @@ func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array {
 	if qe.QBiases != nil && qe.QBiases.Valid() {
 		qbiases = qe.QBiases.TakeAxis(indices, 0)
 	}
-	out := mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode)
-	if qe.GlobalScale != nil {
-		outDType := out.DType()
-		out = mlx.Mul(out, qe.GlobalScale).AsType(outDType)
-	}
-	return out
+	return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode, qe.GlobalScale)
 }
 
 func (qe *QuantizedEmbedding) AsLinear() LinearLayer {

x/models/nn/nn_test.go

@@ -166,7 +166,7 @@ func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) {
 		t.Fatalf("mxfp4 qbiases = %v, want nil", ql.QBiases)
 	}
 
-	dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4")
+	dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4", nil)
 	mlx.Eval(dequantizedWeight)
 
 	qOut := ql.Forward(input).AsType(mlx.DTypeFloat32)

x/models/qwen3_5/qwen3_5.go

@@ -501,11 +501,11 @@ func fuseGateUpProjections(gate, up *stackedExpertWeights) *stackedExpertWeights
 	}
 	gateWeight := gate.Weight
 	if gate.Scales != nil {
-		gateWeight = mlx.Dequantize(gate.Weight, gate.Scales, gate.Biases, gate.GroupSize, gate.Bits, gate.Mode)
+		gateWeight = mlx.Dequantize(gate.Weight, gate.Scales, gate.Biases, gate.GroupSize, gate.Bits, gate.Mode, nil)
 	}
 	upWeight := up.Weight
 	if up.Scales != nil {
-		upWeight = mlx.Dequantize(up.Weight, up.Scales, up.Biases, up.GroupSize, up.Bits, up.Mode)
+		upWeight = mlx.Dequantize(up.Weight, up.Scales, up.Biases, up.GroupSize, up.Bits, up.Mode, nil)
 	}
 	return &stackedExpertWeights{
 		Weight:    mlx.Concatenate([]*mlx.Array{gateWeight, upWeight}, 1),
@@ -552,7 +552,7 @@ func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuanti
 			slog.Warn("dequantizing expert weights: no gather kernel for format", "tensor", key, "mode", mode, "bits", bits)
 		}
 		return &stackedExpertWeights{
-			Weight:    mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode),
+			Weight:    mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode, nil),
 			Bits:      bits,
 			GroupSize: groupSize,
 			Mode:      mode,
@@ -611,7 +611,7 @@ func collectPerExpertProjection(tensors map[string]*mlx.Array, cfg *Config, useQ
 				biases = append(biases, qb)
 			}
 		} else {
-			weights = append(weights, mlx.Dequantize(w, s, qb, gs, b, m))
+			weights = append(weights, mlx.Dequantize(w, s, qb, gs, b, m, nil))
 			numDequantized++
 		}
 	}
@@ -676,7 +676,7 @@ func combinedGateUpProjection(tensors map[string]*mlx.Array, cfg *Config, useQua
 		slog.Warn("dequantizing expert weights: no gather kernel for format", "tensor", key, "mode", mode, "bits", bits)
 	}
 	return &stackedExpertWeights{
-		Weight:    mlx.Dequantize(gateUp, scales, qbiases, groupSize, bits, mode),
+		Weight:    mlx.Dequantize(gateUp, scales, qbiases, groupSize, bits, mode, nil),
 		Bits:      bits,
 		GroupSize: groupSize,
 		Mode:      mode,