ollama source for Momentry Core verification
This commit is contained in:
Accusys
258
x/imagegen/nn/nn.go
Normal file
258
x/imagegen/nn/nn.go
Normal file
@@ -0,0 +1,258 @@
|
||||
// Package nn provides neural network layer types.
|
||||
package nn
|
||||
|
||||
import "github.com/ollama/ollama/x/imagegen/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).
|
||||
// This allows swapping between Linear and QuantizedLinear at runtime.
|
||||
type LinearLayer interface {
|
||||
Forward(x *mlx.Array) *mlx.Array
|
||||
OutputDim() int32 // Returns the output dimension of the layer
|
||||
}
|
||||
|
||||
// Linear applies an affine transformation: y = x @ W.T + b
|
||||
// Weight is stored as [out_features, in_features], matching PyTorch/MLX convention.
|
||||
type Linear struct {
|
||||
Weight *mlx.Array `weight:"weight"` // [out_features, in_features]
|
||||
Bias *mlx.Array `weight:"bias,optional"` // [out_features] or nil
|
||||
}
|
||||
|
||||
// NewLinear creates a linear layer.
|
||||
// Weight should be [out_features, in_features].
|
||||
func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear {
|
||||
return &Linear{Weight: weight, Bias: bias}
|
||||
}
|
||||
|
||||
// NewQuantizedLinear creates a quantized linear layer directly from bf16 weights.
|
||||
// Quantizes the weight immediately and evaluates to break lazy dependencies.
|
||||
// Note: For modes like "nvfp4", qbiases will be nil.
|
||||
func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, mode string) *QuantizedLinear {
|
||||
qw, scales, qbiases := mlx.Quantize(weight, groupSize, bits, mode)
|
||||
// Eval immediately so bf16 weight can be freed
|
||||
// Handle modes that don't return qbiases (e.g., nvfp4)
|
||||
if qbiases != nil {
|
||||
mlx.Eval(qw, scales, qbiases)
|
||||
} else {
|
||||
mlx.Eval(qw, scales)
|
||||
}
|
||||
return &QuantizedLinear{
|
||||
Weight: qw,
|
||||
Scales: scales,
|
||||
QBiases: qbiases,
|
||||
Bias: bias,
|
||||
GroupSize: groupSize,
|
||||
Bits: bits,
|
||||
Mode: mode,
|
||||
}
|
||||
}
|
||||
|
||||
// Forward applies the linear transformation: x @ W.T + bias
|
||||
func (l *Linear) Forward(x *mlx.Array) *mlx.Array {
|
||||
w := mlx.Transpose(l.Weight, 1, 0)
|
||||
if l.Bias != nil {
|
||||
return mlx.AddMM(l.Bias, x, w, 1.0, 1.0)
|
||||
}
|
||||
return mlx.Linear(x, w)
|
||||
}
|
||||
|
||||
// OutputDim returns the output dimension of the linear layer.
|
||||
func (l *Linear) OutputDim() int32 {
|
||||
return l.Weight.Shape()[0]
|
||||
}
|
||||
|
||||
// ToQuantized converts this Linear to a QuantizedLinear.
|
||||
func (l *Linear) ToQuantized(groupSize, bits int, mode string) *QuantizedLinear {
|
||||
qw, scales, qbiases := mlx.Quantize(l.Weight, groupSize, bits, mode)
|
||||
return &QuantizedLinear{
|
||||
Weight: qw,
|
||||
Scales: scales,
|
||||
QBiases: qbiases,
|
||||
Bias: l.Bias,
|
||||
GroupSize: groupSize,
|
||||
Bits: bits,
|
||||
Mode: mode,
|
||||
}
|
||||
}
|
||||
|
||||
// QuantizedLinear applies an affine transformation using quantized weights.
|
||||
// Equivalent to mlx.nn.QuantizedLinear.
|
||||
// Supports multiple quantization modes:
|
||||
// - "affine": scale + zero-point bias (QBiases required)
|
||||
// - "nvfp4": NVIDIA FP4 with E4M3 scales (QBiases nil)
|
||||
type QuantizedLinear struct {
|
||||
Weight *mlx.Array // Quantized weight data
|
||||
Scales *mlx.Array // Scale factors for dequantization
|
||||
QBiases *mlx.Array // Quantization biases (NOT layer bias), nil for nvfp4
|
||||
Bias *mlx.Array // Layer bias [output_dims] or nil
|
||||
GroupSize int
|
||||
Bits int
|
||||
Mode string
|
||||
}
|
||||
|
||||
// Forward applies the quantized linear transformation.
|
||||
func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array {
|
||||
out := mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
|
||||
if ql.Bias != nil {
|
||||
out = mlx.Add(out, ql.Bias)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// OutputDim returns the output dimension of the quantized linear layer.
|
||||
// For mxfp8/mxfp4, quantized weight shape is [out_features, in_features / group_size].
|
||||
// The output dimension is the first dimension of the weight.
|
||||
func (ql *QuantizedLinear) OutputDim() int32 {
|
||||
return ql.Weight.Shape()[0]
|
||||
}
|
||||
|
||||
// RMSNorm represents an RMS normalization layer.
|
||||
type RMSNorm struct {
|
||||
Weight *mlx.Array `weight:"weight"`
|
||||
Eps float32 // optional: used if Forward called with eps=0
|
||||
}
|
||||
|
||||
// NewRMSNorm creates an RMSNorm layer (for models not using weight loader).
|
||||
func NewRMSNorm(weight *mlx.Array, eps float32) *RMSNorm {
|
||||
return &RMSNorm{Weight: weight, Eps: eps}
|
||||
}
|
||||
|
||||
// Forward applies RMS normalization. If eps=0, uses stored Eps.
|
||||
func (rn *RMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array {
|
||||
if eps == 0 {
|
||||
eps = rn.Eps
|
||||
}
|
||||
return mlx.RMSNorm(x, rn.Weight, eps)
|
||||
}
|
||||
|
||||
// Embedding represents an embedding layer.
|
||||
type Embedding struct {
|
||||
Weight *mlx.Array `weight:"weight"`
|
||||
}
|
||||
|
||||
// NewEmbedding creates an embedding layer.
|
||||
func NewEmbedding(weight *mlx.Array) *Embedding {
|
||||
return &Embedding{Weight: weight}
|
||||
}
|
||||
|
||||
// Forward looks up embeddings by indices.
|
||||
func (e *Embedding) Forward(indices *mlx.Array) *mlx.Array {
|
||||
return mlx.Take(e.Weight, indices, 0)
|
||||
}
|
||||
|
||||
// RepeatKV repeats K/V tensors for grouped query attention
|
||||
// x: [B, num_kv_heads, S, head_dim] -> [B, num_heads, S, head_dim]
|
||||
func RepeatKV(x *mlx.Array, repeatFactor int32) *mlx.Array {
|
||||
if repeatFactor == 1 {
|
||||
return x
|
||||
}
|
||||
shape := x.Shape()
|
||||
// [B, num_kv_heads, S, head_dim] -> [B, num_kv_heads, 1, S, head_dim]
|
||||
x = mlx.ExpandDims(x, 2)
|
||||
// Repeat along the new axis
|
||||
reps := []int32{1, 1, repeatFactor, 1, 1}
|
||||
x = mlx.Tile(x, reps)
|
||||
// Reshape: [B, num_kv_heads, repeat, S, head_dim] -> [B, num_kv_heads * repeat, S, head_dim]
|
||||
return mlx.Reshape(x, shape[0], shape[1]*repeatFactor, shape[2], shape[3])
|
||||
}
|
||||
|
||||
// ApplyCausalMask applies causal (lower triangular) mask to attention scores
|
||||
func ApplyCausalMask(scores *mlx.Array) *mlx.Array {
|
||||
// scores: [B, num_heads, S, S]
|
||||
shape := scores.Shape()
|
||||
seqLen := shape[2]
|
||||
|
||||
// Create causal mask: 1 for positions to keep, 0 for positions to mask
|
||||
mask := mlx.Tri(seqLen, seqLen, 0)
|
||||
|
||||
// Where mask is 0, set score to -inf
|
||||
negInf := mlx.NewScalarArray(float32(-1e9))
|
||||
|
||||
// Broadcast mask to match scores shape
|
||||
mask = mlx.ExpandDims(mlx.ExpandDims(mask, 0), 0) // [1, 1, S, S]
|
||||
|
||||
// Use where: if mask > 0, keep scores, else -inf
|
||||
return mlx.Where(mask, scores, negInf)
|
||||
}
|
||||
|
||||
// ApplyCausalMaskWithOffset applies causal mask for cached attention
|
||||
// scores: [B, num_heads, queryLen, keyLen] where keyLen = cacheLen + queryLen
|
||||
// offset: the starting position of the new queries (i.e., cache length)
|
||||
func ApplyCausalMaskWithOffset(scores *mlx.Array, offset int32) *mlx.Array {
|
||||
if offset == 0 {
|
||||
return ApplyCausalMask(scores)
|
||||
}
|
||||
|
||||
shape := scores.Shape()
|
||||
queryLen := shape[2]
|
||||
keyLen := shape[3]
|
||||
|
||||
// For cached attention, new queries can attend to all cached keys plus
|
||||
// new keys up to and including their position.
|
||||
mask := mlx.Tri(queryLen, keyLen, int(offset))
|
||||
|
||||
negInf := mlx.NewScalarArray(float32(-1e9))
|
||||
mask = mlx.ExpandDims(mlx.ExpandDims(mask, 0), 0) // [1, 1, queryLen, keyLen]
|
||||
|
||||
return mlx.Where(mask, scores, negInf)
|
||||
}
|
||||
|
||||
// LayerNorm represents a standard layer normalization layer (with bias).
|
||||
type LayerNorm struct {
|
||||
Weight *mlx.Array `weight:"weight"`
|
||||
Bias *mlx.Array `weight:"bias"`
|
||||
Eps float32
|
||||
}
|
||||
|
||||
// Forward applies layer normalization: (x - mean) / sqrt(var + eps) * weight + bias
|
||||
func (ln *LayerNorm) Forward(x *mlx.Array) *mlx.Array {
|
||||
eps := ln.Eps
|
||||
if eps == 0 {
|
||||
eps = 1e-5
|
||||
}
|
||||
// Compute mean and variance along last dimension
|
||||
mean := mlx.Mean(x, -1, true)
|
||||
centered := mlx.Sub(x, mean)
|
||||
variance := mlx.Mean(mlx.Mul(centered, centered), -1, true)
|
||||
normalized := mlx.Mul(centered, mlx.RSqrt(mlx.AddScalar(variance, eps)))
|
||||
|
||||
// Scale and shift
|
||||
out := mlx.Mul(normalized, ln.Weight)
|
||||
if ln.Bias != nil {
|
||||
out = mlx.Add(out, ln.Bias)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// MultiLinearLayer is an interface for per-head linear layers.
|
||||
// This allows swapping between MultiLinear (bf16) and pre-dequantized weights.
|
||||
type MultiLinearLayer interface {
|
||||
Forward(x *mlx.Array) *mlx.Array
|
||||
}
|
||||
|
||||
// MultiLinear performs per-head linear projections.
|
||||
// Weight shape: [num_heads, output_dims, input_dims]
|
||||
// Input shape: [B, num_heads, L, input_dims]
|
||||
// Output shape: [B, num_heads, L, output_dims]
|
||||
type MultiLinear struct {
|
||||
Weight *mlx.Array `weight:"weight"`
|
||||
}
|
||||
|
||||
// NewMultiLinear creates a MultiLinear layer with the given weight.
|
||||
func NewMultiLinear(weight *mlx.Array) *MultiLinear {
|
||||
return &MultiLinear{Weight: weight}
|
||||
}
|
||||
|
||||
// Forward applies per-head linear transformation: x @ weight.T per head via broadcasting.
|
||||
func (ml *MultiLinear) Forward(x *mlx.Array) *mlx.Array {
|
||||
// Weight: [num_heads, output_dims, input_dims]
|
||||
// x: [B, num_heads, L, input_dims]
|
||||
// wT: [num_heads, input_dims, output_dims]
|
||||
// Result: [B, num_heads, L, output_dims]
|
||||
wT := mlx.Transpose(ml.Weight, 0, 2, 1)
|
||||
return mlx.Matmul(x, wT)
|
||||
}
|
||||
376
x/imagegen/nn/nn_test.go
Normal file
376
x/imagegen/nn/nn_test.go
Normal file
@@ -0,0 +1,376 @@
|
||||
package nn
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/x/imagegen/mlx"
|
||||
)
|
||||
|
||||
// TestMain initializes MLX before running tests.
|
||||
// If MLX libraries are not available, tests are skipped.
|
||||
func TestMain(m *testing.M) {
|
||||
// Change to repo root so ./build/lib/ollama/ path works
|
||||
_, thisFile, _, _ := runtime.Caller(0)
|
||||
repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..")
|
||||
if err := os.Chdir(repoRoot); err != nil {
|
||||
fmt.Printf("Failed to change to repo root: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if err := mlx.InitMLX(); err != nil {
|
||||
fmt.Printf("Skipping nn tests: %v\n", err)
|
||||
os.Exit(0)
|
||||
}
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
// TestLinearNoBias verifies Linear without bias computes x @ w.T correctly.
|
||||
func TestLinearNoBias(t *testing.T) {
|
||||
// Weight: [out=2, in=3] -> transposed at forward time
|
||||
weight := mlx.NewArrayFloat32([]float32{
|
||||
1, 2, 3, // row 0
|
||||
4, 5, 6, // row 1
|
||||
}, []int32{2, 3})
|
||||
mlx.Eval(weight)
|
||||
|
||||
linear := NewLinear(weight, nil)
|
||||
|
||||
// Input: [1, 3]
|
||||
x := mlx.NewArrayFloat32([]float32{1, 1, 1}, []int32{1, 3})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := linear.Forward(x)
|
||||
mlx.Eval(out)
|
||||
|
||||
// Expected: [1,1,1] @ [[1,4],[2,5],[3,6]] = [6, 15]
|
||||
data := out.Data()
|
||||
if len(data) != 2 || data[0] != 6 || data[1] != 15 {
|
||||
t.Errorf("expected [6, 15], got %v", data)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLinearWithBias verifies Linear with bias computes x @ w.T + b correctly.
|
||||
func TestLinearWithBias(t *testing.T) {
|
||||
weight := mlx.NewArrayFloat32([]float32{
|
||||
1, 2, 3,
|
||||
4, 5, 6,
|
||||
}, []int32{2, 3})
|
||||
bias := mlx.NewArrayFloat32([]float32{10, 20}, []int32{2})
|
||||
mlx.Eval(weight, bias)
|
||||
|
||||
linear := NewLinear(weight, bias)
|
||||
|
||||
x := mlx.NewArrayFloat32([]float32{1, 1, 1}, []int32{1, 3})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := linear.Forward(x)
|
||||
mlx.Eval(out)
|
||||
|
||||
// Expected: [6, 15] + [10, 20] = [16, 35]
|
||||
data := out.Data()
|
||||
if len(data) != 2 || data[0] != 16 || data[1] != 35 {
|
||||
t.Errorf("expected [16, 35], got %v", data)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLinearBatched verifies Linear works with batched input.
|
||||
func TestLinearBatched(t *testing.T) {
|
||||
weight := mlx.NewArrayFloat32([]float32{
|
||||
1, 0,
|
||||
0, 1,
|
||||
}, []int32{2, 2}) // Identity
|
||||
mlx.Eval(weight)
|
||||
|
||||
linear := NewLinear(weight, nil)
|
||||
|
||||
// Batch of 3 inputs
|
||||
x := mlx.NewArrayFloat32([]float32{
|
||||
1, 2,
|
||||
3, 4,
|
||||
5, 6,
|
||||
}, []int32{3, 2})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := linear.Forward(x)
|
||||
mlx.Eval(out)
|
||||
|
||||
// Identity should return same values
|
||||
data := out.Data()
|
||||
expected := []float32{1, 2, 3, 4, 5, 6}
|
||||
for i, v := range expected {
|
||||
if data[i] != v {
|
||||
t.Errorf("at %d: expected %f, got %f", i, v, data[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRMSNorm verifies RMSNorm computation.
|
||||
func TestRMSNorm(t *testing.T) {
|
||||
weight := mlx.NewArrayFloat32([]float32{1, 1, 1, 1}, []int32{4})
|
||||
mlx.Eval(weight)
|
||||
|
||||
norm := NewRMSNorm(weight, 1e-5)
|
||||
|
||||
// Input with known RMS
|
||||
x := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{1, 4})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := norm.Forward(x, 0) // eps=0 uses stored Eps
|
||||
mlx.Eval(out)
|
||||
|
||||
// RMS of [2,2,2,2] = 2, so normalized = [1,1,1,1]
|
||||
data := out.Data()
|
||||
for i, v := range data {
|
||||
if math.Abs(float64(v-1.0)) > 1e-4 {
|
||||
t.Errorf("at %d: expected ~1.0, got %f", i, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRMSNormWithScale verifies RMSNorm applies weight scaling.
|
||||
func TestRMSNormWithScale(t *testing.T) {
|
||||
weight := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{4})
|
||||
mlx.Eval(weight)
|
||||
|
||||
norm := NewRMSNorm(weight, 1e-5)
|
||||
|
||||
x := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{1, 4})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := norm.Forward(x, 0) // eps=0 uses stored Eps
|
||||
mlx.Eval(out)
|
||||
|
||||
// Normalized [1,1,1,1] * weight [2,2,2,2] = [2,2,2,2]
|
||||
data := out.Data()
|
||||
for i, v := range data {
|
||||
if math.Abs(float64(v-2.0)) > 1e-4 {
|
||||
t.Errorf("at %d: expected ~2.0, got %f", i, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedding verifies embedding lookup.
|
||||
func TestEmbedding(t *testing.T) {
|
||||
// Embedding table: 4 tokens, dim 3
|
||||
weight := mlx.NewArrayFloat32([]float32{
|
||||
0, 0, 0, // token 0
|
||||
1, 1, 1, // token 1
|
||||
2, 2, 2, // token 2
|
||||
3, 3, 3, // token 3
|
||||
}, []int32{4, 3})
|
||||
mlx.Eval(weight)
|
||||
|
||||
emb := NewEmbedding(weight)
|
||||
|
||||
// Look up tokens [1, 3, 0]
|
||||
indices := mlx.NewArrayInt32([]int32{1, 3, 0}, []int32{3})
|
||||
mlx.Eval(indices)
|
||||
|
||||
out := emb.Forward(indices)
|
||||
mlx.Eval(out)
|
||||
|
||||
data := out.Data()
|
||||
expected := []float32{1, 1, 1, 3, 3, 3, 0, 0, 0}
|
||||
for i, v := range expected {
|
||||
if data[i] != v {
|
||||
t.Errorf("at %d: expected %f, got %f", i, v, data[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRepeatKV verifies K/V repetition for GQA.
|
||||
func TestRepeatKV(t *testing.T) {
|
||||
// [B=1, num_kv_heads=2, S=2, head_dim=2]
|
||||
x := mlx.NewArrayFloat32([]float32{
|
||||
// head 0
|
||||
1, 2, // pos 0
|
||||
3, 4, // pos 1
|
||||
// head 1
|
||||
5, 6, // pos 0
|
||||
7, 8, // pos 1
|
||||
}, []int32{1, 2, 2, 2})
|
||||
mlx.Eval(x)
|
||||
|
||||
// Repeat factor 2: 2 kv heads -> 4 heads
|
||||
out := RepeatKV(x, 2)
|
||||
mlx.Eval(out)
|
||||
|
||||
shape := out.Shape()
|
||||
if shape[0] != 1 || shape[1] != 4 || shape[2] != 2 || shape[3] != 2 {
|
||||
t.Errorf("expected shape [1,4,2,2], got %v", shape)
|
||||
}
|
||||
|
||||
data := out.Data()
|
||||
// After repeat: head0, head0, head1, head1
|
||||
expected := []float32{
|
||||
1, 2, 3, 4, // head 0 (original)
|
||||
1, 2, 3, 4, // head 0 (repeat)
|
||||
5, 6, 7, 8, // head 1 (original)
|
||||
5, 6, 7, 8, // head 1 (repeat)
|
||||
}
|
||||
for i, v := range expected {
|
||||
if data[i] != v {
|
||||
t.Errorf("at %d: expected %f, got %f", i, v, data[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRepeatKVNoOp verifies RepeatKV with factor 1 returns input unchanged.
|
||||
func TestRepeatKVNoOp(t *testing.T) {
|
||||
x := mlx.NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{1, 1, 2, 2})
|
||||
mlx.Eval(x)
|
||||
|
||||
out := RepeatKV(x, 1)
|
||||
// Should return same pointer
|
||||
if out != x {
|
||||
t.Error("RepeatKV with factor 1 should return input unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
// TestApplyCausalMask verifies causal masking.
|
||||
func TestApplyCausalMask(t *testing.T) {
|
||||
// [B=1, heads=1, S=3, S=3] - all ones
|
||||
scores := mlx.Ones(1, 1, 3, 3)
|
||||
mlx.Eval(scores)
|
||||
|
||||
out := ApplyCausalMask(scores)
|
||||
mlx.Eval(out)
|
||||
|
||||
data := out.Data()
|
||||
// Lower triangular should be 1, upper should be -1e9
|
||||
// Row 0: [1, -inf, -inf]
|
||||
// Row 1: [1, 1, -inf]
|
||||
// Row 2: [1, 1, 1]
|
||||
if data[0] != 1 || data[1] >= 0 || data[2] >= 0 {
|
||||
t.Errorf("row 0 wrong: %v", data[0:3])
|
||||
}
|
||||
if data[3] != 1 || data[4] != 1 || data[5] >= 0 {
|
||||
t.Errorf("row 1 wrong: %v", data[3:6])
|
||||
}
|
||||
if data[6] != 1 || data[7] != 1 || data[8] != 1 {
|
||||
t.Errorf("row 2 wrong: %v", data[6:9])
|
||||
}
|
||||
}
|
||||
|
||||
// TestApplyCausalMaskWithOffset verifies causal masking with cache offset.
|
||||
func TestApplyCausalMaskWithOffset(t *testing.T) {
|
||||
// Simulating: cache has 2 tokens, adding 1 new query
|
||||
// scores: [B=1, heads=1, queryLen=1, keyLen=3]
|
||||
scores := mlx.Ones(1, 1, 1, 3)
|
||||
mlx.Eval(scores)
|
||||
|
||||
out := ApplyCausalMaskWithOffset(scores, 2)
|
||||
mlx.Eval(out)
|
||||
|
||||
data := out.Data()
|
||||
// With offset=2, query at position 2 can attend to all 3 positions
|
||||
if data[0] != 1 || data[1] != 1 || data[2] != 1 {
|
||||
t.Errorf("expected [1, 1, 1], got %v", data)
|
||||
}
|
||||
}
|
||||
|
||||
// TestApplyCausalMaskWithOffsetZero verifies offset=0 falls back to regular causal.
|
||||
func TestApplyCausalMaskWithOffsetZero(t *testing.T) {
|
||||
scores := mlx.Ones(1, 1, 2, 2)
|
||||
mlx.Eval(scores)
|
||||
|
||||
out := ApplyCausalMaskWithOffset(scores, 0)
|
||||
mlx.Eval(out)
|
||||
|
||||
data := out.Data()
|
||||
// Standard causal: [1, -inf], [1, 1]
|
||||
if data[0] != 1 || data[1] >= 0 {
|
||||
t.Errorf("row 0 wrong: %v", data[0:2])
|
||||
}
|
||||
if data[2] != 1 || data[3] != 1 {
|
||||
t.Errorf("row 1 wrong: %v", data[2:4])
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkLinearSmall benchmarks small Linear forward pass.
|
||||
func BenchmarkLinearSmall(b *testing.B) {
|
||||
weight := mlx.RandomNormal([]int32{256, 256}, 42)
|
||||
mlx.Eval(weight)
|
||||
|
||||
linear := NewLinear(weight, nil)
|
||||
|
||||
x := mlx.RandomNormal([]int32{1, 256}, 43)
|
||||
mlx.Eval(x)
|
||||
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
out := linear.Forward(x)
|
||||
mlx.Eval(out)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkLinearLarge benchmarks larger Linear forward pass.
|
||||
func BenchmarkLinearLarge(b *testing.B) {
|
||||
weight := mlx.RandomNormal([]int32{4096, 4096}, 42)
|
||||
mlx.Eval(weight)
|
||||
|
||||
linear := NewLinear(weight, nil)
|
||||
|
||||
x := mlx.RandomNormal([]int32{1, 4096}, 43)
|
||||
mlx.Eval(x)
|
||||
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
out := linear.Forward(x)
|
||||
mlx.Eval(out)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkRMSNorm benchmarks RMSNorm forward pass.
|
||||
func BenchmarkRMSNorm(b *testing.B) {
|
||||
weight := mlx.Ones(4096)
|
||||
mlx.Eval(weight)
|
||||
|
||||
norm := NewRMSNorm(weight, 1e-5)
|
||||
|
||||
x := mlx.RandomNormal([]int32{1, 4096}, 42)
|
||||
mlx.Eval(x)
|
||||
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
out := norm.Forward(x, 0)
|
||||
mlx.Eval(out)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkEmbedding benchmarks embedding lookup.
|
||||
func BenchmarkEmbedding(b *testing.B) {
|
||||
// Typical vocab size
|
||||
weight := mlx.RandomNormal([]int32{32000, 4096}, 42)
|
||||
mlx.Eval(weight)
|
||||
|
||||
emb := NewEmbedding(weight)
|
||||
|
||||
// Single token lookup
|
||||
indices := mlx.NewArrayInt32([]int32{1000}, []int32{1})
|
||||
mlx.Eval(indices)
|
||||
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
out := emb.Forward(indices)
|
||||
mlx.Eval(out)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkRepeatKV benchmarks K/V repetition.
|
||||
func BenchmarkRepeatKV(b *testing.B) {
|
||||
// Typical GQA setup: 8 kv heads -> 32 heads
|
||||
x := mlx.RandomNormal([]int32{1, 8, 512, 128}, 42)
|
||||
mlx.Eval(x)
|
||||
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
out := RepeatKV(x, 4)
|
||||
mlx.Eval(out)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user