upstream-archive/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-09-21 13:38:14 -05:00
Code Issues Packages Projects Releases Wiki Activity

mlxrunner: lay out nn one layer kind per file

Browse Source 
nn.go held every layer type in the package apart from attention,
recurrence and rope: the Linear and Embedding interfaces with their dense
and quantized types, Conv1d, RMSNorm, LayerNorm and MultiLinear, with one
test file to match. Finding a layer meant scanning the file named after
the package.

Each layer kind gets its own file: linear.go and embedding.go hold the
interface and the dense and quantized types, conv.go, norm.go and
multilinear.go take the rest, and nn_test.go splits the same way. The
Layer and MultiLinearLayer interfaces go; nothing implemented or accepted
them. No code changes otherwise.
This commit is contained in:
Jesse Gross
2026-09-16 14:06:08 -07:00
parent d27fde67ea
commit 7343e22764
10 changed files with 522 additions and 504 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlxrunner/nn/conv.go
+37
View File
@@ -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)
}
mlxrunner/nn/embedding.go
+61
View File
@@ -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,
}
}
mlxrunner/nn/embedding_test.go
+44
View File
@@ -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")
}
}
mlxrunner/nn/linear.go
+85
View File
@@ -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))
}
mlxrunner/nn/linear_test.go
+102
View File
@@ -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)
}
}
})
}
mlxrunner/nn/multilinear.go
+18
View File
@@ -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)
}
mlxrunner/nn/nn.go
-234
View File
@@ -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)
}
mlxrunner/nn/nn_test.go
-270
View File
@@ -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)
}
}
})
}
mlxrunner/nn/norm.go
+35
View File
@@ -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)
}
mlxrunner/nn/norm_test.go
+140
View File
@@ -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])
}
}
})
}