mlx: scope array lifetimes instead of pinning and sweeping

The bindings freed arrays by sweeping everything not pinned, so freeing
anything required knowing what every other caller still held, and code
that never swept accumulated until memory ran out. The prefix cache's
eviction of a long stored path did exactly that: each merge copied the
KV snapshots and nothing freed the consumed copies until the request
ended, which drove a second long request past physical memory.

Every array now belongs to a scope. A function scope, entered with Scoped
or one of the ScopedEval forms, frees what was created in it when the
function returns; results leave only by being returned. A held scope is
closed by its holder and frees what was attached to it. A graph is
built in a function scope and evaluated after it, so the eval frees each
intermediate as it consumes it. Pin, Unpin, Sweep, and the array list's
mutex are gone.

On an M5 Max with qwen3.8:27b-mlx, the second 84k-token request after a
stored one peaks at 35 GB instead of 57 GB; the cold path is unchanged.
The copies themselves are untouched, so restoring an owned path can still
exceed memory.

x/create/draft.go

@@ -22,7 +22,7 @@ func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, st
 	if configPrefix == "" {
 		return nil, fmt.Errorf("draft config prefix must not be empty")
 	}
-	defer sweepMLX()
+	defer releaseMLXCache()
 
 	inv, err := ReadInventory(modelDir)
 	if err != nil {

x/create/mlxthread.go

@@ -59,14 +59,13 @@ func runOnMLXThread(f func() error) error {
 	return <-done
 }
 
-// sweepMLX releases the MLX buffer cache. It is a no-op if no MLX work has run.
-func sweepMLX() {
+// releaseMLXCache releases the MLX buffer cache. It is a no-op if no MLX work has run.
+func releaseMLXCache() {
 	if !mlxThreadStarted.Load() {
 		return
 	}
 	_ = runOnMLXThread(func() error {
 		mlx.ClearCache()
-		mlx.Sweep()
 		return nil
 	})
 }

x/create/pipeline.go

@@ -14,7 +14,7 @@ import (
 // server-side entry point — the caller supplies blob storage (store) and
 // manifest assembly (writeManifest).
 func Create(modelName, modelDir, quantize string, store BlobStore, writeManifest ManifestWriter, fn func(status string)) error {
-	defer sweepMLX()
+	defer releaseMLXCache()
 
 	inv, err := ReadInventory(modelDir)
 	if err != nil {

x/create/quantize.go

@@ -43,11 +43,8 @@ func quantizeBlob(items []quantizeItem) ([]byte, error) {
 
 func quantizeBlobLocked(items []quantizeItem) ([]byte, error) {
 	allArrays := make(map[string]*mlx.Array)
-	var pinned []*mlx.Array
-	defer func() {
-		mlx.Unpin(pinned...)
-		mlx.Sweep()
-	}()
+	held := mlx.NewScope()
+	defer held.Close()
 
 	tmpDir, err := os.MkdirTemp("", "ollama-quantize-*")
 	if err != nil {
@@ -81,36 +78,21 @@ func quantizeBlobLocked(items []quantizeItem) ([]byte, error) {
 	}
 
 	for _, it := range items {
-		if err := func() error {
-			defer mlx.Sweep()
-			tmpPath, toEval, st, err := loadAndQuantizeArray(it.reader, it.name, it.quantize, it.decodeFP8, allArrays, tmpDir)
-			if tmpPath != "" {
-				defer os.Remove(tmpPath)
-			}
-			if err != nil {
-				return err
-			}
-			if st != nil {
-				defer st.Free()
-			}
-			mlx.Eval(toEval...)
-			final := arraysForItem(allArrays, it)
-			mlx.Pin(final...)
-			pinned = append(pinned, final...)
-
-			if mixed && it.quantize != "" {
-				if gs, _, _ := quant.Params(it.quantize); gs > 0 {
-					if metadata == nil {
-						metadata = make(map[string]string)
-					}
-					metadata[it.name+".quant_type"] = it.quantize
-					metadata[it.name+".group_size"] = strconv.Itoa(gs)
-				}
-			}
-			return nil
-		}(); err != nil {
+		if err := quantizeItemArrays(it, allArrays, tmpDir, held); err != nil {
 			return nil, err
 		}
+		// The item's intermediates are free; hand their buffers back before
+		// the next item, which may never reuse those sizes.
+		mlx.ClearCache()
+		if mixed && it.quantize != "" {
+			if gs, _, _ := quant.Params(it.quantize); gs > 0 {
+				if metadata == nil {
+					metadata = make(map[string]string)
+				}
+				metadata[it.name+".quant_type"] = it.quantize
+				metadata[it.name+".group_size"] = strconv.Itoa(gs)
+			}
+		}
 	}
 
 	outPath := filepath.Join(tmpDir, "blob.safetensors")
@@ -120,18 +102,20 @@ func quantizeBlobLocked(items []quantizeItem) ([]byte, error) {
 	return os.ReadFile(outPath)
 }
 
-func arraysForItem(all map[string]*mlx.Array, it quantizeItem) []*mlx.Array {
-	keys := []string{it.name}
-	if it.quantize != "" {
-		keys = append(keys, it.name+".scale", it.name+".bias")
+// quantizeItemArrays loads and quantizes one item into arrays and holds its
+// finished arrays in held.
+func quantizeItemArrays(it quantizeItem, arrays map[string]*mlx.Array, tmpDir string, held *mlx.Scope) error {
+	tmpPath, toEval, st, err := loadAndQuantizeArray(it.reader, it.name, it.quantize, it.decodeFP8, arrays, tmpDir)
+	if tmpPath != "" {
+		defer os.Remove(tmpPath)
 	}
-	out := make([]*mlx.Array, 0, len(keys))
-	for _, k := range keys {
-		if a := all[k]; a != nil {
-			out = append(out, a)
-		}
+	if err != nil {
+		return err
 	}
-	return out
+	defer st.Free()
+	mlx.Eval(toEval...)
+	held.Attach(toEval...)
+	return nil
 }
 
 // loadAndQuantizeArray writes a safetensors reader to a temp file, loads it
@@ -164,65 +148,69 @@ func loadAndQuantizeArray(r io.Reader, name, quantize string, decodeFP8 bool, ar
 		return tmpPath, nil, nil, fmt.Errorf("failed to load safetensors for %s: %w", name, err)
 	}
 
-	arr := st.Get(name)
-	if arr == nil {
-		st.Free()
-		return tmpPath, nil, nil, fmt.Errorf("tensor %q not found in safetensors", name)
-	}
-
-	// Decode an FP8 source tensor (using its block scale) before quantizing,
-	// so a decode-only request (quantize == "") still yields usable float data.
-	if decodeFP8 {
-		scaleKey := name + ".scale_inv"
-		scaleInv := st.Get(scaleKey)
-		if scaleInv == nil {
-			scaleKey = name + ".scale"
-			scaleInv = st.Get(scaleKey)
+	toEval = mlx.ScopedArrays(func() []*mlx.Array {
+		arr := st.Get(name)
+		if arr == nil {
+			err = fmt.Errorf("tensor %q not found in safetensors", name)
+			return nil
 		}
-		if scaleInv == nil {
-			st.Free()
-			return tmpPath, nil, nil, fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", name+".scale_inv", name+".scale", name)
+
+		// Decode an FP8 source tensor (using its block scale) before quantizing,
+		// so a decode-only request (quantize == "") still yields usable float data.
+		if decodeFP8 {
+			scaleKey := name + ".scale_inv"
+			scaleInv := st.Get(scaleKey)
+			if scaleInv == nil {
+				scaleKey = name + ".scale"
+				scaleInv = st.Get(scaleKey)
+			}
+			if scaleInv == nil {
+				err = fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", name+".scale_inv", name+".scale", name)
+				return nil
+			}
+			arr, err = decodeSourceFP8Tensor(arr, scaleInv)
+			if err != nil {
+				err = fmt.Errorf("failed to decode fp8 tensor %s: %w", name, err)
+				return nil
+			}
 		}
-		arr, err = decodeSourceFP8Tensor(arr, scaleInv)
-		if err != nil {
-			st.Free()
-			return tmpPath, nil, nil, fmt.Errorf("failed to decode fp8 tensor %s: %w", name, err)
+
+		if quantize == "" {
+			arr = mlx.Contiguous(arr, false)
+			arrays[name] = arr
+			return []*mlx.Array{arr}
 		}
-		mlx.Eval(arr)
-	}
 
-	if quantize == "" {
-		arr = mlx.Contiguous(arr, false)
-		arrays[name] = arr
-		return tmpPath, []*mlx.Array{arr}, st, nil
-	}
+		if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 {
+			arr = arr.AsType(mlx.DTypeBFloat16)
+		}
 
-	if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 {
-		arr = arr.AsType(mlx.DTypeBFloat16)
-		mlx.Eval(arr)
-	}
+		groupSize, bits, mode := quant.Params(quantize)
+		qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode)
+		if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 {
+			err = fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode)
+			return nil
+		}
+		if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 {
+			err = fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode)
+			return nil
+		}
 
-	groupSize, bits, mode := quant.Params(quantize)
-	qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode)
-	mlx.Eval(qweight, scales)
-	if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 {
+		qweight = mlx.Contiguous(qweight, false)
+		scales = mlx.Contiguous(scales, false)
+		arrays[name] = qweight
+		arrays[name+".scale"] = scales
+		out := []*mlx.Array{qweight, scales}
+		if qbiases != nil {
+			qbiases = mlx.Contiguous(qbiases, false)
+			arrays[name+".bias"] = qbiases
+			out = append(out, qbiases)
+		}
+		return out
+	})
+	if err != nil {
 		st.Free()
-		return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode)
-	}
-	if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 {
-		st.Free()
-		return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode)
-	}
-
-	qweight = mlx.Contiguous(qweight, false)
-	scales = mlx.Contiguous(scales, false)
-	arrays[name] = qweight
-	arrays[name+".scale"] = scales
-	toEval = append(toEval, qweight, scales)
-	if qbiases != nil {
-		qbiases = mlx.Contiguous(qbiases, false)
-		arrays[name+".bias"] = qbiases
-		toEval = append(toEval, qbiases)
+		return tmpPath, nil, nil, err
 	}
 	return tmpPath, toEval, st, nil
 }

x/mlxrunner/cache/cache.go

@@ -15,7 +15,7 @@ type Cache interface {
 	Offset() int
 
 	// Snapshot copies cache state from fromOffset to current offset into
-	// pinned VRAM arrays. The active cache is unchanged.
+	// owned VRAM arrays. The active cache is unchanged.
 	Snapshot(fromOffset int) Snapshot
 
 	// PrepareSnapshots schedules the cache to capture a snapshot as its
@@ -67,7 +67,7 @@ type Snapshot interface {
 	// never lazy may treat this as a no-op.
 	SetMaterializeHook(func(delta int))
 
-	// Close unpins the snapshot's arrays so they can be freed by Sweep.
+	// Close frees the snapshot's arrays.
 	Close()
 }
 

x/mlxrunner/cache/kvcache.go

@@ -23,6 +23,7 @@ type Attention interface {
 
 type KVCache struct {
 	keys, values *mlx.Array
+	scope        *mlx.Scope
 	offset       int
 	step         int
 
@@ -35,7 +36,7 @@ type KVCache struct {
 }
 
 func NewKVCache() *KVCache {
-	return &KVCache{step: 256}
+	return &KVCache{step: 256, scope: mlx.NewScope()}
 }
 
 // Assumes B = 1; heterogeneous batches are not supported.
@@ -78,7 +79,7 @@ func (c *KVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
 			c.values.Set(c.values.Concatenate(2, newValues))
 		} else {
 			c.keys, c.values = newKeys, newValues
-			mlx.Pin(c.keys, c.values)
+			c.scope.Attach(c.keys, c.values)
 		}
 	}
 
@@ -133,6 +134,7 @@ func (c *KVCache) captureLazySnapshots(start, end int) {
 // and cache is nil.
 type kvSnapshot struct {
 	keys, values         *mlx.Array
+	scope                *mlx.Scope // holds keys and values once copied out
 	fromOffset, toOffset int
 	cache                *KVCache // issuer while lazy; nil once copied out
 
@@ -154,7 +156,7 @@ func (s *kvSnapshot) Size() int {
 func (s *kvSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
 
 func (s *kvSnapshot) Close() {
-	mlx.Unpin(s.keys, s.values)
+	s.scope.Close()
 	if s.cache != nil {
 		s.cache.dropLazySnapshot(s)
 		s.cache = nil
@@ -170,14 +172,15 @@ func (s *kvSnapshot) copyOut() {
 		return
 	}
 	c := s.cache
-	kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
-	vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
-	kCopy := mlx.Contiguous(kSlice, false)
-	vCopy := mlx.Contiguous(vSlice, false)
-	mlx.Pin(kCopy, vCopy)
-	mlx.AsyncEval(kCopy, vCopy)
+	copies := mlx.ScopedAsyncEval(func() []*mlx.Array {
+		kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
+		vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
+		return []*mlx.Array{mlx.Contiguous(kSlice, false), mlx.Contiguous(vSlice, false)}
+	})
+	s.scope = mlx.NewScope()
+	s.scope.Attach(copies...)
 
-	s.keys, s.values = kCopy, vCopy
+	s.keys, s.values = copies[0], copies[1]
 	c.dropLazySnapshot(s)
 	s.cache = nil
 
@@ -249,7 +252,7 @@ func (c *KVCache) Restore(snapshot Snapshot, target int) bool {
 
 	// Rewind to snapshot start, then feed snapshot.
 	c.offset = snap.fromOffset
-	c.appendKV(snap.keys, snap.values)
+	mlx.Scoped(func() { c.appendKV(snap.keys, snap.values) })
 
 	// Clamp to target if needed (target may be less than full snapshot).
 	if target < c.offset {
@@ -287,20 +290,21 @@ func (c *KVCache) Merge(parent, child Snapshot) Snapshot {
 	p.copyOut()
 	ch.copyOut()
 
-	mk := p.keys.Concatenate(2, ch.keys)
-	mv := p.values.Concatenate(2, ch.values)
-	mlx.Pin(mk, mv)
-	mlx.AsyncEval(mk, mv)
-
-	p.Close()
-	ch.Close()
-
-	return &kvSnapshot{
-		keys:       mk,
-		values:     mv,
+	merged := &kvSnapshot{
+		scope:      mlx.NewScope(),
 		fromOffset: p.fromOffset,
 		toOffset:   ch.toOffset,
 	}
+	joined := mlx.ScopedAsyncEval(func() []*mlx.Array {
+		joined := []*mlx.Array{p.keys.Concatenate(2, ch.keys), p.values.Concatenate(2, ch.values)}
+		p.Close()
+		ch.Close()
+		return joined
+	})
+	merged.scope.Attach(joined...)
+	merged.keys, merged.values = joined[0], joined[1]
+
+	return merged
 }
 
 func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
@@ -327,27 +331,23 @@ func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
 		return p, ch
 	}
 
-	pk := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false)
-	pv := mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false)
-	ck := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false)
-	cv := mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false)
-	mlx.Pin(pk, pv, ck, cv)
-	mlx.AsyncEval(pk, pv, ck, cv)
+	p := &kvSnapshot{scope: mlx.NewScope(), fromOffset: snap.fromOffset, toOffset: at}
+	ch := &kvSnapshot{scope: mlx.NewScope(), fromOffset: at, toOffset: snap.toOffset}
+	halves := mlx.ScopedAsyncEval(func() []*mlx.Array {
+		halves := []*mlx.Array{
+			mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false),
+			mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false),
+			mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false),
+			mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false),
+		}
+		snap.Close()
+		return halves
+	})
+	p.scope.Attach(halves[0], halves[1])
+	ch.scope.Attach(halves[2], halves[3])
+	p.keys, p.values = halves[0], halves[1]
+	ch.keys, ch.values = halves[2], halves[3]
 
-	snap.Close()
-
-	p := &kvSnapshot{
-		keys:       pk,
-		values:     pv,
-		fromOffset: snap.fromOffset,
-		toOffset:   at,
-	}
-	ch := &kvSnapshot{
-		keys:       ck,
-		values:     cv,
-		fromOffset: at,
-		toOffset:   snap.toOffset,
-	}
 	return p, ch
 }
 
@@ -358,7 +358,7 @@ func (c *KVCache) Free() {
 	for _, s := range slices.Clone(c.lazySnapshots) {
 		s.copyOut()
 	}
-	mlx.Unpin(c.keys, c.values)
+	c.scope.Close()
 	c.keys, c.values = nil, nil
 	c.offset = 0
 	c.snapshots = pendingSnapshots{}

x/mlxrunner/cache/lazy_test.go

@@ -31,10 +31,10 @@ func firstKeyAt(arr *mlx.Array, p, D int) float32 {
 	return arr.Floats()[p*D]
 }
 
-// settledActiveMemory drains unpinned arrays and the allocator cache, then
-// reports active (allocated, in-use) bytes.
+// settledActiveMemory drains the allocator cache, then reports active
+// (allocated, in-use) bytes. The work between two readings runs in a scope
+// so its intermediates are gone by the second one.
 func settledActiveMemory() int {
-	mlx.Sweep()
 	mlx.ClearCache()
 	return mlx.ActiveMemory()
 }
@@ -62,26 +62,28 @@ func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) {
 
 		baseline := settledActiveMemory()
 
-		snaps := c.TakeSnapshots()
-		// Every captured snapshot is a lazy snapshot (no owned buffer).
-		for i, s := range snaps {
-			if s == nil {
-				continue
+		mlx.Scoped(func() {
+			snaps := c.TakeSnapshots()
+			// Every captured snapshot is a lazy snapshot (no owned buffer).
+			for i, s := range snaps {
+				if s == nil {
+					continue
+				}
+				if ks := s.(*kvSnapshot); ks.keys != nil {
+					t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i)
+				}
 			}
-			if ks := s.(*kvSnapshot); ks.keys != nil {
-				t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i)
-			}
-		}
 
-		// MTP commit: rewind to a partial accept, then discard all snapshots.
-		if !c.Restore(nil, before+draft/2) {
-			t.Fatal("live rewind failed")
-		}
-		for _, s := range snaps {
-			if s != nil {
-				s.Close()
+			// MTP commit: rewind to a partial accept, then discard all snapshots.
+			if !c.Restore(nil, before+draft/2) {
+				t.Fatal("live rewind failed")
 			}
-		}
+			for _, s := range snaps {
+				if s != nil {
+					s.Close()
+				}
+			}
+		})
 
 		after := settledActiveMemory()
 		// Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but
@@ -219,25 +221,28 @@ func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) {
 
 		base := settledActiveMemory()
 
-		// Lazy snapshot [2,10), split at 5.
-		snap := c.Snapshot(2)
-		p, ch := c.Split(snap, 5)
-		ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot)
-		if ps.keys != nil || cs.keys != nil {
-			t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)")
-		}
-		if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 {
-			t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset)
-		}
+		var merged *kvSnapshot
+		mlx.Scoped(func() {
+			// Lazy snapshot [2,10), split at 5.
+			snap := c.Snapshot(2)
+			p, ch := c.Split(snap, 5)
+			ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot)
+			if ps.keys != nil || cs.keys != nil {
+				t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)")
+			}
+			if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 {
+				t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset)
+			}
 
-		// Merge them back into [2,10).
-		merged := c.Merge(p, ch).(*kvSnapshot)
-		if merged.keys != nil {
-			t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)")
-		}
-		if merged.fromOffset != 2 || merged.toOffset != 10 {
-			t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset)
-		}
+			// Merge them back into [2,10).
+			merged = c.Merge(p, ch).(*kvSnapshot)
+			if merged.keys != nil {
+				t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)")
+			}
+			if merged.fromOffset != 2 || merged.toOffset != 10 {
+				t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset)
+			}
+		})
 
 		if after := settledActiveMemory(); after > base {
 			t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base)
@@ -313,15 +318,17 @@ func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) {
 		// Restore the snapshot back to 10. Its slots [5,10) were never overwritten,
 		// so it is still lazy and the data is already in the buffer — a pure offset
 		// move, no allocation.
-		if !c.Restore(snap, 10) {
-			t.Fatal("restore failed")
-		}
-		if snap.keys != nil {
-			t.Fatal("snapshot was copied out; expected the offset-move fast path")
-		}
-		if c.Offset() != 10 {
-			t.Fatalf("offset after restore = %d, want 10", c.Offset())
-		}
+		mlx.Scoped(func() {
+			if !c.Restore(snap, 10) {
+				t.Fatal("restore failed")
+			}
+			if snap.keys != nil {
+				t.Fatal("snapshot was copied out; expected the offset-move fast path")
+			}
+			if c.Offset() != 10 {
+				t.Fatalf("offset after restore = %d, want 10", c.Offset())
+			}
+		})
 		if after := settledActiveMemory(); after > base {
 			t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base)
 		}

