mlx: speed up Qwen 3.8 prompt processing (#18550)

* mlx: speed up Qwen 3.8 prompt processing

Use MLX's gated-delta kernel for long scans and fold dense MLP global scales into SwiGLU.

* address comments

mlx/act.go

@@ -78,6 +78,20 @@ var SwiGLU = Compile2(
 	Shapeless(),
 )
 
+var swiGLUScaled = Compile(
+	"SwiGLUScaled",
+	func(in ...*Array) []*Array {
+		gate, gateFactor, up, upFactor := in[0], in[1], in[2], in[3]
+		// Preserve the standalone scale path's rounding before applying
+		// SwiGLU. Compile still keeps these casts in the fused kernel instead
+		// of materializing two scaled projection outputs.
+		gate = Mul(gate, gateFactor).AsType(gate.DType())
+		up = Mul(up, upFactor).AsType(up.DType())
+		return []*Array{SiLU(gate).Multiply(up)}
+	},
+	Shapeless(),
+)
+
 // GeGLU returns gelu_approx(gate) * up as a fused kernel. Matches mlx_lm's
 // geglu, used by Gemma-family MLP and MoE paths.
 var GeGLU = Compile2(
@@ -88,6 +102,24 @@ var GeGLU = Compile2(
 	Shapeless(),
 )
 
+// SwiGLUScaled applies each non-nil quantization global scale, rounds the
+// projection back to its input dtype, then evaluates SwiGLU. A nil scale means
+// that projection is already complete.
+func SwiGLUScaled(gate, gateScale, up, upScale *Array) *Array {
+	if gateScale == nil && upScale == nil {
+		return SwiGLU(gate, up)
+	}
+	if gateScale == nil {
+		gateScale = identityGlobalScale()
+	}
+	if upScale == nil {
+		upScale = identityGlobalScale()
+	}
+	// Rank 0 or one-dimensional, so convert once here rather than per output
+	// element inside the compiled function.
+	return swiGLUScaled(gate, globalScaleFactor(gateScale), up, globalScaleFactor(upScale))[0]
+}
+
 // LogitSoftcap returns tanh(x / cap) * cap as a fused kernel. Matches
 // mlx_lm's logit_softcap. cap must have the same dtype as x.
 var LogitSoftcap = Compile2(

mlx/act_test.go

@@ -70,6 +70,94 @@ func benchmarkGELU(b *testing.B, fn func(*Array) *Array) {
 	}
 }
 
+func TestSwiGLUScaledMatchesSeparateScaling(t *testing.T) {
+	for _, tt := range []struct {
+		name               string
+		gateScale, upScale []float32
+	}{
+		{name: "no scales"},
+		{name: "scalar scales", gateScale: []float32{0.75}, upScale: []float32{0.75}},
+		{name: "gate scale only", gateScale: []float32{0.75}},
+		{name: "up scale only", upScale: []float32{0.75}},
+		{name: "per-output scales", gateScale: []float32{0.5, 0.75, 1.25, 1.5}, upScale: []float32{1.5, 1.25, 0.75, 0.5}},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			withMLXThread(t, func(t *mlxthreadtest.T) {
+				EnableCompile()
+				gate := FromValues([]float32{-3.25, -1.5, -0.25, 0.5, 1.75, 3, 4.5, 6}, 2, 4).AsType(DTypeBFloat16)
+				up := FromValues([]float32{2.5, -2, 1.25, -0.75, 0.125, 1.5, -3.5, 5}, 2, 4).AsType(DTypeBFloat16)
+				storedScale := func(factors []float32) *Array {
+					if factors == nil {
+						return nil
+					}
+					values := make([]float32, len(factors))
+					for i := range factors {
+						values[i] = factors[i] * float32(Nvfp4MaxProduct)
+					}
+					if len(values) == 1 {
+						return FromValue(values[0])
+					}
+					return FromValues(values, len(values))
+				}
+				gateScale, upScale := storedScale(tt.gateScale), storedScale(tt.upScale)
+
+				wantGate, wantUp := gate, up
+				if gateScale != nil {
+					wantGate = scaleAndCast(wantGate, gateScale)
+				}
+				if upScale != nil {
+					wantUp = scaleAndCast(wantUp, upScale)
+				}
+				want := SwiGLU(wantGate, wantUp)
+				got := SwiGLUScaled(gate, gateScale, up, upScale)
+				wantF32, gotF32 := want.AsType(DTypeFloat32), got.AsType(DTypeFloat32)
+				Eval(wantF32, gotF32)
+
+				wantValues, gotValues := wantF32.Floats(), gotF32.Floats()
+				for i := range wantValues {
+					if gotValues[i] != wantValues[i] {
+						t.Fatalf("SwiGLUScaled()[%d] = %v, want %v", i, gotValues[i], wantValues[i])
+					}
+				}
+			})
+		})
+	}
+}
+
+func BenchmarkSwiGLUSeparateScaling(b *testing.B) {
+	benchmarkSwiGLUScaling(b, func(gate, gateScale, up, upScale *Array) *Array {
+		return SwiGLU(scaleAndCast(gate, gateScale), scaleAndCast(up, upScale))
+	})
+}
+
+func BenchmarkSwiGLUFusedScaling(b *testing.B) {
+	benchmarkSwiGLUScaling(b, SwiGLUScaled)
+}
+
+func benchmarkSwiGLUScaling(b *testing.B, fn func(gate, gateScale, up, upScale *Array) *Array) {
+	thread := mlxTestThread(b)
+	if err := thread.Do(b.Context(), func() error {
+		EnableCompile()
+		// qwen3.8:27b-nvfp4's dense MLP output at a 2048-token prompt.
+		gate := AddScalar(Zeros(DTypeBFloat16, 1, 2048, 17408), 0.5)
+		up := AddScalar(Zeros(DTypeBFloat16, 1, 2048, 17408), 1.5)
+		gateScale := FromValue(float32(Nvfp4MaxProduct) * 0.75)
+		upScale := FromValue(float32(Nvfp4MaxProduct) * 1.25)
+		Eval(gate, up, gateScale, upScale)
+		defer ClearCache()
+
+		Scoped(func() { Eval(fn(gate, gateScale, up, upScale)) })
+
+		b.ResetTimer()
+		for range b.N {
+			Scoped(func() { Eval(fn(gate, gateScale, up, upScale)) })
+		}
+		return nil
+	}); err != nil {
+		b.Fatal(err)
+	}
+}
+
 func TestReLUSquared(t *testing.T) {
 	var got []float32
 	withMLXThread(t, func(t *mlxthreadtest.T) {

mlx/gated_delta.go

@@ -298,6 +298,9 @@ func gatedDeltaRecurrenceGraph(q, k, v, g, beta, state *Array) (y, nextState *Ar
 // directly.
 func gatedDeltaRecurrence(q, k, v, g, beta, state *Array) (y, nextState *Array) {
 	if dims, ok := resolveGatedDeltaRecurrenceDims(q, k, v, g, beta, state); ok {
+		if supportsFastGatedDeltaUpdate(dims) {
+			return fastGatedDeltaUpdate(q, k, v, g, beta, state, nil)
+		}
 		outs := gatedDeltaRecurrenceKernel.run(gpuLaunch{
 			dtypes: []gpuDTypeArg{{"InT", q.DType()}, {"StT", state.DType()}},
 			ints:   []gpuIntArg{{"Dk", dims.Dk}, {"Dv", dims.Dv}, {"Hk", dims.Hk}, {"Hv", dims.Hv}},
@@ -319,6 +322,17 @@ func gatedDeltaRecurrence(q, k, v, g, beta, state *Array) (y, nextState *Array)
 	return y, nextState
 }
 
+// supportsFastGatedDeltaUpdate reports whether MLX's Metal kernel is available
+// for this geometry. Short scans stay on Ollama's whole-step kernel, which also
+// serves MTP.
+func supportsFastGatedDeltaUpdate(dims gatedDeltaRecurrenceDims) bool {
+	if !MetalIsAvailable() || dims.T < 16 || dims.Dk != 128 || dims.Dv != 128 {
+		return false
+	}
+	return (dims.Hk == 16 && (dims.Hv == 16 || dims.Hv == 32 || dims.Hv == 48 || dims.Hv == 64)) ||
+		(dims.Hk == 24 && dims.Hv == 24) || (dims.Hk == 32 && dims.Hv == 32)
+}
+
 // gatedDeltaMaxTokens caps the fused scan length at the current token plus
 // a ten-token draft — several times the depth the EV controller selects in
 // practice. SeqT is a template argument, so every accepted length compiles

mlx/gated_delta_test.go

@@ -76,6 +76,24 @@ func testGatedDeltaMatchesGraph(t *mlxthreadtest.T) {
 	}
 }
 
+func TestGatedDeltaRecurrenceUsesFastUpdate(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		if !MetalIsAvailable() {
+			t.Skip("requires Metal")
+		}
+
+		in := fastGatedDeltaPatternInputs(1, 33, 16, 128, 48, 128, DTypeBFloat16)
+		y, state := gatedDeltaRecurrence(in.q, in.k, in.v, in.gates, in.beta, in.state)
+		wantY, wantState := fastGatedDeltaUpdate(in.q, in.k, in.v, in.gates, in.beta, in.state, nil)
+		if err := requireExact("y", y, wantY); err != nil {
+			t.Fatal(err)
+		}
+		if err := requireExact("state", state, wantState); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
 func TestGatedDeltaGraphRouting(t *testing.T) {
 	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testGatedDeltaGraphRouting(t)

mlx/ops_extra.go

@@ -15,6 +15,17 @@ import (
 // wrappers that apply it themselves divide it back out.
 const Nvfp4MaxProduct = 448 * 6
 
+// globalScaleFactor converts a stored global scale into the multiplier an
+// output needs.
+func globalScaleFactor(scale *Array) *Array {
+	return DivScalar(scale, Nvfp4MaxProduct)
+}
+
+// identityGlobalScale leaves a value unchanged.
+func identityGlobalScale() *Array {
+	return FromValue(float32(Nvfp4MaxProduct))
+}
+
 // scaleAndCast applies a global scale to an output and casts back, fusing the
 // divide, multiply and cast into one kernel. Eagerly this is several: MLX
 // promotes binary operands with an astype, so a bf16 input times a float32
@@ -22,7 +33,7 @@ const Nvfp4MaxProduct = 448 * 6
 var scaleAndCast = Compile2(
 	"GlobalScaleOutput",
 	func(out, scale *Array) *Array {
-		return Mul(out, DivScalar(scale, Nvfp4MaxProduct)).AsType(out.DType())
+		return Mul(out, globalScaleFactor(scale)).AsType(out.DType())
 	},
 	Shapeless(),
 )

mlxrunner/model/qwen3_5/qwen3_5.go

@@ -1197,7 +1197,7 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
 }
 
 func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array {
-	return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x)))
+	return m.DownProj.Forward(nn.SwiGLU(m.GateProj, m.UpProj, x))
 }
 
 func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.Array {

mlxrunner/nn/linear.go

@@ -67,9 +67,7 @@ func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int,
 }
 
 func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array {
-	// Double-scale nvfp4 (e.g., NVIDIA ModelOpt) applies the per-tensor or
-	// per-row global scale inside QuantizedMatmul.
-	out := mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode, ql.GlobalScale)
+	out := ql.matmul(x, ql.GlobalScale)
 	if ql.Bias != nil {
 		bias := ql.Bias
 		if bias.DType() != out.DType() {
@@ -80,6 +78,27 @@ func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array {
 	return out
 }
 
+func (ql *QuantizedLinear) matmul(x, globalScale *mlx.Array) *mlx.Array {
+	return mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true,
+		ql.GroupSize, ql.Bits, ql.Mode, globalScale)
+}
+
+// SwiGLU applies gate and up projections followed by a SwiGLU activation.
+// Quantized projections without bias defer their global scales so the scale
+// and activation operations can be fused.
+func SwiGLU(gate, up LinearLayer, x *mlx.Array) *mlx.Array {
+	gateOut, gateScale := forwardDeferScale(gate, x)
+	upOut, upScale := forwardDeferScale(up, x)
+	return mlx.SwiGLUScaled(gateOut, gateScale, upOut, upScale)
+}
+
+func forwardDeferScale(l LinearLayer, x *mlx.Array) (out, pending *mlx.Array) {
+	if ql, ok := l.(*QuantizedLinear); ok && ql.GlobalScale != nil && ql.Bias == nil {
+		return ql.matmul(x, nil), ql.GlobalScale
+	}
+	return l.Forward(x), nil
+}
+
 func (ql *QuantizedLinear) OutputDim() int32 {
 	return int32(ql.Weight.Dim(0))
 }

mlxrunner/nn/linear_test.go

@@ -2,6 +2,7 @@ package nn
 
 import (
 	"math"
+	"slices"
 	"testing"
 
 	"github.com/ollama/ollama/mlx"
@@ -100,3 +101,65 @@ func TestQuantizedLinearGlobalScaleMatchesDequantized(t *testing.T) {
 		}
 	})
 }
+
+func TestSwiGLU(t *testing.T) {
+	for _, tt := range []struct {
+		name       string
+		quantized  bool
+		withScales bool
+		withBias   bool
+	}{
+		{name: "plain"},
+		{name: "quantized", quantized: true},
+		{name: "quantized with global scales", quantized: true, withScales: true},
+		{name: "quantized with global scales and bias", quantized: true, withScales: true, withBias: true},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			mlxtest.Run(t, func(t *mlxtest.T) {
+				if tt.quantized && !mlx.MetalIsAvailable() && !mlx.CUDAIsAvailable() {
+					t.Skip("nvfp4 quantized_matmul requires a GPU backend")
+				}
+				const width, group = 64, 16
+
+				weightValues := make([]float32, width*width)
+				for i := range weightValues {
+					weightValues[i] = float32((i%23)-11) * 0.011
+				}
+				weight := mlx.FromValues(weightValues, width, width).AsType(mlx.DTypeBFloat16)
+				var bias *mlx.Array
+				if tt.withBias {
+					biasValues := make([]float32, width)
+					for i := range biasValues {
+						biasValues[i] = float32(i%5-2) / 4
+					}
+					bias = mlx.FromValues(biasValues, width).AsType(mlx.DTypeBFloat16)
+				}
+
+				var gate, up LinearLayer
+				if tt.quantized {
+					gateQuantized := NewQuantizedLinear(weight, bias, group, 4, "nvfp4")
+					upQuantized := NewQuantizedLinear(weight, bias, group, 4, "nvfp4")
+					if tt.withScales {
+						gateQuantized.GlobalScale = mlx.FromValue(float32(mlx.Nvfp4MaxProduct) * 0.75)
+						upQuantized.GlobalScale = mlx.FromValue(float32(mlx.Nvfp4MaxProduct) * 1.25)
+					}
+					gate, up = gateQuantized, upQuantized
+				} else {
+					gate, up = NewLinear(weight, bias), NewLinear(weight, bias)
+				}
+
+				xValues := make([]float32, width)
+				for i := range xValues {
+					xValues[i] = float32(i%7-3) / 8
+				}
+				x := mlx.FromValues(xValues, 1, 1, width).AsType(mlx.DTypeBFloat16)
+				got := SwiGLU(gate, up, x).AsType(mlx.DTypeFloat32)
+				want := mlx.SwiGLU(gate.Forward(x), up.Forward(x)).AsType(mlx.DTypeFloat32)
+				mlx.Eval(got, want)
+				if gotValues, wantValues := got.Floats(), want.Floats(); !slices.Equal(gotValues, wantValues) {
+					t.Fatalf("SwiGLU() = %v, want %v", gotValues, wantValues)
+				}
+			})
+		})
+	}
+}