diff --git a/mlx/act.go b/mlx/act.go index 267e9ce63..b5cfb8f22 100644 --- a/mlx/act.go +++ b/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( diff --git a/mlx/act_test.go b/mlx/act_test.go index 68539496c..8173aba94 100644 --- a/mlx/act_test.go +++ b/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) { diff --git a/mlx/gated_delta.go b/mlx/gated_delta.go index c4854a885..09cb05b7d 100644 --- a/mlx/gated_delta.go +++ b/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 diff --git a/mlx/gated_delta_test.go b/mlx/gated_delta_test.go index dc3758d8d..eea40882f 100644 --- a/mlx/gated_delta_test.go +++ b/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) diff --git a/mlx/ops_extra.go b/mlx/ops_extra.go index 63b039e40..9aff32488 100644 --- a/mlx/ops_extra.go +++ b/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(), ) diff --git a/mlxrunner/model/qwen3_5/qwen3_5.go b/mlxrunner/model/qwen3_5/qwen3_5.go index 81dc7a35c..8a85e734f 100644 --- a/mlxrunner/model/qwen3_5/qwen3_5.go +++ b/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 { diff --git a/mlxrunner/nn/linear.go b/mlxrunner/nn/linear.go index 55f1775c6..2ad249200 100644 --- a/mlxrunner/nn/linear.go +++ b/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)) } diff --git a/mlxrunner/nn/linear_test.go b/mlxrunner/nn/linear_test.go index 76f74086b..de00068e3 100644 --- a/mlxrunner/nn/linear_test.go +++ b/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) + } + }) + }) + } +}