mlxrunner: decouple models from attention cache storage layout

Models build their own attention masks and read K/V directly from
the cache's buffers, which ties them to the cache's storage layout.
That blocks multi-sequence batching — right-padded rows need a
query-padding mask composed onto every model — and rules out
variants like paged attention where K/V isn't one contiguous tensor.

Caches now hand back a per-layer KVHistory holding post-update K, V,
and a MaskApplier that merges the cache's storage restrictions into
the model's logical mask. Models describe their mask in logical
terms; SDPA composes model, padding, and applier contributions and
dispatches to the kernel's causal or no-mask fast path when it can.
KVHistory still exposes K, V, and the composed mask for manual
attention paths (e.g. CUDA prefill at head_dim > 128).

Performance for single-sequence inference is unchanged.

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
 }

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)
+}

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())

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 {

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<L, then again on a
+// fresh cache with an unpadded length-qLen batch using the same
+// real prefix. After the call, Offset() must equal qLen (not L),
+// and the resulting cache state must match the unpadded equivalent.
+// Pins the recurrent contract: a forward with padding produces the
+// same end-state as a forward with the real-prefix-only input.
+func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
+	skipIfNoMLX(t)
+	const convTail, convDim = 2, 6
+	const numVHeads, headVDim, headKDim = 1, 4, 6
+	const L = 4
+	const qLen = 2
+
+	// Use distinct values for the real prefix and the padded tail so
+	// we can detect any leak from padded positions into the result.
+	makeQKV := func(seed float32, T int) (q, k, v *mlx.Array) {
+		mkLast := func(off float32, T, n, d int) *mlx.Array {
+			vals := make([]float32, 1*T*n*d)
+			for i := range vals {
+				vals[i] = off + 0.05*float32(i)
+			}
+			return mlx.FromValues(vals, 1, T, n, d)
+		}
+		q = mkLast(seed, T, 1, headKDim)
+		k = mkLast(seed+0.1, T, 1, headKDim)
+		v = mkLast(seed+0.2, T, numVHeads, headVDim)
+		return
+	}
+	makeGB := func(seed float32, T int) (g, beta *mlx.Array) {
+		gVals := make([]float32, 1*T*numVHeads)
+		bVals := make([]float32, 1*T*numVHeads)
+		for i := range gVals {
+			gVals[i] = seed + 0.01*float32(i)
+			bVals[i] = seed - 0.02*float32(i)
+		}
+		g = mlx.FromValues(gVals, 1, T, numVHeads)
+		beta = mlx.FromValues(bVals, 1, T, numVHeads)
+		return
+	}
+	makeQKVPadded := func() (q, k, v *mlx.Array) {
+		qReal, kReal, vReal := makeQKV(0.3, qLen)
+		// Distinct, large junk values in the padded tail to surface
+		// any leak (real outputs are O(1)).
+		qPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 1, headKDim), 99)
+		kPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 1, headKDim), 99)
+		vPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads, headVDim), 99)
+		q = mlx.Concatenate([]*mlx.Array{qReal, qPad}, 1)
+		k = mlx.Concatenate([]*mlx.Array{kReal, kPad}, 1)
+		v = mlx.Concatenate([]*mlx.Array{vReal, vPad}, 1)
+		return
+	}
+	makeGBPadded := func() (g, beta *mlx.Array) {
+		gReal, betaReal := makeGB(0.1, qLen)
+		gPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads), 99)
+		betaPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads), 99)
+		g = mlx.Concatenate([]*mlx.Array{gReal, gPad}, 1)
+		beta = mlx.Concatenate([]*mlx.Array{betaReal, betaPad}, 1)
+		return
+	}
+
+	// The conv input dimension must match the cache's convDim.
+	mkConvInput := func(seed float32, T int) *mlx.Array {
+		vals := make([]float32, 1*T*convDim)
+		for i := range vals {
+			vals[i] = seed + 0.05*float32(i)
+		}
+		return mlx.FromValues(vals, 1, T, convDim)
+	}
+	mkWeight := func(seed float32) *mlx.Array {
+		vals := make([]float32, convDim*(convTail+1))
+		for i := range vals {
+			vals[i] = seed + 0.1*float32(i)
+		}
+		return mlx.FromValues(vals, convDim, convTail+1)
+	}
+	weight := mkWeight(0.2)
+
+	runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) {
+		var convInput *mlx.Array
+		if T == L {
+			realPart := mkConvInput(0.4, qLen)
+			padPart := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, T-qLen, convDim), 99)
+			convInput = mlx.Concatenate([]*mlx.Array{realPart, padPart}, 1)
+		} else {
+			convInput = mkConvInput(0.4, T)
+		}
+
+		history := c.Get(b, mlx.DTypeFloat32)
+		_, nextConv := nn.CausalConv1D(b, convInput, nil, weight, convTail,
+			nn.WithRecurrentHistory(history))
+
+		var q, k, v, g, beta *mlx.Array
+		if T == L {
+			q, k, v = makeQKVPadded()
+			g, beta = makeGBPadded()
+		} else {
+			q, k, v = makeQKV(0.3, T)
+			g, beta = makeGB(0.1, T)
+		}
+		_, newDelta := nn.GatedDelta(b, q, k, v, g, beta,
+			nn.WithRecurrentHistory(history))
+
+		c.Put(b, nextConv, newDelta)
+		return nextConv, newDelta
+	}
+
+	// Padded forward.
+	cPad := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
+	bPad := &batch.Batch{
+		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, L),
+		SeqOffsets:   []int32{0},
+		SeqQueryLens: []int32{int32(qLen)},
+	}
+	nextConvPad, deltaPad := runForward(cPad, bPad, L)
+	mlx.Eval(nextConvPad, deltaPad)
+	if got := cPad.Offset(); got != qLen {
+		t.Fatalf("padded forward: Offset() = %d, want %d (must advance by SeqQueryLens, not L)", got, qLen)
+	}
+
+	// Unpadded reference.
+	cRef := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
+	bRef := &batch.Batch{
+		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, qLen),
+		SeqOffsets:   []int32{0},
+		SeqQueryLens: []int32{int32(qLen)},
+	}
+	nextConvRef, deltaRef := runForward(cRef, bRef, qLen)
+	mlx.Eval(nextConvRef, deltaRef)
+	if got := cRef.Offset(); got != qLen {
+		t.Fatalf("unpadded forward: Offset() = %d, want %d", got, qLen)
+	}
+
+	gp := nextConvPad.Floats()
+	gr := nextConvRef.Floats()
+	if len(gp) != len(gr) {
+		t.Fatalf("nextConv shape mismatch: padded %d vs unpadded %d", len(gp), len(gr))
+	}
+	for i := range gp {
+		if math.Abs(float64(gp[i]-gr[i])) > 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())
+	}
+}

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])
+		}
+	}
+}

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

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
 }
