diff --git a/x/mlxrunner/batch/batch.go b/x/mlxrunner/batch/batch.go index e93754d7b..409f53076 100644 --- a/x/mlxrunner/batch/batch.go +++ b/x/mlxrunner/batch/batch.go @@ -8,7 +8,35 @@ type Batch struct { InputIDs *mlx.Array // SeqOffsets gives each row's current position within its sequence — - // where the chunk in InputIDs starts. Length equals the batch dimension + // where the chunk in InputIDs starts. Length equals the batch dimension // of InputIDs. SeqOffsets []int32 + + // SeqQueryLens is each row's real query length in this forward. Values + // less than L mean the row's tail is padding that must be masked out. + // Length equals the batch dimension of InputIDs. + SeqQueryLens []int32 + + // Memo is per-forward memoization used to cache results, such as masks, + // which are often the same across layers. + Memo Memo +} + +type Memo struct { + entries map[any]any +} + +// Get returns the memoized value for key and true if present, or nil +// and false otherwise. +func (m *Memo) Get(key any) (any, bool) { + v, ok := m.entries[key] + return v, ok +} + +// Put stores value under key, allocating on first use. +func (m *Memo) Put(key, value any) { + if m.entries == nil { + m.entries = map[any]any{} + } + m.entries[key] = value } diff --git a/x/mlxrunner/cache/cache.go b/x/mlxrunner/cache/cache.go index f9202d90f..08d90a211 100644 --- a/x/mlxrunner/cache/cache.go +++ b/x/mlxrunner/cache/cache.go @@ -1,12 +1,17 @@ package cache import ( + "fmt" + "github.com/ollama/ollama/logutil" + "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" ) +// Cache is common state management shared by every cache kind. Writers +// live on the specific caches type Cache interface { - Update(keys, values *mlx.Array) (newKeys, newValues *mlx.Array) // State returns the cache-owned state roots that should be kept/evaluated. State() []*mlx.Array Free() @@ -39,6 +44,16 @@ type Snapshot interface { Close() } +// Attention is the contract for caches that back attention layers +// (KVCache, RotatingKVCache). +type Attention interface { + Cache + + // Update appends (k, v) and returns an opaque nn.KVHistory for + // this layer's SDPA. + Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory +} + type KVCache struct { keys, values *mlx.Array offset int @@ -49,7 +64,14 @@ func NewKVCache() *KVCache { return &KVCache{step: 256} } -func (c *KVCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { +// Assumes B = 1; heterogeneous batches are not supported. +func (c *KVCache) Update(_ *batch.Batch, keys, values *mlx.Array) *nn.KVHistory { + newK, newV := c.appendKV(keys, values) + return nn.NewKVHistory(newK, newV, nil) +} + +// appendKV is the raw write path shared by Update and Restore. +func (c *KVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { B, H, L, Dk, Dv := keys.Dim(0), keys.Dim(1), keys.Dim(2), keys.Dim(3), values.Dim(3) prev := c.offset @@ -141,9 +163,9 @@ func (c *KVCache) Restore(snapshot Snapshot, target int) bool { return false } - // Rewind to snapshot start, then feed snapshot data through Update. + // Rewind to snapshot start, then feed snapshot. c.offset = snap.fromOffset - c.Update(snap.keys, snap.values) + c.appendKV(snap.keys, snap.values) // Clamp to target if needed (target may be less than full snapshot). if target < c.offset { @@ -240,7 +262,22 @@ func NewRotatingKVCache(maxSize int) *RotatingKVCache { return &RotatingKVCache{maxSize: maxSize, KVCache: NewKVCache()} } -func (c *RotatingKVCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { +// Assumes B = 1; heterogeneous batches are not supported. +func (c *RotatingKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory { + newK, newV := c.appendKV(keys, values) + return nn.NewKVHistory(newK, newV, rotatingApplier{ + b: b, + K: newK.Dim(2), + L: keys.Dim(2), + window: c.maxSize, + ringIdx: c.idx, + dtype: keys.DType(), + }) +} + +// appendKV is the raw write path shared by Update and Restore — +// routes to concat for prefill (L > 1) and update for decode. +func (c *RotatingKVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { if keys.Dim(2) > 1 { return c.concat(keys, values) } @@ -444,3 +481,67 @@ func (c *RotatingKVCache) Free() { c.KVCache.Free() c.idx = 0 } + +// rotatingApplier composes the sliding-window storage restriction +// onto the caller's logical mask. +// +// ringIdx is the cache's write cursor at Update time. At L=1 decode +// the ring buffer is not position-ordered — logical col j lives at +// storage slot (ringIdx+j) mod K — so tensor masks built in +// logical space must be gathered into this layout before the kernel +// sees them. At L>1 prefill the concat path has already linearised +// storage, so the gather is identity and ringIdx is unused. +type rotatingApplier struct { + b *batch.Batch + K int + L int + window int + ringIdx int + dtype mlx.DType +} + +func (r rotatingApplier) ApplyMask(logical nn.AttentionMask) nn.AttentionMask { + if r.L == 1 { + // Single-query decode: storage already enforces the window + // (Update keeps the last maxSize tokens, all within + // [absQ-window+1, absQ]), and every stored key's absolute + // position <= absQ. For a zero or plain-causal logical mask + // both constraints reduce to "no mask", so return the zero + // mask and let SDPA dispatch to mode="". + if logical.IsZero() || logical.IsCausal() { + return nn.AttentionMask{} + } + + // Tensor-backed mask (user ArrayMask, causal+Relax, causal + // with accumulated array): materialize in logical-position + // order then gather K cols into ring-slot order so they + // align with the cache output the kernel will index. + arr := logical.AsArray(r.b, r.K, r.dtype) + arr = gatherRingCols(arr, r.ringIdx, r.K) + return nn.ArrayMask(arr) + } + + return logical.Intersect(nn.SlidingWindowMask(r.b, r.K, r.window, r.dtype)) +} + +// gatherRingCols reorders a [B, 1, L, K] mask's K axis from +// logical-position order (col 0 = oldest stored position) into the +// cache's ring-slot order (col 0 = buffer slot 0). Logical col j +// lives at slot (ringIdx+j) mod K, so storage slot s reads from +// logical col (s-ringIdx+K) mod K. Returns arr unchanged when the +// permutation is a no-op: ringIdx % K == 0 (layouts coincide), or +// the K axis broadcasts (dim 3 == 1, i.e. Q-padding-shaped masks +// where every key shares the same value). +func gatherRingCols(arr *mlx.Array, ringIdx, K int) *mlx.Array { + if w := arr.Dim(3); w != 1 && w != K { + panic(fmt.Sprintf("gatherRingCols: K-axis width %d must be 1 or %d", w, K)) + } + ringIdx %= K + if ringIdx == 0 || arr.Dim(3) == 1 { + return arr + } + shift := K - ringIdx + tail := arr.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(shift, K)) + head := arr.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, shift)) + return tail.Concatenate(3, head) +} diff --git a/x/mlxrunner/cache/cache_test.go b/x/mlxrunner/cache/cache_test.go index 86c26004a..22ecbdbfa 100644 --- a/x/mlxrunner/cache/cache_test.go +++ b/x/mlxrunner/cache/cache_test.go @@ -3,6 +3,7 @@ package cache import ( "testing" + "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" ) @@ -13,6 +14,17 @@ func skipIfNoMLX(t *testing.T) { } } +// newKVBatch builds a B=1 batch at SeqOffsets=off with all-real +// queries (SeqQueryLens=L) — the standard single-sequence cache +// test shape. +func newKVBatch(off, L int) *batch.Batch { + return &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, L), + SeqOffsets: []int32{int32(off)}, + SeqQueryLens: []int32{int32(L)}, + } +} + func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) { skipIfNoMLX(t) c := NewKVCache() @@ -20,7 +32,7 @@ func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } // Snapshot [5, 10). @@ -44,7 +56,7 @@ func TestKVCacheDataSurvivesSnapshotRestore(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } snap := c.Snapshot(0) @@ -84,7 +96,7 @@ func TestKVCacheSplitPreservesData(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } snap := c.Snapshot(0) @@ -128,7 +140,7 @@ func TestKVCacheSplitMergeRoundTripData(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } snap := c.Snapshot(0) @@ -163,7 +175,7 @@ func TestRotatingKVCacheRestoreOutsideWindow(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } // Offset 3 is outside the window. @@ -182,7 +194,7 @@ func TestRotatingKVCacheSnapshotPreservesWindow(t *testing.T) { for range 10 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } snap := c.Snapshot(0) @@ -194,7 +206,7 @@ func TestRotatingKVCacheSnapshotPreservesWindow(t *testing.T) { for range 5 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } // Restore to offset 10. @@ -228,7 +240,7 @@ func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) { for range 6 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } if c.Offset() != 6 { t.Fatalf("offset = %d, want 6", c.Offset()) @@ -240,7 +252,7 @@ func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) { for range 3 { k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } // Restore to snapshot state. @@ -255,7 +267,7 @@ func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) { // produce a valid window of size 4 at offset 7. k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) if c.Offset() != 7 { t.Fatalf("offset after post-restore update = %d, want 7", c.Offset()) diff --git a/x/mlxrunner/cache/recurrent.go b/x/mlxrunner/cache/recurrent.go index 0f0016539..cf87953db 100644 --- a/x/mlxrunner/cache/recurrent.go +++ b/x/mlxrunner/cache/recurrent.go @@ -1,6 +1,10 @@ package cache -import "github.com/ollama/ollama/x/mlxrunner/mlx" +import ( + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) // RecurrentCache stores state for linear-recurrent layers. // @@ -65,30 +69,24 @@ func (c *RecurrentCache) ensure(batch int, dtype mlx.DType) { } } -func (c *RecurrentCache) ConvState(batch int, dtype mlx.DType) *mlx.Array { - c.ensure(batch, dtype) - return c.convState +// Get returns the current conv/delta state for the SSM layer's read +// phase. Lazy-initializes zero-filled state tensors using b.InputIDs +// for the batch size; reallocates if the existing state's batch size +// or dtype no longer matches. +func (c *RecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory { + c.ensure(b.InputIDs.Dim(0), dtype) + return nn.NewRecurrentHistory(c.convState, c.deltaState) } -func (c *RecurrentCache) SetConvState(v *mlx.Array) { - c.convState = c.setState(c.convState, v, true) -} - -func (c *RecurrentCache) DeltaState(batch int, dtype mlx.DType) *mlx.Array { - c.ensure(batch, dtype) - return c.deltaState -} - -func (c *RecurrentCache) SetDeltaState(v *mlx.Array) { - c.deltaState = c.setState(c.deltaState, v, false) -} - -func (c *RecurrentCache) Advance(n int) { - c.offset += n -} - -func (c *RecurrentCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { - return keys, values +// Put stores the post-computation conv/delta states for the SSM +// layer's write phase and advances the cache offset by the current +// forward's real token count. +// +// Assumes B = 1; heterogeneous batches are not supported. +func (c *RecurrentCache) Put(b *batch.Batch, newConv, newDelta *mlx.Array) { + c.convState = c.setState(c.convState, newConv, true) + c.deltaState = c.setState(c.deltaState, newDelta, false) + c.offset += int(b.SeqQueryLens[0]) } func (c *RecurrentCache) State() []*mlx.Array { diff --git a/x/mlxrunner/cache/recurrent_test.go b/x/mlxrunner/cache/recurrent_test.go index ef8b7f7a3..31e58dae2 100644 --- a/x/mlxrunner/cache/recurrent_test.go +++ b/x/mlxrunner/cache/recurrent_test.go @@ -1,9 +1,12 @@ package cache import ( + "math" "testing" + "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" ) // TestRecurrentCacheRestoreExactOffset verifies that RecurrentCache restore @@ -12,13 +15,16 @@ import ( func TestRecurrentCacheRestoreExactOffset(t *testing.T) { skipIfNoMLX(t) c := NewRecurrentCache(3, 12, 4, 8, 8) - _ = c.ConvState(1, mlx.DTypeFloat16) - _ = c.DeltaState(1, mlx.DTypeFloat16) - c.Advance(10) + b1 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1)} + c.Get(b1, mlx.DTypeFloat16) // lazy-init + + b10 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 10), SeqQueryLens: []int32{10}} + c.Put(b10, nil, nil) // advance to 10 snap := c.Snapshot(0) // snap.offset == 10 - c.Advance(5) // cache now at 15 + b5 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 5), SeqQueryLens: []int32{5}} + c.Put(b5, nil, nil) // cache now at 15 // target < snap.offset: fails (can't rewind past snapshot) if c.Restore(snap, 5) { @@ -38,3 +44,194 @@ func TestRecurrentCacheRestoreExactOffset(t *testing.T) { t.Fatalf("offset = %d, want 10", c.Offset()) } } + +func TestRecurrentCacheGetLazyInit(t *testing.T) { + skipIfNoMLX(t) + c := NewRecurrentCache(3, 4, 2, 4, 4) + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{1}, + } + h := c.Get(b, mlx.DTypeBFloat16) + if c.Offset() != 0 { + t.Fatalf("Get should not advance; got offset %d", c.Offset()) + } + if h.ConvState() == nil || h.DeltaState() == nil { + t.Fatal("history should expose conv/delta tensors") + } + if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16) + } + if got := h.DeltaState().DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeBFloat16) + } +} + +// TestRecurrentCachePaddedRoundTrip runs Get → CausalConv1D → +// GatedDelta → Put on a B=1 batch with qLen 1e-4 { + t.Fatalf("nextConv[%d]: padded=%v unpadded=%v (padding leaked into conv state)", i, gp[i], gr[i]) + } + } + + dp := deltaPad.Floats() + dr := deltaRef.Floats() + if len(dp) != len(dr) { + t.Fatalf("delta state shape mismatch: padded %d vs unpadded %d", len(dp), len(dr)) + } + for i := range dp { + if math.Abs(float64(dp[i]-dr[i])) > 1e-3 { + t.Fatalf("delta state[%d]: padded=%v unpadded=%v (padding leaked into recurrent state)", i, dp[i], dr[i]) + } + } +} + +func TestRecurrentCachePutAdvances(t *testing.T) { + skipIfNoMLX(t) + c := NewRecurrentCache(3, 4, 2, 4, 4) + b := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 2), SeqQueryLens: []int32{2}} + newConv := mlx.Zeros(mlx.DTypeFloat16, 1, 3, 4) + newDelta := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 4, 4) + c.Put(b, newConv, newDelta) + if c.Offset() != 2 { + t.Fatalf("cache offset not advanced: %d", c.Offset()) + } +} diff --git a/x/mlxrunner/cache/rotating_attention_test.go b/x/mlxrunner/cache/rotating_attention_test.go new file mode 100644 index 000000000..38499d828 --- /dev/null +++ b/x/mlxrunner/cache/rotating_attention_test.go @@ -0,0 +1,237 @@ +package cache + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +// TestRotatingKVCacheDecodeParity drives a rotating cache past its +// wrap point with single-token writes, runs an L=1 decode through +// SDPA, and compares against a reference computed from the same per- +// position K/V in logical-position order with the same caller mask. +// +// Attention is permutation-invariant when K, V, and the mask are +// permuted together, so the cache's storage-order output (with the +// applier's gather composing the caller's logical mask back into +// storage order) must equal the logical-order reference. +func TestRotatingKVCacheDecodeParity(t *testing.T) { + skipIfNoMLX(t) + const H, D = 1, 4 + const window = 4 + const totalWrites = 7 // past wrap (window=4); last write is the L=1 decode + const scale = 1.0 + + // Per-position k, v values. Use distinct seeds so the per-position + // values are clearly distinguishable. + perPosKV := func(pos int) (k, v *mlx.Array) { + kVals := make([]float32, H*D) + vVals := make([]float32, H*D) + for i := range kVals { + kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i) + vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i) + } + k = mlx.FromValues(kVals, 1, H, 1, D) + v = mlx.FromValues(vVals, 1, H, 1, D) + return + } + + q := mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D) + mlx.Eval(q) + + // Drive the cache: write positions 0..totalWrites-2 as a "history", + // then position totalWrites-1 is the actual L=1 decode under test. + c := NewRotatingKVCache(window) + for pos := range totalWrites - 1 { + k, v := perPosKV(pos) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } + + finalPos := totalWrites - 1 + kFinal, vFinal := perPosKV(finalPos) + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{int32(finalPos)}, + SeqQueryLens: []int32{1}, + } + history := c.Update(b, kFinal, vFinal) + + // Reference: the in-window logical-position-ordered K and V are + // the last `window` per-position values (positions + // [finalPos-window+1, finalPos]). Build them in that order. + startPos := max(finalPos-window+1, 0) + logicalKs := make([]*mlx.Array, 0, window) + logicalVs := make([]*mlx.Array, 0, window) + for pos := startPos; pos <= finalPos; pos++ { + kp, vp := perPosKV(pos) + logicalKs = append(logicalKs, kp) + logicalVs = append(logicalVs, vp) + } + kLogical := mlx.Concatenate(logicalKs, 2) + vLogical := mlx.Concatenate(logicalVs, 2) + + // A logical-order ArrayMask with distinct, non-trivial values per + // key column. Picked so each column's contribution to softmax is + // distinct — the test fails if the cache's gather permutes the + // columns wrong before the kernel sees them. + maskVals := []float32{0.1, -0.3, 0.7, -0.2} + logicalMask := mlx.FromValues(maskVals, 1, 1, 1, window) + + cases := []struct { + name string + model nn.AttentionMask + // reference mask uses the same coordinates the model mask + // represents; for ArrayMask it's the same tensor (since the + // reference K/V is in logical order). + refMode string + refMask *mlx.Array + }{ + {"zero", nn.AttentionMask{}, "", nil}, + {"causal-at-L1", nn.CausalMask(), "", nil}, + {"array", nn.ArrayMask(logicalMask), "array", logicalMask}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := nn.ScaledDotProductAttention(b, q, scale, + nn.WithKVHistory(history), + nn.WithMask(tc.model)) + + want := mlx.FastScaledDotProductAttention(q, kLogical, vLogical, scale, + tc.refMode, tc.refMask) + + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if math.Abs(float64(gs[i]-ws[i])) > 1e-5 { + t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i]) + } + } + }) + } +} + +// TestRotatingKVCachePrefillParity drives an L>1 prefill into a +// rotating cache and verifies SDPA output through WithKVHistory +// matches a reference computed from the same K/V with the model mask +// and window restriction composed manually. +func TestRotatingKVCachePrefillParity(t *testing.T) { + skipIfNoMLX(t) + const H, L, D = 1, 6, 4 + const window = 4 + const scale = 1.0 + + qVals := make([]float32, 1*H*L*D) + kVals := make([]float32, 1*H*L*D) + vVals := make([]float32, 1*H*L*D) + for i := range qVals { + qVals[i] = 0.5 + 0.05*float32(i) + kVals[i] = -0.3 + 0.07*float32(i) + vVals[i] = 0.3 + 0.03*float32(i) + } + q := mlx.FromValues(qVals, 1, H, L, D) + k := mlx.FromValues(kVals, 1, H, L, D) + v := mlx.FromValues(vVals, 1, H, L, D) + b := newKVBatch(0, L) + + cases := []struct { + name string + mask nn.AttentionMask + // rect arguments matching nn.AttentionMask.Relax (qLo, qHi, kLo, kHi) + relax [][4]int + causal bool + }{ + {"zero", nn.AttentionMask{}, nil, false}, + {"causal", nn.CausalMask(), nil, true}, + {"causal+relax", nn.CausalMask().Relax(0, 1, 4, 2, 5), [][4]int{{1, 4, 2, 5}}, true}, + } + + negInf := float32(math.Inf(-1)) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c := NewRotatingKVCache(window) + history := c.Update(b, k, v) + + got := nn.ScaledDotProductAttention(b, q, scale, + nn.WithKVHistory(history), + nn.WithMask(tc.mask)) + + // Reference mask: causal blocks k > absQ; relax rectangles + // release causal-blocked cells; window blocks k < absQ - window + 1. + refVals := make([]float32, L*L) + for qi := range L { + absQ := qi + for ki := range L { + blocked := false + if tc.causal && ki > absQ { + blocked = true + } + for _, r := range tc.relax { + qLo, qHi, kLo, kHi := r[0], r[1], r[2], r[3] + if absQ >= qLo && absQ < qHi && ki >= kLo && ki < kHi { + blocked = false + } + } + if window > 0 && ki < absQ-window+1 { + blocked = true + } + if blocked { + refVals[qi*L+ki] = negInf + } + } + } + refMask := mlx.FromValues(refVals, 1, 1, L, L) + want := mlx.FastScaledDotProductAttention(q, k, v, scale, "array", refMask) + + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if math.Abs(float64(gs[i]-ws[i])) > 1e-4 { + t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i]) + } + } + }) + } +} + +// TestRotatingKVCacheMLAParity drives a rotating cache with the MLA +// shape — K = [kvLatent, kPE] concatenated, V = zero-width — then +// uses WithMLAHistory to slice V from K and compares output against +// a manual reference. Pins the cache+MLA integration that +// glm4_moe_lite uses in production. +func TestRotatingKVCacheMLAParity(t *testing.T) { + skipIfNoMLX(t) + const H, L, D, valueDim = 1, 3, 6, 4 + const scale = 1.0 + const window = 8 // window >= L so no window restriction + + kVals := make([]float32, 1*H*L*D) + for i := range kVals { + kVals[i] = 0.1 * float32(i+1) + } + k := mlx.FromValues(kVals, 1, H, L, D) + v := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, 0) + + q := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, D) + b := newKVBatch(0, L) + + c := NewRotatingKVCache(window) + history := c.Update(b, k, v) + got := nn.ScaledDotProductAttention(b, q, scale, + nn.WithMLAHistory(history, valueDim), + nn.WithMask(nn.CausalMask())) + + vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim)) + want := mlx.FastScaledDotProductAttention(q, k, vRef, scale, "causal", nil) + + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if math.Abs(float64(gs[i]-ws[i])) > 1e-5 { + t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i]) + } + } +} diff --git a/x/mlxrunner/cache/rotating_multiturn_test.go b/x/mlxrunner/cache/rotating_multiturn_test.go index 9914e1e86..03991a1b9 100644 --- a/x/mlxrunner/cache/rotating_multiturn_test.go +++ b/x/mlxrunner/cache/rotating_multiturn_test.go @@ -61,13 +61,13 @@ func feedMulti(c *RotatingKVCache, startID float32, n int) float32 { ids[i] = startID + float32(i) } k, v := multiTokenKV(ids) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) return startID + float32(n) } func feedSingle(c *RotatingKVCache, id float32) { k, v := singleTokenKV(id) - c.Update(k, v) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) } // TestRotatingKVCacheConcatMidRotationPreservesContext: after the buffer diff --git a/x/mlxrunner/mlx/fast.go b/x/mlxrunner/mlx/fast.go index 0a2adea41..f38a670df 100644 --- a/x/mlxrunner/mlx/fast.go +++ b/x/mlxrunner/mlx/fast.go @@ -7,19 +7,21 @@ import ( "unsafe" ) -func ScaledDotProductAttention(query, key, value, mask *Array, scale float32) *Array { - if mask == nil { - mask = New("") - } - +func FastScaledDotProductAttention(q, k, v *Array, scale float32, mode string, mask *Array) *Array { sinks := New("") - - mode := "causal" cMode := C.CString(mode) defer C.free(unsafe.Pointer(cMode)) + var maskCtx C.mlx_array + if mask != nil { + maskCtx = mask.ctx + } else { + empty := New("") + maskCtx = empty.ctx + } + out := New("FAST_SDPA") - C.mlx_fast_scaled_dot_product_attention(&out.ctx, query.ctx, key.ctx, value.ctx, C.float(scale), cMode, mask.ctx, sinks.ctx, DefaultStream().ctx) + C.mlx_fast_scaled_dot_product_attention(&out.ctx, q.ctx, k.ctx, v.ctx, C.float(scale), cMode, maskCtx, sinks.ctx, DefaultStream().ctx) return out } @@ -43,4 +45,3 @@ func (r *RMSNorm) Forward(x *Array, eps float32) *Array { C.mlx_fast_rms_norm(&out.ctx, x.ctx, r.Weight.ctx, C.float(eps), DefaultStream().ctx) return out } - diff --git a/x/mlxrunner/mlx/gated_delta.go b/x/mlxrunner/mlx/gated_delta.go index c691f0512..7034c33a6 100644 --- a/x/mlxrunner/mlx/gated_delta.go +++ b/x/mlxrunner/mlx/gated_delta.go @@ -606,11 +606,35 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar return y, nextState, true } -// GatedDelta runs the recurrent update operation. +// FastGatedDelta runs the recurrent update operation. +// +// When mask is non-nil, it must be a [B, T] bool tensor identifying real +// (true) vs. padded (false) positions in q/k/v/g/beta. Padded positions +// are substituted with neutral values (q=k=v=beta=0, g=1) so each padded +// kernel iteration is a no-op — state passes through unchanged and the +// final state equals the state after the last real token of each row. // // It tries the fused CUDA kernel first, then Metal, then falls back to a // backend-agnostic MLX implementation with identical inputs/outputs. -func GatedDelta(q, k, v, g, beta, state *Array) (y, nextState *Array) { +func FastGatedDelta(q, k, v, g, beta, state, mask *Array) (y, nextState *Array) { + // TODO: handle this more efficiently with a masked kernel (MLX-LM has one). + if mask != nil { + B := int32(mask.Dim(0)) + T := int32(mask.Dim(1)) + m4 := Reshape(mask, B, T, 1, 1) + m3 := Reshape(mask, B, T, 1) + zeroQ := FromValue(float32(0)).AsType(q.DType()) + zeroK := FromValue(float32(0)).AsType(k.DType()) + zeroV := FromValue(float32(0)).AsType(v.DType()) + zeroBeta := FromValue(float32(0)).AsType(beta.DType()) + oneG := FromValue(float32(1)).AsType(g.DType()) + q = Where(m4, q, zeroQ) + k = Where(m4, k, zeroK) + v = Where(m4, v, zeroV) + beta = Where(m3, beta, zeroBeta) + g = Where(m3, g, oneG) + } + if y, nextState, ok := gatedDeltaCUDAKernelApply(q, k, v, g, beta, state); ok { return y, nextState } @@ -619,7 +643,7 @@ func GatedDelta(q, k, v, g, beta, state *Array) (y, nextState *Array) { } y, nextState = gatedDeltaFallback(q, k, v, g, beta, state) if y == nil || nextState == nil { - panic("mlx.GatedDelta: fallback failed (invalid inputs or unsupported shapes)") + panic("mlx.FastGatedDelta: fallback failed (invalid inputs or unsupported shapes)") } return y, nextState } diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index ffc718220..623a4f1ce 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/x/mlxrunner/mlx/ops_extra.go @@ -493,35 +493,6 @@ func SoftmaxAxis(a *Array, axis int, precise bool) *Array { return out } -func ScaledDotProductAttentionCausal(q, k, v *Array, scale float32, causalMask bool) *Array { - mask := New("") - sinks := New("") - mode := "" - if causalMask { - mode = "causal" - } - cMode := C.CString(mode) - defer C.free(unsafe.Pointer(cMode)) - - out := New("FAST_SDPA") - C.mlx_fast_scaled_dot_product_attention(&out.ctx, q.ctx, k.ctx, v.ctx, C.float(scale), cMode, mask.ctx, sinks.ctx, DefaultStream().ctx) - return out -} - -// ScaledDotProductAttentionMasked runs the fast SDPA kernel with an explicit -// additive mask. The mask is broadcast to [B, H, Q, K] and added to scores -// before softmax. Pass mode="array" so MLX actually consults mask_arr; the -// empty string is "no mask" and silently ignores the array argument. -func ScaledDotProductAttentionMasked(q, k, v *Array, scale float32, mask *Array) *Array { - sinks := New("") - cMode := C.CString("array") - defer C.free(unsafe.Pointer(cMode)) - - out := New("FAST_SDPA") - C.mlx_fast_scaled_dot_product_attention(&out.ctx, q.ctx, k.ctx, v.ctx, C.float(scale), cMode, mask.ctx, sinks.ctx, DefaultStream().ctx) - return out -} - func LayerNormFn(x, weight, bias *Array, eps float32) *Array { out := New("FAST_LAYERNORM") var w, b C.mlx_array diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 6a9fc05ad..414015f6a 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -124,8 +124,9 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er } r.Model.Forward(&batch.Batch{ - InputIDs: mlx.FromValues(tokens[processed:processed+n], 1, n), - SeqOffsets: []int32{int32(position)}, + InputIDs: mlx.FromValues(tokens[processed:processed+n], 1, n), + SeqOffsets: []int32{int32(position)}, + SeqQueryLens: []int32{int32(n)}, }, caches) mlx.Sweep() materializeCaches() @@ -148,8 +149,9 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er step := func(token *mlx.Array) sampler.Result { fwd := r.Model.Forward(&batch.Batch{ - InputIDs: token, - SeqOffsets: []int32{int32(position)}, + InputIDs: token, + SeqOffsets: []int32{int32(position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, }, caches) position += token.Dim(1) logits := r.Model.Unembed(fwd) diff --git a/x/models/gemma3/gemma3.go b/x/models/gemma3/gemma3.go index 45704195b..6862f18f8 100644 --- a/x/models/gemma3/gemma3.go +++ b/x/models/gemma3/gemma3.go @@ -416,7 +416,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { if caches != nil && i < len(caches) { c = caches[i] } - h = layer.Forward(h, c, positions, B, L, m.TextConfig) + h = layer.Forward(h, b, c, positions, B, L, m.TextConfig) } return mlx.RMSNormFn(h, m.NormScaled, m.RMSNormEps) @@ -456,10 +456,10 @@ func (m *Model) FormatPrompt(prompt string) string { return fmt.Sprintf("user\n%s\nmodel\n", prompt) } -func (l *DecoderLayer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig) *mlx.Array { +func (l *DecoderLayer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig) *mlx.Array { normed := mlx.RMSNormFn(x, l.InputNormScaled, cfg.RMSNormEps) - attnOut := l.Attention.Forward(normed, c, positions, B, L, l.IsSliding, cfg) + attnOut := l.Attention.Forward(normed, b, c, positions, B, L, l.IsSliding, cfg) attnOut = mlx.RMSNormFn(attnOut, l.PostAttnNormScaled, cfg.RMSNormEps) h := mlx.Add(x, attnOut) @@ -471,7 +471,7 @@ func (l *DecoderLayer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array return mlx.Add(h, mlpOut) } -func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig) *mlx.Array { +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig) *mlx.Array { q := a.QProj.Forward(x) k := a.KProj.Forward(x) v := a.VProj.Forward(x) @@ -496,13 +496,16 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B q = mlx.RoPEWithBase(q, int(cfg.HeadDim), false, ropeTheta, 1.0, positions) k = mlx.RoPEWithBase(k, int(cfg.HeadDim), false, ropeTheta, 1.0, positions) - if c != nil { - k, v = c.Update(k, v) - } - // MLX SDPA supports grouped-query attention directly (Q heads can be a // multiple of K/V heads), so avoid materializing repeated K/V tensors. - out := mlx.ScaledDotProductAttentionCausal(q, k, v, cfg.Scale, L > 1) + var kv nn.SDPAOption + if c != nil { + history := c.(cache.Attention).Update(b, k, v) + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) + } + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) return a.OProj.Forward(out) } diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index b8e68522f..df1d1b271 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -88,9 +88,17 @@ type TextConfig struct { KVDonors map[int32]bool `json:"-"` } -// sharedKVEntry stores cached KV state from a donor layer for KV sharing. -type sharedKVEntry struct { - K, V *mlx.Array +// sharedHistory carries a donor layer's K/V to donees that share +// it. Exactly one of history or (k, v) is populated: history when +// the donor had a cache (donees feed it to SDPA via WithKVHistory; +// (k, v) when the donor had no cache (donees feed them via WithKV). +// mask carries any storage-side mask the (k, v) path needs (e.g. +// sliding window) — unused on the history path, where the cache's +// applier supplies the equivalent restriction. +type sharedHistory struct { + history *nn.KVHistory + k, v *mlx.Array + mask nn.AttentionMask } // Attention implements Gemma 4 attention with Q/K normalization and v-norm. @@ -135,34 +143,6 @@ func firstNonNil(tensors map[string]*mlx.Array, keys ...string) *mlx.Array { return nil } -// buildCausalMaskWindow creates a [1, 1, Q, K] additive causal mask with an -// optional sliding-window restriction. When window > 0, positions where -// kv < absQ - window + 1 are also masked (the token can only see the most -// recent `window` keys). When window == 0, only the causal constraint applies. -// -// This is the prefill-time mask for sliding-window attention layers. Without -// the window restriction, a sliding layer would attend to the entire prefix -// during prompt processing and diverge from the reference starting at the -// first position past the window boundary. -func buildCausalMaskWindow(Q, K, window int32) *mlx.Array { - offset := K - Q // cache offset: kv positions before the current query chunk - vals := make([]float32, Q*K) - negInf := float32(math.Inf(-1)) - for q := range Q { - absQ := offset + q - var lo int32 - if window > 0 { - lo = max(absQ-window+1, 0) - } - for kv := range K { - if kv > absQ || kv < lo { - vals[q*K+kv] = negInf - } - } - } - return mlx.FromValues(vals, 1, 1, int(Q), int(K)) -} - // sliceAxis1 slices a tensor along axis 1: a[:, start:stop, ...]. func sliceAxis1(a *mlx.Array, start, stop int32) *mlx.Array { dims := a.Dims() @@ -1026,20 +1006,11 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { perLayerInputs = m.computePLEInputs(b.InputIDs, h) } - var sharedKV map[int32]sharedKVEntry + // KV sharing: each donor layer stores its KVHistory here so later + // shared layers can reuse it in lieu of their own cache update. + var sharedKV map[int32]sharedHistory if len(m.KVShareMap) > 0 { - sharedKV = make(map[int32]sharedKVEntry) - } - - // Per-forward-pass sliding-window mask cache. The first sliding layer - // populates it; every subsequent sliding layer reuses the cached mask. - // Before this, Attention.Forward rebuilt a ~L×kLen mask from scratch on - // the CPU for every sliding layer (~25 rebuilds on a 26B MoE prefill), - // which was the bulk of a ~28% prefill regression vs the pre-SWA-fix - // baseline. - var smc *slidingMaskCache - if L > 1 && m.SlidingWindow > 0 { - smc = &slidingMaskCache{} + sharedKV = make(map[int32]sharedHistory) } for i, layer := range m.Layers { @@ -1055,15 +1026,15 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { } // Get shared KV for this layer if it's a shared layer. - var donorEntry *sharedKVEntry + var donor *sharedHistory if layer.KVShareDonor >= 0 { - if entry, ok := sharedKV[layer.KVShareDonor]; ok { - donorEntry = &entry + if d, ok := sharedKV[layer.KVShareDonor]; ok { + donor = &d } } - var donorKV *sharedKVEntry - h, donorKV = layer.Forward(h, c, positions, B, L, m.TextConfig, pleInput, donorEntry, smc) + var donorKV *sharedHistory + h, donorKV = layer.Forward(h, b, c, positions, B, L, m.TextConfig, pleInput, donor) // If this layer is a donor, store its cached KV for later shared layers. if layer.IsDonor && donorKV != nil { @@ -1074,37 +1045,6 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { return mlx.RMSNormFn(h, m.NormScaled, m.RMSNormEps) } -// slidingMaskCache is a per-forward-pass cache for the sliding-window -// additive mask used by Attention.Forward. All sliding-attention layers in -// a single forward pass see the same (L, kLen, window, dtype) tuple, so -// we can build the mask once on the first sliding layer's call and let -// every subsequent layer reuse it. The cache is instantiated fresh at the -// top of Model.Forward and passed through DecoderLayer → Attention; it is -// nil on paths where no caching is wanted (e.g. L == 1 decode). -type slidingMaskCache struct { - mask *mlx.Array - L int32 - kLen int32 - window int32 -} - -// get returns a cached mask matching (L, kLen, window, dtype), or builds -// and caches a new one. Safe to call with a nil receiver (falls through to -// a direct build without caching). -func (c *slidingMaskCache) get(L, kLen, window int32, dtype mlx.DType) *mlx.Array { - if c == nil { - return buildCausalMaskWindow(L, kLen, window).AsType(dtype) - } - if c.mask != nil && c.L == L && c.kLen == kLen && c.window == window { - return c.mask - } - c.mask = buildCausalMaskWindow(L, kLen, window).AsType(dtype) - c.L = L - c.kLen = kLen - c.window = window - return c.mask -} - func (m *Model) Unembed(x *mlx.Array) *mlx.Array { logits := m.LMHead.Forward(x) @@ -1189,9 +1129,9 @@ func sliceLayerDim(combined *mlx.Array, layerIdx, B, L, pleDim int32) *mlx.Array return mlx.Squeeze(sliced, 2) } -func (l *DecoderLayer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig, pleInput *mlx.Array, donorEntry *sharedKVEntry, slidingMaskCache *slidingMaskCache) (*mlx.Array, *sharedKVEntry) { +func (l *DecoderLayer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig, pleInput *mlx.Array, donor *sharedHistory) (*mlx.Array, *sharedHistory) { normed := mlx.RMSNormFn(x, l.InputNormScaled, cfg.RMSNormEps) - attnOut, donorKV := l.Attention.Forward(normed, c, positions, B, L, l.IsSliding, cfg, donorEntry, slidingMaskCache) + attnOut, kv := l.Attention.Forward(normed, b, c, positions, B, L, l.IsSliding, cfg, donor) attnOut = mlx.RMSNormFn(attnOut, l.PostAttnNormScaled, cfg.RMSNormEps) h := mlx.Add(x, attnOut) @@ -1236,10 +1176,10 @@ func (l *DecoderLayer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array h = mlx.Mul(h, l.LayerScalar) } - return h, donorKV + return h, kv } -func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig, donorEntry *sharedKVEntry, slidingMaskCache *slidingMaskCache) (*mlx.Array, *sharedKVEntry) { +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig, donor *sharedHistory) (*mlx.Array, *sharedHistory) { // Determine head dim and scale based on layer type. headDim := cfg.HeadDim scale := cfg.SlidingScale @@ -1265,25 +1205,19 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B } q = mlx.RoPEWithFreqs(q, ropeDims, false, ropeBase, 1.0, positions, ropeFreqs) - var k, v *mlx.Array - var donorKV *sharedKVEntry - - if donorEntry != nil { - // Shared layer: use donor's cached K/V. - k = donorEntry.K - v = donorEntry.V - } else { + kv := donor + if kv == nil { // Determine KV head count: K=V full-attention layers use NumGlobalKeyValueHeads. kvHeads := cfg.NumKeyValueHeads if a.VProj == nil && !isSliding && cfg.NumGlobalKeyValueHeads > 0 { kvHeads = cfg.NumGlobalKeyValueHeads } - // Non-shared layer: compute K/V. - k = a.KProj.Forward(x) + k := a.KProj.Forward(x) k = mlx.Reshape(k, B, L, kvHeads, headDim) k = mlx.Transpose(k, 0, 2, 1, 3) + var v *mlx.Array if a.VProj != nil { v = a.VProj.Forward(x) v = mlx.Reshape(v, B, L, kvHeads, headDim) @@ -1303,31 +1237,40 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B v = mlx.RMSNormFn(v, nil, cfg.RMSNormEps) // Update cache. + kv = &sharedHistory{} if c != nil { - k, v = c.Update(k, v) - donorKV = &sharedKVEntry{K: k, V: v} + kv.history = c.(cache.Attention).Update(b, k, v) + } else { + kv.k, kv.v = k, v + if isSliding { + kv.mask = nn.SlidingWindowMask(b, k.Dim(2), int(cfg.SlidingWindow), q.DType()) + } } } - // Sliding-window layers must restrict attention to the last `window` keys - // during prefill. The rotating KV cache handles decode, but for L > 1 the - // cache holds all prefix keys so we need an explicit mask. Full-attention - // layers pass window=0 (plain causal). - var window int32 - if isSliding && L > 1 && cfg.SlidingWindow > 0 { - window = cfg.SlidingWindow - } - var out *mlx.Array - switch { - case headDim > 128 && L > 1 && !mlx.MetalIsAvailable(): + if headDim > 128 && L > 1 && !mlx.MetalIsAvailable() { // Manual attention for CUDA prefill with head_dim > 128. // cuDNN SDPA requires head_dim <= 128, and the MLX CUDA SDPA vector // kernel only handles L < 4 (generation). For prefill, we fall back // to explicit matmul+softmax+matmul on CUDA. + var k, v *mlx.Array + mask := nn.CausalMask().Intersect(nn.QPaddingMask(b, q.DType())) + if kv.history != nil { + k, v = kv.history.K(), kv.history.V() + mask = kv.history.Mask(mask) + } else { + k, v = kv.k, kv.v + mask = mask.Intersect(nn.KPaddingMask(b, k.Dim(2), b.SeqQueryLens, q.DType())) + mask = mask.Intersect(kv.mask) + } kvHeads := int32(k.Dim(1)) nRepeats := cfg.NumAttentionHeads / kvHeads kLen := int32(k.Dim(2)) + // AsArray returns [B, 1, L, K]; reshape to rank 5 so that + // right-to-left broadcast against scores [B, kvHeads, + // nRepeats, L, K] aligns the batch dim correctly. + maskArr := mlx.Reshape(mask.AsArray(b, int(kLen), q.DType()), B, 1, 1, L, kLen) q = mlx.MulScalar(q, scale) q = mlx.Reshape(q, B, kvHeads, nRepeats, L, headDim) @@ -1336,22 +1279,20 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B kT := mlx.Transpose(k, 0, 1, 2, 4, 3) scores := mlx.Matmul(q, kT) - mask := buildCausalMaskWindow(L, kLen, window) - scores = mlx.Add(scores, mask) + scores = mlx.Add(scores, maskArr) scores = mlx.SoftmaxAxis(scores, -1, true) out = mlx.Matmul(scores, v) out = mlx.Reshape(out, B, cfg.NumAttentionHeads, L, headDim) - case window > 0: - // Sliding-window prefill: fast SDPA "causal" mode doesn't take a - // window, so supply an explicit additive mask. All sliding layers - // in one forward pass see the same (L, kLen, window, dtype) tuple, - // so we memoize the mask on the first call and reuse it for every - // subsequent sliding layer via slidingMaskCache. - kLen := int32(k.Dim(2)) - mask := slidingMaskCache.get(L, kLen, window, q.DType()) - out = mlx.ScaledDotProductAttentionMasked(q, k, v, scale, mask) - default: - out = mlx.ScaledDotProductAttentionCausal(q, k, v, scale, L > 1) + } else { + var opt nn.SDPAOption + mask := nn.CausalMask() + if kv.history != nil { + opt = nn.WithKVHistory(kv.history) + } else { + opt = nn.WithKV(kv.k, kv.v, b.SeqQueryLens) + mask = mask.Intersect(kv.mask) + } + out = nn.ScaledDotProductAttention(b, q, scale, opt, nn.WithMask(mask)) } out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*headDim) if !mlx.MetalIsAvailable() { @@ -1359,7 +1300,7 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B // strided views differently. Metal handles them natively. out = mlx.Contiguous(out, false) } - return a.OProj.Forward(out), donorKV + return a.OProj.Forward(out), kv } func (m *MLP) Forward(x *mlx.Array) *mlx.Array { diff --git a/x/models/glm4_moe_lite/glm4_moe_lite.go b/x/models/glm4_moe_lite/glm4_moe_lite.go index b2945b2f4..d30875b0d 100644 --- a/x/models/glm4_moe_lite/glm4_moe_lite.go +++ b/x/models/glm4_moe_lite/glm4_moe_lite.go @@ -88,7 +88,7 @@ type MLAAttention struct { } // Forward computes absorbed MLA attention output. -func (a *MLAAttention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { +func (a *MLAAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { q := a.QAProj.Forward(x) q = a.QALayerNorm.Forward(q, cfg.RMSNormEps) q = a.QBProj.Forward(q) @@ -117,18 +117,24 @@ func (a *MLAAttention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array keys := mlx.Concatenate([]*mlx.Array{kvLatent, kPE}, 3) - cachedL := L + // MLA compresses K and V into a single tensor: the cache stores + // [kvLatent, kPE] concatenated along the last dim as its keys, + // and V is the kvLatent prefix (first KVLoraRank positions) of + // that same tensor. WithMLAHistory handles the slice on our + // behalf so the model never touches the history's K/V. + var kv nn.SDPAOption if c != nil { placeholderValues := mlx.ZerosF32([]int32{B, 1, L, 0}) - keys, _ = c.Update(keys, placeholderValues) - cachedL = int32(keys.Dim(2)) + history := c.(cache.Attention).Update(b, keys, placeholderValues) + kv = nn.WithMLAHistory(history, int(cfg.KVLoraRank)) + } else { + values := mlx.SliceStartStop(keys, []int32{0, 0, 0, 0}, []int32{B, 1, L, cfg.KVLoraRank}) + kv = nn.WithKV(keys, values, b.SeqQueryLens) } - values := mlx.SliceStartStop(keys, []int32{0, 0, 0, 0}, []int32{B, 1, cachedL, cfg.KVLoraRank}) - queries := mlx.Concatenate([]*mlx.Array{qLatent, qPE}, 3) - out := mlx.ScaledDotProductAttentionCausal(queries, keys, values, cfg.Scale, L > 1) + out := nn.ScaledDotProductAttention(b, queries, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) out = a.UnembedOut.Forward(out) out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.VHeadDim) @@ -306,8 +312,8 @@ type DenseBlock struct { } // Forward applies the dense block -func (blk *DenseBlock) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { - r := blk.Attention.Forward(blk.InputLayerNorm.Forward(x, cfg.RMSNormEps), c, positions, B, L, cfg) +func (blk *DenseBlock) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + r := blk.Attention.Forward(blk.InputLayerNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg) h := mlx.Add(x, r) r = blk.MLP.Forward(blk.PostAttentionLayerNorm.Forward(h, cfg.RMSNormEps)) @@ -323,8 +329,8 @@ type MoEBlock struct { } // Forward applies the MoE block -func (blk *MoEBlock) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { - r := blk.Attention.Forward(blk.InputLayerNorm.Forward(x, cfg.RMSNormEps), c, positions, B, L, cfg) +func (blk *MoEBlock) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + r := blk.Attention.Forward(blk.InputLayerNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg) h := mlx.Add(x, r) r = blk.MoE.Forward(blk.PostAttentionLayerNorm.Forward(h, cfg.RMSNormEps), cfg) @@ -333,7 +339,7 @@ func (blk *MoEBlock) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, // Block interface for both dense and MoE blocks type Block interface { - Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array + Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array } // Model represents the complete GLM4-MoE-Lite model @@ -707,7 +713,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { if caches != nil { c = caches[i] } - h = layer.Forward(h, c, positions, B, L, m.Config) + h = layer.Forward(h, b, c, positions, B, L, m.Config) } h = m.Norm.Forward(h, m.RMSNormEps) diff --git a/x/models/llama/llama.go b/x/models/llama/llama.go index d28498225..fd47c2b26 100644 --- a/x/models/llama/llama.go +++ b/x/models/llama/llama.go @@ -248,7 +248,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { if caches != nil && i < len(caches) { c = caches[i] } - h = layer.Forward(h, c, positions, B, L, m.Config) + h = layer.Forward(h, b, c, positions, B, L, m.Config) } return m.Norm.Forward(h, m.RMSNormEps) @@ -278,12 +278,12 @@ func (m *Model) NewCaches() []cache.Cache { return caches } -func (l *Layer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { - h := mlx.Add(x, l.Attention.Forward(l.AttentionNorm.Forward(x, cfg.RMSNormEps), c, positions, B, L, cfg)) +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + h := mlx.Add(x, l.Attention.Forward(l.AttentionNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg)) return mlx.Add(h, l.MLP.Forward(l.MLPNorm.Forward(h, cfg.RMSNormEps))) } -func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { q := a.QProj.Forward(x) k := a.KProj.Forward(x) v := a.VProj.Forward(x) @@ -300,13 +300,16 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B q = mlx.RoPEWithBase(q, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) k = mlx.RoPEWithBase(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) - if c != nil { - k, v = c.Update(k, v) - } - // MLX SDPA supports grouped-query attention directly (Q heads can be a // multiple of K/V heads), so avoid materializing repeated K/V tensors. - out := mlx.ScaledDotProductAttentionCausal(q, k, v, cfg.Scale, L > 1) + var kv nn.SDPAOption + if c != nil { + history := c.(cache.Attention).Update(b, k, v) + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) + } + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) return a.OProj.Forward(out) } diff --git a/x/models/nn/recurrent.go b/x/models/nn/recurrent.go new file mode 100644 index 000000000..6f0f7474f --- /dev/null +++ b/x/models/nn/recurrent.go @@ -0,0 +1,259 @@ +package nn + +import ( + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// RecurrentOption configures a call to CausalConv1D or GatedDelta. +type RecurrentOption func(*recurrentConfig) + +// recurrentConfig is the resolved set of inputs supplied via +// RecurrentOption. Exactly one of history or (convState/deltaState) +// must be supplied per call. +type recurrentConfig struct { + history *RecurrentHistory + convState *mlx.Array + deltaState *mlx.Array +} + +// WithRecurrentHistory supplies a cache's per-layer view of conv and +// delta state. The cache hides any storage layout (per-row, paged, +// gather/scatter) behind the history. +func WithRecurrentHistory(h *RecurrentHistory) RecurrentOption { + return func(c *recurrentConfig) { c.history = h } +} + +// WithRecurrentState supplies explicit conv and delta state tensors +// for the no-cache path. Each wrapper consumes one of the two — pass +// nil for the unused slot when calling only one wrapper. +func WithRecurrentState(convState, deltaState *mlx.Array) RecurrentOption { + return func(c *recurrentConfig) { + c.convState = convState + c.deltaState = deltaState + } +} + +// resolve applies opts and panics if WithRecurrentHistory and +// WithRecurrentState were combined or neither was supplied. +func resolveRecurrentConfig(opts []RecurrentOption) recurrentConfig { + var cfg recurrentConfig + for _, opt := range opts { + opt(&cfg) + } + + haveHistory := cfg.history != nil + haveState := cfg.convState != nil || cfg.deltaState != nil + if haveHistory && haveState { + panic("WithRecurrentHistory and WithRecurrentState are mutually exclusive") + } + if !haveHistory && !haveState { + panic("no recurrent state supplied (use WithRecurrentHistory or WithRecurrentState)") + } + + return cfg +} + +// CausalConv1D runs a depthwise causal 1D convolution with recurrent +// state management. Prepends the prior conv state along axis 1, runs +// the conv, and returns (output, nextConv). nextConv is the trailing +// convTail positions of the concat — write it back to the cache via +// Put alongside the scan's new delta state. +// +// Conv selection: when conv is non-nil (a full nn.Conv1d layer), it +// runs through conv.Forward. Otherwise weight is treated as the bare +// depthwise kernel [C, K] and the fallback manual implementation runs. +// Exactly one of conv or weight should be non-nil. +// +// Shapes: input [B, L, D]; prior state [B, convTail, D]; output +// [B, L, D] (the causal conv strips the prepended state). +// +// Prior state comes from exactly one of WithRecurrentHistory (cache +// path) or WithRecurrentState (no-cache path). +func CausalConv1D(b *batch.Batch, input *mlx.Array, conv *Conv1d, weight *mlx.Array, convTail int, opts ...RecurrentOption) (out, nextConv *mlx.Array) { + cfg := resolveRecurrentConfig(opts) + var prior *mlx.Array + if cfg.history != nil { + prior = cfg.history.ConvState() + } else { + prior = cfg.convState + } + + mask := paddingMask(b, int32(input.Dim(1))) + if mask != nil { + zero := mlx.FromValue(float32(0)).AsType(input.DType()) + input = mlx.Where(mlx.ExpandDims(mask, 2), input, zero) + } + + concat := mlx.Concatenate([]*mlx.Array{prior, input}, 1) + if conv != nil { + out = conv.Forward(concat) + } else { + out = depthwiseCausalConv1d(concat, weight, int32(input.Dim(1))) + } + + B := int32(concat.Dim(0)) + total := int32(concat.Dim(1)) + D := int32(concat.Dim(2)) + + // Gather the tail from each of the non-padded sequence ends + if mask != nil && convTail > 0 { + offsets := make([]int32, int(B)*convTail) + + for i := range int(B) { + end := b.SeqQueryLens[i] + + for k := range convTail { + offsets[i*convTail+k] = end + int32(k) + } + } + + positions := mlx.NewArrayInt32(offsets, []int32{B, int32(convTail), 1}) + nextConv = mlx.TakeAlongAxis(concat, positions, 1) + } else { + nextConv = mlx.SliceStartStop(concat, + []int32{0, total - int32(convTail), 0}, + []int32{B, total, D}) + } + + return out, nextConv +} + +// depthwiseCausalConv1d implements a depthwise 1D causal convolution +// manually as a sum of kernel-offset multiplies. x has shape +// [B, inLen, C], weight has shape [C, K]; output has shape [B, outLen, C] +// where outLen = inLen - K + 1 (the caller passes outLen to avoid the +// subtraction). Used as the fallback path in CausalConv1D when no +// full Conv1d layer is configured. +func depthwiseCausalConv1d(x, w *mlx.Array, outLen int32) *mlx.Array { + if x == nil || w == nil { + return nil + } + if w.NumDims() != 2 { + return nil + } + B := int32(x.Dim(0)) + C := int32(w.Dim(0)) + K := int32(w.Dim(1)) + var out *mlx.Array + for i := range K { + seg := mlx.SliceStartStop(x, []int32{0, i, 0}, []int32{B, i + outLen, C}) + wi := mlx.SliceStartStop(w, []int32{0, i}, []int32{C, i + 1}) + wi = mlx.Reshape(wi, 1, 1, C) + term := mlx.Mul(seg, wi) + if out == nil { + out = term + } else { + out = mlx.Add(out, term) + } + } + return out +} + +// GatedDelta wraps mlx.FastGatedDelta with recurrent state management. +// Reads prior delta state from the supplied option and returns +// (output, newDelta). Write newDelta back via the cache's Put +// alongside the conv wrapper's nextConv. +// +// Shape conventions: +// +// q: [B, L, numKeyHeads, headKDim] +// k: [B, L, numKeyHeads, headKDim] +// v: [B, L, numValueHeads, headVDim] +// state: [B, numValueHeads, headVDim, headKDim] +// +// Prior state comes from exactly one of WithRecurrentHistory (cache +// path) or WithRecurrentState (no-cache path). +func GatedDelta(b *batch.Batch, q, k, v, gDecay, beta *mlx.Array, opts ...RecurrentOption) (out, newDelta *mlx.Array) { + cfg := resolveRecurrentConfig(opts) + var state *mlx.Array + if cfg.history != nil { + state = cfg.history.DeltaState() + } else { + state = cfg.deltaState + } + + return mlx.FastGatedDelta(q, k, v, gDecay, beta, state, paddingMask(b, int32(q.Dim(1)))) +} + +// RecurrentHistory is an opaque per-forward view a recurrent cache +// hands to the SSM kernel wrappers — prior conv and delta state +// tensors. Models do not construct this directly; pass it through +// via WithRecurrentHistory, or use WithRecurrentState on the no-cache +// path. +// +// Opaque structure to model code; accessors ConvState/DeltaState +// provide the escape hatch for custom SSM paths. +type RecurrentHistory struct { + convState, deltaState *mlx.Array +} + +// NewRecurrentHistory constructs a RecurrentHistory. Intended for +// cache implementations across packages; model code uses +// WithRecurrentHistory / WithRecurrentState instead. +func NewRecurrentHistory(convState, deltaState *mlx.Array) *RecurrentHistory { + return &RecurrentHistory{convState: convState, deltaState: deltaState} +} + +// ConvState returns the current convolution state tensor. +// +// Last-resort escape hatch for custom SSM paths — may force a slow +// materialization to canonical form depending on the cache's +// internal storage. Prefer CausalConv1D via WithRecurrentHistory. +func (h *RecurrentHistory) ConvState() *mlx.Array { return h.convState } + +// DeltaState returns the current delta state tensor. +// +// Last-resort escape hatch for custom SSM paths — may force a slow +// materialization to canonical form depending on the cache's +// internal storage. Prefer GatedDelta via WithRecurrentHistory. +func (h *RecurrentHistory) DeltaState() *mlx.Array { return h.deltaState } + +type paddingMaskInputs struct { + batch *batch.Batch + L int32 +} + +func (in paddingMaskInputs) build() *mlx.Array { + B := len(in.batch.SeqQueryLens) + + needed := false + for i := range B { + if in.batch.SeqQueryLens[i] < in.L { + needed = true + break + } + } + if !needed { + return nil + } + + L := int(in.L) + vals := make([]bool, B*L) + for i := range B { + n := int(in.batch.SeqQueryLens[i]) + + base := i * L + for j := range n { + vals[base+j] = true + } + } + + return mlx.FromValues(vals, B, L) +} + +// paddingMask derives a [B, L] bool mask from b.SeqQueryLens for +// right-padded inputs (real tokens at [0, len_i), padding at +// [len_i, L)). Returns nil when b has no rows or every row is full — +// the no-padding fast path that costs nothing extra. +func paddingMask(b *batch.Batch, L int32) *mlx.Array { + inputs := paddingMaskInputs{batch: b, L: L} + if cached, ok := b.Memo.Get(inputs); ok { + return cached.(*mlx.Array) + } + + mask := inputs.build() + b.Memo.Put(inputs, mask) + + return mask +} diff --git a/x/models/nn/recurrent_test.go b/x/models/nn/recurrent_test.go new file mode 100644 index 000000000..d6e0d3538 --- /dev/null +++ b/x/models/nn/recurrent_test.go @@ -0,0 +1,340 @@ +package nn + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +func ones(dtype mlx.DType, shape ...int) *mlx.Array { + return mlx.AddScalar(mlx.Zeros(dtype, shape...), 1) +} + +// fromValues builds a tensor with sequentially-numbered float32 +// values so element-by-element parity actually exercises the kernel. +func fromValues(seed float32, shape ...int) *mlx.Array { + n := 1 + for _, d := range shape { + n *= d + } + vals := make([]float32, n) + for i := range vals { + vals[i] = seed + 0.1*float32(i) + } + return mlx.FromValues(vals, shape...) +} + +// depthwiseCausalRef is a Go-side reference for the depthwise causal +// 1D conv fallback. concat is [B, total, C], weight is [C, K], output +// is [B, total-K+1, C]. Used to anchor the wrapper's parity tests. +func depthwiseCausalRef(concat, weight *mlx.Array) []float32 { + mlx.Eval(concat, weight) + cVals := concat.Floats() + wVals := weight.Floats() + B := concat.Dim(0) + total := concat.Dim(1) + C := concat.Dim(2) + K := weight.Dim(1) + outLen := total - K + 1 + out := make([]float32, B*outLen*C) + for bi := range B { + for q := range outLen { + for c := range C { + var sum float32 + for k := range K { + x := cVals[bi*total*C+(q+k)*C+c] + w := wVals[c*K+k] + sum += x * w + } + out[bi*outLen*C+q*C+c] = sum + } + } + } + return out +} + +// TestCausalConv1DParity drives the wrapper with non-trivial prior, +// input, and weight values, then compares against a direct depthwise- +// causal-conv reference. +func TestCausalConv1DParity(t *testing.T) { + skipIfNoMLX(t) + B, L, D, convTail := 1, 4, 3, 2 + K := convTail + 1 + + input := fromValues(0.5, B, L, D) + prior := fromValues(-0.3, B, convTail, D) + weight := fromValues(0.2, D, K) + + out, nextConv := CausalConv1D(&batch.Batch{}, input, nil, weight, convTail, WithRecurrentState(prior, nil)) + mlx.Eval(out, nextConv) + + concat := mlx.Concatenate([]*mlx.Array{prior, input}, 1) + want := depthwiseCausalRef(concat, weight) + got := out.Floats() + if len(got) != len(want) { + t.Fatalf("out len = %d, want %d", len(got), len(want)) + } + for i := range want { + if math.Abs(float64(got[i]-want[i])) > 1e-5 { + t.Fatalf("out[%d]: got %v, want %v", i, got[i], want[i]) + } + } + + // nextConv (no padding) is the trailing convTail rows of concat. + mlx.Eval(concat) + cVals := concat.Floats() + total := concat.Dim(1) + wantTail := make([]float32, B*convTail*D) + for bi := range B { + for k := range convTail { + for d := range D { + wantTail[bi*convTail*D+k*D+d] = cVals[bi*total*D+(total-convTail+k)*D+d] + } + } + } + tail := nextConv.Floats() + if len(tail) != len(wantTail) { + t.Fatalf("nextConv len = %d, want %d", len(tail), len(wantTail)) + } + for i := range wantTail { + if tail[i] != wantTail[i] { + t.Fatalf("nextConv[%d]: got %v, want %v", i, tail[i], wantTail[i]) + } + } +} + +// TestCausalConv1DPaddedRowParity drives a B=2 batch with one short +// row (qLen 1e-5 { + t.Fatalf("row 1 real out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[gotI], wantOut1[refI]) + } + } + } + + // nextConv: row 0 unaffected, row 1 must be the row's real tail + // (positions [qLenShort - convTail, qLenShort) of the per-row + // concat, i.e. the last two real input rows in this setup). + gotTail := nextConv.Floats() + wantTail0 := refNextConv0.Floats() + wantTail1 := refNextConv1.Floats() + for k := range convTail { + for d := range D { + i := k*D + d + if gotTail[i] != wantTail0[i] { + t.Fatalf("row 0 nextConv[k=%d,d=%d]: got %v, want %v", k, d, gotTail[i], wantTail0[i]) + } + } + } + for k := range convTail { + for d := range D { + gotI := convTail*D + k*D + d + refI := k*D + d + if gotTail[gotI] != wantTail1[refI] { + t.Fatalf("row 1 nextConv[k=%d,d=%d]: got %v, want %v (must come from real positions, not the padded tail)", + k, d, gotTail[gotI], wantTail1[refI]) + } + } + } +} + +func TestGatedDeltaZeroFallback(t *testing.T) { + skipIfNoMLX(t) + B, L, nK, nV, dK, dV := 1, 2, 1, 1, 4, 4 + q := ones(mlx.DTypeFloat32, B, L, nK, dK) + k := ones(mlx.DTypeFloat32, B, L, nK, dK) + v := ones(mlx.DTypeFloat32, B, L, nV, dV) + gDecay := ones(mlx.DTypeFloat32, B, L, nV) + beta := ones(mlx.DTypeFloat32, B, L, nV) + + zero := mlx.Zeros(mlx.DTypeFloat32, B, nV, dV, dK) + outA, stateA := GatedDelta(&batch.Batch{}, q, k, v, gDecay, beta, WithRecurrentState(nil, zero)) + outB, stateB := mlx.FastGatedDelta(q, k, v, gDecay, beta, zero, nil) + mlx.Eval(outA, stateA, outB, stateB) + + gotOut, wantOut := outA.Floats(), outB.Floats() + for i := range wantOut { + if gotOut[i] != wantOut[i] { + t.Fatalf("output[%d]: wrapper=%v direct=%v", i, gotOut[i], wantOut[i]) + } + } + gotState, wantState := stateA.Floats(), stateB.Floats() + for i := range wantState { + if gotState[i] != wantState[i] { + t.Fatalf("state[%d]: wrapper=%v direct=%v", i, gotState[i], wantState[i]) + } + } +} + +func TestGatedDeltaUsesPriorState(t *testing.T) { + skipIfNoMLX(t) + B, L, nK, nV, dK, dV := 1, 2, 1, 1, 4, 4 + q := ones(mlx.DTypeFloat32, B, L, nK, dK) + k := ones(mlx.DTypeFloat32, B, L, nK, dK) + v := ones(mlx.DTypeFloat32, B, L, nV, dV) + gDecay := ones(mlx.DTypeFloat32, B, L, nV) + beta := ones(mlx.DTypeFloat32, B, L, nV) + + priorState := mlx.MulScalar(ones(mlx.DTypeFloat32, B, nV, dV, dK), 3) + + outA, _ := GatedDelta(&batch.Batch{}, q, k, v, gDecay, beta, WithRecurrentState(nil, priorState)) + outB, _ := mlx.FastGatedDelta(q, k, v, gDecay, beta, priorState, nil) + mlx.Eval(outA, outB) + + gotOut, wantOut := outA.Floats(), outB.Floats() + for i := range wantOut { + if gotOut[i] != wantOut[i] { + t.Fatalf("output[%d]: wrapper=%v direct=%v", i, gotOut[i], wantOut[i]) + } + } +} + +// TestGatedDeltaPaddedRowParity drives a B=2 batch where row 1 is +// short (qLen < L). The wrapper must substitute neutral values +// (q=k=v=beta=0, g=1) at row 1's padded positions so the recurrence +// is a no-op there — and row 1's final state must equal the state +// after its last real token. Pinned via parity against a B=1 length- +// qLen call on the same row. +func TestGatedDeltaPaddedRowParity(t *testing.T) { + skipIfNoMLX(t) + L, nK, nV, dK, dV := 4, 1, 1, 4, 4 + qLenShort := 2 + + makeRows := func(seedA, seedB float32, shape ...int) *mlx.Array { + // Build a rank-(len(shape)+1) tensor with B=2 rows from two + // distinct seeds so the rows are not accidentally identical. + n := 1 + for _, d := range shape { + n *= d + } + vals := make([]float32, 2*n) + for i := range n { + vals[i] = seedA + 0.1*float32(i) + } + for i := range n { + vals[n+i] = seedB + 0.1*float32(i) + } + full := append([]int{2}, shape...) + return mlx.FromValues(vals, full...) + } + + q := makeRows(0.5, -0.5, L, nK, dK) + k := makeRows(0.7, -0.7, L, nK, dK) + v := makeRows(0.3, -0.3, L, nV, dV) + gDecay := makeRows(0.1, -0.1, L, nV) + beta := makeRows(0.4, -0.4, L, nV) + priorState := makeRows(0.2, -0.2, nV, dV, dK) + + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 2, L), + SeqOffsets: []int32{0, 0}, + SeqQueryLens: []int32{int32(L), int32(qLenShort)}, + } + _, state := GatedDelta(b, q, k, v, gDecay, beta, WithRecurrentState(nil, priorState)) + mlx.Eval(state) + + // Reference for row 1: B=1 length-qLenShort call against the + // row's real prefix and its prior state slice. + row1Slice := func(a *mlx.Array, axisLens ...int32) *mlx.Array { + dims := a.Dims() + start := make([]int32, len(dims)) + stop := make([]int32, len(dims)) + start[0], stop[0] = 1, 2 + for i := 1; i < len(dims); i++ { + stop[i] = int32(dims[i]) + } + // Optionally truncate axis 1 (sequence axis) to qLenShort. + if len(axisLens) >= 1 && len(dims) >= 2 { + stop[1] = axisLens[0] + } + return mlx.SliceStartStop(a, start, stop) + } + q1 := row1Slice(q, int32(qLenShort)) + k1 := row1Slice(k, int32(qLenShort)) + v1 := row1Slice(v, int32(qLenShort)) + gDecay1 := row1Slice(gDecay, int32(qLenShort)) + beta1 := row1Slice(beta, int32(qLenShort)) + priorRow1 := row1Slice(priorState) + + _, refState := mlx.FastGatedDelta(q1, k1, v1, gDecay1, beta1, priorRow1, nil) + mlx.Eval(refState) + + gotState := state.Floats() + wantState := refState.Floats() + row1Stride := nV * dV * dK + for i := range row1Stride { + gotV := gotState[row1Stride+i] + wantV := wantState[i] + if math.Abs(float64(gotV-wantV)) > 1e-4 { + t.Fatalf("row 1 final state[%d]: got %v, want %v", i, gotV, wantV) + } + } +} diff --git a/x/models/nn/sdpa.go b/x/models/nn/sdpa.go new file mode 100644 index 000000000..3aaa2a908 --- /dev/null +++ b/x/models/nn/sdpa.go @@ -0,0 +1,578 @@ +package nn + +import ( + "encoding/binary" + "math" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// SDPAOption configures a call to ScaledDotProductAttention. +type SDPAOption func(*sdpaConfig) + +type sdpaConfig struct { + // Exactly one of (k,v,kLens) or history supplies keys/values. + k, v *mlx.Array + kLens []int32 + history *KVHistory + + // Optional model-supplied logical mask. + mask AttentionMask +} + +// WithKVHistory supplies a cache's per-layer view of K and V. The +// cache hides any storage layout (sliding window, ring buffer, +// k-padding) behind the history. +func WithKVHistory(h *KVHistory) SDPAOption { + return func(c *sdpaConfig) { c.history = h } +} + +// WithMLAHistory supplies a cache's per-layer view for absorbed MLA +// attention, where V is the first valueDim positions of K. +func WithMLAHistory(h *KVHistory, valueDim int) SDPAOption { + v := h.K().Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim)) + return WithKVHistory(&KVHistory{k: h.K(), v: v, applier: h.applier}) +} + +// WithKV supplies explicit K/V tensors for the no-cache path. kLens +// gives per-row real key extents — pass b.SeqQueryLens for self- +// attention, or the caller's own extents for cross-attention. +func WithKV(k, v *mlx.Array, kLens []int32) SDPAOption { + return func(c *sdpaConfig) { c.k = k; c.v = v; c.kLens = kLens } +} + +// WithMask supplies the model's logical-coordinate mask. +func WithMask(m AttentionMask) SDPAOption { + return func(c *sdpaConfig) { c.mask = m } +} + +// ScaledDotProductAttention runs the fast SDPA kernel against q and +// the keys/values supplied via exactly one of WithKV or +// WithKVHistory. Automatically applies any Q/K padding masking required +// for padded batches. +func ScaledDotProductAttention(b *batch.Batch, q *mlx.Array, scale float32, opts ...SDPAOption) *mlx.Array { + var cfg sdpaConfig + for _, opt := range opts { + opt(&cfg) + } + + haveKV := cfg.k != nil || cfg.v != nil + haveHistory := cfg.history != nil + if haveKV && haveHistory { + panic("nn.ScaledDotProductAttention: WithKV and WithKVHistory are mutually exclusive") + } + if !haveKV && !haveHistory { + panic("nn.ScaledDotProductAttention: no keys/values supplied (use WithKV or WithKVHistory)") + } + + k, v := cfg.k, cfg.v + var applier MaskApplier + if cfg.history != nil { + k = cfg.history.K() + v = cfg.history.V() + applier = cfg.history.applier + } + + inputs := dispatchInputs{ + batch: b, + mask: cfg.mask, + applier: applier, + K: k.Dim(2), + dtype: k.DType(), + kLens: newKLensKey(cfg.kLens), + } + + if cached, ok := b.Memo.Get(inputs); ok { + d := cached.(sdpaDispatch) + return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr) + } + + d := inputs.resolve() + b.Memo.Put(inputs, d) + return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr) +} + +// sdpaDispatch is the resolved kernel call for a given SDPA key — +// either a flag-mode fast path (mode "" or "causal", arr nil) or an +// array-mode call with a materialized tensor. Memoized on b.Memo so +// sibling layers skip applier composition, padding build, and AsArray. +type sdpaDispatch struct { + mode string + arr *mlx.Array +} + +// dispatchInputs bundles every value resolve reads and doubles as +// the Memo map key. All fields are comparable: batch is a +// *batch.Batch pointer, the applier interface is comparable when +// its concrete type is, and kLens is a kLensKey string that hashes +// by content. +// +// Making resolve a method on this struct is the enforcement — any +// new dependency must be added as a field, which automatically +// participates in the map key. +// +// applier and kLens are mutually exclusive by construction: +// WithKVHistory sets applier (which owns any K-padding in its output +// space) and leaves kLens ""; WithKV sets kLens and leaves applier nil. +type dispatchInputs struct { + batch *batch.Batch + mask AttentionMask + applier MaskApplier + K int + dtype mlx.DType + kLens kLensKey +} + +// kLensKey is a comparable encoding of an int32 slice (four bytes +// per element, native endian) so it can live in a struct used as a +// map key. Decode back via Int32s. +type kLensKey string + +func newKLensKey(vs []int32) kLensKey { + if len(vs) == 0 { + return "" + } + buf := make([]byte, len(vs)*4) + for i, v := range vs { + binary.NativeEndian.PutUint32(buf[i*4:], uint32(v)) + } + return kLensKey(buf) +} + +// Int32s decodes the key back to a fresh []int32. +func (k kLensKey) Int32s() []int32 { + if len(k) == 0 { + return nil + } + b := []byte(k) + out := make([]int32, len(b)/4) + for i := range out { + out[i] = int32(binary.NativeEndian.Uint32(b[i*4:])) + } + return out +} + +// resolve composes model + padding + storage contributions and +// returns the kernel dispatch decision. Reads only from inputs; any +// new input must be added to dispatchInputs. +// +// Order matters: QPaddingMask is added in logical Q-space before the +// applier runs, so an applier that remaps coordinates receives the +// full logical mask. The applier and KPaddingMask branches are +// mutually exclusive — on the applier path the output may be in a +// remapped K space, so the applier owns any K-padding; on the +// WithKV path kLens describes the direct K tensor, which shares +// logical K space with QPaddingMask. +func (inputs dispatchInputs) resolve() sdpaDispatch { + mask := inputs.mask.Intersect(QPaddingMask(inputs.batch, inputs.dtype)) + + if inputs.applier != nil { + mask = inputs.applier.ApplyMask(mask) + } else if inputs.kLens != "" { + mask = mask.Intersect(KPaddingMask(inputs.batch, inputs.K, inputs.kLens.Int32s(), inputs.dtype)) + } + + switch { + case mask.IsZero(): + return sdpaDispatch{mode: ""} + case mask.IsCausal(): + if inputs.batch.InputIDs.Dim(1) == 1 { + // At L=1 the causal "k > q" constraint is redundant - + // drop it so the kernel dispatches to the no-mask fast path. + return sdpaDispatch{mode: ""} + } else { + return sdpaDispatch{mode: "causal"} + } + default: + return sdpaDispatch{mode: "array", arr: mask.AsArray(inputs.batch, inputs.K, inputs.dtype)} + } +} + +// MaskApplier composes a cache's storage-mask contribution onto a +// fully-composed logical mask. The returned mask may live in the +// applier's own coordinate system (e.g. a rotated or compacted K layout), +// so any addition in logical K space must happen before the applier runs. +// SDPA does not add KPaddingMask on this path — the applier owns any +// K-padding its output needs. +// +// Implementations must be comparable struct values whose fields +// capture everything the composition depends on (no slice, map, or +// func fields); the value doubles as the applier's identity in +// SDPA's dispatch-cache key, where a non-comparable concrete type +// would panic at map insertion. A nil MaskApplier means "no storage +// contribution". +type MaskApplier interface { + ApplyMask(logical AttentionMask) AttentionMask +} + +// KVHistory is the per-forward view a KV cache hands to SDPA: +// post-Update K and V plus an optional MaskApplier that composes +// the cache's storage mask onto the caller's model mask. +type KVHistory struct { + k, v *mlx.Array + applier MaskApplier +} + +// NewKVHistory constructs a KVHistory. Intended for +// cache implementations across packages; model code uses +// WithKVHistory / WithKV instead. +func NewKVHistory(k, v *mlx.Array, applier MaskApplier) *KVHistory { + return &KVHistory{k: k, v: v, applier: applier} +} + +// K returns the post-Update keys tensor. +// +// Last-resort escape hatch for custom attention paths — may force a +// slow materialization to canonical form depending on the cache's +// internal storage. Prefer ScaledDotProductAttention via +// WithKVHistory. +func (h *KVHistory) K() *mlx.Array { return h.k } + +// V returns the post-Update values tensor. +// +// Last-resort escape hatch for custom attention paths — may force a +// slow materialization to canonical form depending on the cache's +// internal storage. Prefer ScaledDotProductAttention via +// WithKVHistory. +func (h *KVHistory) V() *mlx.Array { return h.v } + +// Mask returns the final AttentionMask for this layer's SDPA — +// cache storage restrictions composed onto the caller's fully- +// composed logical mask. +// +// Last-resort escape hatch for custom attention paths — may force a +// slow materialization to canonical form depending on the cache's +// internal storage. Prefer ScaledDotProductAttention via +// WithKVHistory. +func (h *KVHistory) Mask(logical AttentionMask) AttentionMask { + if h.applier == nil { + return logical + } + return h.applier.ApplyMask(logical) +} + +// AttentionMask describes an attention mask in four states: +// - zero value: no mask. +// - flag-form causal (causal=true only): dispatches to the MLX +// kernel's mask_mode="causal" fast path. +// - causal with relaxation rectangles: a causal mask with +// bidirectional attention rectangles, such as for images. +// - additive tensor (array!=nil): broadcast-compatible with +// [B, 1, L, K]; contributed by a custom mask, helpers such as +// QPaddingMask, KPaddingMask, or cache appliers and accumulated +// via Intersect. +// +// The mask is a pure logical description — it carries no batch and +// exists independent of cache storage layout. +// +// All fields are comparable, so AttentionMask values compare with == +// by full identity — SDPA uses this directly as a dispatch-cache key. +type AttentionMask struct { + causal bool + relaxations *relaxNode + array *mlx.Array +} + +type relaxRect struct { + seq, qLo, qHi, kLo, kHi int +} + +// relaxNode is a singly-linked list node holding relaxation +// rectangles. Each AttentionMask must have a fresh set of +// nodes to avoid false sharing between masks. +type relaxNode struct { + rect relaxRect + next *relaxNode +} + +// CausalMask returns a flag-form causal mask. The mask stays +// tensor-free — hitting the kernel's mask_mode="causal" fast path — +// until something composes a relaxation, padding, or applier tensor +// onto it; then SDPA materializes via AsArray. +func CausalMask() AttentionMask { + return AttentionMask{causal: true} +} + +// ArrayMask wraps an explicit additive tensor broadcast-compatible +// with [B, 1, L, K]. +func ArrayMask(a *mlx.Array) AttentionMask { + return AttentionMask{array: a} +} + +// IsZero reports whether the mask is the zero value (no mask at all). +func (m AttentionMask) IsZero() bool { + return !m.causal && m.array == nil && m.relaxations == nil +} + +// IsCausal reports whether the mask is pure flag-form causal — no +// relaxations and no accumulated array. SDPA dispatches to the +// kernel's "causal" fast path on this; any padding, applier +// contribution, or relaxation falls to the array path. +func (m AttentionMask) IsCausal() bool { + return m.causal && m.relaxations == nil && m.array == nil +} + +// Relax records a relaxation rectangle for batch sequence seq — +// positions (q, k) with q in [qLo, qHi) and k in [kLo, kHi) become +// freely attendable regardless of the causal base. Coordinates are +// absolute sequence positions on both axes, matching how causal is +// defined (k <= q). Multiple calls compose as a union per sequence. +// +// Rectangles that cannot change any cell — empty or already fully +// inside causal (kHi-1 <= qLo) — are dropped so IsCausal stays true +// and the mask remains on the kernel's fast path. +// +// Panics on pure ArrayMask (the caller owns the tensor and should +// modify it directly) or on the zero mask (nothing to relax). +func (m AttentionMask) Relax(seq, qLo, qHi, kLo, kHi int) AttentionMask { + if !m.causal { + if m.array != nil { + panic("AttentionMask.Relax: cannot relax a pure ArrayMask; modify the tensor directly") + } + panic("AttentionMask.Relax: cannot relax a zero mask") + } + if qLo >= qHi || kLo >= kHi { + return m + } + if kHi-1 <= qLo { + return m + } + m.relaxations = &relaxNode{ + rect: relaxRect{seq: seq, qLo: qLo, qHi: qHi, kLo: kLo, kHi: kHi}, + next: m.relaxations, + } + return m +} + +// Intersect returns the element-wise sum of this mask and other. Masks are +// additive and apply before softmax, so this is intersection +// semantics — a position is valid only if both sides have 0 there. +// +// At AsArray time a causal+Relax+array mask materializes as: causal +// writes -inf into the upper triangle, Relax overwrites its +// rectangles back to 0, then array is added on top — restricting 0 +// cells further or no-op'ing on -inf cells. +func (m AttentionMask) Intersect(other AttentionMask) AttentionMask { + if m.IsZero() { + return other + } + if other.IsZero() { + return m + } + + result := AttentionMask{ + causal: m.causal || other.causal, + } + + // Relax requires causal, so relaxations != nil implies causal. + switch { + case m.relaxations != nil && other.relaxations != nil: + // Both sides causal+Relax: pairwise rect intersection per sequence. + var list *relaxNode + for a := m.relaxations; a != nil; a = a.next { + for b := other.relaxations; b != nil; b = b.next { + if a.rect.seq != b.rect.seq { + continue + } + qLo := max(a.rect.qLo, b.rect.qLo) + qHi := min(a.rect.qHi, b.rect.qHi) + kLo := max(a.rect.kLo, b.rect.kLo) + kHi := min(a.rect.kHi, b.rect.kHi) + if qHi <= qLo || kHi <= kLo || kHi-1 <= qLo { + continue + } + list = &relaxNode{ + rect: relaxRect{seq: a.rect.seq, qLo: qLo, qHi: qHi, kLo: kLo, kHi: kHi}, + next: list, + } + } + } + result.relaxations = list + case m.relaxations != nil && !other.causal: + result.relaxations = m.relaxations + case other.relaxations != nil && !m.causal: + result.relaxations = other.relaxations + default: + // Implicit: one side causal+Relax, the other plain causal + // (no relaxations). Plain causal blocks every cell Relax + // tried to release, so intersection with its empty release + // set leaves nothing — result.relaxations stays nil and + // collapses to pure causal. + } + + switch { + case m.array != nil && other.array != nil: + result.array = mlx.Add(m.array, other.array) + case m.array != nil: + result.array = m.array + case other.array != nil: + result.array = other.array + } + + return result +} + +// AsArray materializes the mask as a [B, 1, L, K] additive tensor +// (0 where valid, -inf where blocked). B and L come from b; K and +// dtype come from the caller. +// +// Composition order: +// 1. Start from zero. +// 2. If m.causal: -inf where oldestPos+k > SeqOffsets[b] + q per row. +// 3. Apply m.relaxations (qLo/qHi and kLo/kHi are absolute positions). +// 4. Add m.array if present. +func (m AttentionMask) AsArray(b *batch.Batch, K int, dtype mlx.DType) *mlx.Array { + // Pure ArrayMask: caller owns the tensor, nothing to compose. + if !m.causal && m.relaxations == nil && m.array != nil { + if m.array.DType() == dtype { + return m.array + } + return m.array.AsType(dtype) + } + + B := len(b.SeqOffsets) + L := b.InputIDs.Dim(1) + + negInf := float32(math.Inf(-1)) + vals := make([]float32, B*L*K) + if m.causal { + for i := range B { + off := int(b.SeqOffsets[i]) + oldestPos := max(0, off+L-K) + base := i * L * K + for q := range L { + absQ := off + q + row := base + q*K + for k := range K { + if oldestPos+k > absQ { + vals[row+k] = negInf + } + } + } + } + } + + for n := m.relaxations; n != nil; n = n.next { + r := n.rect + if r.seq < 0 || r.seq >= B { + continue + } + off := int(b.SeqOffsets[r.seq]) + oldestPos := max(0, off+L-K) + qLo := min(max(r.qLo-off, 0), L) + qHi := min(max(r.qHi-off, 0), L) + kLo := min(max(r.kLo-oldestPos, 0), K) + kHi := min(max(r.kHi-oldestPos, 0), K) + base := r.seq * L * K + for q := qLo; q < qHi; q++ { + row := base + q*K + for k := kLo; k < kHi; k++ { + vals[row+k] = 0 + } + } + } + + out := mlx.FromValues(vals, B, 1, L, K) + if m.array != nil { + out = mlx.Add(out, m.array) + } + if dtype != mlx.DTypeFloat32 { + out = out.AsType(dtype) + } + return out +} + +// QPaddingMask returns an additive [B, 1, L, 1] mask that blocks +// padded query rows (q >= b.SeqQueryLens[i]) across all keys. It is +// logical — independent of whatever layout the cache uses for K. +// Returns the zero mask when every row is full. +func QPaddingMask(b *batch.Batch, dtype mlx.DType) AttentionMask { + return padTailMask(len(b.SeqOffsets), b.InputIDs.Dim(1), 2, b.SeqQueryLens, dtype) +} + +// KPaddingMask returns an additive [B, 1, 1, K] mask that blocks +// padded key columns (k >= kLens[i]) across all queries. Storage- +// dependent: kLens describes where real content ends in physical K, +// so this is typically used without a cache where the caller knows +// the actual layout. Returns the zero mask when every row is full. +func KPaddingMask(b *batch.Batch, K int, kLens []int32, dtype mlx.DType) AttentionMask { + return padTailMask(len(b.SeqOffsets), K, 3, kLens, dtype) +} + +// SlidingWindowMask returns an additive [B, 1, L, K] mask blocking +// keys outside a per-row window of size `window`: any key whose +// absolute position p < absQ - window + 1 is blocked. Returns the +// zero mask when window <= 0 or no row needs blocking. +// +// Defined in logical position space — the K axis is position-ordered +// with column 0 at oldestPos = max(0, b.SeqOffsets[i]+L-K). +func SlidingWindowMask(b *batch.Batch, K, window int, dtype mlx.DType) AttentionMask { + if window <= 0 { + return AttentionMask{} + } + B := len(b.SeqOffsets) + L := b.InputIDs.Dim(1) + negInf := float32(math.Inf(-1)) + vals := make([]float32, B*L*K) + needed := false + for i := range B { + off := int(b.SeqOffsets[i]) + oldestPos := max(0, off+L-K) + base := i * L * K + for q := range L { + absQ := off + q + lo := absQ - window + 1 + maskCount := lo - oldestPos + if maskCount <= 0 { + continue + } + if maskCount > K { + maskCount = K + } + row := base + q*K + for k := range maskCount { + vals[row+k] = negInf + needed = true + } + } + } + if !needed { + return AttentionMask{} + } + out := mlx.FromValues(vals, B, 1, L, K) + if dtype != mlx.DTypeFloat32 { + out = out.AsType(dtype) + } + return ArrayMask(out) +} + +func padTailMask(B, total, axis int, lens []int32, dtype mlx.DType) AttentionMask { + needed := false + for i := range B { + if int(lens[i]) < total { + needed = true + break + } + } + if !needed { + return AttentionMask{} + } + + negInf := float32(math.Inf(-1)) + vals := make([]float32, B*total) + for i := range B { + n := int(lens[i]) + base := i * total + for j := n; j < total; j++ { + vals[base+j] = negInf + } + } + shape := [4]int{B, 1, 1, 1} + shape[axis] = total + out := mlx.FromValues(vals, shape[0], shape[1], shape[2], shape[3]) + if dtype != mlx.DTypeFloat32 { + out = out.AsType(dtype) + } + return ArrayMask(out) +} diff --git a/x/models/nn/sdpa_test.go b/x/models/nn/sdpa_test.go new file mode 100644 index 000000000..88285473f --- /dev/null +++ b/x/models/nn/sdpa_test.go @@ -0,0 +1,680 @@ +package nn + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// newBatch constructs a synthetic batch for mask/SDPA tests. +// seqOffsets defines B (length of slice) and each row's absolute start; +// L is the padded query length along InputIDs's second axis; +// qLens is per-row real query length (defaults to all L if nil). +func newBatch(seqOffsets []int32, L int, qLens []int32) *batch.Batch { + B := len(seqOffsets) + if qLens == nil { + qLens = make([]int32, B) + for i := range qLens { + qLens[i] = int32(L) + } + } + // InputIDs values don't matter for masking, only the shape. + ids := mlx.FromValues(make([]int32, B*L), B, L) + return &batch.Batch{ + InputIDs: ids, + SeqOffsets: seqOffsets, + SeqQueryLens: qLens, + } +} + +func TestAttentionMaskZero(t *testing.T) { + skipIfNoMLX(t) + var m AttentionMask + if !m.IsZero() { + t.Fatal("zero value should report IsZero") + } + if m.IsCausal() { + t.Fatal("zero value should not report IsCausal") + } + b := newBatch([]int32{0}, 2, nil) + arr := m.AsArray(b, 3, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("zero value AsArray should return a zeros tensor, not nil") + } + mlx.Eval(arr) + got := arr.Floats() + for i, v := range got { + if v != 0 { + t.Fatalf("zero mask should materialize all zeros; got[%d] = %v", i, v) + } + } +} + +func TestAttentionMaskAsArrayCausal(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 6 + b := newBatch([]int32{2}, L, nil) + arr := CausalMask().AsArray(b, K, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("CausalMask AsArray should return a tensor") + } + dims := arr.Dims() + if len(dims) != 4 || dims[0] != 1 || dims[1] != 1 || dims[2] != L || dims[3] != K { + t.Fatalf("want shape [1,1,%d,%d], got %v", L, K, dims) + } + mlx.Eval(arr) + got := arr.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, L*K) + for q := range L { + absQ := int(b.SeqOffsets[0]) + q + for k := range K { + if k > absQ { + want[q*K+k] = negInf + } + } + } + for i := range want { + if !sameF(got[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestAttentionMaskRelaxLazy(t *testing.T) { + skipIfNoMLX(t) + // Relax must not materialize a tensor — the perf invariant the + // causal-flag fast path relies on. Everything else (predicates, + // AsArray contents) is exercised by the materialization tests. + m := CausalMask(). + Relax(0, 1, 3, 2, 5). + Relax(0, 0, 2, 1, 4) + if m.array != nil { + t.Fatal("Relax should not materialize a tensor") + } +} + +// TestAttentionMaskRelaxNoopRectsMatchCausal pins the contract that +// rectangles which can't change any cell — empty in q or k, or fully +// inside the causal triangle — must produce the same materialized +// tensor as plain causal. +func TestAttentionMaskRelaxNoopRectsMatchCausal(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 6 + b := newBatch([]int32{0}, L, nil) + want := CausalMask().AsArray(b, K, mlx.DTypeFloat32) + mlx.Eval(want) + wantF := want.Floats() + + cases := []struct { + name string + qLo, qHi, kLo, kHi int + }{ + {"empty Q rect", 2, 2, 0, 3}, + {"empty K rect", 0, 3, 2, 2}, + {"fully under causal", 5, 7, 0, 3}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + m := CausalMask().Relax(0, tc.qLo, tc.qHi, tc.kLo, tc.kHi) + arr := m.AsArray(b, K, mlx.DTypeFloat32) + mlx.Eval(arr) + got := arr.Floats() + for i := range wantF { + if !sameF(got[i], wantF[i]) { + t.Fatalf("index %d: want %v, got %v", i, wantF[i], got[i]) + } + } + }) + } +} + +func TestAttentionMaskAsArrayWithRelax(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 6 + b := newBatch([]int32{0}, L, nil) + arr := CausalMask().Relax(0, 1, 3, 2, 5).AsArray(b, K, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("expected tensor") + } + mlx.Eval(arr) + got := arr.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, L*K) + for q := range L { + for k := range K { + if k > q { + want[q*K+k] = negInf + } + } + } + for q := 1; q < 3; q++ { + for k := 2; k < 5; k++ { + want[q*K+k] = 0 + } + } + for i := range want { + if !sameF(got[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestAttentionMaskAsArrayPerRow(t *testing.T) { + skipIfNoMLX(t) + L, K := 3, 5 + b := newBatch([]int32{0, 2}, L, nil) + m := CausalMask(). + Relax(0, 0, 2, 0, 3). + Relax(1, 3, 5, 2, 5) + arr := m.AsArray(b, K, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("expected tensor") + } + dims := arr.Dims() + if dims[0] != 2 { + t.Fatalf("expected batch dim 2, got %v", dims) + } + mlx.Eval(arr) + got := arr.Floats() + negInf := float32(math.Inf(-1)) + + want := make([]float32, 2*L*K) + for bi, off := range b.SeqOffsets { + for q := range L { + absQ := int(off) + q + for k := range K { + if k > absQ { + want[bi*L*K+q*K+k] = negInf + } + } + } + } + for q := range 2 { + for k := range 3 { + want[0*L*K+q*K+k] = 0 + } + } + for q := 1; q < 3; q++ { + for k := 2; k < 5; k++ { + want[1*L*K+q*K+k] = 0 + } + } + for i := range want { + if !sameF(got[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestQPaddingMask(t *testing.T) { + skipIfNoMLX(t) + L := 4 + // Row 0 fully real; row 1 has 2 real queries. + b := newBatch([]int32{0, 0}, L, []int32{int32(L), 2}) + m := QPaddingMask(b, mlx.DTypeFloat32) + if m.array == nil { + t.Fatal("expected q-padding tensor") + } + mlx.Eval(m.array) + got := m.array.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, 2*L) + // Row 0: no blocking; row 1: q >= 2 blocked. + for q := 2; q < L; q++ { + want[1*L+q] = negInf + } + for i := range want { + if !sameF(got[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestKPaddingMask(t *testing.T) { + skipIfNoMLX(t) + K := 5 + // Row 0 full keys; row 1 has 3 real keys. + b := newBatch([]int32{0, 0}, 4, nil) + kLens := []int32{int32(K), 3} + m := KPaddingMask(b, K, kLens, mlx.DTypeFloat32) + if m.array == nil { + t.Fatal("expected k-padding tensor") + } + mlx.Eval(m.array) + got := m.array.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, 2*K) + for k := 3; k < K; k++ { + want[1*K+k] = negInf + } + for i := range want { + if !sameF(got[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestQPaddingMaskZeroWhenFull(t *testing.T) { + skipIfNoMLX(t) + b := newBatch([]int32{0}, 4, nil) + m := QPaddingMask(b, mlx.DTypeFloat32) + if !m.IsZero() { + t.Fatal("QPaddingMask at full queries should be zero") + } +} + +func TestKPaddingMaskZeroWhenFull(t *testing.T) { + skipIfNoMLX(t) + K := 4 + b := newBatch([]int32{0}, 4, nil) + kLens := []int32{int32(K)} + m := KPaddingMask(b, K, kLens, mlx.DTypeFloat32) + if !m.IsZero() { + t.Fatal("KPaddingMask at full keys should be zero") + } +} + +func TestAttentionMaskCombineCausal(t *testing.T) { + skipIfNoMLX(t) + var z AttentionMask + got := z.Intersect(CausalMask()) + if !got.IsCausal() { + t.Fatal("zero + CausalMask should be pure causal") + } + got = CausalMask().Intersect(z) + if !got.IsCausal() { + t.Fatal("CausalMask + zero should be pure causal") + } + got = CausalMask().Intersect(CausalMask()) + if !got.IsCausal() { + t.Fatal("causal + causal should stay pure causal") + } +} + +func TestAttentionMaskCombineRelaxDroppedAgainstCausal(t *testing.T) { + skipIfNoMLX(t) + relaxed := CausalMask().Relax(0, 1, 3, 2, 5) + got := relaxed.Intersect(CausalMask()) + if !got.IsCausal() { + t.Fatal("causal-with-Relax + causal should drop relaxations and stay pure causal") + } + got = CausalMask().Intersect(relaxed) + if !got.IsCausal() { + t.Fatal("causal + causal-with-Relax should drop relaxations and stay pure causal") + } + + // Disjoint relaxations on two causals also drop — neither side + // agrees to release the cells the other side relaxed. + got = CausalMask().Relax(0, 1, 3, 2, 5).Intersect(CausalMask().Relax(0, 5, 7, 6, 9)) + if !got.IsCausal() { + t.Fatal("disjoint relaxations on two causals should drop and stay pure causal") + } +} + +func TestAttentionMaskCombineRelaxIntersect(t *testing.T) { + skipIfNoMLX(t) + L, K := 6, 6 + b := newBatch([]int32{0}, L, nil) + + // Overlapping rects on two causals: the surviving relaxation is + // the geometric intersection — q in [1,3) ∩ [2,5) = [2,3), + // k in [2,5) ∩ [3,6) = [3,5). + m := CausalMask().Relax(0, 1, 3, 2, 5).Intersect(CausalMask().Relax(0, 2, 5, 3, 6)) + if m.IsCausal() { + t.Fatal("overlapping relaxations should survive as their intersection, not collapse to pure causal") + } + arr := m.AsArray(b, K, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("expected tensor") + } + mlx.Eval(arr) + vals := arr.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, L*K) + for q := range L { + for k := range K { + if k > q { + want[q*K+k] = negInf + } + } + } + // Intersection rect: q ∈ [2,3), k ∈ [3,5). + for q := 2; q < 3; q++ { + for k := 3; k < 5; k++ { + want[q*K+k] = 0 + } + } + for i := range want { + if !sameF(vals[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], vals[i]) + } + } +} + +func TestAttentionMaskCombineRelaxKeptAgainstNonCausal(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 6 + b := newBatch([]int32{0}, L, nil) + + // Pad q=3 — non-causal additive contribution that should leave + // the relaxation intact (the rect releases above-diagonal cells + // q in [1,3), k in [2,5) where k > q). + pad := QPaddingMask(newBatch([]int32{0}, L, []int32{3}), mlx.DTypeFloat32) + if pad.IsZero() { + t.Fatal("padding mask should be non-zero") + } + got := CausalMask().Relax(0, 1, 3, 2, 5).Intersect(pad) + arr := got.AsArray(b, K, mlx.DTypeFloat32) + if arr == nil { + t.Fatal("expected tensor") + } + mlx.Eval(arr) + vals := arr.Floats() + negInf := float32(math.Inf(-1)) + want := make([]float32, L*K) + for q := range L { + for k := range K { + if k > q { + want[q*K+k] = negInf + } + } + } + for q := 1; q < 3; q++ { + for k := 2; k < 5; k++ { + want[q*K+k] = 0 + } + } + for q := 3; q < L; q++ { + for k := range K { + want[q*K+k] = negInf + } + } + for i := range want { + if !sameF(vals[i], want[i]) { + t.Fatalf("index %d: want %v, got %v", i, want[i], vals[i]) + } + } +} + +func TestAttentionMaskCombineArrays(t *testing.T) { + skipIfNoMLX(t) + a := mlx.FromValues([]float32{0, 0, 0, 0}, 1, 1, 2, 2) + bb := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 1, 2, 2) + sum := ArrayMask(a).Intersect(ArrayMask(bb)) + if sum.array == nil { + t.Fatal("array + array should produce array") + } + mlx.Eval(sum.array) + got := sum.array.Floats() + want := []float32{1, 2, 3, 4} + for i := range want { + if got[i] != want[i] { + t.Fatalf("index %d: want %v, got %v", i, want[i], got[i]) + } + } +} + +func TestAttentionMaskRelaxPanicOnArray(t *testing.T) { + skipIfNoMLX(t) + a := mlx.FromValues([]float32{0}, 1, 1, 1, 1) + defer func() { + if r := recover(); r == nil { + t.Fatal("Relax on ArrayMask should panic") + } + }() + ArrayMask(a).Relax(0, 0, 1, 0, 1) +} + +func TestAttentionMaskRelaxPanicOnZero(t *testing.T) { + skipIfNoMLX(t) + defer func() { + if r := recover(); r == nil { + t.Fatal("Relax on zero mask should panic") + } + }() + var z AttentionMask + z.Relax(0, 0, 1, 0, 1) +} + +func sameF(a, b float32) bool { + if math.IsInf(float64(a), -1) && math.IsInf(float64(b), -1) { + return true + } + return a == b +} + +// sdpaInputs builds non-trivial Q/K/V so masking actually changes the +// kernel output. With zero K/V, SDPA returns zero regardless of mask +// and "parity" tests pass even when the mask path is broken. +func sdpaInputs(L, K int) (q, k, v *mlx.Array) { + const D = 4 + qVals := make([]float32, L*D) + for i := range qVals { + qVals[i] = 0.1 * float32(i+1) + } + kVals := make([]float32, K*D) + for i := range kVals { + kVals[i] = 0.07 * float32(i+1) + } + vVals := make([]float32, K*D) + for i := range vVals { + vVals[i] = float32(i+1) - 0.5*float32(K*D) + } + q = mlx.FromValues(qVals, 1, 1, L, D) + k = mlx.FromValues(kVals, 1, 1, K, D) + v = mlx.FromValues(vVals, 1, 1, K, D) + return +} + +func TestSDPACausalParity(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 4 + q, k, v := sdpaInputs(L, K) + b := newBatch([]int32{int32(K - L)}, L, nil) + got := ScaledDotProductAttention(b, q, 1.0, + WithKV(k, v, []int32{int32(K)}), + WithMask(CausalMask()), + ) + want := mlx.FastScaledDotProductAttention(q, k, v, 1.0, "causal", nil) + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if gs[i] != ws[i] { + t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i]) + } + } +} + +func TestSDPAZeroMaskParity(t *testing.T) { + skipIfNoMLX(t) + L, K := 4, 4 + q, k, v := sdpaInputs(L, K) + b := newBatch([]int32{0}, L, nil) + got := ScaledDotProductAttention(b, q, 1.0, WithKV(k, v, []int32{int32(K)})) + want := mlx.FastScaledDotProductAttention(q, k, v, 1.0, "", nil) + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if gs[i] != ws[i] { + t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i]) + } + } +} + +func TestSDPAArrayMaskParity(t *testing.T) { + skipIfNoMLX(t) + L, K := 3, 3 + q, k, v := sdpaInputs(L, K) + b := newBatch([]int32{0}, L, nil) + mask := mlx.FromValues([]float32{ + 0, -1, -1, + 0, 0, -1, + 0, 0, 0, + }, 1, 1, 3, 3) + got := ScaledDotProductAttention(b, q, 1.0, + WithKV(k, v, []int32{int32(K)}), + WithMask(ArrayMask(mask)), + ) + want := mlx.FastScaledDotProductAttention(q, k, v, 1.0, "array", mask) + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if gs[i] != ws[i] { + t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i]) + } + } +} + +func TestSDPARelaxMaskMaterializes(t *testing.T) { + skipIfNoMLX(t) + L, K := 3, 5 + q, k, v := sdpaInputs(L, K) + b := newBatch([]int32{int32(K - L)}, L, nil) + got := ScaledDotProductAttention(b, q, 1.0, + WithKV(k, v, []int32{int32(K)}), + WithMask(CausalMask().Relax(0, 3, 5, 2, 5)), + ) + ref := CausalMask().Relax(0, 3, 5, 2, 5).AsArray(b, K, k.DType()) + want := mlx.FastScaledDotProductAttention(q, k, v, 1.0, "array", ref) + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if gs[i] != ws[i] { + t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i]) + } + } +} + +func TestSDPAPanicsWithBothKVAndHistory(t *testing.T) { + skipIfNoMLX(t) + L := 3 + q, k, v := sdpaInputs(L, L) + b := newBatch([]int32{0}, L, nil) + history := NewKVHistory(k, v, nil) + defer func() { + if r := recover(); r == nil { + t.Fatal("expected panic when both WithKV and WithKVHistory are supplied") + } + }() + ScaledDotProductAttention(b, q, 1.0, WithKV(k, v, []int32{int32(L)}), WithKVHistory(history)) +} + +func TestSDPAMLAHistorySlicesVFromK(t *testing.T) { + skipIfNoMLX(t) + L, D, valueDim := 2, 5, 3 + kBuf := make([]float32, 1*1*L*D) + for i := range kBuf { + kBuf[i] = float32(i) + 1 + } + k := mlx.FromValues(kBuf, 1, 1, L, D) + v := mlx.Zeros(mlx.DTypeFloat32, 1, 1, L, valueDim) + history := NewKVHistory(k, v, nil) + + q := mlx.Zeros(mlx.DTypeFloat32, 1, 1, L, D) + b := newBatch([]int32{0}, L, nil) + got := ScaledDotProductAttention(b, q, 1.0, + WithMLAHistory(history, valueDim), + ) + vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim)) + want := mlx.FastScaledDotProductAttention(q, k, vRef, 1.0, "", nil) + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if gs[i] != ws[i] { + t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i]) + } + } +} + +func TestSDPAPanicsWithoutKV(t *testing.T) { + skipIfNoMLX(t) + q := mlx.FromValues(make([]float32, 4), 1, 1, 1, 4) + b := newBatch([]int32{0}, 1, nil) + defer func() { + if r := recover(); r == nil { + t.Fatal("expected panic when no K/V supplied") + } + }() + ScaledDotProductAttention(b, q, 1.0) +} + +// fillTensor builds a [B, H, T, D] float32 tensor whose entries are +// distinct, non-zero, and predictable so per-row slices stay distinct. +func fillTensor(seed float32, B, H, T, D int) *mlx.Array { + vals := make([]float32, B*H*T*D) + for i := range vals { + vals[i] = seed + 0.05*float32(i) + } + return mlx.FromValues(vals, B, H, T, D) +} + +// TestSDPAMultiSequenceParity drives a B=2 batch with mixed +// SeqOffsets and SeqQueryLens through ScaledDotProductAttention via +// the no-cache (WithKV) path, then compares each row's real +// positions against a B=1 reference at that row's offset and length. +// Padded-tail outputs are unconstrained and not checked. Pins the +// central multi-sequence contract: right-padded rows must produce +// per-row outputs that don't depend on the padded tails. +func TestSDPAMultiSequenceParity(t *testing.T) { + skipIfNoMLX(t) + const H, D = 1, 4 + const L, K = 4, 6 + const qShort, kShort = 2, 2 + const scale = 1.0 + + q := fillTensor(0.5, 2, H, L, D) + k := fillTensor(-0.3, 2, H, K, D) + v := fillTensor(0.7, 2, H, K, D) + b := newBatch([]int32{2, 0}, L, []int32{int32(L), int32(qShort)}) + + got := ScaledDotProductAttention(b, q, scale, + WithKV(k, v, []int32{int32(K), int32(kShort)}), + WithMask(CausalMask())) + mlx.Eval(got) + gotF := got.Floats() + + // Row 0: full Q at offset 2, full K. B=1 reference. + q0 := mlx.SliceStartStop(q, []int32{0, 0, 0, 0}, []int32{1, H, L, D}) + k0 := mlx.SliceStartStop(k, []int32{0, 0, 0, 0}, []int32{1, H, K, D}) + v0 := mlx.SliceStartStop(v, []int32{0, 0, 0, 0}, []int32{1, H, K, D}) + b0 := newBatch([]int32{2}, L, nil) + ref0 := ScaledDotProductAttention(b0, q0, scale, + WithKV(k0, v0, []int32{int32(K)}), + WithMask(CausalMask())) + mlx.Eval(ref0) + ref0F := ref0.Floats() + + // Row 1: real Q at offset 0, length qShort, with kShort real keys. + q1 := mlx.SliceStartStop(q, []int32{1, 0, 0, 0}, []int32{2, H, int32(qShort), D}) + k1 := mlx.SliceStartStop(k, []int32{1, 0, 0, 0}, []int32{2, H, int32(kShort), D}) + v1 := mlx.SliceStartStop(v, []int32{1, 0, 0, 0}, []int32{2, H, int32(kShort), D}) + b1 := newBatch([]int32{0}, qShort, nil) + ref1 := ScaledDotProductAttention(b1, q1, scale, + WithKV(k1, v1, []int32{int32(kShort)}), + WithMask(CausalMask())) + mlx.Eval(ref1) + ref1F := ref1.Floats() + + // got is [2, H, L, D] = [B=2, 1, 4, 4]. Row 0 is got[0,...] and + // must match ref0 over the full [L, D]. Row 1 is got[1,...] and + // must match ref1 over [qShort, D] only — padded positions are + // unconstrained. + rowStride := H * L * D + for i := range rowStride { + if !approxEqual(gotF[i], ref0F[i], 1e-5) { + t.Fatalf("row 0 [%d]: got %v, want %v", i, gotF[i], ref0F[i]) + } + } + for q := range qShort { + for d := range D { + gotI := rowStride + q*D + d + refI := q*D + d + if !approxEqual(gotF[gotI], ref1F[refI], 1e-5) { + t.Fatalf("row 1 [q=%d,d=%d]: got %v, want %v", q, d, gotF[gotI], ref1F[refI]) + } + } + } +} diff --git a/x/models/qwen3/qwen3.go b/x/models/qwen3/qwen3.go index 1269d635f..8b055a92e 100644 --- a/x/models/qwen3/qwen3.go +++ b/x/models/qwen3/qwen3.go @@ -265,7 +265,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { if caches != nil && i < len(caches) { c = caches[i] } - h = layer.Forward(h, c, positions, B, L, m.Config) + h = layer.Forward(h, b, c, positions, B, L, m.Config) } return m.Norm.Forward(h, m.RMSNormEps) @@ -295,12 +295,12 @@ func (m *Model) NewCaches() []cache.Cache { return caches } -func (l *Layer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { - h := mlx.Add(x, l.Attention.Forward(l.AttentionNorm.Forward(x, cfg.RMSNormEps), c, positions, B, L, cfg)) +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + h := mlx.Add(x, l.Attention.Forward(l.AttentionNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg)) return mlx.Add(h, l.MLP.Forward(l.MLPNorm.Forward(h, cfg.RMSNormEps))) } -func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { q := a.QProj.Forward(x) k := a.KProj.Forward(x) v := a.VProj.Forward(x) @@ -319,13 +319,16 @@ func (a *Attention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B q = mlx.RoPEWithBase(q, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) k = mlx.RoPEWithBase(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) - if c != nil { - k, v = c.Update(k, v) - } - // MLX SDPA supports grouped-query attention directly (Q heads can be a // multiple of K/V heads), so avoid materializing repeated K/V tensors. - out := mlx.ScaledDotProductAttentionCausal(q, k, v, cfg.Scale, L > 1) + var kv nn.SDPAOption + if c != nil { + history := c.(cache.Attention).Update(b, k, v) + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) + } + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) return a.OProj.Forward(out) } diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index b29f44009..47c02ed4a 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -1077,32 +1077,7 @@ func softplus(x *mlx.Array) *mlx.Array { return mlx.Logaddexp(x, mlx.Zeros(x.DType(), x.Dims()...)) } -func depthwiseCausalConv1d(x, w *mlx.Array, outLen int32) *mlx.Array { - if x == nil || w == nil { - return nil - } - if w.NumDims() != 2 { - return nil - } - B := int32(x.Dim(0)) - C := int32(w.Dim(0)) - K := int32(w.Dim(1)) - var out *mlx.Array - for i := int32(0); i < K; i++ { - seg := mlx.SliceStartStop(x, []int32{0, i, 0}, []int32{B, i + outLen, C}) - wi := mlx.SliceStartStop(w, []int32{0, i}, []int32{C, i + 1}) - wi = mlx.Reshape(wi, 1, 1, C) - term := mlx.Mul(seg, wi) - if out == nil { - out = term - } else { - out = mlx.Add(out, term) - } - } - return out -} - -func splitQKVZBA(mixedQKVZ, mixedBA *mlx.Array, cfg *Config, B, L int32) (q, k, v, z, b, a *mlx.Array) { +func splitQKVZBA(mixedQKVZ, mixedBA *mlx.Array, cfg *Config, B, L int32) (q, k, v, z, beta, alpha *mlx.Array) { nk := cfg.LinearNumKeyHeads nv := cfg.LinearNumValueHeads dk := cfg.LinearKeyHeadDim @@ -1119,15 +1094,15 @@ func splitQKVZBA(mixedQKVZ, mixedBA *mlx.Array, cfg *Config, B, L int32) (q, k, z = mlx.Reshape(z, B, L, nv, dv) mixedBA = mlx.Reshape(mixedBA, B, L, nk, 2*vPerK) - b = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, 0}, []int32{B, L, nk, vPerK}) - a = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, vPerK}, []int32{B, L, nk, 2 * vPerK}) - b = mlx.Reshape(b, B, L, nv) - a = mlx.Reshape(a, B, L, nv) + beta = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, 0}, []int32{B, L, nk, vPerK}) + alpha = mlx.SliceStartStop(mixedBA, []int32{0, 0, 0, vPerK}, []int32{B, L, nk, 2 * vPerK}) + beta = mlx.Reshape(beta, B, L, nv) + alpha = mlx.Reshape(alpha, B, L, nv) - return q, k, v, z, b, a + return q, k, v, z, beta, alpha } -func (a *FullAttention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { +func (a *FullAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { qg := a.QProj.Forward(x) qg = mlx.Reshape(qg, B, L, cfg.NumAttentionHeads, cfg.HeadDim*2) q := mlx.SliceStartStop(qg, []int32{0, 0, 0, 0}, []int32{B, L, cfg.NumAttentionHeads, cfg.HeadDim}) @@ -1149,11 +1124,14 @@ func (a *FullAttention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Arra q = mlx.RoPEWithBase(q, int(cfg.RopeDim), false, cfg.RopeTheta, 1.0, positions) k = mlx.RoPEWithBase(k, int(cfg.RopeDim), false, cfg.RopeTheta, 1.0, positions) + var kv nn.SDPAOption if c != nil { - k, v = c.Update(k, v) + history := c.(cache.Attention).Update(b, k, v) + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) } - - out := mlx.ScaledDotProductAttentionCausal(q, k, v, cfg.Scale, L > 1) + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) gateSigmoid := mlx.Sigmoid(gate) out = mlx.Mul(out, gateSigmoid) @@ -1161,20 +1139,20 @@ func (a *FullAttention) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Arra return out } -func (g *GatedDeltaNet) Forward(x *mlx.Array, c cache.Cache, B, L int32, cfg *Config) *mlx.Array { - var qkv, z, b, a *mlx.Array +func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B, L int32, cfg *Config) *mlx.Array { + var qkv, z, beta, alpha *mlx.Array useSplitProj := g.InProjQKV != nil && g.InProjZ != nil && g.InProjB != nil && g.InProjA != nil if useSplitProj { qkv = g.InProjQKV.Forward(x) z = g.InProjZ.Forward(x) z = mlx.Reshape(z, B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim) - b = g.InProjB.Forward(x) - a = g.InProjA.Forward(x) + beta = g.InProjB.Forward(x) + alpha = g.InProjA.Forward(x) } else { mixedQKVZ := g.InProjQKVZ.Forward(x) mixedBA := g.InProjBA.Forward(x) var q, k, v *mlx.Array - q, k, v, z, b, a = splitQKVZBA(mixedQKVZ, mixedBA, cfg, B, L) + q, k, v, z, beta, alpha = splitQKVZBA(mixedQKVZ, mixedBA, cfg, B, L) qkv = mlx.Concatenate([]*mlx.Array{ mlx.Reshape(q, B, L, cfg.LinearNumKeyHeads*cfg.LinearKeyHeadDim), mlx.Reshape(k, B, L, cfg.LinearNumKeyHeads*cfg.LinearKeyHeadDim), @@ -1182,32 +1160,20 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, c cache.Cache, B, L int32, cfg *Co }, -1) } convTail := cfg.LinearConvKernelDim - 1 - var convState *mlx.Array var rc *cache.RecurrentCache - if c != nil { - if typed, ok := c.(*cache.RecurrentCache); ok { - rc = typed - convState = rc.ConvState(int(B), x.DType()) - } - } - if convState == nil { - convState = mlx.Zeros(x.DType(), int(B), int(convTail), int(2*cfg.LinearNumKeyHeads*cfg.LinearKeyHeadDim+cfg.LinearNumValueHeads*cfg.LinearValueHeadDim)) + var rec nn.RecurrentOption + if typed, ok := c.(*cache.RecurrentCache); ok { + rc = typed + rec = nn.WithRecurrentHistory(rc.Get(b, x.DType())) + } else { + rec = nn.WithRecurrentState( + mlx.Zeros(x.DType(), int(B), int(convTail), qkv.Dim(2)), + mlx.Zeros(x.DType(), int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)), + ) } - convInput := mlx.Concatenate([]*mlx.Array{convState, qkv}, 1) - var convOut *mlx.Array - if g.Conv1D != nil { - convOut = g.Conv1D.Forward(convInput) - } else { - convOut = depthwiseCausalConv1d(convInput, g.ConvWeight, L) - } + convOut, nextConv := nn.CausalConv1D(b, qkv, g.Conv1D, g.ConvWeight, int(convTail), rec) convOut = mlx.SiLU(convOut) - if rc != nil { - total := int32(convInput.Dim(1)) - start := total - convTail - nextConv := mlx.SliceStartStop(convInput, []int32{0, start, 0}, []int32{B, total, int32(convInput.Dim(2))}) - rc.SetConvState(nextConv) - } keyDim := cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim valueDim := cfg.LinearNumValueHeads * cfg.LinearValueHeadDim @@ -1221,30 +1187,21 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, c cache.Cache, B, L int32, cfg *Co q = mlx.MulScalar(mlx.RMSNormFn(q, nil, 1e-6), invScale*invScale) k = mlx.MulScalar(mlx.RMSNormFn(k, nil, 1e-6), invScale) - gDecay := softplus(mlx.Add(a, g.DtBias)) + gDecay := softplus(mlx.Add(alpha, g.DtBias)) gDecay = mlx.Mul(gDecay, g.AExp) gDecay = mlx.Exp(mlx.MulScalar(gDecay, -1)) - gDecay = gDecay.AsType(a.DType()) + gDecay = gDecay.AsType(alpha.DType()) - beta := mlx.Sigmoid(b) + betaGate := mlx.Sigmoid(beta) - var state *mlx.Array - if rc != nil { - state = rc.DeltaState(int(B), x.DType()) - } - if state == nil { - state = mlx.Zeros(x.DType(), int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)) - } - - out, state := mlx.GatedDelta(q, k, v, gDecay, beta, state) + out, state := nn.GatedDelta(b, q, k, v, gDecay, betaGate, rec) outDType := out.DType() out = mlx.RMSNormFn(out, g.NormWeight, cfg.RMSNormEps) out = mlx.Mul(out.AsType(mlx.DTypeFloat32), mlx.SiLU(z.AsType(mlx.DTypeFloat32))).AsType(outDType) out = mlx.Reshape(out, B, L, valueDim) out = g.OutProj.Forward(out) if rc != nil { - rc.SetDeltaState(state) - rc.Advance(int(L)) + rc.Put(b, nextConv, state) } return out } @@ -1329,13 +1286,13 @@ func (m *SparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array { return mlx.Reshape(y, B, L, cfg.HiddenSize) } -func (l *Layer) Forward(x *mlx.Array, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { var r *mlx.Array normed := l.InputNorm.Forward(x, cfg.RMSNormEps) if l.IsLinear { - r = l.Linear.Forward(normed, c, B, L, cfg) + r = l.Linear.Forward(normed, b, c, B, L, cfg) } else { - r = l.FullAttn.Forward(normed, c, positions, B, L, cfg) + r = l.FullAttn.Forward(normed, b, c, positions, B, L, cfg) } h := mlx.Add(x, r) r = l.MLP.Forward(l.PostAttentionNorm.Forward(h, cfg.RMSNormEps), cfg) @@ -1353,7 +1310,7 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { if caches != nil && i < len(caches) { c = caches[i] } - h = layer.Forward(h, c, positions, B, L, m.Config) + h = layer.Forward(h, b, c, positions, B, L, m.Config) } out := m.Norm.Forward(h, m.RMSNormEps) return out