From 5db07cad71f5e1b154d10a5810036fb3ad2fcc07 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Mon, 27 Jul 2026 17:26:45 -0700 Subject: [PATCH] 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 | 16 ++++- x/mlxrunner/mlx/ops_extra_test.go | 80 +++++++++++++++++++++++++ x/models/cohere2_moe/cohere2_moe.go | 2 +- x/models/glm4_moe_lite/glm4_moe_lite.go | 6 +- x/models/laguna/laguna.go | 6 +- x/models/nn/nn.go | 7 +-- x/models/nn/nn_test.go | 2 +- x/models/qwen3_5/qwen3_5.go | 10 ++-- 8 files changed, 107 insertions(+), 22 deletions(-) create mode 100644 x/mlxrunner/mlx/ops_extra_test.go diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index 8ac7a5c79..31170ea06 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/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 } diff --git a/x/mlxrunner/mlx/ops_extra_test.go b/x/mlxrunner/mlx/ops_extra_test.go new file mode 100644 index 000000000..8138420d6 --- /dev/null +++ b/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) + } +} diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 5565f9b1f..60925fb3b 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/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, diff --git a/x/models/glm4_moe_lite/glm4_moe_lite.go b/x/models/glm4_moe_lite/glm4_moe_lite.go index 65f84e439..bc625c1f4 100644 --- a/x/models/glm4_moe_lite/glm4_moe_lite.go +++ b/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 diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index 224546d88..494aa030f 100644 --- a/x/models/laguna/laguna.go +++ b/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) } diff --git a/x/models/nn/nn.go b/x/models/nn/nn.go index 7c7461f0a..371af64c3 100644 --- a/x/models/nn/nn.go +++ b/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 { diff --git a/x/models/nn/nn_test.go b/x/models/nn/nn_test.go index 8c475600e..be71a24ab 100644 --- a/x/models/nn/nn_test.go +++ b/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) diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index 158396c6d..e7a411528 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/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,