diff --git a/mlxrunner/nn/conv.go b/mlxrunner/nn/conv.go new file mode 100644 index 000000000..8a4da9c00 --- /dev/null +++ b/mlxrunner/nn/conv.go @@ -0,0 +1,37 @@ +package nn + +import "github.com/ollama/ollama/mlx" + +// Conv1d applies 1D convolution over NLC input. +type Conv1d struct { + Weight *mlx.Array + Bias *mlx.Array + Stride int32 + Padding int32 + Dilation int32 + Groups int32 +} + +func NewConv1d(weight, bias *mlx.Array, stride, padding, dilation, groups int32) *Conv1d { + if stride <= 0 { + stride = 1 + } + if dilation <= 0 { + dilation = 1 + } + if groups <= 0 { + groups = 1 + } + return &Conv1d{ + Weight: weight, + Bias: bias, + Stride: stride, + Padding: padding, + Dilation: dilation, + Groups: groups, + } +} + +func (c *Conv1d) Forward(x *mlx.Array) *mlx.Array { + return mlx.Conv1d(x, c.Weight, c.Bias, c.Stride, c.Padding, c.Dilation, c.Groups) +} diff --git a/mlxrunner/nn/embedding.go b/mlxrunner/nn/embedding.go new file mode 100644 index 000000000..4b382062c --- /dev/null +++ b/mlxrunner/nn/embedding.go @@ -0,0 +1,61 @@ +package nn + +import "github.com/ollama/ollama/mlx" + +// EmbeddingLayer is an interface for embedding layers that can also expose a +// tied-output projection when the model reuses embedding weights as the LM head. +type EmbeddingLayer interface { + Forward(indices *mlx.Array) *mlx.Array + AsLinear() LinearLayer +} + +// Embedding represents an embedding layer. +type Embedding struct { + Weight *mlx.Array +} + +func NewEmbedding(weight *mlx.Array) *Embedding { + return &Embedding{Weight: weight} +} + +func (e *Embedding) Forward(indices *mlx.Array) *mlx.Array { + return e.Weight.TakeAxis(indices, 0) +} + +func (e *Embedding) AsLinear() LinearLayer { + return NewLinear(e.Weight, nil) +} + +// QuantizedEmbedding performs row-wise embedding lookup from affine/nvfp4/etc. +// packed weights and dequantizes only the selected rows. +type QuantizedEmbedding struct { + Weight *mlx.Array + Scales *mlx.Array + QBiases *mlx.Array + GlobalScale *mlx.Array // Per-tensor global scale for double-scale nvfp4 (nil for standard) + GroupSize int + Bits int + Mode string +} + +func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { + weight := qe.Weight.TakeAxis(indices, 0) + scales := qe.Scales.TakeAxis(indices, 0) + var qbiases *mlx.Array + if qe.QBiases != nil { + qbiases = qe.QBiases.TakeAxis(indices, 0) + } + return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode, qe.GlobalScale) +} + +func (qe *QuantizedEmbedding) AsLinear() LinearLayer { + return &QuantizedLinear{ + Weight: qe.Weight, + Scales: qe.Scales, + QBiases: qe.QBiases, + GlobalScale: qe.GlobalScale, + GroupSize: qe.GroupSize, + Bits: qe.Bits, + Mode: qe.Mode, + } +} diff --git a/mlxrunner/nn/embedding_test.go b/mlxrunner/nn/embedding_test.go new file mode 100644 index 000000000..c5eb98465 --- /dev/null +++ b/mlxrunner/nn/embedding_test.go @@ -0,0 +1,44 @@ +package nn + +import ( + "testing" + + "github.com/ollama/ollama/mlx" +) + +func TestQuantizedEmbeddingAsLinearPreservesGlobalScale(t *testing.T) { + weight := &mlx.Array{} + scales := &mlx.Array{} + qbiases := &mlx.Array{} + globalScale := &mlx.Array{} + + embedding := &QuantizedEmbedding{ + Weight: weight, + Scales: scales, + QBiases: qbiases, + GlobalScale: globalScale, + GroupSize: 16, + Bits: 4, + Mode: "nvfp4", + } + + linear, ok := embedding.AsLinear().(*QuantizedLinear) + if !ok { + t.Fatalf("AsLinear type = %T, want *QuantizedLinear", embedding.AsLinear()) + } + if linear.Weight != weight { + t.Fatalf("AsLinear Weight = %p, want %p", linear.Weight, weight) + } + if linear.Scales != scales { + t.Fatalf("AsLinear Scales = %p, want %p", linear.Scales, scales) + } + if linear.QBiases != qbiases { + t.Fatalf("AsLinear QBiases = %p, want %p", linear.QBiases, qbiases) + } + if linear.GlobalScale != globalScale { + t.Fatalf("AsLinear GlobalScale = %p, want %p", linear.GlobalScale, globalScale) + } + if linear.GroupSize != 16 || linear.Bits != 4 || linear.Mode != "nvfp4" { + t.Fatalf("AsLinear quant params = (%d, %d, %q), want (16, 4, %q)", linear.GroupSize, linear.Bits, linear.Mode, "nvfp4") + } +} diff --git a/mlxrunner/nn/linear.go b/mlxrunner/nn/linear.go new file mode 100644 index 000000000..55f1775c6 --- /dev/null +++ b/mlxrunner/nn/linear.go @@ -0,0 +1,85 @@ +package nn + +import "github.com/ollama/ollama/mlx" + +// LinearLayer is an interface for linear layers (both regular and quantized). +type LinearLayer interface { + Forward(x *mlx.Array) *mlx.Array + OutputDim() int32 +} + +// Linear applies an affine transformation: y = x @ W.T + b +type Linear struct { + Weight *mlx.Array + Bias *mlx.Array +} + +func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { + if bias != nil && bias.DType() != weight.DType() { + bias = bias.AsType(weight.DType()) + } + return &Linear{Weight: weight, Bias: bias} +} + +func (l *Linear) Forward(x *mlx.Array) *mlx.Array { + w := l.Weight.Transpose(1, 0) + if l.Bias != nil { + return l.Bias.Addmm(x, w, 1.0, 1.0) + } + return x.Matmul(w) +} + +func (l *Linear) OutputDim() int32 { + return int32(l.Weight.Dim(0)) +} + +// QuantizedLinear applies an affine transformation using quantized weights. +type QuantizedLinear struct { + Weight *mlx.Array // Quantized weight data + Scales *mlx.Array // Scale factors for dequantization + QBiases *mlx.Array // Quantization biases (nil for nvfp4) + Bias *mlx.Array // Layer bias [output_dims] or nil + GlobalScale *mlx.Array // Per-tensor or per-row global scale for double-scale nvfp4 (nil for standard) + GroupSize int + Bits int + Mode string +} + +func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, mode string) *QuantizedLinear { + qw, scales, qbiases := mlx.Quantize(weight, groupSize, bits, mode) + if qbiases != nil { + mlx.Eval(qw, scales, qbiases) + } else { + mlx.Eval(qw, scales) + } + if bias != nil && bias.DType() != weight.DType() { + bias = bias.AsType(weight.DType()) + } + return &QuantizedLinear{ + Weight: qw, + Scales: scales, + QBiases: qbiases, + Bias: bias, + GroupSize: groupSize, + Bits: bits, + Mode: mode, + } +} + +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) + if ql.Bias != nil { + bias := ql.Bias + if bias.DType() != out.DType() { + bias = bias.AsType(out.DType()) + } + out = out.Add(bias) + } + return out +} + +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 new file mode 100644 index 000000000..76f74086b --- /dev/null +++ b/mlxrunner/nn/linear_test.go @@ -0,0 +1,102 @@ +package nn + +import ( + "math" + "testing" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlx/mlxtest" +) + +func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + weightVals := make([]float32, 3*32) + for i := range weightVals { + weightVals[i] = float32((i%11)-5) / 7 + } + inputVals := make([]float32, 2*32) + for i := range inputVals { + inputVals[i] = float32((i%7)-3) / 5 + } + + weight := mlx.FromValues(weightVals, 3, 32).AsType(mlx.DTypeBFloat16) + input := mlx.FromValues(inputVals, 2, 32).AsType(mlx.DTypeBFloat16) + mlx.Eval(weight, input) + + ql := NewQuantizedLinear(weight, nil, 32, 4, "mxfp4") + if ql.QBiases != nil { + t.Fatalf("mxfp4 qbiases = %v, want nil", ql.QBiases) + } + + dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4", nil) + mlx.Eval(dequantizedWeight) + + qOut := ql.Forward(input).AsType(mlx.DTypeFloat32) + dOut := NewLinear(dequantizedWeight, nil).Forward(input).AsType(mlx.DTypeFloat32) + mlx.Eval(qOut, dOut) + + got := qOut.Floats() + want := dOut.Floats() + if len(got) != len(want) { + t.Fatalf("output length = %d, want %d", len(got), len(want)) + } + + for i := range got { + if !approxEqual(got[i], want[i], 1e-3) { + t.Fatalf("output[%d] = %.6f, want %.6f", i, got[i], want[i]) + } + } + }) +} + +// A dense nvfp4 projection carries the checkpoint's global scale through +// QuantizedMatmul, which applies it to the output in a single fused kernel. +// The dequantized weights are the reference. +func TestQuantizedLinearGlobalScaleMatchesDequantized(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + if !mlx.MetalIsAvailable() && !mlx.CUDAIsAvailable() { + t.Skip("nvfp4 quantized_matmul requires a GPU backend") + } + const rows, cols, group = 64, 64, 16 + + weightValues := make([]float32, rows*cols) + for i := range weightValues { + weightValues[i] = float32((i%23)-11) * 0.011 + } + weight := mlx.FromValues(weightValues, rows, cols).AsType(mlx.DTypeBFloat16) + packed, scales, _ := mlx.Quantize(weight, group, 4, "nvfp4") + mlx.Eval(packed, scales) + + globalScale := mlx.FromValues([]float32{0.375}, 1) + linear := &QuantizedLinear{ + Weight: packed, Scales: scales, GlobalScale: globalScale, + GroupSize: group, Bits: 4, Mode: "nvfp4", + } + + xValues := make([]float32, cols) + for i := range xValues { + xValues[i] = float32(i%7-3) / 8 + } + x := mlx.FromValues(xValues, 1, cols).AsType(mlx.DTypeBFloat16) + + got := linear.Forward(x).AsType(mlx.DTypeFloat32) + dense := mlx.Dequantize(packed, scales, nil, group, 4, "nvfp4", globalScale) + want := mlx.Matmul(x.AsType(mlx.DTypeFloat32), mlx.Transpose(dense.AsType(mlx.DTypeFloat32), 1, 0)) + mlx.Eval(got, want) + + gotValues, wantValues := got.Floats(), want.Floats() + if len(gotValues) != len(wantValues) { + t.Fatalf("output length = %d, want %d", len(gotValues), len(wantValues)) + } + for i := range gotValues { + if math.IsNaN(float64(gotValues[i])) || math.IsInf(float64(gotValues[i]), 0) { + t.Fatalf("output[%d] = %v, want finite", i, gotValues[i]) + } + delta := math.Abs(float64(gotValues[i] - wantValues[i])) + tolerance := 0.02 * math.Max(math.Abs(float64(wantValues[i])), 1) + if delta > tolerance { + t.Fatalf("output[%d] = %v, want %v (delta %v > %v)", i, gotValues[i], wantValues[i], delta, tolerance) + } + } + }) +} diff --git a/mlxrunner/nn/multilinear.go b/mlxrunner/nn/multilinear.go new file mode 100644 index 000000000..1b0e2ed5e --- /dev/null +++ b/mlxrunner/nn/multilinear.go @@ -0,0 +1,18 @@ +package nn + +import "github.com/ollama/ollama/mlx" + +// MultiLinear performs per-head linear projections. +// Weight shape: [num_heads, output_dims, input_dims] +type MultiLinear struct { + Weight *mlx.Array +} + +func NewMultiLinear(weight *mlx.Array) *MultiLinear { + return &MultiLinear{Weight: weight} +} + +func (ml *MultiLinear) Forward(x *mlx.Array) *mlx.Array { + wT := ml.Weight.Transpose(0, 2, 1) + return x.Matmul(wT) +} diff --git a/mlxrunner/nn/nn.go b/mlxrunner/nn/nn.go deleted file mode 100644 index 90f918faf..000000000 --- a/mlxrunner/nn/nn.go +++ /dev/null @@ -1,234 +0,0 @@ -package nn - -import "github.com/ollama/ollama/mlx" - -// Layer is the interface for neural network layers with a Forward method. -type Layer interface { - Forward(x *mlx.Array) *mlx.Array -} - -// LinearLayer is an interface for linear layers (both regular and quantized). -type LinearLayer interface { - Forward(x *mlx.Array) *mlx.Array - OutputDim() int32 -} - -// EmbeddingLayer is an interface for embedding layers that can also expose a -// tied-output projection when the model reuses embedding weights as the LM head. -type EmbeddingLayer interface { - Forward(indices *mlx.Array) *mlx.Array - AsLinear() LinearLayer -} - -// Conv1d applies 1D convolution over NLC input. -type Conv1d struct { - Weight *mlx.Array - Bias *mlx.Array - Stride int32 - Padding int32 - Dilation int32 - Groups int32 -} - -func NewConv1d(weight, bias *mlx.Array, stride, padding, dilation, groups int32) *Conv1d { - if stride <= 0 { - stride = 1 - } - if dilation <= 0 { - dilation = 1 - } - if groups <= 0 { - groups = 1 - } - return &Conv1d{ - Weight: weight, - Bias: bias, - Stride: stride, - Padding: padding, - Dilation: dilation, - Groups: groups, - } -} - -func (c *Conv1d) Forward(x *mlx.Array) *mlx.Array { - return mlx.Conv1d(x, c.Weight, c.Bias, c.Stride, c.Padding, c.Dilation, c.Groups) -} - -// Linear applies an affine transformation: y = x @ W.T + b -type Linear struct { - Weight *mlx.Array - Bias *mlx.Array -} - -func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { - if bias != nil && bias.DType() != weight.DType() { - bias = bias.AsType(weight.DType()) - } - return &Linear{Weight: weight, Bias: bias} -} - -func (l *Linear) Forward(x *mlx.Array) *mlx.Array { - w := l.Weight.Transpose(1, 0) - if l.Bias != nil { - return l.Bias.Addmm(x, w, 1.0, 1.0) - } - return x.Matmul(w) -} - -func (l *Linear) OutputDim() int32 { - return int32(l.Weight.Dim(0)) -} - -// QuantizedLinear applies an affine transformation using quantized weights. -type QuantizedLinear struct { - Weight *mlx.Array // Quantized weight data - Scales *mlx.Array // Scale factors for dequantization - QBiases *mlx.Array // Quantization biases (nil for nvfp4) - Bias *mlx.Array // Layer bias [output_dims] or nil - GlobalScale *mlx.Array // Per-tensor or per-row global scale for double-scale nvfp4 (nil for standard) - GroupSize int - Bits int - Mode string -} - -func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, mode string) *QuantizedLinear { - qw, scales, qbiases := mlx.Quantize(weight, groupSize, bits, mode) - if qbiases != nil { - mlx.Eval(qw, scales, qbiases) - } else { - mlx.Eval(qw, scales) - } - if bias != nil && bias.DType() != weight.DType() { - bias = bias.AsType(weight.DType()) - } - return &QuantizedLinear{ - Weight: qw, - Scales: scales, - QBiases: qbiases, - Bias: bias, - GroupSize: groupSize, - Bits: bits, - Mode: mode, - } -} - -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) - if ql.Bias != nil { - bias := ql.Bias - if bias.DType() != out.DType() { - bias = bias.AsType(out.DType()) - } - out = out.Add(bias) - } - return out -} - -func (ql *QuantizedLinear) OutputDim() int32 { - return int32(ql.Weight.Dim(0)) -} - -// RMSNorm represents an RMS normalization layer. -type RMSNorm struct { - Weight *mlx.Array - Eps float32 -} - -func NewRMSNorm(weight *mlx.Array, eps float32) *RMSNorm { - return &RMSNorm{Weight: weight, Eps: eps} -} - -func (rn *RMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array { - if eps == 0 { - eps = rn.Eps - } - return mlx.RMSNormFn(x, rn.Weight, eps) -} - -// Embedding represents an embedding layer. -type Embedding struct { - Weight *mlx.Array -} - -func NewEmbedding(weight *mlx.Array) *Embedding { - return &Embedding{Weight: weight} -} - -func (e *Embedding) Forward(indices *mlx.Array) *mlx.Array { - return e.Weight.TakeAxis(indices, 0) -} - -func (e *Embedding) AsLinear() LinearLayer { - return NewLinear(e.Weight, nil) -} - -// QuantizedEmbedding performs row-wise embedding lookup from affine/nvfp4/etc. -// packed weights and dequantizes only the selected rows. -type QuantizedEmbedding struct { - Weight *mlx.Array - Scales *mlx.Array - QBiases *mlx.Array - GlobalScale *mlx.Array // Per-tensor global scale for double-scale nvfp4 (nil for standard) - GroupSize int - Bits int - Mode string -} - -func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { - weight := qe.Weight.TakeAxis(indices, 0) - scales := qe.Scales.TakeAxis(indices, 0) - var qbiases *mlx.Array - if qe.QBiases != nil { - qbiases = qe.QBiases.TakeAxis(indices, 0) - } - return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode, qe.GlobalScale) -} - -func (qe *QuantizedEmbedding) AsLinear() LinearLayer { - return &QuantizedLinear{ - Weight: qe.Weight, - Scales: qe.Scales, - QBiases: qe.QBiases, - GlobalScale: qe.GlobalScale, - GroupSize: qe.GroupSize, - Bits: qe.Bits, - Mode: qe.Mode, - } -} - -// LayerNorm represents a standard layer normalization layer (with bias). -type LayerNorm struct { - Weight *mlx.Array - Bias *mlx.Array - Eps float32 -} - -func (ln *LayerNorm) Forward(x *mlx.Array) *mlx.Array { - eps := ln.Eps - if eps == 0 { - eps = 1e-5 - } - return mlx.LayerNormFn(x, ln.Weight, ln.Bias, eps) -} - -// MultiLinearLayer is an interface for per-head linear layers. -type MultiLinearLayer interface { - Forward(x *mlx.Array) *mlx.Array -} - -// MultiLinear performs per-head linear projections. -// Weight shape: [num_heads, output_dims, input_dims] -type MultiLinear struct { - Weight *mlx.Array -} - -func NewMultiLinear(weight *mlx.Array) *MultiLinear { - return &MultiLinear{Weight: weight} -} - -func (ml *MultiLinear) Forward(x *mlx.Array) *mlx.Array { - wT := ml.Weight.Transpose(0, 2, 1) - return x.Matmul(wT) -} diff --git a/mlxrunner/nn/nn_test.go b/mlxrunner/nn/nn_test.go deleted file mode 100644 index c31cb2db6..000000000 --- a/mlxrunner/nn/nn_test.go +++ /dev/null @@ -1,270 +0,0 @@ -package nn - -import ( - "math" - "testing" - - "github.com/ollama/ollama/mlx" - "github.com/ollama/ollama/mlx/mlxtest" -) - -func approxEqual(a, b, tol float32) bool { - return float32(math.Abs(float64(a-b))) < tol -} - -// TestLayerNormNoBias verifies LayerNorm without bias against manual computation. -func TestLayerNormNoBias(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - // Input: [1, 4] — single row, 4 features - x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) - weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4) - mlx.Eval(x, weight) - - ln := &LayerNorm{Weight: weight, Eps: 1e-5} - out := ln.Forward(x) - mlx.Eval(out) - - data := out.Floats() - if len(data) != 4 { - t.Fatalf("expected 4 values, got %d", len(data)) - } - - // Manual LayerNorm: mean=2.5, var=1.25, std=sqrt(1.25+1e-5) - // normalized = (x - mean) / std - mean := float32(2.5) - variance := float32(1.25) - std := float32(math.Sqrt(float64(variance + 1e-5))) - for i, v := range []float32{1, 2, 3, 4} { - expected := (v - mean) / std - if !approxEqual(data[i], expected, 1e-4) { - t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i]) - } - } - }) -} - -// TestLayerNormWithBias verifies LayerNorm with weight and bias. -func TestLayerNormWithBias(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) - weight := mlx.FromValues([]float32{2, 2, 2, 2}, 4) - bias := mlx.FromValues([]float32{10, 20, 30, 40}, 4) - mlx.Eval(x, weight, bias) - - ln := &LayerNorm{Weight: weight, Bias: bias, Eps: 1e-5} - out := ln.Forward(x) - mlx.Eval(out) - - data := out.Floats() - if len(data) != 4 { - t.Fatalf("expected 4 values, got %d", len(data)) - } - - mean := float32(2.5) - variance := float32(1.25) - std := float32(math.Sqrt(float64(variance + 1e-5))) - biases := []float32{10, 20, 30, 40} - for i, v := range []float32{1, 2, 3, 4} { - expected := ((v-mean)/std)*2 + biases[i] - if !approxEqual(data[i], expected, 1e-4) { - t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i]) - } - } - }) -} - -// TestLayerNormBatched verifies LayerNorm normalizes each row independently. -func TestLayerNormBatched(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - // Input: [2, 3] — two rows - x := mlx.FromValues([]float32{ - 1, 2, 3, - 10, 20, 30, - }, 2, 3) - weight := mlx.FromValues([]float32{1, 1, 1}, 3) - mlx.Eval(x, weight) - - ln := &LayerNorm{Weight: weight, Eps: 1e-5} - out := ln.Forward(x) - mlx.Eval(out) - - data := out.Floats() - if len(data) != 6 { - t.Fatalf("expected 6 values, got %d", len(data)) - } - - // Each row should be independently normalized. - // Row 0: [1,2,3] mean=2, var=2/3 - // Row 1: [10,20,30] mean=20, var=200/3 - // After normalization both rows should have the same pattern - // since [10,20,30] = 10*[1,2,3], the normalized values are identical. - for i := range 3 { - if !approxEqual(data[i], data[i+3], 1e-4) { - t.Errorf("row 0 elem %d (%.6f) != row 1 elem %d (%.6f); expected identical normalized values", - i, data[i], i, data[i+3]) - } - } - - // Verify the normalized values sum to ~0 (mean-centered) - sum := data[0] + data[1] + data[2] - if !approxEqual(sum, 0, 1e-4) { - t.Errorf("normalized row sum should be ~0, got %.6f", sum) - } - }) -} - -// TestLayerNormDefaultEps verifies the default epsilon of 1e-5 is used when Eps is 0. -func TestLayerNormDefaultEps(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) - weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4) - mlx.Eval(x, weight) - - // Eps=0 should use default 1e-5 - ln0 := &LayerNorm{Weight: weight, Eps: 0} - out0 := ln0.Forward(x) - mlx.Eval(out0) - - lnExplicit := &LayerNorm{Weight: weight, Eps: 1e-5} - outExplicit := lnExplicit.Forward(x) - mlx.Eval(outExplicit) - - d0 := out0.Floats() - dE := outExplicit.Floats() - for i := range d0 { - if !approxEqual(d0[i], dE[i], 1e-6) { - t.Errorf("index %d: Eps=0 gave %.6f, Eps=1e-5 gave %.6f", i, d0[i], dE[i]) - } - } - }) -} - -func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - weightVals := make([]float32, 3*32) - for i := range weightVals { - weightVals[i] = float32((i%11)-5) / 7 - } - inputVals := make([]float32, 2*32) - for i := range inputVals { - inputVals[i] = float32((i%7)-3) / 5 - } - - weight := mlx.FromValues(weightVals, 3, 32).AsType(mlx.DTypeBFloat16) - input := mlx.FromValues(inputVals, 2, 32).AsType(mlx.DTypeBFloat16) - mlx.Eval(weight, input) - - ql := NewQuantizedLinear(weight, nil, 32, 4, "mxfp4") - if ql.QBiases != nil { - t.Fatalf("mxfp4 qbiases = %v, want nil", ql.QBiases) - } - - dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4", nil) - mlx.Eval(dequantizedWeight) - - qOut := ql.Forward(input).AsType(mlx.DTypeFloat32) - dOut := NewLinear(dequantizedWeight, nil).Forward(input).AsType(mlx.DTypeFloat32) - mlx.Eval(qOut, dOut) - - got := qOut.Floats() - want := dOut.Floats() - if len(got) != len(want) { - t.Fatalf("output length = %d, want %d", len(got), len(want)) - } - - for i := range got { - if !approxEqual(got[i], want[i], 1e-3) { - t.Fatalf("output[%d] = %.6f, want %.6f", i, got[i], want[i]) - } - } - }) -} - -func TestQuantizedEmbeddingAsLinearPreservesGlobalScale(t *testing.T) { - weight := &mlx.Array{} - scales := &mlx.Array{} - qbiases := &mlx.Array{} - globalScale := &mlx.Array{} - - embedding := &QuantizedEmbedding{ - Weight: weight, - Scales: scales, - QBiases: qbiases, - GlobalScale: globalScale, - GroupSize: 16, - Bits: 4, - Mode: "nvfp4", - } - - linear, ok := embedding.AsLinear().(*QuantizedLinear) - if !ok { - t.Fatalf("AsLinear type = %T, want *QuantizedLinear", embedding.AsLinear()) - } - if linear.Weight != weight { - t.Fatalf("AsLinear Weight = %p, want %p", linear.Weight, weight) - } - if linear.Scales != scales { - t.Fatalf("AsLinear Scales = %p, want %p", linear.Scales, scales) - } - if linear.QBiases != qbiases { - t.Fatalf("AsLinear QBiases = %p, want %p", linear.QBiases, qbiases) - } - if linear.GlobalScale != globalScale { - t.Fatalf("AsLinear GlobalScale = %p, want %p", linear.GlobalScale, globalScale) - } - if linear.GroupSize != 16 || linear.Bits != 4 || linear.Mode != "nvfp4" { - t.Fatalf("AsLinear quant params = (%d, %d, %q), want (16, 4, %q)", linear.GroupSize, linear.Bits, linear.Mode, "nvfp4") - } -} - -// A dense nvfp4 projection carries the checkpoint's global scale through -// QuantizedMatmul, which applies it to the output in a single fused kernel. -// The dequantized weights are the reference. -func TestQuantizedLinearGlobalScaleMatchesDequantized(t *testing.T) { - mlxtest.Run(t, func(t *mlxtest.T) { - if !mlx.MetalIsAvailable() && !mlx.CUDAIsAvailable() { - t.Skip("nvfp4 quantized_matmul requires a GPU backend") - } - const rows, cols, group = 64, 64, 16 - - weightValues := make([]float32, rows*cols) - for i := range weightValues { - weightValues[i] = float32((i%23)-11) * 0.011 - } - weight := mlx.FromValues(weightValues, rows, cols).AsType(mlx.DTypeBFloat16) - packed, scales, _ := mlx.Quantize(weight, group, 4, "nvfp4") - mlx.Eval(packed, scales) - - globalScale := mlx.FromValues([]float32{0.375}, 1) - linear := &QuantizedLinear{ - Weight: packed, Scales: scales, GlobalScale: globalScale, - GroupSize: group, Bits: 4, Mode: "nvfp4", - } - - xValues := make([]float32, cols) - for i := range xValues { - xValues[i] = float32(i%7-3) / 8 - } - x := mlx.FromValues(xValues, 1, cols).AsType(mlx.DTypeBFloat16) - - got := linear.Forward(x).AsType(mlx.DTypeFloat32) - dense := mlx.Dequantize(packed, scales, nil, group, 4, "nvfp4", globalScale) - want := mlx.Matmul(x.AsType(mlx.DTypeFloat32), mlx.Transpose(dense.AsType(mlx.DTypeFloat32), 1, 0)) - mlx.Eval(got, want) - - gotValues, wantValues := got.Floats(), want.Floats() - if len(gotValues) != len(wantValues) { - t.Fatalf("output length = %d, want %d", len(gotValues), len(wantValues)) - } - for i := range gotValues { - if math.IsNaN(float64(gotValues[i])) || math.IsInf(float64(gotValues[i]), 0) { - t.Fatalf("output[%d] = %v, want finite", i, gotValues[i]) - } - delta := math.Abs(float64(gotValues[i] - wantValues[i])) - tolerance := 0.02 * math.Max(math.Abs(float64(wantValues[i])), 1) - if delta > tolerance { - t.Fatalf("output[%d] = %v, want %v (delta %v > %v)", i, gotValues[i], wantValues[i], delta, tolerance) - } - } - }) -} diff --git a/mlxrunner/nn/norm.go b/mlxrunner/nn/norm.go new file mode 100644 index 000000000..126c447b7 --- /dev/null +++ b/mlxrunner/nn/norm.go @@ -0,0 +1,35 @@ +package nn + +import "github.com/ollama/ollama/mlx" + +// RMSNorm represents an RMS normalization layer. +type RMSNorm struct { + Weight *mlx.Array + Eps float32 +} + +func NewRMSNorm(weight *mlx.Array, eps float32) *RMSNorm { + return &RMSNorm{Weight: weight, Eps: eps} +} + +func (rn *RMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array { + if eps == 0 { + eps = rn.Eps + } + return mlx.RMSNormFn(x, rn.Weight, eps) +} + +// LayerNorm represents a standard layer normalization layer (with bias). +type LayerNorm struct { + Weight *mlx.Array + Bias *mlx.Array + Eps float32 +} + +func (ln *LayerNorm) Forward(x *mlx.Array) *mlx.Array { + eps := ln.Eps + if eps == 0 { + eps = 1e-5 + } + return mlx.LayerNormFn(x, ln.Weight, ln.Bias, eps) +} diff --git a/mlxrunner/nn/norm_test.go b/mlxrunner/nn/norm_test.go new file mode 100644 index 000000000..40c37b2ce --- /dev/null +++ b/mlxrunner/nn/norm_test.go @@ -0,0 +1,140 @@ +package nn + +import ( + "math" + "testing" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlx/mlxtest" +) + +func approxEqual(a, b, tol float32) bool { + return float32(math.Abs(float64(a-b))) < tol +} + +// TestLayerNormNoBias verifies LayerNorm without bias against manual computation. +func TestLayerNormNoBias(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + // Input: [1, 4] — single row, 4 features + x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) + weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4) + mlx.Eval(x, weight) + + ln := &LayerNorm{Weight: weight, Eps: 1e-5} + out := ln.Forward(x) + mlx.Eval(out) + + data := out.Floats() + if len(data) != 4 { + t.Fatalf("expected 4 values, got %d", len(data)) + } + + // Manual LayerNorm: mean=2.5, var=1.25, std=sqrt(1.25+1e-5) + // normalized = (x - mean) / std + mean := float32(2.5) + variance := float32(1.25) + std := float32(math.Sqrt(float64(variance + 1e-5))) + for i, v := range []float32{1, 2, 3, 4} { + expected := (v - mean) / std + if !approxEqual(data[i], expected, 1e-4) { + t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i]) + } + } + }) +} + +// TestLayerNormWithBias verifies LayerNorm with weight and bias. +func TestLayerNormWithBias(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) + weight := mlx.FromValues([]float32{2, 2, 2, 2}, 4) + bias := mlx.FromValues([]float32{10, 20, 30, 40}, 4) + mlx.Eval(x, weight, bias) + + ln := &LayerNorm{Weight: weight, Bias: bias, Eps: 1e-5} + out := ln.Forward(x) + mlx.Eval(out) + + data := out.Floats() + if len(data) != 4 { + t.Fatalf("expected 4 values, got %d", len(data)) + } + + mean := float32(2.5) + variance := float32(1.25) + std := float32(math.Sqrt(float64(variance + 1e-5))) + biases := []float32{10, 20, 30, 40} + for i, v := range []float32{1, 2, 3, 4} { + expected := ((v-mean)/std)*2 + biases[i] + if !approxEqual(data[i], expected, 1e-4) { + t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i]) + } + } + }) +} + +// TestLayerNormBatched verifies LayerNorm normalizes each row independently. +func TestLayerNormBatched(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + // Input: [2, 3] — two rows + x := mlx.FromValues([]float32{ + 1, 2, 3, + 10, 20, 30, + }, 2, 3) + weight := mlx.FromValues([]float32{1, 1, 1}, 3) + mlx.Eval(x, weight) + + ln := &LayerNorm{Weight: weight, Eps: 1e-5} + out := ln.Forward(x) + mlx.Eval(out) + + data := out.Floats() + if len(data) != 6 { + t.Fatalf("expected 6 values, got %d", len(data)) + } + + // Each row should be independently normalized. + // Row 0: [1,2,3] mean=2, var=2/3 + // Row 1: [10,20,30] mean=20, var=200/3 + // After normalization both rows should have the same pattern + // since [10,20,30] = 10*[1,2,3], the normalized values are identical. + for i := range 3 { + if !approxEqual(data[i], data[i+3], 1e-4) { + t.Errorf("row 0 elem %d (%.6f) != row 1 elem %d (%.6f); expected identical normalized values", + i, data[i], i, data[i+3]) + } + } + + // Verify the normalized values sum to ~0 (mean-centered) + sum := data[0] + data[1] + data[2] + if !approxEqual(sum, 0, 1e-4) { + t.Errorf("normalized row sum should be ~0, got %.6f", sum) + } + }) +} + +// TestLayerNormDefaultEps verifies the default epsilon of 1e-5 is used when Eps is 0. +func TestLayerNormDefaultEps(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4) + weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4) + mlx.Eval(x, weight) + + // Eps=0 should use default 1e-5 + ln0 := &LayerNorm{Weight: weight, Eps: 0} + out0 := ln0.Forward(x) + mlx.Eval(out0) + + lnExplicit := &LayerNorm{Weight: weight, Eps: 1e-5} + outExplicit := lnExplicit.Forward(x) + mlx.Eval(outExplicit) + + d0 := out0.Floats() + dE := outExplicit.Floats() + for i := range d0 { + if !approxEqual(d0[i], dE[i], 1e-6) { + t.Errorf("index %d: Eps=0 gave %.6f, Eps=1e-5 gave %.6f", i, d0[i], dE[i]) + } + } + }) +}