x/mlxrunner/cache/recurrent.go

@@ -20,6 +20,7 @@ import (
 type RecurrentCache struct {
 	convState  *mlx.Array
 	deltaState *mlx.Array
+	scope      *mlx.Scope
 	offset     int
 
 	convTail  int
@@ -77,13 +78,14 @@ func (c *RecurrentCache) captureBoundary(reached int, conv, delta *mlx.Array) {
 
 func (c *RecurrentCache) setState(old, v *mlx.Array) *mlx.Array {
 	v = v.Clone()
-	mlx.Pin(v)
-	mlx.Unpin(old)
+	c.scope.Attach(v)
+	c.scope.Discard(old)
 	return v
 }
 
 func NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim int32) *RecurrentCache {
 	return &RecurrentCache{
+		scope:     mlx.NewScope(),
 		convTail:  int(convTail),
 		convDim:   int(convDim),
 		numVHeads: int(numVHeads),
@@ -174,26 +176,28 @@ func (c *RecurrentCache) State() []*mlx.Array {
 // does not depend on any parent state.
 type recurrentSnapshot struct {
 	convState, deltaState *mlx.Array
+	scope                 *mlx.Scope
 	offset                int
 }
 
 func (s *recurrentSnapshot) Size() int { return s.convState.NumBytes() + s.deltaState.NumBytes() }
-func (s *recurrentSnapshot) Close()    { mlx.Unpin(s.convState, s.deltaState) }
+func (s *recurrentSnapshot) Close()    { s.scope.Close() }
 
 // SetMaterializeHook is a no-op: recurrent snapshots own their compact copy from
 // construction.
 func (s *recurrentSnapshot) SetMaterializeHook(func(int)) {}
 
-// newRecurrentSnapshot clones and pins conv/delta into an owned snapshot at
+// newRecurrentSnapshot clones conv/delta into an owned snapshot at
 // offset. It does not schedule the eval — capture-path snapshots ride the
 // cache's State into the caller's batched eval.
 func newRecurrentSnapshot(conv, delta *mlx.Array, offset int) *recurrentSnapshot {
 	snap := &recurrentSnapshot{
 		convState:  conv.Clone(),
 		deltaState: delta.Clone(),
+		scope:      mlx.NewScope(),
 		offset:     offset,
 	}
-	mlx.Pin(snap.convState, snap.deltaState)
+	snap.scope.Attach(snap.convState, snap.deltaState)
 	return snap
 }
 
@@ -226,8 +230,10 @@ func (c *RecurrentCache) Restore(snapshot Snapshot, target int) bool {
 		return false
 	}
 
-	c.convState = c.setState(c.convState, snap.convState)
-	c.deltaState = c.setState(c.deltaState, snap.deltaState)
+	mlx.Scoped(func() {
+		c.convState = c.setState(c.convState, snap.convState)
+		c.deltaState = c.setState(c.deltaState, snap.deltaState)
+	})
 	c.offset = snap.offset
 
 	return true
@@ -248,7 +254,7 @@ func (c *RecurrentCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
 }
 
 func (c *RecurrentCache) Free() {
-	mlx.Unpin(c.convState, c.deltaState)
+	c.scope.Close()
 	c.convState, c.deltaState = nil, nil
 	c.offset = 0
 	c.snapshots = pendingSnapshots{}

x/mlxrunner/cache/rotating.go

@@ -13,6 +13,7 @@ import (
 // RotatingKVCache implements sliding window attention with bounded memory.
 type RotatingKVCache struct {
 	keys, values *mlx.Array
+	scope        *mlx.Scope
 	offset       int
 	step         int
 	maxSize      int
@@ -29,7 +30,7 @@ type RotatingKVCache struct {
 }
 
 func NewRotatingKVCache(maxSize int) *RotatingKVCache {
-	return &RotatingKVCache{maxSize: maxSize, step: 256}
+	return &RotatingKVCache{maxSize: maxSize, step: 256, scope: mlx.NewScope()}
 }
 
 // Assumes B = 1; heterogeneous batches are not supported.
@@ -65,7 +66,7 @@ func (c *RotatingKVCache) concat(keys, values *mlx.Array) (newK *mlx.Array, newV
 
 	if c.keys == nil {
 		c.keys, c.values = keys.Clone(), values.Clone()
-		mlx.Pin(c.keys, c.values)
+		c.scope.Attach(c.keys, c.values)
 	} else {
 		if c.idx < c.keys.Dim(2) {
 			if c.offset <= c.maxSize {
@@ -123,7 +124,7 @@ func (c *RotatingKVCache) update(keys, values *mlx.Array) (*mlx.Array, *mlx.Arra
 			c.values.Set(c.values.Concatenate(2, newValues))
 		} else {
 			c.keys, c.values = newKeys, newValues
-			mlx.Pin(c.keys, c.values)
+			c.scope.Attach(c.keys, c.values)
 		}
 		c.idx = prev
 	}
@@ -194,18 +195,18 @@ func (c *RotatingKVCache) State() []*mlx.Array {
 	}
 }
 
-// replaceBuffer swaps in newK/newV as the cache's keys/values, unpinning the old
-// buffer and pinning the new one.
+// replaceBuffer swaps in newK/newV as the cache's keys/values, releasing the
+// old buffer and holding the new one.
 func (c *RotatingKVCache) replaceBuffer(newK, newV *mlx.Array) {
-	mlx.Unpin(c.keys, c.values)
+	c.scope.Discard(c.keys, c.values)
 	c.keys, c.values = newK, newV
-	mlx.Pin(c.keys, c.values)
+	c.scope.Attach(c.keys, c.values)
 }
 
 func (c *RotatingKVCache) Free() {
 	// Freeing drops the buffer lazy snapshots index into; copy them out first.
 	c.copyOutLazySnapshots()
-	mlx.Unpin(c.keys, c.values)
+	c.scope.Close()
 	c.keys, c.values = nil, nil
 	c.offset = 0
 	c.idx = 0
@@ -283,6 +284,7 @@ func (c *RotatingKVCache) lazyRotatingSnapshot(o int) Snapshot {
 // reorders or drops those slots.
 type rotatingSnapshot struct {
 	keys, values         *mlx.Array // owned window once copied out; nil while lazy
+	scope                *mlx.Scope // holds keys and values once copied out
 	fromOffset, toOffset int        // absolute offset range the window covers
 	idx                  int        // buffer write position a restore installs
 
@@ -307,7 +309,7 @@ func (s *rotatingSnapshot) Size() int {
 func (s *rotatingSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
 
 func (s *rotatingSnapshot) Close() {
-	mlx.Unpin(s.keys, s.values)
+	s.scope.Close()
 	if s.cache != nil {
 		s.cache.dropLazySnapshot(s)
 		s.cache = nil
@@ -322,14 +324,15 @@ func (s *rotatingSnapshot) copyOut() {
 		return
 	}
 	c := s.cache
-	kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
-	vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
-	k := mlx.Contiguous(kSlice, false)
-	v := mlx.Contiguous(vSlice, false)
-	mlx.Pin(k, v)
-	mlx.AsyncEval(k, v)
+	copies := mlx.ScopedAsyncEval(func() []*mlx.Array {
+		kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
+		vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
+		return []*mlx.Array{mlx.Contiguous(kSlice, false), mlx.Contiguous(vSlice, false)}
+	})
+	s.scope = mlx.NewScope()
+	s.scope.Attach(copies...)
 
-	s.keys, s.values = k, v
+	s.keys, s.values = copies[0], copies[1]
 	c.dropLazySnapshot(s)
 	s.cache = nil
 
@@ -363,11 +366,13 @@ func (c *RotatingKVCache) Snapshot(fromOffset int) Snapshot {
 	state := c.State()
 	k := state[0].Clone()
 	v := state[1].Clone()
-	mlx.Pin(k, v)
+	scope := mlx.NewScope()
+	scope.Attach(k, v)
 
 	return &rotatingSnapshot{
 		keys:       k,
 		values:     v,
+		scope:      scope,
 		fromOffset: fromOffset,
 		toOffset:   c.offset,
 		idx:        c.idx,
@@ -414,10 +419,12 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool {
 		c.dropLazySnapshot(snap)
 		c.copyOutLazySnapshots()
 		liveLen := snap.sliceEnd - snap.sliceStart
-		c.replaceBuffer(
-			c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
-			c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
-		)
+		mlx.Scoped(func() {
+			c.replaceBuffer(
+				c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
+				c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
+			)
+		})
 		snap.sliceStart, snap.sliceEnd = 0, liveLen
 		c.lazySnapshots = append(c.lazySnapshots, snap)
 		c.offset = snap.toOffset
@@ -435,7 +442,7 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool {
 	snap.copyOut()
 	c.copyOutLazySnapshots()
 
-	c.replaceBuffer(snap.keys.Clone(), snap.values.Clone())
+	mlx.Scoped(func() { c.replaceBuffer(snap.keys.Clone(), snap.values.Clone()) })
 	c.offset = snap.toOffset
 	c.idx = snap.idx
 

x/mlxrunner/dflash.go

@@ -8,7 +8,7 @@ import (
 	"github.com/ollama/ollama/x/mlxrunner/model/base"
 )
 
-// dflashPendingFlushTokens bounds the pinned feature rows between flushes.
+// dflashPendingFlushTokens bounds the held feature rows between flushes.
 const dflashPendingFlushTokens = 256
 
 // dflashDrafter drafts with a block-diffusion draft model (DFlash): one
@@ -48,6 +48,7 @@ type dflashDraftSession struct {
 	ctxOffset       int
 	pendingFeatures []*mlx.Array
 	pendingCount    int
+	pending         *mlx.Scope // holds the rows until the flush
 
 	// blockOutstanding tracks the proposal's scheduled rollback point, which
 	// commitBlock has to drain even when it needs no rewind.
@@ -66,7 +67,10 @@ func (d *dflashDraftSession) committed(tokens, features *mlx.Array, position int
 	}
 	if start < n {
 		f := features.Slice(mlx.Slice(), mlx.Slice(start, n), mlx.Slice())
-		mlx.Pin(f)
+		if d.pending == nil {
+			d.pending = mlx.NewScope()
+		}
+		d.pending.Attach(f)
 		d.pendingFeatures = append(d.pendingFeatures, f)
 		d.pendingCount += n - start
 		if d.pendingCount >= dflashPendingFlushTokens {
@@ -91,7 +95,8 @@ func (d *dflashDraftSession) takePending() *mlx.Array {
 		return nil
 	}
 	features := mlx.Concatenate(d.pendingFeatures, 1)
-	mlx.Unpin(d.pendingFeatures...)
+	d.pending.Close()
+	d.pending = nil
 	d.pendingFeatures = nil
 	d.ctxOffset += d.pendingCount
 	d.pendingCount = 0
@@ -116,24 +121,25 @@ func (d *dflashDraftSession) flush() {
 	spec := d.drafter.spec
 	d.commitBlock()
 
-	offset := d.ctxOffset
-	features := d.takePending()
-	if features == nil {
+	if len(d.pendingFeatures) == 0 {
 		return
 	}
-	spec.draft.Forward(&batch.Batch{
-		SeqOffsets: []int32{int32(offset)},
-		Hidden:     features,
-		Layout:     d.layout,
-	}, spec.targets, spec.draftKV)
-
-	// Force the cache writes: a session that never drafts would otherwise
-	// leave the flush chain unevaluated, pinning every feature until close.
-	state := make([]*mlx.Array, 0, 2*len(spec.draftKV))
-	for _, c := range spec.draftKV {
-		state = append(state, c.State()...)
-	}
-	mlx.AsyncEval(state...)
+	offset := d.ctxOffset
+	// Evaluating the cache state forces the writes: a session that never
+	// drafts would otherwise leave the flush chain unevaluated, holding
+	// every feature until close.
+	mlx.ScopedAsyncEval(func() []*mlx.Array {
+		spec.draft.Forward(&batch.Batch{
+			SeqOffsets: []int32{int32(offset)},
+			Hidden:     d.takePending(),
+			Layout:     d.layout,
+		}, spec.targets, spec.draftKV)
+		state := make([]*mlx.Array, 0, 2*len(spec.draftKV))
+		for _, c := range spec.draftKV {
+			state = append(state, c.State()...)
+		}
+		return state
+	})
 }
 
 // propose drafts a block after the not-yet-validated current token, one
@@ -152,34 +158,39 @@ func (d *dflashDraftSession) propose(current *mlx.Array, maxTokens int) *draftCa
 	// Send only the anchor and the rows being sampled, not the full trained
 	// block. Exact for causal layers, and measured as free for bidirectional
 	// ones.
-	masks := make([]int32, n)
-	for i := range masks {
-		masks[i] = d.drafter.maskToken
-	}
-	block := current.ExpandDims(-1).Concatenate(1, mlx.FromValues(masks, 1, len(masks)))
+	var candidates *draftCandidates
+	mlx.ScopedArrays(func() []*mlx.Array {
+		masks := make([]int32, n)
+		for i := range masks {
+			masks[i] = d.drafter.maskToken
+		}
+		block := current.ExpandDims(-1).Concatenate(1, mlx.FromValues(masks, 1, len(masks)))
 
-	offset := d.ctxOffset
-	features := d.takePending()
+		offset := d.ctxOffset
+		features := d.takePending()
 
-	scheduleSpeculation(spec.draftKV, d.ctxOffset, 1)
-	d.blockOutstanding = true
+		scheduleSpeculation(spec.draftKV, d.ctxOffset, 1)
+		d.blockOutstanding = true
 
-	hidden, _ := spec.draft.Forward(&batch.Batch{
-		InputIDs:     block,
-		SeqOffsets:   []int32{int32(offset)},
-		SeqQueryLens: []int32{int32(n + 1)},
-		Hidden:       features,
-		Layout:       d.layout,
-	}, spec.targets, spec.draftKV)
+		hidden, _ := spec.draft.Forward(&batch.Batch{
+			InputIDs:     block,
+			SeqOffsets:   []int32{int32(offset)},
+			SeqQueryLens: []int32{int32(n + 1)},
+			Hidden:       features,
+			Layout:       d.layout,
+		}, spec.targets, spec.draftKV)
 
-	// Row i predicts the token at its own position, so the anchor row is
-	// unused. Rows 1..n are sampled from one batched distribution; penalties
-	// see only the committed history, not the other rows of the block.
-	logits := spec.draft.Unembed(hidden.Slice(mlx.Slice(), mlx.Slice(1, n+1), mlx.Slice()))
-	dist := r.Sampler.Distribution(pipelineSlot, logits, nil)
-	tokens := r.Sampler.SampleDistribution(pipelineSlot, dist)
-	return &draftCandidates{
-		tokens: tokens.ExpandDims(0),
-		dist:   dist,
-	}
+		// Row i predicts the token at its own position, so the anchor row is
+		// unused. Rows 1..n are sampled from one batched distribution; penalties
+		// see only the committed history, not the other rows of the block.
+		logits := spec.draft.Unembed(hidden.Slice(mlx.Slice(), mlx.Slice(1, n+1), mlx.Slice()))
+		dist := r.Sampler.Distribution(pipelineSlot, logits, nil)
+		tokens := r.Sampler.SampleDistribution(pipelineSlot, dist)
+		candidates = &draftCandidates{
+			tokens: tokens.ExpandDims(0),
+			dist:   dist,
+		}
+		return candidates.Arrays()
+	})
+	return candidates
 }

x/mlxrunner/grammar.go

@@ -32,7 +32,7 @@ const (
 )
 
 // grammarEngine is the runner's structured-output subsystem: xgrammar bound
-// to the model's vocabulary, plus the pinned lookup table for expanding
+// to the model's vocabulary, plus the held lookup table for expanding
 // packed token masks on the device.
 type grammarEngine struct {
 	// compileMu is the single compile slot: one native compile at a time
@@ -46,6 +46,7 @@ type grammarEngine struct {
 
 	maskTable  *mlx.Array
 	byteShifts *mlx.Array
+	scope      *mlx.Scope
 }
 
 func newGrammarEngine(logitsWidth int, tokenizer *tokenizer.Tokenizer) *grammarEngine {
@@ -91,7 +92,7 @@ func validateGrammarVocab(logitsWidth, tokenizerSize int) error {
 }
 
 // initMask builds the byte-to-mask lookup table for expanding packed token
-// masks on the device, pinned for the runner's lifetime: row v holds, for
+// masks on the device, held for the runner's lifetime: row v holds, for
 // each of the byte value v's eight bits low to high, 0 where the bit is set
 // (token allowed) and -inf where it is clear.
 func (e *grammarEngine) initMask(vocabSize int) {
@@ -106,7 +107,8 @@ func (e *grammarEngine) initMask(vocabSize int) {
 	}
 	e.maskTable = mlx.FromValues(vals, 256, 8)
 	e.byteShifts = mlx.FromValues([]int32{0, 8, 16, 24}, 4)
-	mlx.Pin(e.maskTable, e.byteShifts)
+	e.scope = mlx.NewScope()
+	e.scope.Attach(e.maskTable, e.byteShifts)
 }
 
 func (e *grammarEngine) close() {
@@ -116,7 +118,7 @@ func (e *grammarEngine) close() {
 		e.compiler.Close()
 		e.compiler = nil
 	}
-	mlx.Unpin(e.maskTable, e.byteShifts)
+	e.scope.Close()
 	e.maskTable, e.byteShifts = nil, nil
 }
 

x/mlxrunner/media.go

@@ -54,6 +54,7 @@ type requestMedia struct {
 	// toggled in place so every batch shares the same slice.
 	manifest []batch.MediaItem
 	features []*mlx.Array // parallel to items; nil until encoded
+	scope    *mlx.Scope
 
 	// layout is the request's one-row Batch.Layout, shared by every batch
 	// like the manifest; nil when the model returned no layout.
@@ -70,6 +71,7 @@ func (r *Runner) openMedia(request Request) *requestMedia {
 		inputLen: len(request.Tokens),
 		manifest: make([]batch.MediaItem, len(request.MediaItems)),
 		features: make([]*mlx.Array, len(request.MediaItems)),
+		scope:    mlx.NewScope(),
 	}
 	if request.Layout != nil {
 		m.layout = []any{request.Layout}
@@ -114,7 +116,7 @@ func (m *requestMedia) extendChunk(pos, n int) int {
 }
 
 // batchMedia returns the manifest for chunk [pos, pos+n), encoding and
-// pinning each item's features on first overlap; nothing evaluates here —
+// holding each item's features on first overlap; nothing evaluates here —
 // the consuming forward pulls the encoder.
 func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem {
 	if m == nil {
@@ -125,9 +127,11 @@ func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem {
 			continue
 		}
 		if m.features[i] == nil {
-			data := mlx.FromValues(item.item.MediaData, item.item.Dims...)
-			m.features[i] = m.model.EncodeMedia(item.item, data)
-			mlx.Pin(m.features[i])
+			m.features[i] = mlx.ScopedArrays(func() []*mlx.Array {
+				data := mlx.FromValues(item.item.MediaData, item.item.Dims...)
+				return []*mlx.Array{m.model.EncodeMedia(item.item, data)}
+			})[0]
+			m.scope.Attach(m.features[i])
 			// The upload copied the pixels; free them here — release never
 			// passes the end of an expansion reaching the prompt's last token.
 			item.item.MediaData = nil
@@ -137,9 +141,9 @@ func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem {
 	return m.manifest
 }
 
-// release frees what items fully evaluated or restored at position pos no
-// longer need: the pinned features and the preprocessed pixel buffer.
-func (m *requestMedia) release(pos int) {
+// free frees what items fully evaluated or restored at position pos no
+// longer need: the held features and the preprocessed pixel buffer.
+func (m *requestMedia) free(pos int) {
 	if m == nil {
 		return
 	}
@@ -147,7 +151,7 @@ func (m *requestMedia) release(pos int) {
 		if item.pos+item.length <= pos {
 			item.item.MediaData = nil
 			if m.features[i] != nil {
-				mlx.Unpin(m.features[i])
+				m.scope.Discard(m.features[i])
 				m.features[i] = nil
 				m.manifest[i].Features = nil
 			}
@@ -155,18 +159,16 @@ func (m *requestMedia) release(pos int) {
 	}
 }
 
-// close unpins whatever remains when the pipeline exits.
+// close frees whatever remains when the pipeline exits.
 func (m *requestMedia) close() {
 	if m == nil {
 		return
 	}
-	for i, f := range m.features {
-		if f != nil {
-			mlx.Unpin(f)
-			m.features[i] = nil
-			m.manifest[i].Features = nil
-		}
+	for i := range m.features {
+		m.features[i] = nil
+		m.manifest[i].Features = nil
 	}
+	m.scope.Close()
 }
 
 // expandMedia tokenizes the [img-N]-tagged prompt into segments, expands

x/mlxrunner/media_test.go

@@ -105,11 +105,11 @@ func TestBatchMediaLifecycle(t *testing.T) {
 			t.Fatalf("second overlap re-encoded (calls=%d)", calls)
 		}
 
-		m.release(4)
+		m.free(4)
 		if m.manifest[0].Features == nil {
 			t.Fatal("release dropped features before the expansion was evaluated")
 		}
-		m.release(6)
+		m.free(6)
 		if m.manifest[0].Features != nil {
 			t.Fatal("release kept features past the expansion end")
 		}

x/mlxrunner/mlx/act_test.go

@@ -23,7 +23,6 @@ func TestGELUCompiledMatchesEager(t *testing.T) {
 			withMLXThread(t, func(t *mlxthreadtest.T) {
 				EnableCompile()
 				input := FromValues(values, len(values)).AsType(tt.dtype)
-				Pin(input)
 
 				want := gelu(input)
 				got := GELU(input)
@@ -38,7 +37,6 @@ func TestGELUCompiledMatchesEager(t *testing.T) {
 						t.Fatalf("%s GELU[%d] = %v, want %v (delta %v)", tt.name, i, gotValues[i], wantValues[i], delta)
 					}
 				}
-				Unpin(input)
 			})
 		})
 	}
@@ -58,22 +56,13 @@ func benchmarkGELU(b *testing.B, fn func(*Array) *Array) {
 		EnableCompile()
 		input := AddScalar(Zeros(DTypeBFloat16, 1, 4096, 8192), 1)
 		Eval(input)
-		Pin(input)
-		defer func() {
-			Unpin(input)
-			Sweep()
-			ClearCache()
-		}()
+		defer ClearCache()
 
-		warmup := fn(input)
-		Eval(warmup)
-		Sweep()
+		Scoped(func() { Eval(fn(input)) })
 
 		b.ResetTimer()
 		for range b.N {
-			output := fn(input)
-			Eval(output)
-			Sweep()
+			Scoped(func() { Eval(fn(input)) })
 		}
 		return nil
 	}); err != nil {
@@ -85,8 +74,6 @@ func TestReLUSquared(t *testing.T) {
 	var got []float32
 	withMLXThread(t, func(t *mlxthreadtest.T) {
 		x := FromValues([]float32{-2, -0, 0.5, 2}, 4)
-		Pin(x)
-		defer Unpin(x)
 
 		y := ReLUSquared(x)
 		Eval(y)

x/mlxrunner/mlx/array.go

@@ -8,40 +8,22 @@ import (
 	"fmt"
 	"log/slog"
 	"reflect"
-	"sort"
 	"strings"
-	"sync"
-	"sync/atomic"
 	"unsafe"
-
-	"github.com/ollama/ollama/logutil"
 )
 
+// An Array's lifetime is governed by the scope it belongs to; see scope.go.
 type Array struct {
-	ctx    C.mlx_array
-	name   string
-	pinned atomic.Int32
+	ctx   C.mlx_array
+	name  string
+	scope *Scope
 }
 
-var (
-	arrays   []*Array
-	arraysMu sync.Mutex
-)
-
 // constructor utilities
 
 func New(name string) *Array {
 	t := &Array{name: name}
-
-	if tracing {
-		traceScratch = append(traceScratch, t)
-	} else {
-		arraysMu.Lock()
-		defer arraysMu.Unlock()
-
-		arrays = append(arrays, t)
-	}
-
+	currentScope.take(t)
 	return t
 }
 
@@ -135,52 +117,17 @@ func (t *Array) Clone() *Array {
 	return tt
 }
 
-// lifecycle utilities
-
-// Pin marks arrays as in-use so they are retained during Sweep.
-func Pin(s ...*Array) {
-	for _, t := range s {
-		if t != nil {
-			t.pinned.Add(1)
-		}
-	}
-}
-
-// Unpin marks arrays as no longer in-use, allowing Sweep to free them.
-func Unpin(s ...*Array) {
-	for _, t := range s {
-		if t != nil {
-			if t.pinned.Add(-1) < 0 {
-				panic(fmt.Sprintf("mlx.Unpin: negative pin count on array %q", t.name))
-			}
-		}
-	}
-}
-
-// Sweep releases all unpinned arrays, primarily intermediate tensors. MLX will truly
-// free them when there are no other references, including dependencies in the graph.
-func Sweep() {
-	arraysMu.Lock()
-	defer arraysMu.Unlock()
-	n := 0
-	for _, t := range arrays {
-		if t.pinned.Load() > 0 && t.Valid() {
-			arrays[n] = t
-			n++
-		} else if t.Valid() {
-			mlxCheck(C.mlx_array_free(t.ctx))
-			t.ctx.ctx = nil
-		}
-	}
-	arrays = arrays[:n]
-}
-
 // misc. utilities
 
 func (t *Array) Valid() bool {
 	return t.ctx.ctx != nil
 }
 
+func (t *Array) free() {
+	mlxCheck(C.mlx_array_free(t.ctx))
+	t.ctx.ctx = nil
+}
+
 func (t *Array) String() string {
 	str := mlxCheck(C.mlx_string_new())
 	mlxCheck(C.mlx_array_tostring(&str, t.ctx))
@@ -191,7 +138,6 @@ func (t *Array) String() string {
 func (t *Array) LogValue() slog.Value {
 	attrs := []slog.Attr{
 		slog.String("name", t.name),
-		slog.Int("pinned", int(t.pinned.Load())),
 	}
 	if t.Valid() {
 		attrs = append(attrs,
@@ -288,20 +234,3 @@ func (t *Array) Save(name string) error {
 	}
 	return nil
 }
-
-// LogArrays logs all live arrays, sorted by size
-func LogArrays() {
-	arraysMu.Lock()
-	defer arraysMu.Unlock()
-	sort.Slice(arrays, func(i, j int) bool {
-		return arrays[i].NumBytes() > arrays[j].NumBytes()
-	})
-
-	var total int
-	for _, t := range arrays {
-		nb := t.NumBytes()
-		total += nb
-		logutil.Trace(fmt.Sprintf("tensor %-60s %5s %5s pinned=%d %v", t.name, t.DType(), PrettyBytes(nb), t.pinned.Load(), t.Dims()))
-	}
-	logutil.Trace(fmt.Sprintf("tensors total: %d, size: %s, active: %s", len(arrays), PrettyBytes(total), PrettyBytes(ActiveMemory())))
-}

x/mlxrunner/mlx/compile.go

@@ -120,10 +120,6 @@ func Compile3(name string, fn func(*Array, *Array, *Array) *Array, opts ...Compi
 // single-threaded at this level a plain Go bool suffices.
 var tracing bool
 
-// traceScratch collects arrays created during a compile trace so they can be
-// freed as a group when the callback returns.
-var traceScratch []*Array
-
 //export closureCallback
 func closureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload unsafe.Pointer) (rc C.int) {
 	defer func() {
@@ -136,26 +132,19 @@ func closureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload
 	handle := *(*cgo.Handle)(payload)
 	fn := handle.Value().(CompileFunc)
 
-	// When tracing, we track all of the intermediates that are created and free them separately at the end of
-	// the process. This will give the effect of a single op - inputs are owned by the original caller (via
-	// the MLX layer) and outputs are transferred back to MLX to create a new Go side tensor.
+	// The trace runs in its own scope so its intermediates are freed as a
+	// group when the callback returns. This gives the effect of a single op:
+	// inputs are owned by the original caller (via the MLX layer) and outputs
+	// are transferred back to MLX to create a new Go side tensor.
 	if tracing {
 		panic("mlx: nested compile trace")
 	}
 	tracing = true
-	traceScratch = nil
+	s := enterScope()
+	s.noEscape = true
 	defer func() {
-		for _, a := range traceScratch {
-			if a.pinned.Load() > 0 {
-				panic("mlx: traced array was pinned during compilation")
-			}
-			if a.Valid() {
-				mlxCheck(C.mlx_array_free(a.ctx))
-				a.ctx.ctx = nil
-			}
-		}
 		tracing = false
-		traceScratch = nil
+		exitScope(s)
 	}()
 
 	n := int(mlxCheck(C.mlx_vector_array_size(input)))

x/mlxrunner/mlx/compile_test.go

@@ -33,29 +33,27 @@ func testCompileFusion(t *mlxthreadtest.T) {
 
 	a := FromValues(data, n)
 	b := FromValues(data, n)
-	Pin(a, b)
-	defer Unpin(a, b)
 
 	// Compiled: ops fused into a single kernel.
 	EnableCompile()
 	fn := Compile2("diamond", body, Shapeless())
-	warm := fn(a, b)
-	Eval(warm)
-	Sweep()
+	Scoped(func() { Eval(fn(a, b)) })
 	ClearCache()
 	ResetPeakMemory()
-	y := fn(a, b)
-	Eval(y)
-	compiledPeak := PeakMemory()
-	Sweep()
+	var compiledPeak int
+	Scoped(func() {
+		Eval(fn(a, b))
+		compiledPeak = PeakMemory()
+	})
 
 	// Uncompiled: ops evaluated individually, intermediates materialized.
 	ClearCache()
 	ResetPeakMemory()
-	z := body(a, b)
-	Eval(z)
-	uncompiledPeak := PeakMemory()
-	Sweep()
+	var uncompiledPeak int
+	Scoped(func() {
+		Eval(body(a, b))
+		uncompiledPeak = PeakMemory()
+	})
 
 	if compiledPeak == 0 && uncompiledPeak == 0 {
 		t.Skip("peak memory tracking not available")
@@ -88,8 +86,6 @@ func testCompileNested(t *mlxthreadtest.T) {
 
 	gate := FromValues([]float32{0, 1, 2}, 3)
 	up := FromValues([]float32{1, 1, 1}, 3)
-	Pin(gate, up)
-	defer Unpin(gate, up)
 
 	y := outer(gate, up)
 	Eval(y)
@@ -116,8 +112,6 @@ func testCompileCallbackPanicRecovers(t *mlxthreadtest.T) {
 	})
 
 	x := FromValues([]float32{1}, 1)
-	Pin(x)
-	defer Unpin(x)
 
 	defer func() {
 		r := recover()
@@ -139,26 +133,22 @@ func TestCompileNoTrackingGrowth(t *testing.T) {
 
 func testCompileNoTrackingGrowth(t *mlxthreadtest.T) {
 	// Repeated invocations of a compiled kernel should not grow the
-	// tracked-arrays list; the callback's traceScratch collects
-	// intermediates during tracing and frees them when the callback returns.
+	// tracked-arrays list; the callback's scope collects intermediates
+	// during tracing and frees them when the callback returns.
 	fn := Compile2("mul_add", func(a, b *Array) *Array {
 		return a.Multiply(b).Add(b)
 	})
 
 	a := FromValues([]float32{1, 2}, 2)
 	b := FromValues([]float32{3, 4}, 2)
-	Pin(a, b)
-	defer Unpin(a, b)
 
-	Sweep()
-	before := len(arrays)
+	before := len(currentScope.arrays)
 
 	for range 100 {
-		_ = fn(a, b)
-		Sweep()
+		Scoped(func() { _ = fn(a, b) })
 	}
 
-	after := len(arrays)
+	after := len(currentScope.arrays)
 	if after > before+2 {
 		t.Fatalf("tracked arrays grew from %d to %d across 100 calls (includes initial trace)", before, after)
 	}

x/mlxrunner/mlx/scope.go

@@ -0,0 +1,203 @@
+package mlx
+
+import (
+	"fmt"
+)
+
+// Array lifetimes
+//
+// Every handle belongs to a scope, a set of arrays freed together. A
+// function scope is entered with Scoped, ScopedArrays, ScopedEval, or
+// ScopedAsyncEval and ends when the function returns; a held scope is
+// created with NewScope and ends when its holder closes it. A function
+// scope frees what was created in it or detached into it; a held scope
+// frees what was attached to it. An array moves between scopes in three
+// ways only: by being returned from a function scope to the caller's scope,
+// by Attach into a held scope, or by Detach from a held scope back into the
+// current one for a caller that still reads them.
+//
+// MLX frees a buffer once no handle and no queued graph refers to it, so a
+// graph is built in one function scope and evaluated after that scope ends,
+// and the eval frees each intermediate as it consumes it. The function that
+// finishes a graph opens the scope, returns the arrays that leave it, and
+// discards or closes inside it whatever the graph consumed. An operation
+// that only adds to its caller's graph, a model forward, a cache update, a
+// sampling distribution, builds into the open scope and never opens its
+// own; one that finishes a graph of its own, a cache copy, a sample, opens
+// one like any other finisher. Whoever needs the values evaluates them once
+// the scope has ended: the finisher itself with ScopedEval or
+// ScopedAsyncEval, or its caller with Eval or AsyncEval. A holder holds
+// what outlives the function that produced it, and only a holder gives its
+// arrays up.
+
+type Scope struct {
+	arrays []*Array
+	parent *Scope
+	// noEscape refuses to let an array move out of the scope.
+	noEscape bool
+}
+
+// Function scopes
+
+// Scoped runs fn in a function scope. Arrays created or released inside it
+// are freed when fn returns.
+func Scoped(fn func()) {
+	s := enterScope()
+	defer exitScope(s)
+	fn()
+}
+
+// ScopedArrays runs fn in a function scope and moves the arrays it returns to
+// the caller's scope.
+func ScopedArrays(fn func() []*Array) []*Array {
+	s := enterScope()
+	defer exitScope(s)
+	ts := fn()
+	escape(ts...)
+	return ts
+}
+
+// ScopedEval runs fn in a function scope, moves the arrays it returns to the
+// caller's scope, ends the scope, and then evaluates them.
+func ScopedEval(fn func() []*Array) []*Array {
+	ts := ScopedArrays(fn)
+	Eval(ts...)
+	return ts
+}
+
+// ScopedAsyncEval is ScopedEval with an asynchronous evaluation.
+func ScopedAsyncEval(fn func() []*Array) []*Array {
+	ts := ScopedArrays(fn)
+	AsyncEval(ts...)
+	return ts
+}
+
+// Held scopes
+
+func NewScope() *Scope {
+	return &Scope{}
+}
+
+// Attach takes arrays from the function scope that built them or from the
+// root. An array some held scope already holds, this one included, is that
+// holder's to discard or detach first: a second Attach means two owners.
+func (s *Scope) Attach(arrays ...*Array) {
+	for _, t := range arrays {
+		if t == nil {
+			continue
+		}
+		if t.scope != rootScope && t.scope.parent == nil {
+			panic(fmt.Sprintf("mlx: array %q is already held", t.name))
+		}
+		s.take(t)
+	}
+}
+
+// Detach moves arrays back to the current scope, for a caller that still
+// reads them.
+func (s *Scope) Detach(arrays ...*Array) {
+	for _, t := range arrays {
+		if t == nil {
+			continue
+		}
+		if t.scope != s {
+			panic(fmt.Sprintf("mlx: array %q is not held by this scope", t.name))
+		}
+		currentScope.take(t)
+	}
+}
+
+func (s *Scope) Discard(arrays ...*Array) {
+	for _, t := range arrays {
+		if t == nil {
+			continue
+		}
+		if t.scope != s {
+			panic(fmt.Sprintf("mlx: array %q is not held by this scope", t.name))
+		}
+		s.remove(t)
+		t.free()
+	}
+}
+
+// Close frees whatever the scope still holds. A nil *Scope holds nothing.
+func (s *Scope) Close() {
+	if s == nil {
+		return
+	}
+	s.end()
+}
+
+// Internals
+
+var (
+	rootScope    = &Scope{}
+	currentScope = rootScope
+)
+
+func enterScope() *Scope {
+	s := &Scope{parent: currentScope}
+	currentScope = s
+	return s
+}
+
+// exitScope ends s.
+func exitScope(s *Scope) {
+	if currentScope != s {
+		panic("mlx: scope exited out of order")
+	}
+	currentScope = s.parent
+	s.end()
+}
+
+// escape moves arrays of the current scope to the caller's scope so they
+// outlive the current one. Arrays held elsewhere are left where they are.
+func escape(arrays ...*Array) {
+	if currentScope.parent == nil {
+		return
+	}
+	for _, t := range arrays {
+		if t != nil && t.scope == currentScope {
+			currentScope.parent.take(t)
+		}
+	}
+}
+
+// end frees the arrays in s.
+func (s *Scope) end() {
+	for _, t := range s.arrays {
+		t.free()
+	}
+	s.arrays = nil
+}
+
+// take moves t into s, out of the scope it was in.
+func (s *Scope) take(t *Array) {
+	if from := t.scope; from != nil {
+		if from == s {
+			return
+		}
+		if !t.Valid() {
+			panic(fmt.Sprintf("mlx: array %q used after its scope ended", t.name))
+		}
+		if from.noEscape {
+			panic(fmt.Sprintf("mlx: array %q escaped a scope that allows no escape", t.name))
+		}
+		from.remove(t)
+	}
+	t.scope = s
+	s.arrays = append(s.arrays, t)
+}
+
+// remove drops t from s's list. An array usually leaves the scope that just
+// built it, so the search runs from the end; order in the list is free.
+func (s *Scope) remove(t *Array) {
+	for i := len(s.arrays) - 1; i >= 0; i-- {
+		if s.arrays[i] == t {
+			last := len(s.arrays) - 1
+			s.arrays[i] = s.arrays[last]
+			s.arrays = s.arrays[:last]
+			return
+		}
+	}
+}

x/mlxrunner/mlx/scope_test.go

@@ -0,0 +1,151 @@
+package mlx
+
+import (
+	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
+)
+
+// A function scope frees what was created in it. What fn returns moves to the
+// caller's scope instead, and a returned array the scope does not own stays
+// where it is.
+func TestScopeFreesWhatIsNotReturned(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		held := NewScope()
+		defer held.Close()
+		var kept, returned, dropped *Array
+		Scoped(func() {
+			kept = FromValue(1)
+			held.Attach(kept)
+			out := ScopedArrays(func() []*Array {
+				dropped = FromValue(2)
+				return []*Array{FromValue(3), nil, kept}
+			})
+			returned = out[0]
+			if !returned.Valid() {
+				t.Fatal("returned array was freed with the scope that created it")
+			}
+			if dropped.Valid() {
+				t.Fatal("array not returned survived its scope")
+			}
+		})
+		if returned.Valid() {
+			t.Fatal("returned array survived the scope it was returned into")
+		}
+		if !kept.Valid() {
+			t.Fatal("returning a held array moved it out of its scope")
+		}
+	})
+}
+
+// ScopedEval ends the build scope before it evaluates, so the intermediates
+// are gone by then and the returned arrays come back evaluated.
+func TestScopedEvalEvaluatesAfterBuild(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		Scoped(func() {
+			var tmp *Array
+			out := ScopedEval(func() []*Array {
+				tmp = FromValue(2)
+				return []*Array{FromValue(1).Add(tmp)}
+			})
+			if tmp.Valid() {
+				t.Fatal("intermediate survived the build scope")
+			}
+			if !out[0].Valid() || out[0].Int() != 3 {
+				t.Fatal("returned array was not evaluated after the build scope")
+			}
+		})
+	})
+}
+
+// A held scope keeps arrays past the function scope that created them.
+// Discard frees one now, Detach hands one to the current scope, and Close
+// frees what remains. A scope refuses to discard or detach an array it does
+// not hold, and to attach one that is already held, by itself or another
+// scope.
+func TestHeldScope(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		held, other := NewScope(), NewScope()
+		defer other.Close()
+		var kept, discarded, detached *Array
+		Scoped(func() {
+			kept, discarded, detached = FromValue(1), FromValue(2), FromValue(3)
+			held.Attach(kept, discarded, detached)
+			held.Discard(discarded)
+			if discarded.Valid() {
+				t.Fatal("discarded array survived")
+			}
+			Scoped(func() { held.Detach(detached) })
+			if detached.Valid() {
+				t.Fatal("detached array survived the scope it was detached into")
+			}
+			if !panics(func() { other.Discard(kept) }) {
+				t.Fatal("no panic discarding an array the scope does not hold")
+			}
+			if !panics(func() { other.Detach(kept) }) {
+				t.Fatal("no panic detaching an array the scope does not hold")
+			}
+			if !panics(func() { other.Attach(kept) }) {
+				t.Fatal("no panic holding an array another scope holds")
+			}
+			if !panics(func() { held.Attach(kept) }) {
+				t.Fatal("no panic holding an array twice")
+			}
+		})
+		if !kept.Valid() {
+			t.Fatal("held array was freed with the scope that created it")
+		}
+		held.Close()
+		if kept.Valid() {
+			t.Fatal("held array survived its scope's close")
+		}
+	})
+}
+
+// A scope ends when its function panics, so whatever recovers is back in
+// the scope it started from.
+func TestScopeEndsOnPanic(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		start := currentScope
+		var a *Array
+		func() {
+			defer func() { _ = recover() }()
+			Scoped(func() {
+				a = FromValue(1)
+				panic("build failed")
+			})
+		}()
+		if a.Valid() {
+			t.Fatal("array survived the scope that panicked")
+		}
+		if currentScope != start {
+			t.Fatal("registry not back in the caller's scope after a panic")
+		}
+	})
+}
+
+// Nothing built in a compile trace may outlive it: holding a trace array
+// fails the compiled call.
+func TestCompileTraceRefusesEscapes(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		held := NewScope()
+		defer held.Close()
+		double := Compile("scope_test_escape", func(in ...*Array) []*Array {
+			out := in[0].Add(in[0])
+			held.Attach(out)
+			return []*Array{out}
+		})
+		defer func() {
+			if recover() == nil {
+				t.Fatal("no panic for an array held out of a compile trace")
+			}
+		}()
+		Scoped(func() { double(FromValue(1)) })
+	})
+}
+
+func panics(fn func()) (panicked bool) {
+	defer func() { panicked = recover() != nil }()
+	fn()
+	return false
+}

x/mlxrunner/mlx/thread_test.go

@@ -55,11 +55,12 @@ func TestThreadedMLXOperations(t *testing.T) {
 
 			for range iterations {
 				if err := thread.Do(context.Background(), func() error {
-					a := FromValues([]float32{1, 2, 3, 4}, 2, 2)
-					b := Matmul(a, a)
-					AsyncEval(b)
-					Eval(b)
-					Sweep()
+					Scoped(func() {
+						a := FromValues([]float32{1, 2, 3, 4}, 2, 2)
+						b := Matmul(a, a)
+						AsyncEval(b)
+						Eval(b)
+					})
 					ClearCache()
 					return nil
 				}); err != nil {

x/mlxrunner/model/base/base.go

@@ -176,22 +176,3 @@ func NewDraft(root *model.Root, target Model) (DraftModel, error) {
 
 	return fn(root, target)
 }
-
-// Weights returns a function that loads model weights, then pins all
-// arrays reachable from the model struct and sweeps everything else.
-func Weights(m Model) func(map[string]*mlx.Array) error {
-	return func(tensors map[string]*mlx.Array) error {
-		if err := m.LoadWeights(tensors); err != nil {
-			return err
-		}
-
-		collected := mlx.Collect(m)
-		for _, arr := range collected {
-			mlx.Pin(arr)
-		}
-		mlx.Sweep()
-		mlx.Eval(collected...)
-
-		return nil
-	}
-}

x/mlxrunner/mtp.go

@@ -10,7 +10,7 @@ import (
 )
 
 // mtpPendingFlushTokens caps how many committed look-ahead tokens wait in the
-// pending buffer before a batched flush, bounding the pinned hidden states
+// pending buffer before a batched flush, bounding the held hidden states
 // regardless of what else triggers a flush.
 const mtpPendingFlushTokens = 256
 
@@ -37,7 +37,7 @@ func (d *mtpDrafter) draftLimit() int { return 0 }
 // open returns the drafting session for one request, its pairing frontier
 // synced to the draft caches' restored offset.
 func (d *mtpDrafter) open(layout []any) draftSession {
-	s := &mtpDraftSession{drafter: d, layout: layout}
+	s := &mtpDraftSession{drafter: d, layout: layout, scope: mlx.NewScope()}
 	if kv := d.spec.draftKV; len(kv) > 0 {
 		// A restored prefix arrives with the draft caches already written;
 		// pairing resumes from their absolute offset.
@@ -54,20 +54,22 @@ func (d *mtpDrafter) open(layout []any) draftSession {
 type mtpDraftSession struct {
 	drafter *mtpDrafter
 	layout  []any
+	scope   *mlx.Scope
 
 	// frontier is the slot after the last reported token; frontierHidden is
-	// the pinned target hidden at frontier-1, fused into the next pair.
+	// the held target hidden at frontier-1, fused into the next pair.
 	frontier       int
 	frontierHidden *mlx.Array
 
 	// committedDraftOffset is the slot after the last pair written to the
-	// draft caches; later pairs wait pinned in the pending lists until
+	// draft caches; later pairs wait held in the pending lists until
 	// flushed. pendingCount is the look-ahead tokens those lists hold, summed
 	// across the buffered runs.
 	committedDraftOffset int
 	pendingTokens        []*mlx.Array
 	pendingHiddens       []*mlx.Array
 	pendingCount         int
+	pending              *mlx.Scope // holds the lists' arrays until the flush
 
 	// heldHidden is the frontier row's pre-unembed hidden and heldAuxHidden
 	// its fusion hidden, carried from the last flush so the first proposal
@@ -76,7 +78,7 @@ type mtpDraftSession struct {
 	heldAuxHidden *mlx.Array
 
 	// pendingMedia holds manifest rows the deferred flush may still embed,
-	// pinned since prefill releases them after the target's chunk;
+	// on their own handles since prefill releases them after the target's chunk;
 	// lastDelivered marks each row's newest delivered end.
 	pendingMedia  map[int]batch.MediaItem
 	lastDelivered map[int]int
@@ -113,7 +115,7 @@ func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int, me
 	d.setFrontierHidden(lastHiddenRow(hiddens))
 }
 
-// captureMedia pins the run's feature-bearing rows for the deferred
+// captureMedia holds the run's feature-bearing rows for the deferred
 // flush, which embeds them after prefill has released the features. A row
 // spanning chunks arrives once per chunk.
 func (d *mtpDraftSession) captureMedia(media []batch.MediaItem, end int) {
@@ -129,7 +131,8 @@ func (d *mtpDraftSession) captureMedia(media []batch.MediaItem, end int) {
 				d.pendingMedia = make(map[int]batch.MediaItem)
 				d.lastDelivered = make(map[int]int)
 			}
-			mlx.Pin(item.Features)
+			item.Features = item.Features.Clone()
+			d.scope.Attach(item.Features)
 			d.pendingMedia[item.Pos] = item
 		}
 		d.lastDelivered[item.Pos] = end
@@ -150,7 +153,8 @@ func (d *mtpDraftSession) flushMedia(embedEnd int) []batch.MediaItem {
 	slices.SortFunc(manifest, func(a, b batch.MediaItem) int { return a.Pos - b.Pos })
 	for pos, last := range d.lastDelivered {
 		if last <= embedEnd {
-			mlx.Unpin(d.pendingMedia[pos].Features)
+			// The flush's forward still reads the row; it dies with the build.
+			d.scope.Detach(d.pendingMedia[pos].Features)
 			delete(d.pendingMedia, pos)
 			delete(d.lastDelivered, pos)
 		}
@@ -160,7 +164,7 @@ func (d *mtpDraftSession) flushMedia(embedEnd int) []batch.MediaItem {
 
 func (d *mtpDraftSession) closeMedia() {
 	for pos, item := range d.pendingMedia {
-		mlx.Unpin(item.Features)
+		d.scope.Discard(item.Features)
 		delete(d.pendingMedia, pos)
 		delete(d.lastDelivered, pos)
 	}
@@ -174,7 +178,7 @@ func (d *mtpDraftSession) settle(next *mlx.Array) {
 		return
 	}
 	if d.frontierHidden != nil && d.frontier-1 == d.committedDraftOffset+d.pendingCount {
-		d.queueCacheWrites(next.ExpandDims(-1), d.frontierHidden)
+		d.queueCacheWrites(next.ExpandDims(-1), d.frontierHidden.Clone())
 	}
 	d.flush()
 }
@@ -184,14 +188,18 @@ func (d *mtpDraftSession) close() {
 	d.closeMedia()
 	d.setFrontierHidden(nil)
 	d.setHeld(nil, nil)
+	d.scope.Close()
 }
 
 // queueCacheWrites buffers completed draft-cache writes — look-ahead tokens
 // fused with their target hiddens — flushing once the buffer reaches the token
-// cap so the pinned hiddens stay bounded. flush coalesces the buffered writes
+// cap so the held hiddens stay bounded. flush coalesces the buffered writes
 // into one head forward, so a contiguous run lands in a single draft-cache extend.
 func (d *mtpDraftSession) queueCacheWrites(tokens, hiddens *mlx.Array) {
-	mlx.Pin(tokens, hiddens)
+	if d.pending == nil {
+		d.pending = mlx.NewScope()
+	}
+	d.pending.Attach(tokens, hiddens)
 	d.pendingTokens = append(d.pendingTokens, tokens)
 	d.pendingHiddens = append(d.pendingHiddens, hiddens)
 	d.pendingCount += tokens.Dim(1)
@@ -216,45 +224,51 @@ func (d *mtpDraftSession) flush() {
 		}
 	}
 
-	ids := mlx.Concatenate(d.pendingTokens, 1)
-	hiddens := mlx.Concatenate(d.pendingHiddens, 1)
-	// The pair at slot S embeds the look-ahead token S+1, so this flush
-	// embeds prompt tokens up to committedDraftOffset+len+1.
-	hidden, auxHidden := spec.draft.Forward(&batch.Batch{
-		InputIDs:     ids,
-		SeqOffsets:   []int32{int32(d.committedDraftOffset)},
-		SeqQueryLens: []int32{int32(ids.Dim(1))},
-		Hidden:       hiddens,
-		Media:        d.flushMedia(d.committedDraftOffset + ids.Dim(1) + 1),
-		Layout:       d.layout,
-	}, spec.targets, spec.draftKV)
-	d.setHeld(lastHiddenRow(hidden), lastHiddenRow(auxHidden))
-	d.committedDraftOffset += ids.Dim(1)
+	// Evaluating the state forces the draft writes: a session that never
+	// drafts would otherwise leave the flush chain unevaluated, holding
+	// every hidden until close.
+	n := d.pendingCount
+	out := mlx.ScopedAsyncEval(func() []*mlx.Array {
+		ids, hiddens := d.takePending()
+		// The pair at slot S embeds the look-ahead token S+1, so this flush
+		// embeds prompt tokens up to committedDraftOffset+len+1.
+		hidden, auxHidden := spec.draft.Forward(&batch.Batch{
+			InputIDs:     ids,
+			SeqOffsets:   []int32{int32(d.committedDraftOffset)},
+			SeqQueryLens: []int32{int32(n)},
+			Hidden:       hiddens,
+			Media:        d.flushMedia(d.committedDraftOffset + n + 1),
+			Layout:       d.layout,
+		}, spec.targets, spec.draftKV)
+		out := []*mlx.Array{lastHiddenRow(hidden), lastHiddenRow(auxHidden)}
+		for _, c := range spec.draftKV {
+			out = append(out, c.State()...)
+		}
+		return out
+	})
+	d.setHeld(out[0], out[1])
+	d.committedDraftOffset += n
+}
 
-	// Force the draft writes: a session that never drafts would otherwise
-	// leave the flush chain unevaluated, pinning every hidden until close.
-	state := make([]*mlx.Array, 0, 2*len(spec.draftKV))
-	for _, c := range spec.draftKV {
-		state = append(state, c.State()...)
-	}
-	mlx.AsyncEval(state...)
-
-	mlx.Unpin(d.pendingTokens...)
-	mlx.Unpin(d.pendingHiddens...)
-	d.pendingTokens, d.pendingHiddens = nil, nil
-	d.pendingCount = 0
+// takePending returns the buffered pairs as one batch, advancing past them.
+func (d *mtpDraftSession) takePending() (ids, hiddens *mlx.Array) {
+	ids = mlx.Concatenate(d.pendingTokens, 1)
+	hiddens = mlx.Concatenate(d.pendingHiddens, 1)
+	d.pending.Close()
+	d.pending, d.pendingTokens, d.pendingHiddens, d.pendingCount = nil, nil, nil, 0
+	return ids, hiddens
 }
 
 func (d *mtpDraftSession) setFrontierHidden(h *mlx.Array) {
-	mlx.Pin(h)
-	mlx.Unpin(d.frontierHidden)
+	d.scope.Attach(h)
+	d.scope.Discard(d.frontierHidden)
 	d.frontierHidden = h
 }
 
-// setHeld replaces the held flush outputs, pinned until the next flush or close.
+// setHeld replaces the held flush outputs, kept until the next flush or close.
 func (d *mtpDraftSession) setHeld(hidden, auxHidden *mlx.Array) {
-	mlx.Pin(hidden, auxHidden)
-	mlx.Unpin(d.heldHidden, d.heldAuxHidden)
+	d.scope.Attach(hidden, auxHidden)
+	d.scope.Discard(d.heldHidden, d.heldAuxHidden)
 	d.heldHidden, d.heldAuxHidden = hidden, auxHidden
 }
 
@@ -276,55 +290,60 @@ func (d *mtpDraftSession) propose(current *mlx.Array, maxTokens int) *draftCandi
 		}
 	}
 
-	lastToken := current.ExpandDims(-1)
-	lastHidden := d.frontierHidden
-	draftDists := make([]sampler.Distribution, 0, maxTokens)
-	var prefix *mlx.Array
+	var candidates *draftCandidates
+	mlx.ScopedArrays(func() []*mlx.Array {
+		lastToken := current.ExpandDims(-1)
+		lastHidden := d.frontierHidden
+		draftDists := make([]sampler.Distribution, 0, maxTokens)
+		var prefix *mlx.Array
 
-	for i := range maxTokens {
-		var hidden, auxHidden *mlx.Array
-		if i == 0 && len(spec.draftKV) > 0 {
-			// The settle flush already produced the frontier row; reuse it
-			// instead of re-running the head.
-			hidden, auxHidden = d.heldHidden, d.heldAuxHidden
-		} else {
-			// A head with draft caches writes each draft token to the next
-			// draft-cache slot, advancing one per step from the last committed
-			// slot (the held i==0 step stands in for that slot). A cacheless
-			// head stays at the last committed slot every step, re-attending
-			// the committed prefix read-only ("single-position").
-			pos := d.frontier - 1
-			if len(spec.draftKV) > 0 {
-				pos = d.frontier - 1 + i
+		for i := range maxTokens {
+			var hidden, auxHidden *mlx.Array
+			if i == 0 && len(spec.draftKV) > 0 {
+				// The settle flush already produced the frontier row; reuse it
+				// instead of re-running the head.
+				hidden, auxHidden = d.heldHidden, d.heldAuxHidden
+			} else {
+				// A head with draft caches writes each draft token to the next
+				// draft-cache slot, advancing one per step from the last committed
+				// slot (the held i==0 step stands in for that slot). A cacheless
+				// head stays at the last committed slot every step, re-attending
+				// the committed prefix read-only ("single-position").
+				pos := d.frontier - 1
+				if len(spec.draftKV) > 0 {
+					pos = d.frontier - 1 + i
+				}
+				hidden, auxHidden = spec.draft.Forward(&batch.Batch{
+					InputIDs:     lastToken,
+					SeqOffsets:   []int32{int32(pos)},
+					SeqQueryLens: []int32{1},
+					Hidden:       lastHidden,
+					Layout:       d.layout,
+				}, spec.targets, spec.draftKV)
 			}
-			hidden, auxHidden = spec.draft.Forward(&batch.Batch{
-				InputIDs:     lastToken,
-				SeqOffsets:   []int32{int32(pos)},
-				SeqQueryLens: []int32{1},
-				Hidden:       lastHidden,
-				Layout:       d.layout,
-			}, spec.targets, spec.draftKV)
-		}
-		// Unembed only the row being sampled, never the batch.
-		stepLogits := spec.draft.Unembed(hidden).Squeeze(1)
-		lastHidden = auxHidden
-		// The chain's earlier drafts ride along as the row's history, so
-		// penalties shape proposals the same way they shape validation.
-		dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix)
-		nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist)
+			// Unembed only the row being sampled, never the batch.
+			stepLogits := spec.draft.Unembed(hidden).Squeeze(1)
+			lastHidden = auxHidden
+			// The chain's earlier drafts ride along as the row's history, so
+			// penalties shape proposals the same way they shape validation.
+			dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix)
+			nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist)
 
-		lastToken = nextToken.ExpandDims(-1)
-		draftDists = append(draftDists, dist)
-		if prefix == nil {
-			prefix = lastToken
-		} else {
-			prefix = prefix.Concatenate(1, lastToken)
+			lastToken = nextToken.ExpandDims(-1)
+			draftDists = append(draftDists, dist)
+			if prefix == nil {
+				prefix = lastToken
+			} else {
+				prefix = prefix.Concatenate(1, lastToken)
+			}
 		}
-	}
-	return &draftCandidates{
-		tokens: prefix,
-		dist:   sampler.ConcatenateDistributions(draftDists),
-	}
+		candidates = &draftCandidates{
+			tokens: prefix,
+			dist:   sampler.ConcatenateDistributions(draftDists),
+		}
+		return candidates.Arrays()
+	})
+	return candidates
 }
 
 func lastHiddenRow(hidden *mlx.Array) *mlx.Array {

x/mlxrunner/mtp_test.go

@@ -274,8 +274,6 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) {
 
 		spec := testSpeculationSession(r, caches)
 		current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-		unpin := pinAcceptInputs(current, candidates)
-		defer unpin()
 		results, accepted, observed, err := spec.accept(&position, current, candidates, nil)
 		if err != nil {
 			t.Fatalf("accept: %v", err)
@@ -311,8 +309,6 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) {
 
 		spec := testSpeculationSession(r, caches)
 		current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-		unpin := pinAcceptInputs(current, candidates)
-		defer unpin()
 		results, accepted, observed, err := spec.accept(&position, current, candidates, nil)
 		if err != nil {
 			t.Fatalf("accept: %v", err)
@@ -350,8 +346,6 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) {
 
 		spec := testSpeculationSession(r, caches)
 		current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-		unpin := pinAcceptInputs(current, candidates)
-		defer unpin()
 		results, accepted, observed, err := spec.accept(&position, current, candidates, nil)
 		if err != nil {
 			t.Fatalf("accept: %v", err)
@@ -1377,12 +1371,3 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates {
 	d.committed(mlx.FromValues([]int32{0}, 1, 1), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab), 0, nil)
 	return d.propose(mlx.FromValues([]int32{0}, 1), len(tokens))
 }
-
-// pinAcceptInputs pins the arrays accept's caller must keep alive across
-// accept's internal sweep — current and the candidate tokens — as the decoder
-// and next do in the live engine. It returns the matching unpin.
-func pinAcceptInputs(current sampler.Result, candidates *draftCandidates) func() {
-	arrays := append(current.Arrays(), candidates.tokens)
-	mlx.Pin(arrays...)
-	return func() { mlx.Unpin(arrays...) }
-}

x/mlxrunner/pipeline.go

@@ -90,21 +90,19 @@ func (r *Runner) Prepare(request *Request) (err error) {
 // The runner serializes requests today so we just use a fixed slot ID.
 const pipelineSlot = 0
 
-func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) error {
+func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) (err error) {
 	mlx.ResetPeakMemory()
+	mlx.Scoped(func() { err = r.generate(ctx, request) })
+	mlx.ClearCache()
 
-	defer func() {
-		r.Sampler.Remove(pipelineSlot)
-		mlx.Sweep()
-		mlx.ClearCache()
-
-		if slog.Default().Enabled(context.TODO(), logutil.LevelTrace) {
-			mlx.LogArrays()
-			r.cache.dumpTree()
-		}
-		slog.Info("peak memory", "size", mlx.PrettyBytes(mlx.PeakMemory()))
-	}()
+	if slog.Default().Enabled(context.TODO(), logutil.LevelTrace) {
+		r.cache.dumpTree()
+	}
+	slog.Info("memory", "peak", mlx.PrettyBytes(mlx.PeakMemory()), "held", mlx.PrettyBytes(mlx.ActiveMemory()))
+	return err
+}
 
+func (r *Runner) generate(ctx context.Context, request Request) error {
 	inputs := request.Tokens
 
 	session := r.cache.begin(inputs, request.MediaItems)
@@ -126,6 +124,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 
 	// Register the sampler after prefill completes.
 	r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs)
+	defer r.Sampler.Remove(pipelineSlot)
 
 	grammar, err := request.Grammar.resolve(ctx)
 	if err != nil {
@@ -166,7 +165,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
 		snapshotOffsets = append(snapshotOffsets, end)
 	}
 
-	materializeCaches := func() {
+	cacheState := func() []*mlx.Array {
 		state := make([]*mlx.Array, 0, 2*len(caches))
 		for _, c := range caches {
 			if c == nil {
@@ -174,10 +173,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
 			}
 			state = append(state, c.State()...)
 		}
-		if len(state) == 0 {
-			return
-		}
-		mlx.Eval(state...)
+		return state
 	}
 
 	session.schedulePrefillSnapshots(snapshotOffsets)
@@ -185,7 +181,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
 	total, processed := len(tokens), 0
 	position := len(inputs) - len(tokens)
 	// Free restored items' buffers now: on a full cache hit the loop never runs.
-	media.release(position)
+	media.free(position)
 	for total-processed > 1 {
 		if err := ctx.Err(); err != nil {
 			// Settle the drafter with the next prompt token so the caches
@@ -198,26 +194,27 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
 		n := min(prefillChunk, total-processed-1)
 		n = media.extendChunk(position, n)
 
-		chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n)
-		manifest := media.batchMedia(position, n)
-		_, auxHidden := r.Model.Forward(&batch.Batch{
-			InputIDs:     chunkIDs,
-			SeqOffsets:   []int32{int32(position)},
-			SeqQueryLens: []int32{int32(n)},
-			Media:        manifest,
-			Layout:       media.rowLayout(),
-		}, caches)
-		// Report to the drafter only after the chunk's eval: a draft flush
-		// evaluates, and an eval before the sweep cannot free any buffer the
-		// chunk's live handles retain — on media chunks, the whole vision tower.
-		mlx.Pin(chunkIDs, auxHidden)
-		mlx.Sweep()
-		materializeCaches()
-		spec.committed(chunkIDs, auxHidden, position, manifest)
-		mlx.Unpin(chunkIDs, auxHidden)
-		// Released after committed so the drafter can capture rows its
-		// deferred flush still embeds.
-		media.release(position + n)
+		mlx.Scoped(func() {
+			chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n)
+			chunkMedia := media.batchMedia(position, n)
+			auxHidden := mlx.ScopedArrays(func() []*mlx.Array {
+				_, auxHidden := r.Model.Forward(&batch.Batch{
+					InputIDs:     chunkIDs,
+					SeqOffsets:   []int32{int32(position)},
+					SeqQueryLens: []int32{int32(n)},
+					Media:        chunkMedia,
+					Layout:       media.rowLayout(),
+				}, caches)
+				return []*mlx.Array{auxHidden}
+			})[0]
+			mlx.Eval(cacheState()...)
+			// Report to the drafter only after the chunk's eval: a draft
+			// flush evaluates.
+			spec.committed(chunkIDs, auxHidden, position, chunkMedia)
+			// Freed after committed so the drafter can capture rows its
+			// deferred flush still embeds.
+			media.free(position + n)
+		})
 		processed += n
 		position += n
 		slog.Info("Prompt processing progress", "processed", processed, "total", total)
@@ -290,55 +287,63 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess
 			return err
 		}
 
-		results, err := d.next(request.Options.NumPredict - generated)
+		var done bool
+		var err error
+		mlx.Scoped(func() {
+			var results []sampler.Result
+			results, err = d.next(request.Options.NumPredict - generated)
+			if err != nil {
+				return
+			}
+
+			// Record the whole run before streaming any of it: a cancelled
+			// stream returns early and must not leave the caches ahead of
+			// session.outputs.
+			stream := len(results)
+			for i, res := range results {
+				id := res.Token.Int()
+				session.outputs = append(session.outputs, id)
+				if done {
+					continue
+				}
+				if r.Tokenizer.IsEOS(id) {
+					final.DoneReason = 0
+					done = true
+					stream = i
+					continue
+				}
+				generated++
+				if generated >= request.Options.NumPredict {
+					done = true
+					stream = i + 1
+				}
+			}
+
+			for _, res := range results[:stream] {
+				resp, ok := detok.detokenize(res)
+				if !ok {
+					continue
+				}
+				// Two-pass structured output cancels the first pass before its final response.
+				if request.IncludeIntermediateMetrics {
+					resp.PromptEvalCount = len(request.Tokens)
+					resp.PromptEvalCachedCount = final.PromptEvalCachedCount
+					resp.PromptEvalDuration = promptEval
+					resp.EvalCount = generated
+					resp.EvalDuration = time.Since(now)
+				}
+				select {
+				case <-ctx.Done():
+					err = ctx.Err()
+					return
+				case request.Responses <- resp:
+				}
+			}
+		})
 		if err != nil {
 			return err
 		}
 
-		// Record the whole run before streaming any of it: a cancelled
-		// stream returns early and must not leave the caches ahead of
-		// session.outputs.
-		done := false
-		stream := len(results)
-		for i, res := range results {
-			id := res.Token.Int()
-			session.outputs = append(session.outputs, id)
-			if done {
-				continue
-			}
-			if r.Tokenizer.IsEOS(id) {
-				final.DoneReason = 0
-				done = true
-				stream = i
-				continue
-			}
-			generated++
-			if generated >= request.Options.NumPredict {
-				done = true
-				stream = i + 1
-			}
-		}
-
-		for _, res := range results[:stream] {
-			resp, ok := detok.detokenize(res)
-			if !ok {
-				continue
-			}
-			// Two-pass structured output cancels the first pass before its final response.
-			if request.IncludeIntermediateMetrics {
-				resp.PromptEvalCount = len(request.Tokens)
-				resp.PromptEvalCachedCount = final.PromptEvalCachedCount
-				resp.PromptEvalDuration = promptEval
-				resp.EvalCount = generated
-				resp.EvalDuration = time.Since(now)
-			}
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			case request.Responses <- resp:
-			}
-		}
-
 		if done {
 			break
 		}
@@ -375,6 +380,7 @@ type pipelinedDecoder struct {
 	grammars []*grammar // row i's grammar; nil rows are unconstrained
 	position int
 	pending  sampler.Result // in flight: sampled, not yet forwarded
+	scope    *mlx.Scope     // holds pending across steps
 	// Steps run ahead asynchronously: when one faults, its token is already
 	// forwarded and still has to be returned, so err waits for the next call.
 	err error
@@ -383,18 +389,15 @@ type pipelinedDecoder struct {
 func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int, layout []any, g *grammar) *pipelinedDecoder {
 	t := &pipelinedDecoder{
 		r: r, spec: spec, caches: caches, layout: layout, position: position,
-		grammars: []*grammar{g},
+		grammars: []*grammar{g}, scope: mlx.NewScope(),
 	}
 	logits := t.forward(seed)
-	mlx.Pin(logits)
-	defer mlx.Unpin(logits)
 
 	if r.grammarEngine.hasGrammar(t.grammars) {
 		// Dispatch the forward before the host builds the first masks. The
 		// first sample commits nothing, so there is nothing to accept. A mask
 		// fault here is a step fault like any other: the seed is already
 		// forwarded, so the error waits for the first call.
-		mlx.Sweep()
 		mlx.AsyncEval(logits)
 		var errs []error
 		logits, errs = r.grammarEngine.mask(t.grammars, logits, nil)
@@ -421,12 +424,9 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
 	}
 	out := t.pending
 	logits := t.forward(out.Token.ExpandDims(-1))
-	mlx.Pin(logits)
-	defer mlx.Unpin(logits)
 
 	if t.r.grammarEngine.hasGrammar(t.grammars) {
 		// Dispatch the forward before the host's grammar work.
-		mlx.Sweep()
 		mlx.AsyncEval(logits)
 
 		err := t.failRows(t.r.grammarEngine.accept(t.grammars, out.Token.Ints()))
@@ -438,23 +438,25 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
 
 	t.pending = t.sample(logits)
 
-	mlx.Unpin(out.Arrays()...)
+	t.scope.Detach(out.Arrays()...)
 	return []sampler.Result{out}, nil
 }
 
 // forward runs the model one step over token, shaped [B, L], and returns the
 // final position's [B, 1, V] logits, still lazy.
 func (t *pipelinedDecoder) forward(token *mlx.Array) *mlx.Array {
-	hidden, auxHidden := t.r.Model.Forward(&batch.Batch{
-		InputIDs:     token,
-		SeqOffsets:   []int32{int32(t.position)},
-		SeqQueryLens: []int32{int32(token.Dim(1))},
-		Layout:       t.layout,
-	}, t.caches)
-	t.spec.committed(token, auxHidden, t.position, nil)
-	t.position += token.Dim(1)
-	logits := t.r.Model.Unembed(hidden)
-	return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice())
+	return mlx.ScopedArrays(func() []*mlx.Array {
+		hidden, auxHidden := t.r.Model.Forward(&batch.Batch{
+			InputIDs:     token,
+			SeqOffsets:   []int32{int32(t.position)},
+			SeqQueryLens: []int32{int32(token.Dim(1))},
+			Layout:       t.layout,
+		}, t.caches)
+		t.spec.committed(token, auxHidden, t.position, nil)
+		t.position += token.Dim(1)
+		logits := t.r.Model.Unembed(hidden)
+		return []*mlx.Array{logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice())}
+	})[0]
 }
 
 // sample dispatches the batched sample over the decoder's rows. On an
@@ -462,9 +464,8 @@ func (t *pipelinedDecoder) forward(token *mlx.Array) *mlx.Array {
 // is in flight before the previous tokens are synchronized.
 func (t *pipelinedDecoder) sample(logits *mlx.Array) sampler.Result {
 	next := t.r.Sampler.Sample([]int{pipelineSlot}, logits.Squeeze(1))
-	mlx.Pin(next.Arrays()...)
-	mlx.Sweep()
 	mlx.AsyncEval(next.Arrays()...)
+	t.scope.Attach(next.Arrays()...)
 	return next
 }
 
@@ -477,6 +478,9 @@ func (t *pipelinedDecoder) drain() ([]sampler.Result, int, error) {
 		// The sample leaves without its forward, so its accept runs here.
 		err = t.failRows(t.r.grammarEngine.accept(t.grammars, t.pending.Token.Ints()))
 	}
+	if err == nil {
+		t.scope.Detach(t.pending.Arrays()...)
+	}
 	return []sampler.Result{t.pending}, t.position, err
 }
 
@@ -484,7 +488,7 @@ func (t *pipelinedDecoder) close() {
 	// The in-flight sample's forward was never dispatched; its report settles
 	// the drafter level with the caches' resting offset.
 	t.spec.settle(t.pending.Token)
-	mlx.Unpin(t.pending.Arrays()...)
+	t.scope.Close()
 }
 
 // detokenizer serializes sampled tokens into response chunks, holding bytes

x/mlxrunner/prefix_cache.go

@@ -656,16 +656,18 @@ func (c *prefixCache) evictNode(node *trieNode) {
 func (c *prefixCache) dumpTree() {
 	// Summary stats
 	var cacheBytes int
-	for _, kv := range c.caches {
-		if kv == nil {
-			continue
-		}
-		for _, a := range kv.State() {
-			if a != nil {
-				cacheBytes += a.NumBytes()
+	mlx.Scoped(func() {
+		for _, kv := range c.caches {
+			if kv == nil {
+				continue
+			}
+			for _, a := range kv.State() {
+				if a != nil {
+					cacheBytes += a.NumBytes()
+				}
 			}
 		}
-	}
+	})
 
 	// Build active path set for marking.
 	active := make(map[*trieNode]bool, len(c.activePath))

x/mlxrunner/runner.go

@@ -42,6 +42,7 @@ type Request struct {
 
 type Runner struct {
 	Model         base.Model
+	weights       *mlx.Scope
 	Tokenizer     *tokenizer.Tokenizer
 	Requests      chan Request
 	Sampler       *sample.Sampler
@@ -57,89 +58,104 @@ type Runner struct {
 }
 
 func (r *Runner) Load(modelName string) error {
-	root, err := model.Open(modelName)
+	weights, err := r.loadModel(modelName)
 	if err != nil {
 		return err
 	}
-	defer root.Close()
-
-	m, err := base.New(root)
-	if err != nil {
-		return err
-	}
-
-	// Load all tensor blobs from manifest
-	tensors, err := loadTensorsFromManifest(root)
-	if err != nil {
-		return err
-	}
-
-	// On Metal, materialize the loaded tensors with CPU reads before any
-	// weight graph exists, so the weight eval never commits a command buffer
-	// that waits on file data. CUDA loads read at dispatch and need no pre-pass.
-	if mlx.MetalIsAvailable() {
-		mlx.Eval(slices.Collect(maps.Values(tensors))...)
-	}
-
-	// Assign weights to model (model-specific logic). Target and draft weights
-	// must be loaded before sweeping so tensors from a combined manifest are
-	// not discarded before the draft model can retain them.
-	if err := m.LoadWeights(tensors); err != nil {
-		return err
-	}
-
-	var draftModel base.DraftModel
-	draft, err := base.NewDraft(root, m)
-	if err != nil {
-		return err
-	}
-	if draft != nil {
-		if err := draft.LoadWeights(tensors); err != nil {
-			return err
-		}
-		draftModel = draft
-	} else if sd, ok := m.(base.SelfDraft); ok {
-		// Inline draft head: already loaded with the target; nil if none shipped.
-		draftModel = sd.SelfDraft()
-	}
-
-	collected := mlx.Collect(m)
-	if draft != nil {
-		draftArrays := mlx.Collect(draft)
-		collected = append(collected, draftArrays...)
-		if root.Draft != nil {
-			slog.Info("Loaded draft model", "tensor_prefix", root.Draft.TensorPrefix, "config", root.Draft.Config, "arrays", len(draftArrays))
-		} else {
-			slog.Info("Loaded draft model", "arrays", len(draftArrays))
-		}
-	}
-	for _, arr := range collected {
-		mlx.Pin(arr)
-	}
-	mlx.Sweep()
-	mlx.Eval(collected...)
+	mlx.Eval(weights...)
+	r.weights = mlx.NewScope()
+	r.weights.Attach(weights...)
 	configureWiredMemory()
-
-	r.Model = m
-	r.Tokenizer = m.Tokenizer()
-	r.contextLength = m.MaxContextLength()
-	caches := m.NewCaches()
-	draftCaches := newDraftCaches(draftModel)
-	r.cache = newPrefixCache(slices.Concat(caches, draftCaches))
-	r.Sampler = sample.New(r.contextLength)
-	r.spec = newSpeculation(r, draftModel, caches, draftCaches)
-	r.grammarEngine = newGrammarEngine(logitsWidth(m), r.Tokenizer)
-
-	mlx.EnableCompile()
-
 	return nil
 }
 
+func (r *Runner) loadModel(modelName string) (weights []*mlx.Array, err error) {
+	weights = mlx.ScopedArrays(func() []*mlx.Array {
+		root, e := model.Open(modelName)
+		if e != nil {
+			err = e
+			return nil
+		}
+		defer root.Close()
+
+		m, e := base.New(root)
+		if e != nil {
+			err = e
+			return nil
+		}
+
+		// Load all tensor blobs from manifest
+		tensors, e := loadTensorsFromManifest(root)
+		if e != nil {
+			err = e
+			return nil
+		}
+
+		// On Metal, materialize the loaded tensors with CPU reads before any
+		// weight graph exists, so the weight eval never commits a command buffer
+		// that waits on file data. CUDA loads read at dispatch and need no pre-pass.
+		if mlx.MetalIsAvailable() {
+			mlx.Eval(slices.Collect(maps.Values(tensors))...)
+		}
+
+		// Assign weights to model (model-specific logic). Target and draft weights
+		// must be loaded before the load scope ends so tensors from a combined
+		// manifest are not discarded before the draft model can retain them.
+		if err = m.LoadWeights(tensors); err != nil {
+			return nil
+		}
+
+		var draftModel base.DraftModel
+		draft, e := base.NewDraft(root, m)
+		if e != nil {
+			err = e
+			return nil
+		}
+		if draft != nil {
+			if err = draft.LoadWeights(tensors); err != nil {
+				return nil
+			}
+			draftModel = draft
+		} else if sd, ok := m.(base.SelfDraft); ok {
+			// Inline draft head: already loaded with the target; nil if none shipped.
+			draftModel = sd.SelfDraft()
+		}
+
+		w := mlx.Collect(m)
+		if draft != nil {
+			draftArrays := mlx.Collect(draft)
+			w = append(w, draftArrays...)
+			if root.Draft != nil {
+				slog.Info("Loaded draft model", "tensor_prefix", root.Draft.TensorPrefix, "config", root.Draft.Config, "arrays", len(draftArrays))
+			} else {
+				slog.Info("Loaded draft model", "arrays", len(draftArrays))
+			}
+		}
+
+		r.Model = m
+		r.Tokenizer = m.Tokenizer()
+		r.contextLength = m.MaxContextLength()
+		caches := m.NewCaches()
+		draftCaches := newDraftCaches(draftModel)
+		r.cache = newPrefixCache(slices.Concat(caches, draftCaches))
+		r.Sampler = sample.New(r.contextLength)
+		r.spec = newSpeculation(r, draftModel, caches, draftCaches)
+		r.grammarEngine = newGrammarEngine(logitsWidth(m), r.Tokenizer)
+
+		mlx.EnableCompile()
+
+		return w
+	})
+	return weights, err
+}
+
 func (r *Runner) Close() {
 	if r.grammarEngine != nil {
 		r.grammarEngine.close()
 		r.grammarEngine = nil
 	}
+	r.weights.Close()
+	r.weights = nil
 }
 
 // newDraftCaches returns nil when the model ships no draft.
@@ -152,24 +168,23 @@ func newDraftCaches(draft base.DraftModel) []cache.Cache {
 
 // logitsWidth reads a model's logits width off a one-token forward's static
 // shape — the same Forward and Unembed path decode logits take. Nothing is
-// evaluated, and the probe's caches and graph are released before returning,
-// which sweeps every unpinned array: call this only at load, after the
-// model's weights are pinned.
-func logitsWidth(m base.Model) int {
-	caches := m.NewCaches()
-	hidden, _ := m.Forward(&batch.Batch{
-		InputIDs:     mlx.FromValues([]int32{0}, 1, 1),
-		SeqOffsets:   []int32{0},
-		SeqQueryLens: []int32{1},
-	}, caches)
-	logits := m.Unembed(hidden)
-	width := logits.Dim(logits.NumDims() - 1)
-	for _, c := range caches {
-		if c != nil {
-			c.Free()
+// evaluated.
+func logitsWidth(m base.Model) (width int) {
+	mlx.Scoped(func() {
+		caches := m.NewCaches()
+		hidden, _ := m.Forward(&batch.Batch{
+			InputIDs:     mlx.FromValues([]int32{0}, 1, 1),
+			SeqOffsets:   []int32{0},
+			SeqQueryLens: []int32{1},
+		}, caches)
+		logits := m.Unembed(hidden)
+		width = logits.Dim(logits.NumDims() - 1)
+		for _, c := range caches {
+			if c != nil {
+				c.Free()
+			}
 		}
-	}
-	mlx.Sweep()
+	})
 	return width
 }
 

x/mlxrunner/sample/logprob_test.go

@@ -26,16 +26,15 @@ func runSampleLogprobs(t *mlxtest.T, logits []float32, topK int) (int32, float64
 	s := New(128)
 	t.Cleanup(func() {
 		s.Free()
-		mlx.Sweep()
 	})
 	s.Add(0, Options{Logprobs: true, TopLogprobs: topK}, nil)
 
 	tensor := mlx.FromValues(logits, 1, len(logits))
-	res := s.Sample([]int{0}, tensor)
-
-	mlx.Pin(res.Arrays()...)
-	t.Cleanup(func() { mlx.Unpin(res.Arrays()...) })
-	mlx.Sweep()
+	var res Result
+	mlx.ScopedArrays(func() []*mlx.Array {
+		res = s.Sample([]int{0}, tensor)
+		return res.Arrays()
+	})
 	mlx.Eval(res.Arrays()...)
 
 	selected := res.Token.Int()
@@ -249,15 +248,12 @@ func TestBatchedLogprobsPerRow(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 		s.Add(1, Options{Logprobs: true}, nil)
 		s.Add(2, Options{Logprobs: true}, nil)
 
 		logits := mlx.FromValues(append(append([]float32{}, rowA...), rowB...), 2, 3)
 		res := s.Sample([]int{1, 2}, logits)
-		mlx.Pin(res.Arrays()...)
-		t.Cleanup(func() { mlx.Unpin(res.Arrays()...) })
 		mlx.Eval(res.Arrays()...)
 
 		got := res.Logprob.Floats()

x/mlxrunner/sample/sample.go

@@ -38,8 +38,8 @@ type Result struct {
 }
 
 // Arrays returns the tensor fields as a slice so callers can drive the mlx
-// lifecycle verbs (Pin, Unpin, Eval, AsyncEval) over the whole group. Unset
-// fields stay nil; the mlx helpers skip them.
+// lifecycle verbs (Eval, AsyncEval, a held scope's Attach) over the whole
+// group. Unset fields stay nil; the mlx helpers skip them.
 func (r Result) Arrays() []*mlx.Array {
 	return []*mlx.Array{r.Token, r.Logprob, r.TopTokens, r.TopLogprobs}
 }
@@ -174,6 +174,7 @@ type Sampler struct {
 	// belongs to slots[i]; W is max(RepeatLastN) across penalty slots.
 	// Allocated on the first penalty slot, rebuilt only in Add/Remove.
 	history *mlx.Array
+	scope   *mlx.Scope
 
 	// allSameOpts: every registered slot shares Options. When true the
 	// canonical shared value is s.slots[0].opts.
@@ -209,6 +210,7 @@ func New(numCtx int) *Sampler {
 		byID:        make(map[int]*slotState),
 		allSameOpts: true,
 		numCtx:      numCtx,
+		scope:       mlx.NewScope(),
 	}
 }
 
@@ -270,29 +272,32 @@ func (s *Sampler) Add(seqID int, opts Options, priorTokens []int32) {
 	// Grow the pool to hold this slot's row. The pool is lazy — the first
 	// penalty slot allocates it — and thereafter every registered slot
 	// gets a row (rows for non-penalty slots are zero and never read).
-	// Invariant: s.history is pinned whenever non-nil.
 	if s.history != nil || opts.usesHistory() {
-		targetWidth := max(opts.RepeatLastN, s.historyWidth())
-		newRow := makeHistoryRow(priorTokens, opts.RepeatLastN, targetWidth)
+		pool := mlx.ScopedArrays(func() []*mlx.Array {
+			targetWidth := max(opts.RepeatLastN, s.historyWidth())
+			newRow := makeHistoryRow(priorTokens, opts.RepeatLastN, targetWidth)
 
-		var pool *mlx.Array
-		switch {
-		case s.history == nil && len(s.slots) == 0:
-			pool = newRow
-		case s.history == nil:
-			// First penalty slot with non-penalty slots already registered;
-			// seed zero rows so s.slots and pool row indices stay aligned.
-			zeros := mlx.Zeros(mlx.DTypeInt32, len(s.slots), targetWidth)
-			pool = zeros.Concatenate(0, newRow)
-		case targetWidth > s.historyWidth():
-			pad := mlx.Zeros(mlx.DTypeInt32, s.history.Dim(0), targetWidth-s.historyWidth())
-			pool = s.history.Concatenate(1, pad).Concatenate(0, newRow)
-		default:
-			pool = s.history.Concatenate(0, newRow)
-		}
+			var pool *mlx.Array
+			switch {
+			case s.history == nil && len(s.slots) == 0:
+				pool = newRow
+			case s.history == nil:
+				// First penalty slot with non-penalty slots already registered;
+				// seed zero rows so s.slots and pool row indices stay aligned.
+				zeros := mlx.Zeros(mlx.DTypeInt32, len(s.slots), targetWidth)
+				pool = zeros.Concatenate(0, newRow)
+			case targetWidth > s.historyWidth():
+				pad := mlx.Zeros(mlx.DTypeInt32, s.history.Dim(0), targetWidth-s.historyWidth())
+				pool = s.history.Concatenate(1, pad).Concatenate(0, newRow)
+			default:
+				pool = s.history.Concatenate(0, newRow)
+			}
 
-		mlx.Pin(pool)
-		mlx.Unpin(s.history)
+			// The concatenation still reads the old pool.
+			s.scope.Discard(s.history)
+			return []*mlx.Array{pool}
+		})[0]
+		s.scope.Attach(pool)
 		s.history = pool
 
 		if opts.usesHistory() {
@@ -363,34 +368,38 @@ func (s *Sampler) Remove(seqID int) {
 		return
 	}
 
-	n := s.history.Dim(0)
-	var newHistory *mlx.Array
-	switch {
-	case n == 1:
-		newHistory = nil
-	case row == 0:
-		newHistory = s.history.Slice(mlx.Slice(1, n), mlx.Slice())
-	case row == n-1:
-		newHistory = s.history.Slice(mlx.Slice(0, row), mlx.Slice())
-	default:
-		before := s.history.Slice(mlx.Slice(0, row), mlx.Slice())
-		after := s.history.Slice(mlx.Slice(row+1, n), mlx.Slice())
-		newHistory = before.Concatenate(0, after)
-	}
+	newHistory := mlx.ScopedArrays(func() []*mlx.Array {
+		n := s.history.Dim(0)
+		var newHistory *mlx.Array
+		switch {
+		case n == 1:
+			newHistory = nil
+		case row == 0:
+			newHistory = s.history.Slice(mlx.Slice(1, n), mlx.Slice())
+		case row == n-1:
+			newHistory = s.history.Slice(mlx.Slice(0, row), mlx.Slice())
+		default:
+			before := s.history.Slice(mlx.Slice(0, row), mlx.Slice())
+			after := s.history.Slice(mlx.Slice(row+1, n), mlx.Slice())
+			newHistory = before.Concatenate(0, after)
+		}
 
-	mlx.Pin(newHistory)
-	mlx.Unpin(s.history)
+		s.scope.Discard(s.history)
+		return []*mlx.Array{newHistory}
+	})[0]
+	s.scope.Attach(newHistory)
 	s.history = newHistory
 }
 
 // Free releases the pooled history tensor and resets the sampler to the
 // New-equivalent state so it may be reused.
 func (s *Sampler) Free() {
-	mlx.Unpin(s.history)
+	s.scope.Close()
 	*s = Sampler{
 		byID:        make(map[int]*slotState),
 		allSameOpts: true,
 		numCtx:      s.numCtx,
+		scope:       mlx.NewScope(),
 	}
 }
 
@@ -411,34 +420,38 @@ func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result {
 		slots[i] = slot
 	}
 
-	var token *mlx.Array
-	if opts0, ok := s.canBatch(slots); ok {
-		token = s.sampleTokensUniform(slots, opts0, logits)
-	} else {
-		token = s.sampleTokensSerial(slots, logits)
-	}
-
-	res := Result{Token: token}
-	if s.anyLogprobs {
-		// Log-softmax over original logits so every row holds a truthful
-		// value (compute-for-all; consumers filter per-slot). Subtract
-		// max first for numerical stability in the logsumexp.
-		lp := logits.AsType(mlx.DTypeFloat32)
-		lp = lp.Subtract(lp.MaxAxis(-1, true))
-		lp = lp.Subtract(lp.LogsumexpAxis(-1, true))
-		res.Logprob = lp.TakeAlongAxis(token.ExpandDims(-1), -1)
-		if s.maxTopLogprobs > 0 {
-			k := s.maxTopLogprobs
-			if vocab := lp.Dim(lp.NumDims() - 1); k > vocab {
-				k = vocab
-			}
-			// Argpartition on the negated values places the K largest
-			// (unsorted) in positions [0:K].
-			idx := lp.Negative().ArgpartitionAxis(k-1, -1).Slice(mlx.Slice(), mlx.Slice(0, k))
-			res.TopTokens = idx.AsType(mlx.DTypeInt32)
-			res.TopLogprobs = lp.TakeAlongAxis(idx, -1)
+	var res Result
+	mlx.ScopedArrays(func() []*mlx.Array {
+		var token *mlx.Array
+		if opts0, ok := s.canBatch(slots); ok {
+			token = s.sampleTokensUniform(slots, opts0, logits)
+		} else {
+			token = s.sampleTokensSerial(slots, logits)
 		}
-	}
+
+		res = Result{Token: token}
+		if s.anyLogprobs {
+			// Log-softmax over original logits so every row holds a truthful
+			// value (compute-for-all; consumers filter per-slot). Subtract
+			// max first for numerical stability in the logsumexp.
+			lp := logits.AsType(mlx.DTypeFloat32)
+			lp = lp.Subtract(lp.MaxAxis(-1, true))
+			lp = lp.Subtract(lp.LogsumexpAxis(-1, true))
+			res.Logprob = lp.TakeAlongAxis(token.ExpandDims(-1), -1)
+			if s.maxTopLogprobs > 0 {
+				k := s.maxTopLogprobs
+				if vocab := lp.Dim(lp.NumDims() - 1); k > vocab {
+					k = vocab
+				}
+				// Argpartition on the negated values places the K largest
+				// (unsorted) in positions [0:K].
+				idx := lp.Negative().ArgpartitionAxis(k-1, -1).Slice(mlx.Slice(), mlx.Slice(0, k))
+				res.TopTokens = idx.AsType(mlx.DTypeInt32)
+				res.TopLogprobs = lp.TakeAlongAxis(idx, -1)
+			}
+		}
+		return res.Arrays()
+	})
 	return res
 }
 

x/mlxrunner/sample/sample_test.go

@@ -36,7 +36,6 @@ func sampleOne(t *mlxtest.T, opts Options, priorTokens []int32, values []float32
 	s := New(128)
 	t.Cleanup(func() {
 		s.Free()
-		mlx.Sweep()
 	})
 	s.Add(0, opts, priorTokens)
 
@@ -139,7 +138,6 @@ func TestDistributionAppliesTopKBeforeTopP(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 		s.Add(0, Options{Temperature: 1, TopK: 2, TopP: 0.7}, nil)
 
@@ -210,7 +208,6 @@ func TestSeededSamplingIsReproducible(t *testing.T) {
 			s := New(128)
 			t.Cleanup(func() {
 				s.Free()
-				mlx.Sweep()
 			})
 			s.Add(0, Options{Temperature: 1, TopK: 4, Seed: seed, UseSeed: true}, nil)
 
@@ -243,7 +240,6 @@ func TestSeededBernoulliIsReproducible(t *testing.T) {
 			s := New(128)
 			t.Cleanup(func() {
 				s.Free()
-				mlx.Sweep()
 			})
 			s.Add(0, Options{Seed: 99, UseSeed: true}, nil)
 
@@ -269,7 +265,6 @@ func TestSampleHistoryWindow(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 
 		// RepeatLastN=2 with priors {1, 2, 3}: makeHistoryRow keeps only
@@ -300,7 +295,6 @@ func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 
 		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{1, 2})
@@ -333,7 +327,6 @@ func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 
 		// A proposal step passes one logits row with the chain's earlier drafts:
@@ -361,7 +354,6 @@ func TestDistributionMultiRowWithoutChain(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 
 		// A block drafter's proposal batch samples every row from one call with
@@ -394,7 +386,6 @@ func TestCommitBatchesRingWrites(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 
 		s.Add(0, Options{RepeatLastN: 4, RepeatPenalty: 1.1}, []int32{10, 11, 12})
@@ -492,7 +483,6 @@ func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) {
 			s := New(128)
 			t.Cleanup(func() {
 				s.Free()
-				mlx.Sweep()
 			})
 			for _, spec := range tc.slots {
 				s.Add(spec.id, spec.opts, spec.priors)
@@ -519,7 +509,6 @@ func TestRemoveDoesNotLeakHistory(t *testing.T) {
 		s := New(128)
 		t.Cleanup(func() {
 			s.Free()
-			mlx.Sweep()
 		})
 		s.Add(1, opts, []int32{1})
 		s.Add(2, opts, []int32{2})

x/mlxrunner/server.go

@@ -53,10 +53,7 @@ func Execute(args []string) error {
 	if err != nil {
 		return err
 	}
-	defer worker.Stop(context.Background(), func() {
-		mlx.Sweep()
-		mlx.ClearCache()
-	})
+	defer worker.Stop(context.Background(), mlx.ClearCache)
 	runnerCtx, cancelRunner := context.WithCancel(context.Background())
 	defer cancelRunner()
 

x/mlxrunner/speculate.go

@@ -24,7 +24,7 @@ type draftSession interface {
 	// hidden state at that slot. Runs arrive in slot order — prefill
 	// chunks, the decode seed, then each round's validated tokens. media is
 	// the run's manifest (feature-bearing for items the run overlaps), valid
-	// only for the call; a session that defers its forward pins what it
+	// only for the call; a session that defers its forward holds what it
 	// keeps. Nil outside prefill.
 	committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem)
 
@@ -201,6 +201,7 @@ type speculativeDecoder struct {
 	current  sampler.Result    // emitted (or the seed), not yet forwarded
 	inner    *pipelinedDecoder // pipelines plain tokens while parked; nil while drafting
 	grammar  *grammar
+	scope    *mlx.Scope // holds current across rounds
 }
 
 // decoder returns the decoder for this engine's session. A speculationSession that
@@ -208,8 +209,9 @@ type speculativeDecoder struct {
 // running the inner pipelined decoder whose reports keep the draft KV level.
 func (s *speculationSession) decoder(seed *mlx.Array, position int, grammar *grammar) decoder {
 	current := sampler.Result{Token: seed}
-	mlx.Pin(current.Arrays()...)
-	return &speculativeDecoder{s: s, position: position, current: current, grammar: grammar}
+	scope := mlx.NewScope()
+	scope.Attach(current.Arrays()...)
+	return &speculativeDecoder{s: s, position: position, current: current, grammar: grammar, scope: scope}
 }
 
 func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) {
@@ -231,16 +233,13 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) {
 			// land that extra token within it rather than overshooting. At
 			// remaining 1 the cap is 0 and the last token decodes plainly.
 			candidates = s.drafter.propose(st.current.Token, min(s.limit, remaining-1))
+			mlx.AsyncEval(candidates.Arrays()...)
 		}
 		var accepted, observed int
 		var err error
 		if candidates == nil {
 			results, err = st.park(remaining)
 		} else {
-			// candidates stays pinned across accept's internal sweep and the
-			// draft-count read below; accept pins only its own intermediates.
-			mlx.Pin(candidates.tokens)
-			defer mlx.Unpin(candidates.tokens)
 			results, accepted, observed, err = st.s.accept(&st.position, st.current, candidates, st.grammar)
 		}
 		if err != nil {
@@ -257,11 +256,11 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) {
 	return results, nil
 }
 
-// advance retires the last returned token as the next call's current, pinned
-// across the sweeps the next call runs before reading it. Nothing is forced here.
+// advance retires the last returned token as the next call's current, held
+// until the next call reads it. Nothing is forced here.
 func (st *speculativeDecoder) advance(next sampler.Result) {
-	mlx.Pin(next.Arrays()...)
-	mlx.Unpin(st.current.Arrays()...)
+	st.scope.Attach(next.Arrays()...)
+	st.scope.Discard(st.current.Arrays()...)
 	st.current = next
 }
 
@@ -314,7 +313,7 @@ func (st *speculativeDecoder) close() {
 		// the drafter level with the caches' resting offset.
 		st.s.settle(st.current.Token)
 	}
-	mlx.Unpin(st.current.Arrays()...)
+	st.scope.Close()
 	st.s.logStats()
 }
 
@@ -388,10 +387,6 @@ func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) {
 // acceptance model learns from, capped at the EOS (a terminator, not a target
 // rejection). NumPredict is the decode loop's to enforce, so a token past the
 // budget is left for decode to drop, not cut here.
-//
-// The caller keeps current and the candidate tokens pinned across the call,
-// since accept sweeps before its eval and reads both afterward; accept pins
-// only the intermediates it produces.
 func (s *speculationSession) accept(position *int, current sampler.Result, candidates *draftCandidates, g *grammar) (results []sampler.Result, accepted, observed int, err error) {
 	r := s.spec.r
 	before := *position
@@ -412,55 +407,52 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	}
 	defer commit(0)
 
-	dist := candidates.dist.Arrays()
-	mlx.Pin(dist...)
-	mlx.Sweep()
-	mlx.AsyncEval(candidates.tokens)
-	mlx.Unpin(dist...)
+	var auxHiddenSeq, acceptedMask, residualTokens, bonusToken *mlx.Array
+	var draftIDs []int32
+	var constrained bool
+	var maskErr error
+	mlx.ScopedEval(func() []*mlx.Array {
+		var hiddenSeq *mlx.Array
+		hiddenSeq, auxHiddenSeq = r.Model.Forward(&batch.Batch{
+			InputIDs:     current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens),
+			SeqOffsets:   []int32{int32(before)},
+			SeqQueryLens: []int32{int32(draftCount + 1)},
+			Layout:       s.layout,
+		}, s.spec.targets)
 
-	hiddenSeq, auxHiddenSeq := r.Model.Forward(&batch.Batch{
-		InputIDs:     current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens),
-		SeqOffsets:   []int32{int32(before)},
-		SeqQueryLens: []int32{int32(draftCount + 1)},
-		Layout:       s.layout,
-	}, s.spec.targets)
+		// Row i of the fused hidden is the state after the token at before+i, so
+		// the rows already line up with the drafts: row 0 (current's state)
+		// predicts draft 0, and the row after the last accepted draft is the
+		// bonus row. No separate base-logits forward exists on this path.
+		logits := r.Model.Unembed(hiddenSeq)
 
-	// Row i of the fused hidden is the state after the token at before+i, so
-	// the rows already line up with the drafts: row 0 (current's state)
-	// predicts draft 0, and the row after the last accepted draft is the
-	// bonus row. No separate base-logits forward exists on this path.
-	logits := r.Model.Unembed(hiddenSeq)
-
-	draftIDs := candidates.tokens.Ints()
-	constrained := g.constraining()
-	if constrained {
-		var errs []error
-		logits, errs = r.grammarEngine.mask([]*grammar{g}, logits, [][]int32{draftIDs})
-		if err := errors.Join(errs...); err != nil {
-			return nil, 0, 0, err
+		draftIDs = candidates.tokens.Ints()
+		constrained = g.constraining()
+		if constrained {
+			var errs []error
+			logits, errs = r.grammarEngine.mask([]*grammar{g}, logits, [][]int32{draftIDs})
+			if maskErr = errors.Join(errs...); maskErr != nil {
+				return nil
+			}
 		}
+
+		targetDist := r.Sampler.Distribution(pipelineSlot, logits, candidates.tokens)
+		draftDist := candidates.dist
+		acceptedMask = r.sampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens)
+
+		// The next token is sampled for every possible outcome before anything
+		// is evaluated — the residual at each rejection point in one batched
+		// draw, plus the bonus row — so a single Eval covers acceptance and the
+		// next token instead of a second host round trip after the rejection
+		// point is known.
+		residualTokens = r.Sampler.SampleDistribution(pipelineSlot, targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist))
+		bonusToken = r.sampleTokenAt(targetDist, draftCount)
+		return []*mlx.Array{auxHiddenSeq, acceptedMask, residualTokens, bonusToken}
+	})
+	if maskErr != nil {
+		return nil, 0, 0, maskErr
 	}
 
-	targetDist := r.Sampler.Distribution(pipelineSlot, logits, candidates.tokens)
-	draftDist := candidates.dist
-	acceptedMask := r.sampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens)
-
-	// The next token is sampled for every possible outcome before anything
-	// is evaluated — the residual at each rejection point in one batched
-	// draw, plus the bonus row — so a single Eval covers acceptance and the
-	// next token instead of a second host round trip after the rejection
-	// point is known.
-	residualTokens := r.Sampler.SampleDistribution(pipelineSlot, targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist))
-	bonusToken := r.sampleTokenAt(targetDist, draftCount)
-
-	// Pin the arrays read after the eval; current and the candidate tokens
-	// stay pinned by the caller across the call.
-	live := []*mlx.Array{hiddenSeq, auxHiddenSeq, acceptedMask, residualTokens, bonusToken}
-	mlx.Pin(live...)
-	defer mlx.Unpin(live...)
-	mlx.Sweep()
-	mlx.Eval(candidates.tokens, acceptedMask, residualTokens, bonusToken)
-
 	acceptedFlags := acceptedMask.Ints()
 	for _, ok := range acceptedFlags {
 		if ok == 0 {

x/models/qwen3_5/vision.go

@@ -97,7 +97,7 @@ type visionLayout struct {
 // The publisher uses the same tensor layout, image preprocessing, and MRoPE
 // layout for both families.
 type VisionAdapter struct {
-	// Model is exported so mlx.Collect traverses and pins every tower weight.
+	// Model is exported so mlx.Collect reaches every tower weight.
 	// An unexported wrapper field is invisible to the reflection collector.
 	Model *Model
 }

x/models/qwen4_exp/engram_cache.go

@@ -14,6 +14,7 @@ import (
 type engramCache struct {
 	history     *mlx.Array
 	convHistory *mlx.Array
+	scope       *mlx.Scope
 	offset      int
 	eosID       int64
 	width       int
@@ -27,24 +28,25 @@ type engramCache struct {
 type engramSnapshot struct {
 	history     *mlx.Array
 	convHistory *mlx.Array
+	scope       *mlx.Scope
 	offset      int
 }
 
 func newEngramCache(width, convTail, convDim int, eosID int64) *engramCache {
-	return &engramCache{width: width, convTail: convTail, convDim: convDim, eosID: eosID}
+	return &engramCache{width: width, convTail: convTail, convDim: convDim, eosID: eosID, scope: mlx.NewScope()}
 }
 
 func (c *engramCache) setHistory(value *mlx.Array) {
 	value = value.Clone()
-	mlx.Pin(value)
-	mlx.Unpin(c.history)
+	c.scope.Attach(value)
+	c.scope.Discard(c.history)
 	c.history = value
 }
 
 func (c *engramCache) setConvHistory(value *mlx.Array) {
 	value = value.Clone()
-	mlx.Pin(value)
-	mlx.Unpin(c.convHistory)
+	c.scope.Attach(value)
+	c.scope.Discard(c.convHistory)
 	c.convHistory = value
 }
 
@@ -126,7 +128,7 @@ func (c *engramCache) State() []*mlx.Array {
 }
 
 func (c *engramCache) Free() {
-	mlx.Unpin(c.history, c.convHistory)
+	c.scope.Close()
 	c.history = nil
 	c.convHistory = nil
 	c.offset = 0
@@ -176,8 +178,10 @@ func (c *engramCache) Restore(snapshot cache.Snapshot, target int) bool {
 	if !ok || value.offset != target {
 		return false
 	}
-	c.setHistory(value.history)
-	c.setConvHistory(value.convHistory)
+	mlx.Scoped(func() {
+		c.setHistory(value.history)
+		c.setConvHistory(value.convHistory)
+	})
 	c.offset = target
 	return true
 }
@@ -194,12 +198,12 @@ func (c *engramCache) Split(snapshot cache.Snapshot, _ int) (cache.Snapshot, cac
 }
 
 func newEngramSnapshot(history, convHistory *mlx.Array, offset int) *engramSnapshot {
-	snapshot := &engramSnapshot{history: history.Clone(), convHistory: convHistory.Clone(), offset: offset}
-	mlx.Pin(snapshot.history, snapshot.convHistory)
+	snapshot := &engramSnapshot{history: history.Clone(), convHistory: convHistory.Clone(), scope: mlx.NewScope(), offset: offset}
+	snapshot.scope.Attach(snapshot.history, snapshot.convHistory)
 	mlx.AsyncEval(snapshot.history, snapshot.convHistory)
 	return snapshot
 }
 
 func (s *engramSnapshot) Size() int                    { return s.history.NumBytes() + s.convHistory.NumBytes() }
 func (s *engramSnapshot) SetMaterializeHook(func(int)) {}
-func (s *engramSnapshot) Close()                       { mlx.Unpin(s.history, s.convHistory) }
+func (s *engramSnapshot) Close()                       { s.scope.Close() }