-

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
 }

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

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)

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("<start_of_turn>user\n%s<end_of_turn>\n<start_of_turn>model\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)
 }

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 {

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)

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)
 }

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
+}

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<L). For the short row, (a) `out` positions [0..qLen)
+// must equal a B=1 reference at length qLen, (b) `nextConv` for the
+// short row must be the row's last convTail real positions (not the
+// padded tail), (c) the full row must be unaffected.
+func TestCausalConv1DPaddedRowParity(t *testing.T) {
+	skipIfNoMLX(t)
+	L, D, convTail := 4, 3, 2
+	qLenShort := 2
+	K := convTail + 1
+
+	weight := fromValues(0.2, D, K)
+	priorFull := fromValues(0.5, 2, convTail, D)
+	priorShort := mlx.SliceStartStop(priorFull,
+		[]int32{1, 0, 0},
+		[]int32{2, int32(convTail), int32(D)})
+
+	// Pad row 1 with arbitrary values past qLenShort — the wrapper
+	// must zero them before convolving. Distinct values let us catch
+	// any leak.
+	inputFull := fromValues(1.0, 1, L, D)
+	inputShortReal := mlx.FromValues([]float32{
+		2.0, 2.1, 2.2,
+		2.3, 2.4, 2.5,
+	}, 1, qLenShort, D)
+	inputShortPad := mlx.FromValues([]float32{
+		99, 99, 99,
+		99, 99, 99,
+	}, 1, L-qLenShort, D)
+	inputShortFull := mlx.Concatenate([]*mlx.Array{inputShortReal, inputShortPad}, 1)
+	input := mlx.Concatenate([]*mlx.Array{inputFull, inputShortFull}, 0)
+
+	b := &batch.Batch{
+		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
+		SeqOffsets:   []int32{0, 0},
+		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
+	}
+
+	out, nextConv := CausalConv1D(b, input, nil, weight, convTail, WithRecurrentState(priorFull, nil))
+	mlx.Eval(out, nextConv)
+
+	// Reference for row 0: B=1 unpadded length-L call.
+	refOut0, refNextConv0 := CausalConv1D(&batch.Batch{},
+		inputFull, nil, weight, convTail,
+		WithRecurrentState(mlx.SliceStartStop(priorFull,
+			[]int32{0, 0, 0},
+			[]int32{1, int32(convTail), int32(D)}), nil))
+	// Reference for row 1: B=1 unpadded length-qLenShort call.
+	refOut1, refNextConv1 := CausalConv1D(&batch.Batch{},
+		inputShortReal, nil, weight, convTail,
+		WithRecurrentState(priorShort, nil))
+	mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)
+
+	gotOut := out.Floats()
+	wantOut0 := refOut0.Floats()
+	wantOut1 := refOut1.Floats()
+
+	for q := range L {
+		for d := range D {
+			i := q*D + d
+			if gotOut[i] != wantOut0[i] {
+				t.Fatalf("row 0 out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[i], wantOut0[i])
+			}
+		}
+	}
+	for q := range qLenShort {
+		for d := range D {
+			gotI := L*D + q*D + d
+			refI := q*D + d
+			if math.Abs(float64(gotOut[gotI]-wantOut1[refI])) > 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)
+		}
+	}
+}

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)
+}

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])
+			}
+		}
+	}
+}

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)
 }

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