ci: wire up MLX unit tests for PR runs (#17022)

* ci: wire up MLX unit tests for PR runs

Download the latest Darwin release payload matching the current MLX and MLX-C revisions so macOS PR tests can exercise MLX without rebuilding it. If no matching release exists after a pin bump, leave MLX tests skipped until the next release.

Add whole-tree race coverage and smoke-run committed benchmarks. Verify generated UI types, and stabilize tests exposed by the broader CI coverage.

* review comments

* mlx: run tests on one pinned worker

Keep MLX tests and benchmarks on a shared pinned thread while preserving Fatal, Skip, and Cleanup semantics. Also clean stale CI payloads and ensure updater workers shut down cleanly.

* addres comments

.github/scripts/prepare_mlx_darwin.sh

@@ -0,0 +1,159 @@
+#!/usr/bin/env bash
+
+# Prepare prebuilt MLX Metal runtime libraries for macOS CI unit tests.
+#
+# Building MLX is expensive, so to enable the MLX-specific unit tests this
+# helper finds the newest Ollama release whose MLX_VERSION and MLX_C_VERSION
+# match the current checkout, downloads that release's ollama-darwin.tgz, and
+# extracts only mlx_metal_v* into build/lib/ollama
+#
+# If no matching release artifact exists, the helper only emits a warning. That
+# covers the expected window after an MLX update lands and before the next release
+# publishes a matching Darwin payload.
+
+set -euo pipefail
+
+repo="${OLLAMA_MLX_RELEASE_REPO:-ollama/ollama}"
+scan_limit="${OLLAMA_MLX_RELEASE_SCAN_LIMIT:-50}"
+cache_dir="${OLLAMA_MLX_DARWIN_CACHE:-.cache/mlx-darwin-release}"
+target_dir="${OLLAMA_MLX_DARWIN_TARGET:-build/lib/ollama}"
+tarball="${cache_dir}/ollama-darwin.tgz"
+tag_file="${cache_dir}/matched-tag"
+pins_file="${cache_dir}/matched-pins"
+target_pins_file="${target_dir}/.mlx-release-pins"
+tmpdir=""
+tmp_tarball=""
+
+cleanup() {
+  [ -z "${tmpdir}" ] || rm -rf "${tmpdir}"
+  [ -z "${tmp_tarball}" ] || rm -f "${tmp_tarball}"
+}
+trap cleanup EXIT
+
+warn() {
+  if [ -n "${GITHUB_ACTIONS:-}" ]; then
+    echo "::warning::$*"
+  else
+    echo "warning: $*" >&2
+  fi
+}
+
+read_pin() {
+  tr -d '[:space:]' <"$1"
+}
+
+has_payload() {
+  local variant
+  for variant in "${target_dir}"/mlx_metal_v*; do
+    [ -d "${variant}" ] || continue
+    [ -f "${variant}/libmlx.dylib" ] && [ -f "${variant}/libmlxc.dylib" ] && return 0
+  done
+  return 1
+}
+
+has_matching_payload() {
+  [ -f "${target_pins_file}" ] || return 1
+  [ "$(cat "${target_pins_file}")" = "${current_pins}" ] || return 1
+  has_payload
+}
+
+extract_payload() {
+  local tag="$1"
+  tmpdir="$(mktemp -d)"
+
+  tar -xzf "${tarball}" -C "${tmpdir}"
+  mkdir -p "${target_dir}"
+
+  rm -rf "${target_dir}"/mlx_metal_v*
+
+  local found=false
+  local src dest
+  for src in "${tmpdir}"/mlx_metal_v*; do
+    [ -d "${src}" ] || continue
+    found=true
+    dest="${target_dir}/$(basename "${src}")"
+    rm -rf "${dest}"
+    cp -R "${src}" "${dest}"
+  done
+
+  if [ "${found}" != true ] || ! has_payload; then
+    echo "Downloaded ${tarball} did not contain a usable MLX Metal payload" >&2
+    exit 1
+  fi
+
+  echo "${current_pins}" >"${target_pins_file}"
+  echo "Prepared MLX Darwin payload from ${repo} ${tag}:"
+  find "${target_dir}" -maxdepth 2 -type f \( -name 'libmlx.dylib' -o -name 'libmlxc.dylib' -o -name '*.metallib' \) -print
+
+  rm -rf "${tmpdir}"
+  tmpdir=""
+}
+
+if [ "$(uname -s)" != "Darwin" ]; then
+  warn "MLX Darwin payload setup is only supported on macOS"
+  exit 0
+fi
+
+current_mlx="$(read_pin MLX_VERSION)"
+current_mlxc="$(read_pin MLX_C_VERSION)"
+current_pins="${current_mlx} ${current_mlxc}"
+
+if has_matching_payload; then
+  echo "MLX Darwin payload already present in ${target_dir}"
+  exit 0
+fi
+
+mkdir -p "${cache_dir}"
+
+if [ -s "${tarball}" ] && [ -f "${tag_file}" ] && [ "$(cat "${pins_file}" 2>/dev/null || true)" = "${current_pins}" ]; then
+  extract_payload "$(cat "${tag_file}")"
+  exit 0
+fi
+
+matched_tag=""
+matched_url=""
+
+while read -r tag; do
+  [ -n "${tag}" ] || continue
+
+  if ! tag_mlx="$(curl -fsSL "https://raw.githubusercontent.com/${repo}/${tag}/MLX_VERSION" | tr -d '[:space:]')"; then
+    continue
+  fi
+  if [ "${tag_mlx}" != "${current_mlx}" ]; then
+    continue
+  fi
+
+  if ! tag_mlxc="$(curl -fsSL "https://raw.githubusercontent.com/${repo}/${tag}/MLX_C_VERSION" | tr -d '[:space:]')"; then
+    continue
+  fi
+  if [ "${tag_mlxc}" != "${current_mlxc}" ]; then
+    continue
+  fi
+
+  url="https://github.com/${repo}/releases/download/${tag}/ollama-darwin.tgz"
+  if curl -fsIL "${url}" >/dev/null; then
+    matched_tag="${tag}"
+    matched_url="${url}"
+    break
+  fi
+
+  echo "MLX pins match ${tag}, but ${url} is not available"
+done < <(
+  git ls-remote --tags --refs --sort=-version:refname "https://github.com/${repo}.git" 'v*' |
+    awk -v limit="${scan_limit}" '{ sub("refs/tags/", "", $2); print $2; if (limit > 0 && NR >= limit) exit }'
+)
+
+if [ -z "${matched_tag}" ]; then
+  warn "MLX unit tests are temporarily disabled until a release build publishes ollama-darwin.tgz for MLX_VERSION=${current_mlx} MLX_C_VERSION=${current_mlxc}"
+  exit 0
+fi
+
+tmp_tarball="${tarball}.tmp"
+rm -f "${tmp_tarball}"
+curl -fL --retry 3 --retry-delay 2 -o "${tmp_tarball}" "${matched_url}"
+mv "${tmp_tarball}" "${tarball}"
+tmp_tarball=""
+echo "${matched_tag}" >"${tag_file}"
+echo "${current_pins}" >"${pins_file}"
+
+extract_payload "${matched_tag}"

.github/workflows/test.yaml

@@ -414,6 +414,15 @@ jobs:
       - uses: actions/setup-node@v4
         with:
           node-version: '20'
+      - name: Cache MLX Darwin release payload
+        if: ${{ startsWith(matrix.os, 'macos') }}
+        uses: actions/cache@v4
+        with:
+          path: .cache/mlx-darwin-release
+          key: mlx-darwin-${{ hashFiles('MLX_VERSION', 'MLX_C_VERSION') }}
+      - name: Prepare MLX Darwin release payload
+        if: ${{ startsWith(matrix.os, 'macos') }}
+        run: .github/scripts/prepare_mlx_darwin.sh
       - name: Install UI dependencies
         working-directory: ./app/ui/app
         run: npm ci
@@ -438,12 +447,57 @@ jobs:
       - name: Run go generate
         run: go generate ./...
 
+      - name: Verify UI generated types are current
+        if: ${{ startsWith(matrix.os, 'ubuntu') }}
+        run: git diff --exit-code -- app/ui/app/codegen/gotypes.gen.ts
+
       - name: go test
         if: always()
-        run: go test -count=1 -benchtime=1x ./...
+        # Smoke-run each benchmark once to catch panics and bit rot; this does
+        # not assert timings. -benchtime without -bench is inert.
+        run: go test -count=1 -bench=. -benchtime=1x ./...
 
       - name: go test app with live updater tag
         if: ${{ needs.changes.outputs.app_changed == 'True' && contains(fromJSON('["macos-latest","windows-latest"]'), matrix.os) }}
         run: go test -count=1 -tags updater_live ./app/...
 
       - uses: golangci/golangci-lint-action@v9
+
+  race:
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+    runs-on: ${{ matrix.os }}
+    env:
+      CGO_ENABLED: '1'
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: 'go.mod'
+          cache-dependency-path: |
+            go.sum
+            LLAMA_CPP_VERSION
+            MLX_VERSION
+            MLX_C_VERSION
+      - name: Cache MLX Darwin release payload
+        if: ${{ startsWith(matrix.os, 'macos') }}
+        uses: actions/cache@v4
+        with:
+          path: .cache/mlx-darwin-release
+          key: mlx-darwin-${{ hashFiles('MLX_VERSION', 'MLX_C_VERSION') }}
+      - name: Prepare MLX Darwin release payload
+        if: ${{ startsWith(matrix.os, 'macos') }}
+        run: .github/scripts/prepare_mlx_darwin.sh
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '20'
+      # app/ui embeds app/dist, so the UI has to be built before app/... will
+      # even compile.
+      - name: Build UI
+        working-directory: ./app/ui/app
+        run: |
+          npm ci
+          npm run build
+      - name: go test -race
+        run: go test -race -count=1 ./...

app/server/server_test.go

@@ -365,7 +365,7 @@ time=2025-06-30T09:25:56.197-07:00 level=DEBUG source=ggml.go:155 msg="key not f
 			if err != nil {
 				t.Fatalf("failed to write log file %s: %s", serverLogPath, err)
 			}
-			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Millisecond)
+			ctx, cancel := context.WithTimeout(t.Context(), time.Second)
 			defer cancel()
 			info, err := GetInferenceInfo(ctx)
 			if err != nil {

app/ui/ui_test.go

@@ -997,7 +997,7 @@ func TestSettingsToggleAutoUpdateOn_WithPendingUpdate_ShowsNotification(t *testi
 	}
 }
 
-func TestSettingsToggleAutoUpdateOn_NoPendingUpdate_TriggersCheck(t *testing.T) {
+func TestSettingsToggleAutoUpdateOn_NoPendingUpdate_DoesNotNotify(t *testing.T) {
 	testStore := &store.Store{
 		DBPath: filepath.Join(t.TempDir(), "db.sqlite"),
 	}
@@ -1027,12 +1027,6 @@ func TestSettingsToggleAutoUpdateOn_NoPendingUpdate_TriggersCheck(t *testing.T)
 	}}
 	defer upd.Store.Close()
 
-	// Initialize the checkNow channel by starting (and immediately stopping) the checker
-	// so TriggerImmediateCheck doesn't panic on nil channel
-	ctx, cancel := context.WithCancel(t.Context())
-	upd.StartBackgroundUpdaterChecker(ctx, func(string) error { return nil })
-	defer cancel()
-
 	var notificationCalled atomic.Bool
 	server := &Server{
 		Store:   testStore,

app/updater/updater.go

@@ -353,11 +353,23 @@ func (u *Updater) TriggerImmediateCheck() {
 }
 
 func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) {
+	u.startBackgroundUpdaterChecker(ctx, cb)
+}
+
+func (u *Updater) startBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) <-chan struct{} {
 	u.checkNow = make(chan struct{}, 1)
 	u.checkNow <- struct{}{} // Trigger first check after initial delay
+	done := make(chan struct{})
 	go func() {
+		defer close(done)
 		// Don't blast an update message immediately after startup
-		time.Sleep(UpdateCheckInitialDelay)
+		initialDelay := time.NewTimer(UpdateCheckInitialDelay)
+		defer initialDelay.Stop()
+		select {
+		case <-ctx.Done():
+			return
+		case <-initialDelay.C:
+		}
 		slog.Info("beginning update checker", "interval", UpdateCheckInterval)
 		ticker := time.NewTicker(UpdateCheckInterval)
 		defer ticker.Stop()
@@ -406,4 +418,5 @@ func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(str
 			}
 		}
 	}()
+	return done
 }

app/updater/updater_test.go

@@ -190,6 +190,20 @@ func TestDownloadNewReleaseDoesNotUseRawETagAsPathComponent(t *testing.T) {
 	}
 }
 
+// stopChecker cancels the background update checker and waits for its
+// goroutine to return. Tests must join it before returning: the goroutine
+// reads package-level knobs (UpdateCheckURLBase, UpdateCheckInterval, ...)
+// that the next test rewrites.
+func stopChecker(t *testing.T, cancel context.CancelFunc, done <-chan struct{}) {
+	t.Helper()
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(10 * time.Second):
+		t.Error("background update checker did not stop")
+	}
+}
+
 // waitDownloadIdle blocks until no download is in flight, so staged-file
 // handles close before t.TempDir cleanup removes the stage directory. After
 // the context is cancelled a new download can't write (it aborts at the HEAD
@@ -289,11 +303,12 @@ func TestBackgroundCheckerSkipsAlreadyStagedETagDownload(t *testing.T) {
 	defer cancel()
 
 	callbacks := make(chan string, 4)
-	updater.StartBackgroundUpdaterChecker(ctx, func(ver string) error {
+	checkerDone := updater.startBackgroundUpdaterChecker(ctx, func(ver string) error {
 		callbacks <- ver
 		return nil
 	})
-	t.Cleanup(updater.waitDownloadIdle)
+	defer updater.waitDownloadIdle()
+	defer stopChecker(t, cancel, checkerDone)
 
 	for range 2 {
 		select {
@@ -334,10 +349,16 @@ func TestBackgoundChecker(t *testing.T) {
 	UpdateStageDir = t.TempDir()
 	haveUpdate := false
 	verified := false
-	done := make(chan int)
+	// Buffered + non-blocking send: the checker keeps calling cb every
+	// UpdateCheckInterval, and a blocking send would wedge its goroutine once
+	// the test stops receiving.
+	done := make(chan int, 1)
 	cb := func(ver string) error {
 		haveUpdate = true
-		done <- 0
+		select {
+		case done <- 0:
+		default:
+		}
 		return nil
 	}
 	stallTimer := time.NewTimer(5 * time.Second)
@@ -381,8 +402,9 @@ func TestBackgoundChecker(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
-	t.Cleanup(updater.waitDownloadIdle)
+	checkerDone := updater.startBackgroundUpdaterChecker(ctx, cb)
+	defer updater.waitDownloadIdle()
+	defer stopChecker(t, cancel, checkerDone)
 	select {
 	case <-stallTimer.C:
 		t.Fatal("stalled")
@@ -440,12 +462,13 @@ func TestAutoUpdateDisabledSkipsDownload(t *testing.T) {
 	}
 
 	cb := func(ver string) error {
-		t.Fatal("callback should not be called when auto-update is disabled")
+		t.Error("callback should not be called when auto-update is disabled")
 		return nil
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
-	t.Cleanup(updater.waitDownloadIdle)
+	checkerDone := updater.startBackgroundUpdaterChecker(ctx, cb)
+	defer updater.waitDownloadIdle()
+	defer stopChecker(t, cancel, checkerDone)
 
 	// Wait enough time for multiple check cycles
 	time.Sleep(50 * time.Millisecond)
@@ -507,8 +530,9 @@ func TestAutoUpdateReenabledDownloadsUpdate(t *testing.T) {
 		return nil
 	}
 
-	upd.StartBackgroundUpdaterChecker(ctx, cb)
-	t.Cleanup(upd.waitDownloadIdle)
+	checkerDone := upd.startBackgroundUpdaterChecker(ctx, cb)
+	defer upd.waitDownloadIdle()
+	defer stopChecker(t, cancel, checkerDone)
 
 	// Wait for a few cycles with auto-update disabled - no download should happen
 	time.Sleep(50 * time.Millisecond)
@@ -641,8 +665,9 @@ func TestTriggerImmediateCheck(t *testing.T) {
 		return nil
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
-	t.Cleanup(updater.waitDownloadIdle)
+	checkerDone := updater.startBackgroundUpdaterChecker(ctx, cb)
+	defer updater.waitDownloadIdle()
+	defer stopChecker(t, cancel, checkerDone)
 
 	// Wait for the initial check that fires after the initial delay
 	select {

x/internal/mlxtest/mlxtest.go

@@ -3,39 +3,51 @@
 package mlxtest
 
 import (
-	"runtime"
+	"sync"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
+var testThread = sync.OnceValues(func() (*mlxthreadtest.Thread, error) {
+	return mlxthreadtest.Start("mlx-test", func() error {
+		if err := mlx.CheckInit(); err != nil {
+			return err
+		}
+		if mlx.GPUIsAvailable() {
+			mlx.SetDefaultDeviceGPU()
+		}
+		return nil
+	})
+})
+
+// T is the test state available to callbacks running on the MLX thread.
+type T = mlxthreadtest.T
+
 // SkipIfUnavailable skips the test when the MLX dynamic library cannot be
 // loaded (e.g. no MLX backend built for this platform).
 func SkipIfUnavailable(t *testing.T) {
 	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
+	if _, err := testThread(); err != nil {
 		t.Skipf("MLX not available: %v", err)
 	}
 }
 
-// Setup prepares a test that calls into MLX natively: it skips when MLX is
-// unavailable and pins the test goroutine to its OS thread for the duration
-// of the test.
-//
-// The thread pin is load-bearing, not defensive: MLX's default stream cache
-// is thread-local, and anything that migrates the goroutine mid-test (the
-// race detector's scheduler in particular) otherwise panics with
-// "There is no Stream(gpu, 0) in current thread".
-//
-// Setup deliberately does not switch devices or sweep caches: switching the
-// default device re-creates the process-wide default stream, and sweeping the
-// allocator cache between tests changes allocator reuse — both perturbed
-// tests that share lazy arrays with subtests running on other threads.
-func Setup(t *testing.T) {
+// Run executes fn on the MLX thread shared by the package's test binary.
+func Run(t *testing.T, fn func(*T)) {
 	t.Helper()
 
-	SkipIfUnavailable(t)
+	thread, err := testThread()
+	if err != nil {
+		t.Skipf("MLX not available: %v", err)
+	}
 
-	runtime.LockOSThread()
-	t.Cleanup(runtime.UnlockOSThread)
+	mlxthreadtest.Run(t, thread, fn)
+}
+
+// RunSubtest runs a named subtest on the shared MLX test thread.
+func RunSubtest(t *testing.T, name string, fn func(*T)) {
+	t.Helper()
+	t.Run(name, func(t *testing.T) { Run(t, fn) })
 }

x/internal/mlxthreadtest/mlxthreadtest.go

@@ -0,0 +1,211 @@
+// Package mlxthreadtest runs tests on a persistent MLX worker thread.
+package mlxthreadtest
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthread"
+)
+
+// Thread is a pinned worker used by MLX tests.
+type Thread struct {
+	worker *mlxthread.Thread
+	id     uint64
+}
+
+// Start creates a pinned test worker.
+func Start(name string, init func() error) (*Thread, error) {
+	t := &Thread{}
+	thread, err := mlxthread.Start(name, func() error {
+		t.id = currentThreadID()
+		if init != nil {
+			return init()
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	t.worker = thread
+	return t, nil
+}
+
+// Do runs fn on the pinned worker.
+func (t *Thread) Do(ctx context.Context, fn func() error) error {
+	if t.onWorker() {
+		panic("mlxthreadtest.Thread.Do called from its pinned worker")
+	}
+	return t.worker.Do(ctx, fn)
+}
+
+// Stop shuts down the pinned worker after running cleanup on it.
+func (t *Thread) Stop(ctx context.Context, cleanup func()) error {
+	if t.onWorker() {
+		panic("mlxthreadtest.Thread.Stop called from its pinned worker")
+	}
+	return t.worker.Stop(ctx, cleanup)
+}
+
+func (t *Thread) onWorker() bool {
+	id := currentThreadID()
+	return id != 0 && id == t.id
+}
+
+// T is the subset of testing.T supported by MLX test bodies. Operations that
+// end a test are replayed by Run on the test goroutine so the MLX worker remains
+// alive.
+type T struct {
+	testReporter
+	cleanups []func()
+	skipped  bool
+	aborted  bool
+}
+
+// abortPanic unwinds a test body without terminating the worker goroutine.
+var abortPanic = new(struct{ marker byte })
+
+type testReporter interface {
+	Error(...any)
+	Errorf(string, ...any)
+	Fail()
+	Failed() bool
+	Helper()
+	Log(...any)
+	Logf(string, ...any)
+}
+
+// Run executes fn on thread. The callback must use its T argument; calling
+// FailNow or SkipNow on a captured *testing.T terminates the pinned worker.
+func Run(t *testing.T, thread *Thread, fn func(*T)) {
+	t.Helper()
+	if thread.onWorker() {
+		panic("mlxthreadtest.Run called recursively from its pinned worker")
+	}
+
+	mt := &T{testReporter: t}
+	result := make(chan runResult, 1)
+	go func() {
+		defer func() {
+			if v := recover(); v != nil {
+				result <- runResult{panicValue: v}
+			}
+		}()
+		err := thread.Do(context.Background(), func() error {
+			returned := false
+			defer func() {
+				if v := recover(); v != nil {
+					panic(v)
+				}
+				if !returned {
+					result <- runResult{goexit: true}
+				}
+			}()
+
+			mt.run(fn)
+			returned = true
+			return nil
+		})
+		result <- runResult{err: err}
+	}()
+
+	res := <-result
+	if res.goexit {
+		panic("pinned test body called runtime.Goexit; use the test value passed to the callback")
+	}
+	if res.panicValue != nil {
+		panic(res.panicValue)
+	}
+	if res.err != nil {
+		t.Fatal(res.err)
+	}
+	if mt.skipped {
+		t.SkipNow()
+	}
+	if mt.aborted {
+		t.FailNow()
+	}
+}
+
+type runResult struct {
+	err        error
+	panicValue any
+	goexit     bool
+}
+
+// Cleanup registers fn to run on the MLX thread after the current body.
+func (t *T) Cleanup(fn func()) {
+	t.cleanups = append(t.cleanups, fn)
+}
+
+func (t *T) FailNow() {
+	t.Fail()
+	t.aborted = true
+	panic(abortPanic)
+}
+
+func (t *T) Fatal(args ...any) {
+	t.Helper()
+	t.Error(args...)
+	t.aborted = true
+	panic(abortPanic)
+}
+
+func (t *T) Fatalf(format string, args ...any) {
+	t.Helper()
+	t.Errorf(format, args...)
+	t.aborted = true
+	panic(abortPanic)
+}
+
+func (t *T) Skip(args ...any) {
+	t.Helper()
+	t.Log(args...)
+	t.SkipNow()
+}
+
+func (t *T) Skipf(format string, args ...any) {
+	t.Helper()
+	t.Logf(format, args...)
+	t.SkipNow()
+}
+
+func (t *T) SkipNow() {
+	t.skipped = true
+	t.aborted = true
+	panic(abortPanic)
+}
+
+func (t *T) Skipped() bool {
+	return t.skipped
+}
+
+func (t *T) run(fn func(*T)) {
+	defer func() {
+		if v := recover(); v != nil && v != abortPanic {
+			panic(v)
+		}
+	}()
+	defer t.runCleanups()
+	fn(t)
+}
+
+func (t *T) runCleanups() {
+	var panicValue any
+	for len(t.cleanups) > 0 {
+		last := len(t.cleanups) - 1
+		cleanup := t.cleanups[last]
+		t.cleanups = t.cleanups[:last]
+		func() {
+			defer func() {
+				if v := recover(); v != nil && v != abortPanic && panicValue == nil {
+					panicValue = v
+				}
+			}()
+			cleanup()
+		}()
+	}
+	if panicValue != nil {
+		panic(panicValue)
+	}
+}

x/internal/mlxthreadtest/mlxthreadtest_test.go

@@ -0,0 +1,379 @@
+package mlxthreadtest
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+)
+
+type fakeReporter struct{ failed bool }
+
+func (t *fakeReporter) Error(...any)          { t.failed = true }
+func (t *fakeReporter) Errorf(string, ...any) { t.failed = true }
+func (t *fakeReporter) Fail()                 { t.failed = true }
+func (t *fakeReporter) Failed() bool          { return t.failed }
+func (*fakeReporter) Helper()                 {}
+func (*fakeReporter) Log(...any)              {}
+func (*fakeReporter) Logf(string, ...any)     {}
+
+func TestControlsDoNotStopWorker(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := thread.Stop(context.Background(), nil); err != nil {
+			t.Error(err)
+		}
+	})
+
+	tests := []struct {
+		name        string
+		body        func(*T)
+		wantFailed  bool
+		wantSkipped bool
+	}{
+		{name: "fail now", body: func(t *T) { t.FailNow() }, wantFailed: true},
+		{name: "fatal", body: func(t *T) { t.Fatal("stop") }, wantFailed: true},
+		{name: "fatalf", body: func(t *T) { t.Fatalf("%s", "stop") }, wantFailed: true},
+		{name: "skip now", body: func(t *T) { t.SkipNow() }, wantSkipped: true},
+		{name: "skip", body: func(t *T) { t.Skip("stop") }, wantSkipped: true},
+		{name: "skipf", body: func(t *T) { t.Skipf("%s", "stop") }, wantSkipped: true},
+		{
+			name: "failed skip",
+			body: func(t *T) {
+				t.Error("failed")
+				t.Skip("stop")
+			},
+			wantFailed:  true,
+			wantSkipped: true,
+		},
+		{
+			name: "cleanup fail now",
+			body: func(t *T) {
+				t.Cleanup(func() { t.FailNow() })
+			},
+			wantFailed: true,
+		},
+		{
+			name: "cleanup fatal",
+			body: func(t *T) {
+				t.Cleanup(func() { t.Fatal("stop") })
+			},
+			wantFailed: true,
+		},
+		{
+			name: "cleanup skip now",
+			body: func(t *T) {
+				t.Cleanup(func() { t.SkipNow() })
+			},
+			wantSkipped: true,
+		},
+		{
+			name: "cleanup skip",
+			body: func(t *T) {
+				t.Cleanup(func() { t.Skip("stop") })
+			},
+			wantSkipped: true,
+		},
+	}
+	for _, tt := range tests {
+		reporter := &fakeReporter{}
+		mt := &T{testReporter: reporter}
+		if err := thread.Do(context.Background(), func() error {
+			mt.run(tt.body)
+			return nil
+		}); err != nil {
+			t.Fatalf("%s: %v", tt.name, err)
+		}
+		if mt.Failed() != tt.wantFailed || mt.skipped != tt.wantSkipped {
+			t.Fatalf("%s: failed:%v skipped:%v", tt.name, mt.Failed(), mt.skipped)
+		}
+
+		ran := false
+		if err := thread.Do(context.Background(), func() error {
+			ran = true
+			return nil
+		}); err != nil {
+			t.Fatalf("%s follow-up: %v", tt.name, err)
+		}
+		if !ran {
+			t.Fatalf("worker did not run after %s", tt.name)
+		}
+	}
+
+	reached := false
+	t.Run("replay skip", func(t *testing.T) {
+		Run(t, thread, func(t *T) { t.Skip("stop") })
+		reached = true
+	})
+	if reached {
+		t.Fatal("Run returned after Skip")
+	}
+	if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+		t.Fatalf("worker did not survive replayed Skip: %v", err)
+	}
+}
+
+func TestRunReplaysFatalWithoutStoppingWorker(t *testing.T) {
+	if marker := os.Getenv("OLLAMA_TEST_MLX_FATAL_MARKER"); marker != "" {
+		thread, err := Start("test", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() {
+			if err := thread.Stop(context.Background(), nil); err != nil {
+				t.Error(err)
+			}
+		})
+		t.Cleanup(func() {
+			if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+				t.Errorf("worker did not survive replayed Fatal: %v", err)
+				return
+			}
+			if err := os.WriteFile(marker, nil, 0o600); err != nil {
+				t.Error(err)
+			}
+		})
+
+		Run(t, thread, func(t *T) { t.Fatal("stop") })
+		return
+	}
+
+	marker := filepath.Join(t.TempDir(), "worker-survived")
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+	defer cancel()
+	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestRunReplaysFatalWithoutStoppingWorker$", "-test.timeout=2s")
+	cmd.Env = append(os.Environ(), "OLLAMA_TEST_MLX_FATAL_MARKER="+marker)
+	output, err := cmd.CombinedOutput()
+	if err == nil {
+		t.Fatal("subprocess succeeded, want Fatal failure")
+	}
+	if ctx.Err() != nil {
+		t.Fatalf("subprocess hung: %v", ctx.Err())
+	}
+	if !strings.Contains(string(output), "stop") {
+		t.Fatalf("missing Fatal diagnostic:\n%s", output)
+	}
+	if _, err := os.Stat(marker); err != nil {
+		t.Fatalf("worker survival check did not complete: %v\n%s", err, output)
+	}
+}
+
+func TestRunRejectsRecursiveDispatch(t *testing.T) {
+	if currentThreadID() == 0 {
+		t.Skip("OS thread IDs are not available on this platform")
+	}
+
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := thread.Stop(context.Background(), nil); err != nil {
+			t.Error(err)
+		}
+	})
+
+	tests := []struct {
+		name string
+		want string
+		body func(*testing.T)
+	}{
+		{
+			name: "Run",
+			want: "called recursively",
+			body: func(t *testing.T) { Run(t, thread, func(*T) {}) },
+		},
+		{
+			name: "Do",
+			want: "Thread.Do called from its pinned worker",
+			body: func(*testing.T) {
+				_ = thread.Do(context.Background(), func() error { return nil })
+			},
+		},
+		{
+			name: "Stop",
+			want: "Thread.Stop called from its pinned worker",
+			body: func(*testing.T) { _ = thread.Stop(context.Background(), nil) },
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var got any
+			func() {
+				defer func() { got = recover() }()
+				Run(t, thread, func(*T) { tt.body(t) })
+			}()
+			if !strings.Contains(fmt.Sprint(got), tt.want) {
+				t.Fatalf("got panic %v, want %q", got, tt.want)
+			}
+			if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+				t.Fatalf("worker did not survive recursive %s: %v", tt.name, err)
+			}
+		})
+	}
+}
+
+func TestRunRejectsTestingTGoexit(t *testing.T) {
+	if method := os.Getenv("OLLAMA_TEST_MLX_GOEXIT"); method != "" {
+		thread, err := Start("test", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		Run(t, thread, func(*T) {
+			switch method {
+			case "fail":
+				t.FailNow()
+			case "skip":
+				t.SkipNow()
+			default:
+				panic("unknown Goexit test method")
+			}
+		})
+		return
+	}
+
+	for _, method := range []string{"fail", "skip"} {
+		t.Run(method, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+			defer cancel()
+			cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestRunRejectsTestingTGoexit$", "-test.timeout=2s")
+			cmd.Env = append(os.Environ(), "OLLAMA_TEST_MLX_GOEXIT="+method)
+			output, err := cmd.CombinedOutput()
+			if err == nil {
+				t.Fatal("subprocess succeeded, want misuse failure")
+			}
+			if ctx.Err() != nil {
+				t.Fatalf("subprocess hung: %v", ctx.Err())
+			}
+			if !strings.Contains(string(output), "pinned test body called runtime.Goexit") {
+				t.Fatalf("missing Goexit diagnostic:\n%s", output)
+			}
+		})
+	}
+}
+
+func TestCleanupOrder(t *testing.T) {
+	reporter := &fakeReporter{}
+	mt := &T{testReporter: reporter}
+	var got []string
+
+	mt.run(func(t *T) {
+		t.Cleanup(func() { got = append(got, "cleanup 1") })
+		t.Cleanup(func() {
+			got = append(got, "cleanup 2")
+			t.Cleanup(func() { got = append(got, "cleanup 3") })
+		})
+		got = append(got, "body")
+	})
+
+	want := []string{"body", "cleanup 2", "cleanup 3", "cleanup 1"}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("got %v, want %v", got, want)
+	}
+}
+
+func TestCleanupFailureDoesNotHideBodyPanic(t *testing.T) {
+	reporter := &fakeReporter{}
+	mt := &T{testReporter: reporter}
+	var got any
+	func() {
+		defer func() { got = recover() }()
+		mt.run(func(t *T) {
+			t.Cleanup(func() { t.Fatal("cleanup failure") })
+			panic("body panic")
+		})
+	}()
+	if got != "body panic" {
+		t.Fatalf("got panic %v, want body panic", got)
+	}
+	if !reporter.failed {
+		t.Fatal("cleanup Fatal did not fail the test")
+	}
+}
+
+func TestCleanupPanicRunsRemainingCleanups(t *testing.T) {
+	reporter := &fakeReporter{}
+	mt := &T{testReporter: reporter}
+	var cleanups []string
+	var got any
+	func() {
+		defer func() { got = recover() }()
+		mt.run(func(t *T) {
+			t.Cleanup(func() { cleanups = append(cleanups, "first") })
+			t.Cleanup(func() {
+				cleanups = append(cleanups, "second")
+				panic("cleanup panic")
+			})
+		})
+	}()
+	if got != "cleanup panic" {
+		t.Fatalf("got panic %v, want cleanup panic", got)
+	}
+	if want := []string{"second", "first"}; !reflect.DeepEqual(cleanups, want) {
+		t.Fatalf("got cleanups %v, want %v", cleanups, want)
+	}
+}
+
+func TestCleanupRunsOnWorker(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := thread.Stop(context.Background(), nil); err != nil {
+			t.Error(err)
+		}
+	})
+
+	var bodyID, cleanupID uint64
+	Run(t, thread, func(t *T) {
+		bodyID = currentThreadID()
+		t.Cleanup(func() { cleanupID = currentThreadID() })
+	})
+	if bodyID != thread.id || cleanupID != thread.id {
+		t.Fatalf("body thread %d, cleanup thread %d, want worker thread %d", bodyID, cleanupID, thread.id)
+	}
+}
+
+func panicAtTestBody() {
+	panic("test panic")
+}
+
+func TestPanicIncludesTestBodyStack(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := thread.Stop(context.Background(), nil); err != nil {
+			t.Error(err)
+		}
+	})
+
+	var cleanupRan bool
+	var got any
+	func() {
+		defer func() { got = recover() }()
+		Run(t, thread, func(t *T) {
+			t.Cleanup(func() { cleanupRan = true })
+			panicAtTestBody()
+		})
+	}()
+	if !cleanupRan {
+		t.Fatal("cleanup did not run after panic")
+	}
+	if !strings.Contains(fmt.Sprint(got), "panicAtTestBody") {
+		t.Fatalf("panic stack does not include test body:\n%v", got)
+	}
+	if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+		t.Fatalf("worker did not survive panic: %v", err)
+	}
+}

x/internal/mlxthreadtest/threadid_darwin.go

@@ -0,0 +1,10 @@
+//go:build darwin
+
+package mlxthreadtest
+
+import "syscall"
+
+func currentThreadID() uint64 {
+	id, _, _ := syscall.RawSyscall(syscall.SYS_THREAD_SELFID, 0, 0, 0)
+	return uint64(id)
+}

x/internal/mlxthreadtest/threadid_linux.go

@@ -0,0 +1,9 @@
+//go:build linux
+
+package mlxthreadtest
+
+import "syscall"
+
+func currentThreadID() uint64 {
+	return uint64(syscall.Gettid())
+}

x/internal/mlxthreadtest/threadid_other.go

@@ -0,0 +1,7 @@
+//go:build !darwin && !linux && !windows
+
+package mlxthreadtest
+
+func currentThreadID() uint64 {
+	return 0
+}

x/internal/mlxthreadtest/threadid_windows.go

@@ -0,0 +1,9 @@
+//go:build windows
+
+package mlxthreadtest
+
+import "golang.org/x/sys/windows"
+
+func currentThreadID() uint64 {
+	return uint64(windows.GetCurrentThreadId())
+}

x/mlxrunner/cache/cache_test.go

@@ -3,17 +3,11 @@ package cache
 import (
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}
-
 // newKVBatch builds a B=1 batch at SeqOffsets=off with all-real
 // queries (SeqQueryLens=L) — the standard single-sequence cache
 // test shape.
@@ -26,258 +20,265 @@ func newKVBatch(off, L int) *batch.Batch {
 }
 
 func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewKVCache()
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
 
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	// Snapshot [5, 10).
-	snap := c.Snapshot(5)
+		// Snapshot [5, 10).
+		snap := c.Snapshot(5)
 
-	// Free the cache completely — offset is now 0.
-	c.Free()
+		// Free the cache completely — offset is now 0.
+		c.Free()
 
-	// Restore should fail because cache doesn't have data up to fromOffset=5.
-	if c.Restore(snap, 10) {
-		t.Fatal("expected Restore to fail with no base data")
-	}
+		// Restore should fail because cache doesn't have data up to fromOffset=5.
+		if c.Restore(snap, 10) {
+			t.Fatal("expected Restore to fail with no base data")
+		}
+	})
 }
 
 // TestKVCacheDataSurvivesSnapshotRestore verifies that actual array data
 // is preserved through a snapshot→free→restore cycle.
 func TestKVCacheDataSurvivesSnapshotRestore(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewKVCache()
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
 
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	snap := c.Snapshot(0)
-	if snap == nil {
-		t.Fatal("Snapshot returned nil")
-	}
+		snap := c.Snapshot(0)
+		if snap == nil {
+			t.Fatal("Snapshot returned nil")
+		}
 
-	// Free and restore to a fresh cache.
-	c2 := NewKVCache()
-	if !c2.Restore(snap, 10) {
-		t.Fatal("Restore failed")
-	}
-	if c2.Offset() != 10 {
-		t.Fatalf("offset = %d, want 10", c2.Offset())
-	}
+		// Free and restore to a fresh cache.
+		c2 := NewKVCache()
+		if !c2.Restore(snap, 10) {
+			t.Fatal("Restore failed")
+		}
+		if c2.Offset() != 10 {
+			t.Fatalf("offset = %d, want 10", c2.Offset())
+		}
 
-	// Verify State() returns arrays with correct sequence dimension.
-	state := c2.State()
-	if len(state) != 2 {
-		t.Fatalf("State() returned %d arrays, want 2", len(state))
-	}
-	// keys shape: [B, H, seqLen, Dk]
-	if state[0].Dim(2) != 10 {
-		t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2))
-	}
-	if state[1].Dim(2) != 10 {
-		t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2))
-	}
+		// Verify State() returns arrays with correct sequence dimension.
+		state := c2.State()
+		if len(state) != 2 {
+			t.Fatalf("State() returned %d arrays, want 2", len(state))
+		}
+		// keys shape: [B, H, seqLen, Dk]
+		if state[0].Dim(2) != 10 {
+			t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2))
+		}
+		if state[1].Dim(2) != 10 {
+			t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2))
+		}
+	})
 }
 
 // TestKVCacheSplitPreservesData verifies that split produces two snapshots
 // that can be sequentially restored to rebuild the original cache state.
 func TestKVCacheSplitPreservesData(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewKVCache()
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
 
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	snap := c.Snapshot(0)
-	parent, child := c.Split(snap, 5)
-	if parent == nil || child == nil {
-		t.Fatal("Split returned nil")
-	}
+		snap := c.Snapshot(0)
+		parent, child := c.Split(snap, 5)
+		if parent == nil || child == nil {
+			t.Fatal("Split returned nil")
+		}
 
-	// Restore parent → offset=5, seq dim=5.
-	c2 := NewKVCache()
-	if !c2.Restore(parent, 5) {
-		t.Fatal("Restore(parent) failed")
-	}
-	if c2.Offset() != 5 {
-		t.Fatalf("offset after parent = %d, want 5", c2.Offset())
-	}
-	state := c2.State()
-	if state[0].Dim(2) != 5 {
-		t.Fatalf("keys seq dim after parent = %d, want 5", state[0].Dim(2))
-	}
+		// Restore parent → offset=5, seq dim=5.
+		c2 := NewKVCache()
+		if !c2.Restore(parent, 5) {
+			t.Fatal("Restore(parent) failed")
+		}
+		if c2.Offset() != 5 {
+			t.Fatalf("offset after parent = %d, want 5", c2.Offset())
+		}
+		state := c2.State()
+		if state[0].Dim(2) != 5 {
+			t.Fatalf("keys seq dim after parent = %d, want 5", state[0].Dim(2))
+		}
 
-	// Restore child on top → offset=10, seq dim=10.
-	if !c2.Restore(child, 10) {
-		t.Fatal("Restore(child) failed")
-	}
-	if c2.Offset() != 10 {
-		t.Fatalf("offset after child = %d, want 10", c2.Offset())
-	}
-	state = c2.State()
-	if state[0].Dim(2) != 10 {
-		t.Fatalf("keys seq dim after child = %d, want 10", state[0].Dim(2))
-	}
+		// Restore child on top → offset=10, seq dim=10.
+		if !c2.Restore(child, 10) {
+			t.Fatal("Restore(child) failed")
+		}
+		if c2.Offset() != 10 {
+			t.Fatalf("offset after child = %d, want 10", c2.Offset())
+		}
+		state = c2.State()
+		if state[0].Dim(2) != 10 {
+			t.Fatalf("keys seq dim after child = %d, want 10", state[0].Dim(2))
+		}
+	})
 }
 
 // TestKVCacheSplitMergeRoundTripData verifies that splitting and merging back
 // produces a snapshot equivalent to the original.
 func TestKVCacheSplitMergeRoundTripData(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewKVCache()
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
 
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	snap := c.Snapshot(0)
-	parent, child := c.Split(snap, 6)
-	merged := c.Merge(parent, child)
-	if merged == nil {
-		t.Fatal("Merge returned nil")
-	}
+		snap := c.Snapshot(0)
+		parent, child := c.Split(snap, 6)
+		merged := c.Merge(parent, child)
+		if merged == nil {
+			t.Fatal("Merge returned nil")
+		}
 
-	c2 := NewKVCache()
-	if !c2.Restore(merged, 10) {
-		t.Fatal("Restore(merged) failed")
-	}
-	if c2.Offset() != 10 {
-		t.Fatalf("offset = %d, want 10", c2.Offset())
-	}
+		c2 := NewKVCache()
+		if !c2.Restore(merged, 10) {
+			t.Fatal("Restore(merged) failed")
+		}
+		if c2.Offset() != 10 {
+			t.Fatalf("offset = %d, want 10", c2.Offset())
+		}
 
-	state := c2.State()
-	if state[0].Dim(2) != 10 {
-		t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2))
-	}
-	if state[1].Dim(2) != 10 {
-		t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2))
-	}
+		state := c2.State()
+		if state[0].Dim(2) != 10 {
+			t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2))
+		}
+		if state[1].Dim(2) != 10 {
+			t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2))
+		}
+	})
 }
 
 func TestRotatingKVCacheRestoreOutsideWindow(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRotatingKVCache(4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRotatingKVCache(4)
 
-	// Feed 10 tokens (window size 4, so positions 0-5 are evicted).
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		// Feed 10 tokens (window size 4, so positions 0-5 are evicted).
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	// Offset 3 is outside the window.
-	if c.Restore(nil, 3) {
-		t.Fatal("Restore(nil, 3) should fail when outside window")
-	}
+		// Offset 3 is outside the window.
+		if c.Restore(nil, 3) {
+			t.Fatal("Restore(nil, 3) should fail when outside window")
+		}
+	})
 }
 
 // TestRotatingKVCacheSnapshotPreservesWindow verifies that after restoring
 // from a snapshot, the rotating cache has the correct window of data.
 func TestRotatingKVCacheSnapshotPreservesWindow(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRotatingKVCache(4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRotatingKVCache(4)
 
-	// Feed 10 tokens one at a time. Window size 4, so only last 4 are kept.
-	for range 10 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		// Feed 10 tokens one at a time. Window size 4, so only last 4 are kept.
+		for range 10 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	snap := c.Snapshot(0)
-	if snap == nil {
-		t.Fatal("Snapshot returned nil")
-	}
+		snap := c.Snapshot(0)
+		if snap == nil {
+			t.Fatal("Snapshot returned nil")
+		}
 
-	// Feed 5 more tokens.
-	for range 5 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		// Feed 5 more tokens.
+		for range 5 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	// Restore to offset 10.
-	if !c.Restore(snap, 10) {
-		t.Fatal("Restore failed")
-	}
-	if c.Offset() != 10 {
-		t.Fatalf("offset = %d, want 10", c.Offset())
-	}
+		// Restore to offset 10.
+		if !c.Restore(snap, 10) {
+			t.Fatal("Restore failed")
+		}
+		if c.Offset() != 10 {
+			t.Fatalf("offset = %d, want 10", c.Offset())
+		}
 
-	state := c.State()
-	if len(state) != 2 {
-		t.Fatalf("State() returned %d arrays, want 2", len(state))
-	}
-	// Seq dim should be min(offset, maxSize) = min(10, 4) = 4.
-	seqDim := state[0].Dim(2)
-	if seqDim != 4 {
-		t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim)
-	}
+		state := c.State()
+		if len(state) != 2 {
+			t.Fatalf("State() returned %d arrays, want 2", len(state))
+		}
+		// Seq dim should be min(offset, maxSize) = min(10, 4) = 4.
+		seqDim := state[0].Dim(2)
+		if seqDim != 4 {
+			t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim)
+		}
+	})
 }
 
 // TestRotatingKVCacheRestoreFromSnapshot verifies that restoring from a
 // snapshot correctly preserves the write position (idx), so subsequent
 // single-token updates land in the right buffer slot.
 func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRotatingKVCache(4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRotatingKVCache(4)
 
-	// Fill the window: 6 tokens into a size-4 window.
-	// After this, idx has wrapped and the buffer has rotated.
-	for range 6 {
+		// Fill the window: 6 tokens into a size-4 window.
+		// After this, idx has wrapped and the buffer has rotated.
+		for range 6 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
+		if c.Offset() != 6 {
+			t.Fatalf("offset = %d, want 6", c.Offset())
+		}
+
+		snap := c.Snapshot(0)
+
+		// Mutate the cache further so live state diverges from snapshot.
+		for range 3 {
+			k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
+
+		// Restore to snapshot state.
+		if !c.Restore(snap, 6) {
+			t.Fatal("Restore failed")
+		}
+		if c.Offset() != 6 {
+			t.Fatalf("offset after restore = %d, want 6", c.Offset())
+		}
+
+		// Feed one more token. If idx was restored correctly, this should
+		// produce a valid window of size 4 at offset 7.
 		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
 		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
 		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
-	if c.Offset() != 6 {
-		t.Fatalf("offset = %d, want 6", c.Offset())
-	}
 
-	snap := c.Snapshot(0)
-
-	// Mutate the cache further so live state diverges from snapshot.
-	for range 3 {
-		k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
-
-	// Restore to snapshot state.
-	if !c.Restore(snap, 6) {
-		t.Fatal("Restore failed")
-	}
-	if c.Offset() != 6 {
-		t.Fatalf("offset after restore = %d, want 6", c.Offset())
-	}
-
-	// Feed one more token. If idx was restored correctly, this should
-	// produce a valid window of size 4 at offset 7.
-	k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-	v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
-	c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-
-	if c.Offset() != 7 {
-		t.Fatalf("offset after post-restore update = %d, want 7", c.Offset())
-	}
-	state := c.State()
-	if len(state) != 2 {
-		t.Fatalf("State() returned %d arrays, want 2", len(state))
-	}
-	seqDim := state[0].Dim(2)
-	if seqDim != 4 {
-		t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim)
-	}
+		if c.Offset() != 7 {
+			t.Fatalf("offset after post-restore update = %d, want 7", c.Offset())
+		}
+		state := c.State()
+		if len(state) != 2 {
+			t.Fatalf("State() returned %d arrays, want 2", len(state))
+		}
+		seqDim := state[0].Dim(2)
+		if seqDim != 4 {
+			t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim)
+		}
+	})
 }

x/mlxrunner/cache/lazy_test.go

@@ -3,6 +3,7 @@ package cache
 import (
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
@@ -44,53 +45,53 @@ func settledActiveMemory() int {
 // snapshot stays lazy and is discarded before any overwrite. Compare against the
 // bytes an eager per-token copy would cost.
 func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const before, draft, H, D = 16, 8, 4, 8
 
-	const before, draft, H, D = 16, 8, 4, 8
+		c := NewKVCache()
+		fillKV(c, before)
 
-	c := NewKVCache()
-	fillKV(c, before)
-
-	offsets := make([]int, draft)
-	for i := range offsets {
-		offsets[i] = before + i
-	}
-	c.PrepareSnapshots(offsets)
-
-	k, v := batchKV(draft)
-	c.Update(newKVBatch(before, draft), k, v)
-
-	baseline := settledActiveMemory()
-
-	snaps := c.TakeSnapshots()
-	// Every captured snapshot is a lazy snapshot (no owned buffer).
-	for i, s := range snaps {
-		if s == nil {
-			continue
+		offsets := make([]int, draft)
+		for i := range offsets {
+			offsets[i] = before + i
 		}
-		if ks := s.(*kvSnapshot); ks.keys != nil {
-			t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i)
-		}
-	}
+		c.PrepareSnapshots(offsets)
 
-	// 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()
-		}
-	}
+		k, v := batchKV(draft)
+		c.Update(newKVBatch(before, draft), k, v)
 
-	after := settledActiveMemory()
-	// Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but
-	// well under one per-token copy (draft tokens * keys+values).
-	perToken := (c.keys.NumBytes() + c.values.NumBytes()) / c.keys.Dim(2)
-	if after-baseline > perToken {
-		t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing",
-			after-baseline, perToken)
-	}
+		baseline := settledActiveMemory()
+
+		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)
+			}
+		}
+
+		// 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
+		// well under one per-token copy (draft tokens * keys+values).
+		perToken := (c.keys.NumBytes() + c.values.NumBytes()) / c.keys.Dim(2)
+		if after-baseline > perToken {
+			t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing",
+				after-baseline, perToken)
+		}
+	})
 }
 
 // TestKVLazySnapshotSizeZeroUntilMaterialized verifies the accounting contract:
@@ -98,40 +99,40 @@ func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) {
 // destructive write triggers copyOut the materialize hook fires with the
 // newly-allocated bytes and Size() reports the owned arrays.
 func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
+		snap := c.Snapshot(5).(*kvSnapshot)
+		defer snap.Close()
+		if snap.Size() != 0 {
+			t.Fatalf("lazy snapshot Size = %d, want 0", snap.Size())
+		}
 
-	snap := c.Snapshot(5).(*kvSnapshot)
-	defer snap.Close()
-	if snap.Size() != 0 {
-		t.Fatalf("lazy snapshot Size = %d, want 0", snap.Size())
-	}
+		var hookDelta int
+		snap.SetMaterializeHook(func(delta int) { hookDelta = delta })
 
-	var hookDelta int
-	snap.SetMaterializeHook(func(delta int) { hookDelta = delta })
+		// Rewind and overwrite to force copyOut.
+		if !c.Restore(nil, 5) {
+			t.Fatal("rewind failed")
+		}
+		nk, nv := distinctKV(100, 3, H, D)
+		c.Update(newKVBatch(5, 3), nk, nv)
 
-	// Rewind and overwrite to force copyOut.
-	if !c.Restore(nil, 5) {
-		t.Fatal("rewind failed")
-	}
-	nk, nv := distinctKV(100, 3, H, D)
-	c.Update(newKVBatch(5, 3), nk, nv)
-
-	if snap.keys == nil {
-		t.Fatal("snapshot was not materialized by overwriting write")
-	}
-	want := snap.keys.NumBytes() + snap.values.NumBytes()
-	if hookDelta != want {
-		t.Fatalf("hook fired with delta %d, want %d (owned bytes)", hookDelta, want)
-	}
-	if snap.Size() != want {
-		t.Fatalf("materialized Size = %d, want %d", snap.Size(), want)
-	}
+		if snap.keys == nil {
+			t.Fatal("snapshot was not materialized by overwriting write")
+		}
+		want := snap.keys.NumBytes() + snap.values.NumBytes()
+		if hookDelta != want {
+			t.Fatalf("hook fired with delta %d, want %d (owned bytes)", hookDelta, want)
+		}
+		if snap.Size() != want {
+			t.Fatalf("materialized Size = %d, want %d", snap.Size(), want)
+		}
+	})
 }
 
 // TestKVLazySnapshotCopiedOutOnOverwrite verifies that after a rewind, a write
@@ -139,110 +140,110 @@ func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
 // still reads the pre-overwrite data — both keys and values, across the whole
 // captured range (guarding the Slice+Contiguous copy-out representation).
 func TestKVLazySnapshotCopiedOutOnOverwrite(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		// Fill [0,10) with position-encoded values (keys p, values -p).
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	// Fill [0,10) with position-encoded values (keys p, values -p).
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
-
-	// Lazy snapshot [5,10).
-	snap := c.Snapshot(5).(*kvSnapshot)
-	if snap.keys != nil {
-		t.Fatal("Snapshot should return a lazy snapshot")
-	}
-
-	// Rewind to 5 and overwrite [5,8) with different data.
-	if !c.Restore(nil, 5) {
-		t.Fatal("rewind failed")
-	}
-	nk, nv := distinctKV(100, 3, H, D) // positions 100,101,102
-	c.Update(newKVBatch(5, 3), nk, nv)
-
-	// The overwrite must have copied the lazy snapshot out beforehand, preserving
-	// the pre-overwrite keys and values across the whole [5,10) range.
-	if snap.keys == nil {
-		t.Fatal("lazy snapshot was not copied out before the overwriting write")
-	}
-	mlx.Eval(snap.keys, snap.values)
-	keys := snap.keys.Floats()
-	vals := snap.values.Floats()
-	for l := range snap.toOffset - snap.fromOffset {
-		wantK := float32(snap.fromOffset + l)
-		if got := keys[l*D]; got != wantK {
-			t.Fatalf("snapshot keys[%d] = %v, want %v (pre-overwrite data)", l, got, wantK)
+		// Lazy snapshot [5,10).
+		snap := c.Snapshot(5).(*kvSnapshot)
+		if snap.keys != nil {
+			t.Fatal("Snapshot should return a lazy snapshot")
 		}
-		if got := vals[l*D]; got != -wantK {
-			t.Fatalf("snapshot values[%d] = %v, want %v (pre-overwrite data)", l, got, -wantK)
+
+		// Rewind to 5 and overwrite [5,8) with different data.
+		if !c.Restore(nil, 5) {
+			t.Fatal("rewind failed")
 		}
-	}
-	snap.Close()
+		nk, nv := distinctKV(100, 3, H, D) // positions 100,101,102
+		c.Update(newKVBatch(5, 3), nk, nv)
+
+		// The overwrite must have copied the lazy snapshot out beforehand, preserving
+		// the pre-overwrite keys and values across the whole [5,10) range.
+		if snap.keys == nil {
+			t.Fatal("lazy snapshot was not copied out before the overwriting write")
+		}
+		mlx.Eval(snap.keys, snap.values)
+		keys := snap.keys.Floats()
+		vals := snap.values.Floats()
+		for l := range snap.toOffset - snap.fromOffset {
+			wantK := float32(snap.fromOffset + l)
+			if got := keys[l*D]; got != wantK {
+				t.Fatalf("snapshot keys[%d] = %v, want %v (pre-overwrite data)", l, got, wantK)
+			}
+			if got := vals[l*D]; got != -wantK {
+				t.Fatalf("snapshot values[%d] = %v, want %v (pre-overwrite data)", l, got, -wantK)
+			}
+		}
+		snap.Close()
+	})
 }
 
 // TestKVLazySnapshotCopiedOutOnFree verifies Free copies out outstanding lazy snapshots
 // so they survive after the cache buffer is gone.
 func TestKVLazySnapshotCopiedOutOnFree(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
+		snap := c.Snapshot(5).(*kvSnapshot)
+		c.Free()
 
-	snap := c.Snapshot(5).(*kvSnapshot)
-	c.Free()
-
-	if snap.keys == nil {
-		t.Fatal("lazy snapshot was not copied out on Free")
-	}
-	mlx.Eval(snap.keys)
-	if got := firstKeyAt(snap.keys, 0, D); got != 5 {
-		t.Fatalf("snapshot[0] key = %v, want 5 (data preserved through Free)", got)
-	}
-	snap.Close()
+		if snap.keys == nil {
+			t.Fatal("lazy snapshot was not copied out on Free")
+		}
+		mlx.Eval(snap.keys)
+		if got := firstKeyAt(snap.keys, 0, D); got != 5 {
+			t.Fatalf("snapshot[0] key = %v, want 5 (data preserved through Free)", got)
+		}
+		snap.Close()
+	})
 }
 
 // TestKVLazySnapshotSplitMergeNoCopy verifies Split of a lazy snapshot and Merge of two
 // adjacent lazy snapshots are pure arithmetic — they produce lazy snapshots and allocate
 // nothing — while still tracking the correct offsets and data.
 func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
+		base := settledActiveMemory()
 
-	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)
+		}
 
-	// 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)
-	}
-	merged.Close()
+		if after := settledActiveMemory(); after > base {
+			t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base)
+		}
+		merged.Close()
+	})
 }
 
 // TestKVLazySnapshotSurvivesPathSwitch reproduces the switchToPath sequence that
@@ -251,41 +252,41 @@ func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) {
 // different path (Restore feeds new data via appendKV, overwriting the old
 // leaf's slots). The paged-out snapshot must still hold the original tokens.
 func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		// Active path tokens [0,10): positions 0..9.
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	// Active path tokens [0,10): positions 0..9.
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
+		// Page out the diverging leaf [4,10) as a lazy snapshot (what switchToPath does
+		// before rewinding).
+		leaf := c.Snapshot(4).(*kvSnapshot)
 
-	// Page out the diverging leaf [4,10) as a lazy snapshot (what switchToPath does
-	// before rewinding).
-	leaf := c.Snapshot(4).(*kvSnapshot)
-
-	// Rewind to the common ancestor at 4 (Restore(nil): offset move only).
-	if !c.Restore(nil, 4) {
-		t.Fatal("rewind to ancestor failed")
-	}
-
-	// Page in the new path [4,9): positions 200..204, overwriting the old leaf's
-	// slots. appendKV must copy the leaf lazy snapshot out first.
-	nk, nv := distinctKV(200, 5, H, D)
-	c.Update(newKVBatch(4, 5), nk, nv)
-
-	if leaf.keys == nil {
-		t.Fatal("leaf lazy snapshot was not copied out during page-in")
-	}
-	mlx.Eval(leaf.keys)
-	// leaf covers [4,10): its keys are the original positions 4..9.
-	keys := leaf.keys.Floats()
-	for l := range leaf.toOffset - leaf.fromOffset {
-		if got, want := keys[l*D], float32(leaf.fromOffset+l); got != want {
-			t.Fatalf("paged-out leaf key[%d] = %v, want %v (original path data)", l, got, want)
+		// Rewind to the common ancestor at 4 (Restore(nil): offset move only).
+		if !c.Restore(nil, 4) {
+			t.Fatal("rewind to ancestor failed")
 		}
-	}
-	leaf.Close()
+
+		// Page in the new path [4,9): positions 200..204, overwriting the old leaf's
+		// slots. appendKV must copy the leaf lazy snapshot out first.
+		nk, nv := distinctKV(200, 5, H, D)
+		c.Update(newKVBatch(4, 5), nk, nv)
+
+		if leaf.keys == nil {
+			t.Fatal("leaf lazy snapshot was not copied out during page-in")
+		}
+		mlx.Eval(leaf.keys)
+		// leaf covers [4,10): its keys are the original positions 4..9.
+		keys := leaf.keys.Floats()
+		for l := range leaf.toOffset - leaf.fromOffset {
+			if got, want := keys[l*D], float32(leaf.fromOffset+l); got != want {
+				t.Fatalf("paged-out leaf key[%d] = %v, want %v (original path data)", l, got, want)
+			}
+		}
+		leaf.Close()
+	})
 }
 
 // TestKVRestoreLiveLazySnapshotIsOffsetMove verifies the same-path rewind/rematch
@@ -294,46 +295,46 @@ func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) {
 // the offset without cloning or replaying. This is the switchToPath sequence
 // where a paged-out leaf is restored before any write displaced it.
 func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 4, 8
 
-	const H, D = 4, 8
+		c := NewKVCache()
+		k, v := distinctKV(0, 10, H, D)
+		c.Update(newKVBatch(0, 10), k, v)
 
-	c := NewKVCache()
-	k, v := distinctKV(0, 10, H, D)
-	c.Update(newKVBatch(0, 10), k, v)
-
-	// Page out the leaf [5,10) as a lazy snapshot, then rewind (offset move only).
-	snap := c.Snapshot(5).(*kvSnapshot)
-	if !c.Restore(nil, 5) {
-		t.Fatal("rewind failed")
-	}
-
-	base := settledActiveMemory()
-
-	// 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())
-	}
-	if after := settledActiveMemory(); after > base {
-		t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base)
-	}
-
-	// The buffer still holds the original positions 0..9.
-	st := c.State()
-	mlx.Eval(st[0])
-	keys := st[0].Floats()
-	for l := range 10 {
-		if got := keys[l*D]; got != float32(l) {
-			t.Fatalf("restored key[%d] = %v, want %v", l, got, float32(l))
+		// Page out the leaf [5,10) as a lazy snapshot, then rewind (offset move only).
+		snap := c.Snapshot(5).(*kvSnapshot)
+		if !c.Restore(nil, 5) {
+			t.Fatal("rewind failed")
 		}
-	}
-	snap.Close()
+
+		base := settledActiveMemory()
+
+		// 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())
+		}
+		if after := settledActiveMemory(); after > base {
+			t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base)
+		}
+
+		// The buffer still holds the original positions 0..9.
+		st := c.State()
+		mlx.Eval(st[0])
+		keys := st[0].Floats()
+		for l := range 10 {
+			if got := keys[l*D]; got != float32(l) {
+				t.Fatalf("restored key[%d] = %v, want %v", l, got, float32(l))
+			}
+		}
+		snap.Close()
+	})
 }

x/mlxrunner/cache/recurrent_test.go

@@ -4,6 +4,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/models/nn"
@@ -13,66 +14,68 @@ import (
 // only succeeds when target exactly matches the snapshot's offset. Recurrent
 // state is cumulative, so it can't be rewound or fast-forwarded.
 func TestRecurrentCacheRestoreExactOffset(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRecurrentCache(3, 12, 4, 8, 8)
-	b1 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1)}
-	c.Get(b1, mlx.DTypeFloat16) // lazy-init
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRecurrentCache(3, 12, 4, 8, 8)
+		b1 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1)}
+		c.Get(b1, mlx.DTypeFloat16) // lazy-init
 
-	keep := func() ([]*mlx.Array, []*mlx.Array) {
-		s := c.State()
-		return []*mlx.Array{s[0]}, []*mlx.Array{s[1]}
-	}
+		keep := func() ([]*mlx.Array, []*mlx.Array) {
+			s := c.State()
+			return []*mlx.Array{s[0]}, []*mlx.Array{s[1]}
+		}
 
-	b10 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 10), SeqQueryLens: []int32{10}}
-	cs, ds := keep()
-	c.Put(b10, cs, ds) // advance to 10
+		b10 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 10), SeqQueryLens: []int32{10}}
+		cs, ds := keep()
+		c.Put(b10, cs, ds) // advance to 10
 
-	snap := c.Snapshot(0) // snap.offset == 10
+		snap := c.Snapshot(0) // snap.offset == 10
 
-	b5 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 5), SeqQueryLens: []int32{5}}
-	cs, ds = keep()
-	c.Put(b5, cs, ds) // cache now at 15
+		b5 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 5), SeqQueryLens: []int32{5}}
+		cs, ds = keep()
+		c.Put(b5, cs, ds) // cache now at 15
 
-	// target < snap.offset: fails (can't rewind past snapshot)
-	if c.Restore(snap, 5) {
-		t.Fatal("Restore(snap, 5) should fail — target != snap.offset")
-	}
+		// target < snap.offset: fails (can't rewind past snapshot)
+		if c.Restore(snap, 5) {
+			t.Fatal("Restore(snap, 5) should fail — target != snap.offset")
+		}
 
-	// target > snap.offset: fails (can't advance without feeding tokens)
-	if c.Restore(snap, 15) {
-		t.Fatal("Restore(snap, 15) should fail — target != snap.offset")
-	}
+		// target > snap.offset: fails (can't advance without feeding tokens)
+		if c.Restore(snap, 15) {
+			t.Fatal("Restore(snap, 15) should fail — target != snap.offset")
+		}
 
-	// target == snap.offset: succeeds
-	if !c.Restore(snap, 10) {
-		t.Fatal("Restore(snap, 10) should succeed — target == snap.offset")
-	}
-	if c.Offset() != 10 {
-		t.Fatalf("offset = %d, want 10", c.Offset())
-	}
+		// target == snap.offset: succeeds
+		if !c.Restore(snap, 10) {
+			t.Fatal("Restore(snap, 10) should succeed — target == snap.offset")
+		}
+		if c.Offset() != 10 {
+			t.Fatalf("offset = %d, want 10", c.Offset())
+		}
+	})
 }
 
 func TestRecurrentCacheGetLazyInit(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRecurrentCache(3, 4, 2, 4, 4)
-	b := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
-		SeqOffsets:   []int32{0},
-		SeqQueryLens: []int32{1},
-	}
-	h := c.Get(b, mlx.DTypeBFloat16)
-	if c.Offset() != 0 {
-		t.Fatalf("Get should not advance; got offset %d", c.Offset())
-	}
-	if h.ConvState() == nil || h.DeltaState() == nil {
-		t.Fatal("history should expose conv/delta tensors")
-	}
-	if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 {
-		t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16)
-	}
-	if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 {
-		t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32)
-	}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRecurrentCache(3, 4, 2, 4, 4)
+		b := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
+			SeqOffsets:   []int32{0},
+			SeqQueryLens: []int32{1},
+		}
+		h := c.Get(b, mlx.DTypeBFloat16)
+		if c.Offset() != 0 {
+			t.Fatalf("Get should not advance; got offset %d", c.Offset())
+		}
+		if h.ConvState() == nil || h.DeltaState() == nil {
+			t.Fatal("history should expose conv/delta tensors")
+		}
+		if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 {
+			t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16)
+		}
+		if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 {
+			t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32)
+		}
+	})
 }
 
 // TestRecurrentCachePaddedRoundTrip runs Get → CausalConv1D →
@@ -83,138 +86,140 @@ func TestRecurrentCacheGetLazyInit(t *testing.T) {
 // Pins the recurrent contract: a forward with padding produces the
 // same end-state as a forward with the real-prefix-only input.
 func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
-	skipIfNoMLX(t)
-	const convTail, convDim = 2, 6
-	const numVHeads, headVDim, headKDim = 1, 4, 6
-	const L = 4
-	const qLen = 2
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const convTail, convDim = 2, 6
+		const numVHeads, headVDim, headKDim = 1, 4, 6
+		const L = 4
+		const qLen = 2
 
-	// Distinct values for the real prefix and large junk in the padded
-	// tail so any leak from padded positions is visible.
-	const packedDim = 2*headKDim + numVHeads*headVDim
-	mkPacked := func(seed float32, T int) (packed, ba *mlx.Array) {
-		pv := make([]float32, T*packedDim)
-		bv := make([]float32, T*2*numVHeads)
-		for i := range pv {
-			pv[i] = seed + 0.05*float32(i)
+		// Distinct values for the real prefix and large junk in the padded
+		// tail so any leak from padded positions is visible.
+		const packedDim = 2*headKDim + numVHeads*headVDim
+		mkPacked := func(seed float32, T int) (packed, ba *mlx.Array) {
+			pv := make([]float32, T*packedDim)
+			bv := make([]float32, T*2*numVHeads)
+			for i := range pv {
+				pv[i] = seed + 0.05*float32(i)
+			}
+			for i := range bv {
+				bv[i] = seed - 0.02*float32(i)
+			}
+			return mlx.FromValues(pv, 1, T, packedDim), mlx.FromValues(bv, 1, T, 2*numVHeads)
 		}
-		for i := range bv {
-			bv[i] = seed - 0.02*float32(i)
+		mkPackedPadded := func() (packed, ba *mlx.Array) {
+			pReal, baReal := mkPacked(0.3, qLen)
+			pPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, packedDim), 99)
+			baPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 2*numVHeads), 99)
+			return mlx.Concatenate([]*mlx.Array{pReal, pPad}, 1), mlx.Concatenate([]*mlx.Array{baReal, baPad}, 1)
 		}
-		return mlx.FromValues(pv, 1, T, packedDim), mlx.FromValues(bv, 1, T, 2*numVHeads)
-	}
-	mkPackedPadded := func() (packed, ba *mlx.Array) {
-		pReal, baReal := mkPacked(0.3, qLen)
-		pPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, packedDim), 99)
-		baPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 2*numVHeads), 99)
-		return mlx.Concatenate([]*mlx.Array{pReal, pPad}, 1), mlx.Concatenate([]*mlx.Array{baReal, baPad}, 1)
-	}
-	dtBias := mlx.FromValues([]float32{0.3}, numVHeads)
-	aExp := mlx.FromValues([]float32{0.12}, numVHeads)
+		dtBias := mlx.FromValues([]float32{0.3}, numVHeads)
+		aExp := mlx.FromValues([]float32{0.12}, numVHeads)
 
-	// The conv input dimension must match the cache's convDim.
-	mkConvInput := func(seed float32, T int) *mlx.Array {
-		vals := make([]float32, 1*T*convDim)
-		for i := range vals {
-			vals[i] = seed + 0.05*float32(i)
+		// The conv input dimension must match the cache's convDim.
+		mkConvInput := func(seed float32, T int) *mlx.Array {
+			vals := make([]float32, 1*T*convDim)
+			for i := range vals {
+				vals[i] = seed + 0.05*float32(i)
+			}
+			return mlx.FromValues(vals, 1, T, convDim)
 		}
-		return mlx.FromValues(vals, 1, T, convDim)
-	}
-	mkWeight := func(seed float32) *mlx.Array {
-		vals := make([]float32, convDim*(convTail+1))
-		for i := range vals {
-			vals[i] = seed + 0.1*float32(i)
+		mkWeight := func(seed float32) *mlx.Array {
+			vals := make([]float32, convDim*(convTail+1))
+			for i := range vals {
+				vals[i] = seed + 0.1*float32(i)
+			}
+			return mlx.FromValues(vals, convDim, convTail+1)
 		}
-		return mlx.FromValues(vals, convDim, convTail+1)
-	}
-	weight := mkWeight(0.2)
-	// Build the depthwise causal Conv1d as the model does at load time: the
-	// [C, K] kernel becomes [C, K, 1] and the conv is grouped per channel.
-	conv := nn.NewConv1d(mlx.ExpandDims(weight, 2), nil, 1, 0, 1, convDim)
+		weight := mkWeight(0.2)
+		// Build the depthwise causal Conv1d as the model does at load time: the
+		// [C, K] kernel becomes [C, K, 1] and the conv is grouped per channel.
+		conv := nn.NewConv1d(mlx.ExpandDims(weight, 2), nil, 1, 0, 1, convDim)
 
-	runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) {
-		var convInput *mlx.Array
-		if T == L {
-			realPart := mkConvInput(0.4, qLen)
-			padPart := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, T-qLen, convDim), 99)
-			convInput = mlx.Concatenate([]*mlx.Array{realPart, padPart}, 1)
-		} else {
-			convInput = mkConvInput(0.4, T)
+		runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) {
+			var convInput *mlx.Array
+			if T == L {
+				realPart := mkConvInput(0.4, qLen)
+				padPart := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, T-qLen, convDim), 99)
+				convInput = mlx.Concatenate([]*mlx.Array{realPart, padPart}, 1)
+			} else {
+				convInput = mkConvInput(0.4, T)
+			}
+
+			history := c.Get(b, mlx.DTypeFloat32)
+			_, convStates := nn.CausalConv1D(b, convInput, conv, convTail,
+				nn.WithRecurrentHistory(history))
+
+			var packed, ba *mlx.Array
+			if T == L {
+				packed, ba = mkPackedPadded()
+			} else {
+				packed, ba = mkPacked(0.3, T)
+			}
+			_, deltaStates := nn.GatedDelta(b, packed, ba, dtBias, aExp, nn.WithRecurrentHistory(history))
+
+			c.Put(b, convStates, deltaStates)
+			return convStates[len(convStates)-1], deltaStates[len(deltaStates)-1]
 		}
 
-		history := c.Get(b, mlx.DTypeFloat32)
-		_, convStates := nn.CausalConv1D(b, convInput, conv, convTail,
-			nn.WithRecurrentHistory(history))
-
-		var packed, ba *mlx.Array
-		if T == L {
-			packed, ba = mkPackedPadded()
-		} else {
-			packed, ba = mkPacked(0.3, T)
+		// Padded forward.
+		cPad := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
+		bPad := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, L),
+			SeqOffsets:   []int32{0},
+			SeqQueryLens: []int32{int32(qLen)},
 		}
-		_, deltaStates := nn.GatedDelta(b, packed, ba, dtBias, aExp, nn.WithRecurrentHistory(history))
-
-		c.Put(b, convStates, deltaStates)
-		return convStates[len(convStates)-1], deltaStates[len(deltaStates)-1]
-	}
-
-	// Padded forward.
-	cPad := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
-	bPad := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, L),
-		SeqOffsets:   []int32{0},
-		SeqQueryLens: []int32{int32(qLen)},
-	}
-	nextConvPad, deltaPad := runForward(cPad, bPad, L)
-	mlx.Eval(nextConvPad, deltaPad)
-	if got := cPad.Offset(); got != qLen {
-		t.Fatalf("padded forward: Offset() = %d, want %d (must advance by SeqQueryLens, not L)", got, qLen)
-	}
-
-	// Unpadded reference.
-	cRef := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
-	bRef := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, qLen),
-		SeqOffsets:   []int32{0},
-		SeqQueryLens: []int32{int32(qLen)},
-	}
-	nextConvRef, deltaRef := runForward(cRef, bRef, qLen)
-	mlx.Eval(nextConvRef, deltaRef)
-	if got := cRef.Offset(); got != qLen {
-		t.Fatalf("unpadded forward: Offset() = %d, want %d", got, qLen)
-	}
-
-	gp := nextConvPad.Floats()
-	gr := nextConvRef.Floats()
-	if len(gp) != len(gr) {
-		t.Fatalf("nextConv shape mismatch: padded %d vs unpadded %d", len(gp), len(gr))
-	}
-	for i := range gp {
-		if math.Abs(float64(gp[i]-gr[i])) > 1e-4 {
-			t.Fatalf("nextConv[%d]: padded=%v unpadded=%v (padding leaked into conv state)", i, gp[i], gr[i])
+		nextConvPad, deltaPad := runForward(cPad, bPad, L)
+		mlx.Eval(nextConvPad, deltaPad)
+		if got := cPad.Offset(); got != qLen {
+			t.Fatalf("padded forward: Offset() = %d, want %d (must advance by SeqQueryLens, not L)", got, qLen)
 		}
-	}
 
-	dp := deltaPad.Floats()
-	dr := deltaRef.Floats()
-	if len(dp) != len(dr) {
-		t.Fatalf("delta state shape mismatch: padded %d vs unpadded %d", len(dp), len(dr))
-	}
-	for i := range dp {
-		if math.Abs(float64(dp[i]-dr[i])) > 1e-3 {
-			t.Fatalf("delta state[%d]: padded=%v unpadded=%v (padding leaked into recurrent state)", i, dp[i], dr[i])
+		// Unpadded reference.
+		cRef := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim)
+		bRef := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, qLen),
+			SeqOffsets:   []int32{0},
+			SeqQueryLens: []int32{int32(qLen)},
 		}
-	}
+		nextConvRef, deltaRef := runForward(cRef, bRef, qLen)
+		mlx.Eval(nextConvRef, deltaRef)
+		if got := cRef.Offset(); got != qLen {
+			t.Fatalf("unpadded forward: Offset() = %d, want %d", got, qLen)
+		}
+
+		gp := nextConvPad.Floats()
+		gr := nextConvRef.Floats()
+		if len(gp) != len(gr) {
+			t.Fatalf("nextConv shape mismatch: padded %d vs unpadded %d", len(gp), len(gr))
+		}
+		for i := range gp {
+			if math.Abs(float64(gp[i]-gr[i])) > 1e-4 {
+				t.Fatalf("nextConv[%d]: padded=%v unpadded=%v (padding leaked into conv state)", i, gp[i], gr[i])
+			}
+		}
+
+		dp := deltaPad.Floats()
+		dr := deltaRef.Floats()
+		if len(dp) != len(dr) {
+			t.Fatalf("delta state shape mismatch: padded %d vs unpadded %d", len(dp), len(dr))
+		}
+		for i := range dp {
+			if math.Abs(float64(dp[i]-dr[i])) > 1e-3 {
+				t.Fatalf("delta state[%d]: padded=%v unpadded=%v (padding leaked into recurrent state)", i, dp[i], dr[i])
+			}
+		}
+	})
 }
 
 func TestRecurrentCachePutAdvances(t *testing.T) {
-	skipIfNoMLX(t)
-	c := NewRecurrentCache(3, 4, 2, 4, 4)
-	b := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 2), SeqQueryLens: []int32{2}}
-	newConv := mlx.Zeros(mlx.DTypeFloat16, 1, 3, 4)
-	newDelta := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 4, 4)
-	c.Put(b, []*mlx.Array{newConv}, []*mlx.Array{newDelta})
-	if c.Offset() != 2 {
-		t.Fatalf("cache offset not advanced: %d", c.Offset())
-	}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRecurrentCache(3, 4, 2, 4, 4)
+		b := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 2), SeqQueryLens: []int32{2}}
+		newConv := mlx.Zeros(mlx.DTypeFloat16, 1, 3, 4)
+		newDelta := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 4, 4)
+		c.Put(b, []*mlx.Array{newConv}, []*mlx.Array{newDelta})
+		if c.Offset() != 2 {
+			t.Fatalf("cache offset not advanced: %d", c.Offset())
+		}
+	})
 }

x/mlxrunner/cache/rotating_attention_test.go

@@ -4,6 +4,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/models/nn"
@@ -19,7 +20,6 @@ import (
 // applier's gather composing the caller's logical mask back into
 // storage order) must equal the logical-order reference.
 func TestRotatingKVCacheDecodeParity(t *testing.T) {
-	skipIfNoMLX(t)
 	const H, D = 1, 4
 	const window = 4
 	const totalWrites = 7 // past wrap (window=4); last write is the L=1 decode
@@ -39,46 +39,49 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) {
 		return
 	}
 
-	q := mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D)
-	mlx.Eval(q)
+	var q, kLogical, vLogical, logicalMask *mlx.Array
+	var b *batch.Batch
+	var history *nn.KVHistory
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		q = mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D)
 
-	// Drive the cache: write positions 0..totalWrites-2 as a "history",
-	// then position totalWrites-1 is the actual L=1 decode under test.
-	c := NewRotatingKVCache(window)
-	for pos := range totalWrites - 1 {
-		k, v := perPosKV(pos)
-		c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
-	}
+		// Drive the cache: write positions 0..totalWrites-2 as a "history",
+		// then position totalWrites-1 is the actual L=1 decode under test.
+		c := NewRotatingKVCache(window)
+		for pos := range totalWrites - 1 {
+			k, v := perPosKV(pos)
+			c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
+		}
 
-	finalPos := totalWrites - 1
-	kFinal, vFinal := perPosKV(finalPos)
-	b := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
-		SeqOffsets:   []int32{int32(finalPos)},
-		SeqQueryLens: []int32{1},
-	}
-	history := c.Update(b, kFinal, vFinal)
+		finalPos := totalWrites - 1
+		kFinal, vFinal := perPosKV(finalPos)
+		b = &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
+			SeqOffsets:   []int32{int32(finalPos)},
+			SeqQueryLens: []int32{1},
+		}
+		history = c.Update(b, kFinal, vFinal)
 
-	// Reference: the in-window logical-position-ordered K and V are
-	// the last `window` per-position values (positions
-	// [finalPos-window+1, finalPos]). Build them in that order.
-	startPos := max(finalPos-window+1, 0)
-	logicalKs := make([]*mlx.Array, 0, window)
-	logicalVs := make([]*mlx.Array, 0, window)
-	for pos := startPos; pos <= finalPos; pos++ {
-		kp, vp := perPosKV(pos)
-		logicalKs = append(logicalKs, kp)
-		logicalVs = append(logicalVs, vp)
-	}
-	kLogical := mlx.Concatenate(logicalKs, 2)
-	vLogical := mlx.Concatenate(logicalVs, 2)
+		// Reference: the in-window logical-position-ordered K and V are
+		// the last `window` per-position values (positions
+		// [finalPos-window+1, finalPos]). Build them in that order.
+		startPos := max(finalPos-window+1, 0)
+		logicalKs := make([]*mlx.Array, 0, window)
+		logicalVs := make([]*mlx.Array, 0, window)
+		for pos := startPos; pos <= finalPos; pos++ {
+			kp, vp := perPosKV(pos)
+			logicalKs = append(logicalKs, kp)
+			logicalVs = append(logicalVs, vp)
+		}
+		kLogical = mlx.Concatenate(logicalKs, 2)
+		vLogical = mlx.Concatenate(logicalVs, 2)
 
-	// A logical-order ArrayMask with distinct, non-trivial values per
-	// key column. Picked so each column's contribution to softmax is
-	// distinct — the test fails if the cache's gather permutes the
-	// columns wrong before the kernel sees them.
-	maskVals := []float32{0.1, -0.3, 0.7, -0.2}
-	logicalMask := mlx.FromValues(maskVals, 1, 1, 1, window)
+		// A logical-order ArrayMask with distinct, non-trivial values per
+		// key column. Picked so each column's contribution to softmax is
+		// distinct — the test fails if the cache's gather permutes the
+		// columns wrong before the kernel sees them.
+		logicalMask = mlx.FromValues([]float32{0.1, -0.3, 0.7, -0.2}, 1, 1, 1, window)
+	})
 
 	cases := []struct {
 		name  string
@@ -95,7 +98,7 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			got := nn.ScaledDotProductAttention(b, q, scale,
 				nn.WithKVHistory(history),
 				nn.WithMask(tc.model))
@@ -115,45 +118,57 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) {
 }
 
 func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) {
-	skipIfNoMLX(t)
-	if !mlx.MetalIsAvailable() {
-		t.Skip("MLX Metal not available")
-	}
 	const H, D = 1, 4
 	const window = 4
 	const total = 7
 	const scale = 1.0
 
-	q := mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D)
-	mlx.Eval(q)
-
-	full := NewKVCache()
-	sliding := NewRotatingKVCache(window)
-	for pos := range total {
-		kVals := make([]float32, H*D)
-		vVals := make([]float32, H*D)
-		for i := range kVals {
-			kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i)
-			vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i)
-		}
-		k := mlx.FromValues(kVals, 1, H, 1, D)
-		v := mlx.FromValues(vVals, 1, H, 1, D)
-		full.Update(newKVBatch(full.Offset(), 1), k, v)
-		sliding.Update(newKVBatch(sliding.Offset(), 1), k, v)
-	}
-
-	b := newKVBatch(total-1, 1)
-	cases := []struct {
+	var available bool
+	var q *mlx.Array
+	var b *batch.Batch
+	var cases []struct {
 		name string
 		h    *nn.KVHistory
 		mask nn.AttentionMask
-	}{
-		{name: "full", h: full.View(b), mask: nn.CausalMask()},
-		{name: "sliding", h: sliding.View(b), mask: nn.CausalMask()},
+	}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		available = mlx.MetalIsAvailable()
+		if !available {
+			return
+		}
+
+		q = mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D)
+		full := NewKVCache()
+		sliding := NewRotatingKVCache(window)
+		for pos := range total {
+			kVals := make([]float32, H*D)
+			vVals := make([]float32, H*D)
+			for i := range kVals {
+				kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i)
+				vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i)
+			}
+			k := mlx.FromValues(kVals, 1, H, 1, D)
+			v := mlx.FromValues(vVals, 1, H, 1, D)
+			full.Update(newKVBatch(full.Offset(), 1), k, v)
+			sliding.Update(newKVBatch(sliding.Offset(), 1), k, v)
+		}
+
+		b = newKVBatch(total-1, 1)
+		cases = []struct {
+			name string
+			h    *nn.KVHistory
+			mask nn.AttentionMask
+		}{
+			{name: "full", h: full.View(b), mask: nn.CausalMask()},
+			{name: "sliding", h: sliding.View(b), mask: nn.CausalMask()},
+		}
+	})
+	if !available {
+		t.Skip("MLX Metal not available")
 	}
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			got := nn.ScaledDotProductAttention(b, q, scale, nn.WithKVHistory(tc.h), nn.WithMask(tc.mask))
 			want := mlx.FastScaledDotProductAttention(q, tc.h.K(), tc.h.V(), scale, "", nil)
 
@@ -173,23 +188,26 @@ func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) {
 // matches a reference computed from the same K/V with the model mask
 // and window restriction composed manually.
 func TestRotatingKVCachePrefillParity(t *testing.T) {
-	skipIfNoMLX(t)
 	const H, L, D = 1, 6, 4
 	const window = 4
 	const scale = 1.0
 
-	qVals := make([]float32, 1*H*L*D)
-	kVals := make([]float32, 1*H*L*D)
-	vVals := make([]float32, 1*H*L*D)
-	for i := range qVals {
-		qVals[i] = 0.5 + 0.05*float32(i)
-		kVals[i] = -0.3 + 0.07*float32(i)
-		vVals[i] = 0.3 + 0.03*float32(i)
-	}
-	q := mlx.FromValues(qVals, 1, H, L, D)
-	k := mlx.FromValues(kVals, 1, H, L, D)
-	v := mlx.FromValues(vVals, 1, H, L, D)
-	b := newKVBatch(0, L)
+	var q, k, v *mlx.Array
+	var b *batch.Batch
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		qVals := make([]float32, 1*H*L*D)
+		kVals := make([]float32, 1*H*L*D)
+		vVals := make([]float32, 1*H*L*D)
+		for i := range qVals {
+			qVals[i] = 0.5 + 0.05*float32(i)
+			kVals[i] = -0.3 + 0.07*float32(i)
+			vVals[i] = 0.3 + 0.03*float32(i)
+		}
+		q = mlx.FromValues(qVals, 1, H, L, D)
+		k = mlx.FromValues(kVals, 1, H, L, D)
+		v = mlx.FromValues(vVals, 1, H, L, D)
+		b = newKVBatch(0, L)
+	})
 
 	cases := []struct {
 		name string
@@ -205,7 +223,7 @@ func TestRotatingKVCachePrefillParity(t *testing.T) {
 
 	negInf := float32(math.Inf(-1))
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			c := NewRotatingKVCache(window)
 			history := c.Update(b, k, v)
 
@@ -257,78 +275,79 @@ func TestRotatingKVCachePrefillParity(t *testing.T) {
 // an identical write with no snapshots scheduled. Capture happens after the
 // write via lazy snapshots, so it must not perturb the write itself.
 func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) {
-	skipIfNoMLX(t)
-	const H, D = 1, 4
-	const window = 4
-	const before = 5 // past wrap before the batched write
-	const draft = 4
-	const scale = 1.0
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, D = 1, 4
+		const window = 4
+		const before = 5 // past wrap before the batched write
+		const draft = 4
+		const scale = 1.0
 
-	perPosKV := func(pos int) (k, v *mlx.Array) {
-		kVals := make([]float32, H*D)
-		vVals := make([]float32, H*D)
-		for i := range kVals {
-			kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i)
-			vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i)
-		}
-		return mlx.FromValues(kVals, 1, H, 1, D), mlx.FromValues(vVals, 1, H, 1, D)
-	}
-
-	// Build the batched K/V for offsets [before, before+draft).
-	batchK := make([]*mlx.Array, draft)
-	batchV := make([]*mlx.Array, draft)
-	for i := range draft {
-		batchK[i], batchV[i] = perPosKV(before + i)
-	}
-	kBatch := mlx.Concatenate(batchK, 2)
-	vBatch := mlx.Concatenate(batchV, 2)
-
-	qVals := make([]float32, H*draft*D)
-	for i := range qVals {
-		qVals[i] = 0.5 + 0.05*float32(i)
-	}
-	q := mlx.FromValues(qVals, 1, H, draft, D)
-	b := newKVBatch(before, draft)
-
-	// Run the same write twice: once with snapshots scheduled, once without.
-	run := func(schedule bool) []float32 {
-		c := NewRotatingKVCache(window)
-		for pos := range before {
-			k, v := perPosKV(pos)
-			c.Update(newKVBatch(c.Offset(), 1), k, v)
-		}
-		if schedule {
-			offsets := make([]int, draft)
-			for i := range offsets {
-				offsets[i] = before + i
+		perPosKV := func(pos int) (k, v *mlx.Array) {
+			kVals := make([]float32, H*D)
+			vVals := make([]float32, H*D)
+			for i := range kVals {
+				kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i)
+				vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i)
 			}
-			c.PrepareSnapshots(offsets)
+			return mlx.FromValues(kVals, 1, H, 1, D), mlx.FromValues(vVals, 1, H, 1, D)
 		}
-		history := c.Update(b, kBatch, vBatch)
-		out := nn.ScaledDotProductAttention(b, q, scale,
-			nn.WithKVHistory(history),
-			nn.WithMask(nn.CausalMask()))
-		if schedule {
-			for _, s := range c.TakeSnapshots() {
-				if s != nil {
-					s.Close()
+
+		// Build the batched K/V for offsets [before, before+draft).
+		batchK := make([]*mlx.Array, draft)
+		batchV := make([]*mlx.Array, draft)
+		for i := range draft {
+			batchK[i], batchV[i] = perPosKV(before + i)
+		}
+		kBatch := mlx.Concatenate(batchK, 2)
+		vBatch := mlx.Concatenate(batchV, 2)
+
+		qVals := make([]float32, H*draft*D)
+		for i := range qVals {
+			qVals[i] = 0.5 + 0.05*float32(i)
+		}
+		q := mlx.FromValues(qVals, 1, H, draft, D)
+		b := newKVBatch(before, draft)
+
+		// Run the same write twice: once with snapshots scheduled, once without.
+		run := func(schedule bool) []float32 {
+			c := NewRotatingKVCache(window)
+			for pos := range before {
+				k, v := perPosKV(pos)
+				c.Update(newKVBatch(c.Offset(), 1), k, v)
+			}
+			if schedule {
+				offsets := make([]int, draft)
+				for i := range offsets {
+					offsets[i] = before + i
+				}
+				c.PrepareSnapshots(offsets)
+			}
+			history := c.Update(b, kBatch, vBatch)
+			out := nn.ScaledDotProductAttention(b, q, scale,
+				nn.WithKVHistory(history),
+				nn.WithMask(nn.CausalMask()))
+			if schedule {
+				for _, s := range c.TakeSnapshots() {
+					if s != nil {
+						s.Close()
+					}
 				}
 			}
+			mlx.Eval(out)
+			return out.Floats()
 		}
-		mlx.Eval(out)
-		return out.Floats()
-	}
 
-	withSnap := run(true)
-	noSnap := run(false)
-	if len(withSnap) != len(noSnap) {
-		t.Fatalf("output length %d vs %d", len(withSnap), len(noSnap))
-	}
-	for i := range noSnap {
-		if math.Abs(float64(withSnap[i]-noSnap[i])) > 1e-5 {
-			t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i])
+		withSnap := run(true)
+		noSnap := run(false)
+		if len(withSnap) != len(noSnap) {
+			t.Fatalf("output length %d vs %d", len(withSnap), len(noSnap))
 		}
-	}
+		for i := range noSnap {
+			if math.Abs(float64(withSnap[i]-noSnap[i])) > 1e-5 {
+				t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i])
+			}
+		}
+	})
 }
 
 // TestRotatingKVCacheMLAParity drives a rotating cache with the MLA
@@ -337,35 +356,36 @@ func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) {
 // a manual reference. Pins the cache+MLA integration that
 // glm4_moe_lite uses in production.
 func TestRotatingKVCacheMLAParity(t *testing.T) {
-	skipIfNoMLX(t)
-	const H, L, D, valueDim = 1, 3, 6, 4
-	const scale = 1.0
-	const window = 8 // window >= L so no window restriction
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const H, L, D, valueDim = 1, 3, 6, 4
+		const scale = 1.0
+		const window = 8 // window >= L so no window restriction
 
-	kVals := make([]float32, 1*H*L*D)
-	for i := range kVals {
-		kVals[i] = 0.1 * float32(i+1)
-	}
-	k := mlx.FromValues(kVals, 1, H, L, D)
-	v := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, 0)
-
-	q := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, D)
-	b := newKVBatch(0, L)
-
-	c := NewRotatingKVCache(window)
-	history := c.Update(b, k, v)
-	got := nn.ScaledDotProductAttention(b, q, scale,
-		nn.WithMLAHistory(history, valueDim),
-		nn.WithMask(nn.CausalMask()))
-
-	vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim))
-	want := mlx.FastScaledDotProductAttention(q, k, vRef, scale, "causal", nil)
-
-	mlx.Eval(got, want)
-	gs, ws := got.Floats(), want.Floats()
-	for i := range ws {
-		if math.Abs(float64(gs[i]-ws[i])) > 1e-5 {
-			t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i])
+		kVals := make([]float32, 1*H*L*D)
+		for i := range kVals {
+			kVals[i] = 0.1 * float32(i+1)
 		}
-	}
+		k := mlx.FromValues(kVals, 1, H, L, D)
+		v := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, 0)
+
+		q := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, D)
+		b := newKVBatch(0, L)
+
+		c := NewRotatingKVCache(window)
+		history := c.Update(b, k, v)
+		got := nn.ScaledDotProductAttention(b, q, scale,
+			nn.WithMLAHistory(history, valueDim),
+			nn.WithMask(nn.CausalMask()))
+
+		vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim))
+		want := mlx.FastScaledDotProductAttention(q, k, vRef, scale, "causal", nil)
+
+		mlx.Eval(got, want)
+		gs, ws := got.Floats(), want.Floats()
+		for i := range ws {
+			if math.Abs(float64(gs[i]-ws[i])) > 1e-5 {
+				t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i])
+			}
+		}
+	})
 }

x/mlxrunner/cache/rotating_multiturn_test.go

@@ -3,6 +3,7 @@ package cache
 import (
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
@@ -27,7 +28,7 @@ func multiTokenKV(ids []float32) (*mlx.Array, *mlx.Array) {
 
 // stateIDs returns the ids currently in the cache in slot order (logical
 // after a concat, physical/rotated after a single-token update).
-func stateIDs(t *testing.T, c *RotatingKVCache) []float32 {
+func stateIDs(t *mlxtest.T, c *RotatingKVCache) []float32 {
 	t.Helper()
 	state := c.State()
 	if state == nil {
@@ -75,32 +76,32 @@ func feedSingle(c *RotatingKVCache, id float32) {
 // pre-existing tokens in logical order so the first Q of the new batch
 // has a full sliding window.
 func TestRotatingKVCacheConcatMidRotationPreservesContext(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		c := NewRotatingKVCache(window)
 
-	const window = 4
-	c := NewRotatingKVCache(window)
+		nextID := feedMulti(c, 1, 3)
+		for range 6 {
+			feedSingle(c, nextID)
+			nextID++
+		}
+		if c.Offset() != 9 {
+			t.Fatalf("setup: offset=%d want 9", c.Offset())
+		}
+		if c.idx >= c.maxSize {
+			t.Fatalf("setup: expected mid-rotation idx (<%d), got %d", c.maxSize, c.idx)
+		}
 
-	nextID := feedMulti(c, 1, 3)
-	for range 6 {
-		feedSingle(c, nextID)
-		nextID++
-	}
-	if c.Offset() != 9 {
-		t.Fatalf("setup: offset=%d want 9", c.Offset())
-	}
-	if c.idx >= c.maxSize {
-		t.Fatalf("setup: expected mid-rotation idx (<%d), got %d", c.maxSize, c.idx)
-	}
-
-	feedMulti(c, 10, 2)
-	got := stateIDs(t, c)
-	want := []float32{7, 8, 9, 10, 11}
-	if !equalSlice(got, want) {
-		t.Fatalf("post-concat window=%v want %v", got, want)
-	}
-	if c.Offset() != 11 {
-		t.Fatalf("offset=%d want 11", c.Offset())
-	}
+		feedMulti(c, 10, 2)
+		got := stateIDs(t, c)
+		want := []float32{7, 8, 9, 10, 11}
+		if !equalSlice(got, want) {
+			t.Fatalf("post-concat window=%v want %v", got, want)
+		}
+		if c.Offset() != 11 {
+			t.Fatalf("offset=%d want 11", c.Offset())
+		}
+	})
 }
 
 // TestRotatingKVCacheConcatAlignedInvariant: with an aligned buffer
@@ -108,38 +109,38 @@ func TestRotatingKVCacheConcatMidRotationPreservesContext(t *testing.T) {
 // tokens plus the full new batch. This is the chunked-prefill contract
 // x/mlxrunner/pipeline.go relies on.
 func TestRotatingKVCacheConcatAlignedInvariant(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		c := NewRotatingKVCache(window)
 
-	const window = 4
-	c := NewRotatingKVCache(window)
-
-	// Chunk 1 fills past maxSize, leaving Dim == maxSize aligned.
-	feedMulti(c, 1, 6)
-	// Chunk 2: the buffer is intentionally oversized to (maxSize-1) + L
-	// so the first new Q has its full window in scope for this forward.
-	feedMulti(c, 7, 3)
-	got := stateIDs(t, c)
-	want := []float32{4, 5, 6, 7, 8, 9}
-	if !equalSlice(got, want) {
-		t.Fatalf("post-chunk-2 buffer=%v want %v", got, want)
-	}
-
-	// The next decode trims oversize back to maxSize; order may be
-	// physical (rotated), so check as a set.
-	feedSingle(c, 10)
-	got = stateIDs(t, c)
-	if len(got) != window {
-		t.Fatalf("post-decode Dim=%d want %d", len(got), window)
-	}
-	seen := map[float32]bool{}
-	for _, v := range got {
-		seen[v] = true
-	}
-	for _, w := range []float32{7, 8, 9, 10} {
-		if !seen[w] {
-			t.Fatalf("post-decode window missing %v (got %v)", w, got)
+		// Chunk 1 fills past maxSize, leaving Dim == maxSize aligned.
+		feedMulti(c, 1, 6)
+		// Chunk 2: the buffer is intentionally oversized to (maxSize-1) + L
+		// so the first new Q has its full window in scope for this forward.
+		feedMulti(c, 7, 3)
+		got := stateIDs(t, c)
+		want := []float32{4, 5, 6, 7, 8, 9}
+		if !equalSlice(got, want) {
+			t.Fatalf("post-chunk-2 buffer=%v want %v", got, want)
 		}
-	}
+
+		// The next decode trims oversize back to maxSize; order may be
+		// physical (rotated), so check as a set.
+		feedSingle(c, 10)
+		got = stateIDs(t, c)
+		if len(got) != window {
+			t.Fatalf("post-decode Dim=%d want %d", len(got), window)
+		}
+		seen := map[float32]bool{}
+		for _, v := range got {
+			seen[v] = true
+		}
+		for _, w := range []float32{7, 8, 9, 10} {
+			if !seen[w] {
+				t.Fatalf("post-decode window missing %v (got %v)", w, got)
+			}
+		}
+	})
 }
 
 // TestRotatingKVCacheConcatAfterDecodeGrowsBuffer: update() grows the
@@ -148,20 +149,20 @@ func TestRotatingKVCacheConcatAlignedInvariant(t *testing.T) {
 // has not wrapped. Those trailing slots are zero padding and must not
 // be pulled back into the live window on the next concat.
 func TestRotatingKVCacheConcatAfterDecodeGrowsBuffer(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 512
+		c := NewRotatingKVCache(window)
 
-	const window = 512
-	c := NewRotatingKVCache(window)
+		feedMulti(c, 1, 3)
+		feedSingle(c, 4)
+		feedMulti(c, 5, 3)
 
-	feedMulti(c, 1, 3)
-	feedSingle(c, 4)
-	feedMulti(c, 5, 3)
-
-	got := stateIDs(t, c)
-	want := []float32{1, 2, 3, 4, 5, 6, 7}
-	if !equalSlice(got, want) {
-		t.Fatalf("growing-buffer concat=%v want %v", got, want)
-	}
+		got := stateIDs(t, c)
+		want := []float32{1, 2, 3, 4, 5, 6, 7}
+		if !equalSlice(got, want) {
+			t.Fatalf("growing-buffer concat=%v want %v", got, want)
+		}
+	})
 }
 
 // TestRotatingKVCacheConcatAfterLiveRewind: x/mlxrunner/cache.go calls
@@ -171,53 +172,53 @@ func TestRotatingKVCacheConcatAfterDecodeGrowsBuffer(t *testing.T) {
 // tokens. A subsequent concat must drop those, not treat them as wrapped
 // window content.
 func TestRotatingKVCacheConcatAfterLiveRewind(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 8
+		c := NewRotatingKVCache(window)
 
-	const window = 8
-	c := NewRotatingKVCache(window)
+		// Grow the buffer to exactly maxSize without wrapping.
+		feedMulti(c, 1, 2)
+		for id := float32(3); id <= 8; id++ {
+			feedSingle(c, id)
+		}
+		if c.Offset() != window {
+			t.Fatalf("setup: offset=%d want %d", c.Offset(), window)
+		}
 
-	// Grow the buffer to exactly maxSize without wrapping.
-	feedMulti(c, 1, 2)
-	for id := float32(3); id <= 8; id++ {
-		feedSingle(c, id)
-	}
-	if c.Offset() != window {
-		t.Fatalf("setup: offset=%d want %d", c.Offset(), window)
-	}
+		if !c.Restore(nil, 2) {
+			t.Fatalf("live rewind to 2 failed")
+		}
+		if c.Offset() != 2 {
+			t.Fatalf("post-rewind offset=%d want 2", c.Offset())
+		}
 
-	if !c.Restore(nil, 2) {
-		t.Fatalf("live rewind to 2 failed")
-	}
-	if c.Offset() != 2 {
-		t.Fatalf("post-rewind offset=%d want 2", c.Offset())
-	}
-
-	feedMulti(c, 9, 3)
-	got := stateIDs(t, c)
-	want := []float32{1, 2, 9, 10, 11}
-	if !equalSlice(got, want) {
-		t.Fatalf("post-rewind concat=%v want %v", got, want)
-	}
-	if c.Offset() != 5 {
-		t.Fatalf("offset=%d want 5", c.Offset())
-	}
+		feedMulti(c, 9, 3)
+		got := stateIDs(t, c)
+		want := []float32{1, 2, 9, 10, 11}
+		if !equalSlice(got, want) {
+			t.Fatalf("post-rewind concat=%v want %v", got, want)
+		}
+		if c.Offset() != 5 {
+			t.Fatalf("offset=%d want 5", c.Offset())
+		}
+	})
 }
 
 // TestRotatingKVCacheConcatGrowingBuffer: when oldLen < maxSize the trim
 // formula drops to non-positive and all pre-existing tokens are kept.
 func TestRotatingKVCacheConcatGrowingBuffer(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		c := NewRotatingKVCache(window)
 
-	const window = 4
-	c := NewRotatingKVCache(window)
-
-	feedMulti(c, 1, 2)
-	feedMulti(c, 3, 2)
-	got := stateIDs(t, c)
-	want := []float32{1, 2, 3, 4}
-	if !equalSlice(got, want) {
-		t.Fatalf("growing buffer=%v want %v", got, want)
-	}
+		feedMulti(c, 1, 2)
+		feedMulti(c, 3, 2)
+		got := stateIDs(t, c)
+		want := []float32{1, 2, 3, 4}
+		if !equalSlice(got, want) {
+			t.Fatalf("growing buffer=%v want %v", got, want)
+		}
+	})
 }
 
 // TestRotatingKVCacheRunnerChunkedPrefill mirrors the
@@ -225,114 +226,114 @@ func TestRotatingKVCacheConcatGrowingBuffer(t *testing.T) {
 // repeated L>1 Update() calls on a single cache. Scaled-down proxy for
 // the Gemma 4 26B case (sliding_window=1024, prefillChunkSize=2048).
 func TestRotatingKVCacheRunnerChunkedPrefill(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		c := NewRotatingKVCache(window)
 
-	const window = 4
-	c := NewRotatingKVCache(window)
-
-	feedMulti(c, 1, 8)
-	if c.Offset() != 8 {
-		t.Fatalf("chunk 1: offset=%d want 8", c.Offset())
-	}
-
-	feedMulti(c, 9, 8)
-	got := stateIDs(t, c)
-	want := []float32{6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
-	if !equalSlice(got, want) {
-		t.Fatalf("chunk 2: buffer=%v want %v", got, want)
-	}
-
-	feedMulti(c, 17, 4)
-	got = stateIDs(t, c)
-	want = []float32{14, 15, 16, 17, 18, 19, 20}
-	if !equalSlice(got, want) {
-		t.Fatalf("chunk 3: buffer=%v want %v", got, want)
-	}
-
-	// Decode trims oversize back to maxSize; order may be physical.
-	feedSingle(c, 21)
-	got = stateIDs(t, c)
-	if len(got) != window {
-		t.Fatalf("post-decode Dim=%d want %d", len(got), window)
-	}
-	seen := map[float32]bool{}
-	for _, v := range got {
-		seen[v] = true
-	}
-	for _, w := range []float32{18, 19, 20, 21} {
-		if !seen[w] {
-			t.Fatalf("post-decode window missing %v (got %v)", w, got)
+		feedMulti(c, 1, 8)
+		if c.Offset() != 8 {
+			t.Fatalf("chunk 1: offset=%d want 8", c.Offset())
 		}
-	}
+
+		feedMulti(c, 9, 8)
+		got := stateIDs(t, c)
+		want := []float32{6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
+		if !equalSlice(got, want) {
+			t.Fatalf("chunk 2: buffer=%v want %v", got, want)
+		}
+
+		feedMulti(c, 17, 4)
+		got = stateIDs(t, c)
+		want = []float32{14, 15, 16, 17, 18, 19, 20}
+		if !equalSlice(got, want) {
+			t.Fatalf("chunk 3: buffer=%v want %v", got, want)
+		}
+
+		// Decode trims oversize back to maxSize; order may be physical.
+		feedSingle(c, 21)
+		got = stateIDs(t, c)
+		if len(got) != window {
+			t.Fatalf("post-decode Dim=%d want %d", len(got), window)
+		}
+		seen := map[float32]bool{}
+		for _, v := range got {
+			seen[v] = true
+		}
+		for _, w := range []float32{18, 19, 20, 21} {
+			if !seen[w] {
+				t.Fatalf("post-decode window missing %v (got %v)", w, got)
+			}
+		}
+	})
 }
 
 // TestRotatingKVCacheMultiTurnChatSimulation walks a prefill → decode →
 // prefill sequence and checks that each new prefill retains the last
 // (maxSize-1) pre-existing tokens in logical order.
 func TestRotatingKVCacheMultiTurnChatSimulation(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		c := NewRotatingKVCache(window)
 
-	const window = 4
-	c := NewRotatingKVCache(window)
+		nextID := feedMulti(c, 1, 2)
+		for range 5 {
+			feedSingle(c, nextID)
+			nextID++
+		}
+		if c.Offset() != 7 {
+			t.Fatalf("turn 1: offset=%d want 7", c.Offset())
+		}
 
-	nextID := feedMulti(c, 1, 2)
-	for range 5 {
-		feedSingle(c, nextID)
-		nextID++
-	}
-	if c.Offset() != 7 {
-		t.Fatalf("turn 1: offset=%d want 7", c.Offset())
-	}
+		feedMulti(c, nextID, 3)
+		nextID += 3
+		got := stateIDs(t, c)
+		want := []float32{5, 6, 7, 8, 9, 10}
+		if !equalSlice(got, want) {
+			t.Fatalf("turn 2 prefill buffer=%v want %v", got, want)
+		}
 
-	feedMulti(c, nextID, 3)
-	nextID += 3
-	got := stateIDs(t, c)
-	want := []float32{5, 6, 7, 8, 9, 10}
-	if !equalSlice(got, want) {
-		t.Fatalf("turn 2 prefill buffer=%v want %v", got, want)
-	}
+		for range 4 {
+			feedSingle(c, nextID)
+			nextID++
+		}
+		if c.Offset() != 14 {
+			t.Fatalf("turn 2 decode: offset=%d want 14", c.Offset())
+		}
 
-	for range 4 {
-		feedSingle(c, nextID)
-		nextID++
-	}
-	if c.Offset() != 14 {
-		t.Fatalf("turn 2 decode: offset=%d want 14", c.Offset())
-	}
-
-	feedMulti(c, nextID, 2)
-	got = stateIDs(t, c)
-	want = []float32{12, 13, 14, 15, 16}
-	if !equalSlice(got, want) {
-		t.Fatalf("turn 3 prefill buffer=%v want %v", got, want)
-	}
+		feedMulti(c, nextID, 2)
+		got = stateIDs(t, c)
+		want = []float32{12, 13, 14, 15, 16}
+		if !equalSlice(got, want) {
+			t.Fatalf("turn 3 prefill buffer=%v want %v", got, want)
+		}
+	})
 }
 
 // TestRotatingKVCacheOffsetTracking: Offset() is the monotonic logical
 // token count through any mix of Update() calls — Gemma 4 uses
 // donorEntry.Offset - L for the consumer's RoPE offset.
 func TestRotatingKVCacheOffsetTracking(t *testing.T) {
-	skipIfNoMLX(t)
-
-	c := NewRotatingKVCache(4)
-	nextID := feedMulti(c, 1, 3)
-	if c.Offset() != 3 {
-		t.Fatalf("after prefill 3: offset=%d want 3", c.Offset())
-	}
-	for i := range 5 {
-		feedSingle(c, nextID)
-		nextID++
-		if c.Offset() != 3+i+1 {
-			t.Fatalf("after decode %d: offset=%d want %d", i, c.Offset(), 3+i+1)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewRotatingKVCache(4)
+		nextID := feedMulti(c, 1, 3)
+		if c.Offset() != 3 {
+			t.Fatalf("after prefill 3: offset=%d want 3", c.Offset())
 		}
-	}
-	nextID = feedMulti(c, nextID, 2)
-	if c.Offset() != 10 {
-		t.Fatalf("after turn-2 prefill: offset=%d want 10", c.Offset())
-	}
-	// L > maxSize concat.
-	feedMulti(c, nextID, 7)
-	if c.Offset() != 17 {
-		t.Fatalf("after large prefill: offset=%d want 17", c.Offset())
-	}
+		for i := range 5 {
+			feedSingle(c, nextID)
+			nextID++
+			if c.Offset() != 3+i+1 {
+				t.Fatalf("after decode %d: offset=%d want %d", i, c.Offset(), 3+i+1)
+			}
+		}
+		nextID = feedMulti(c, nextID, 2)
+		if c.Offset() != 10 {
+			t.Fatalf("after turn-2 prefill: offset=%d want 10", c.Offset())
+		}
+		// L > maxSize concat.
+		feedMulti(c, nextID, 7)
+		if c.Offset() != 17 {
+			t.Fatalf("after large prefill: offset=%d want 17", c.Offset())
+		}
+	})
 }

x/mlxrunner/cache/snapshot_capture_test.go

@@ -4,6 +4,7 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
@@ -55,7 +56,7 @@ func fillTagged(c Attention, n int) {
 // windowTags reads c's logical window and returns the per-position tag (the
 // absolute offset each slot holds, recovered as value-1). It uses element 0 of
 // each token, which taggedKV set uniformly. Returns nil if the window is empty.
-func windowTags(t *testing.T, c *RotatingKVCache) []int {
+func windowTags(t *mlxtest.T, c *RotatingKVCache) []int {
 	t.Helper()
 	state := c.State()
 	if len(state) == 0 {
@@ -100,68 +101,68 @@ func wantWindowTags(offset, window int) []int {
 // speculation commit path — a live rewind, since KV is append-only — restores
 // the cache to each accepted offset with the prefix intact.
 func TestKVCachePerTokenSnapshotRestore(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const before = 6
+		const draft = 4
 
-	const before = 6
-	const draft = 4
+		for accepted := 0; accepted <= draft; accepted++ {
+			c := NewKVCache()
+			fillKV(c, before)
 
-	for accepted := 0; accepted <= draft; accepted++ {
-		c := NewKVCache()
-		fillKV(c, before)
+			offsets := make([]int, draft)
+			for i := range offsets {
+				offsets[i] = before + i
+			}
+			c.PrepareSnapshots(offsets)
 
-		offsets := make([]int, draft)
-		for i := range offsets {
-			offsets[i] = before + i
-		}
-		c.PrepareSnapshots(offsets)
+			k, v := batchKV(draft)
+			c.Update(newKVBatch(before, draft), k, v)
+			if c.Offset() != before+draft {
+				t.Fatalf("accepted=%d: offset after write = %d, want %d", accepted, c.Offset(), before+draft)
+			}
 
-		k, v := batchKV(draft)
-		c.Update(newKVBatch(before, draft), k, v)
-		if c.Offset() != before+draft {
-			t.Fatalf("accepted=%d: offset after write = %d, want %d", accepted, c.Offset(), before+draft)
-		}
+			snaps := c.TakeSnapshots()
+			if len(snaps) != draft {
+				t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft)
+			}
 
-		snaps := c.TakeSnapshots()
-		if len(snaps) != draft {
-			t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft)
-		}
+			// Captures are edge-local: offset before is zero-width (nil), and each
+			// later offset holds exactly the single token [before+i-1, before+i).
+			if snaps[0] != nil {
+				t.Fatalf("accepted=%d: snaps[0] = %v, want nil (zero-width base)", accepted, snaps[0])
+			}
+			for i := 1; i < draft; i++ {
+				ks := snaps[i].(*kvSnapshot)
+				if ks.fromOffset != before+i-1 || ks.toOffset != before+i {
+					t.Fatalf("accepted=%d: snaps[%d] = [%d,%d), want [%d,%d)", accepted, i, ks.fromOffset, ks.toOffset, before+i-1, before+i)
+				}
+			}
 
-		// Captures are edge-local: offset before is zero-width (nil), and each
-		// later offset holds exactly the single token [before+i-1, before+i).
-		if snaps[0] != nil {
-			t.Fatalf("accepted=%d: snaps[0] = %v, want nil (zero-width base)", accepted, snaps[0])
-		}
-		for i := 1; i < draft; i++ {
-			ks := snaps[i].(*kvSnapshot)
-			if ks.fromOffset != before+i-1 || ks.toOffset != before+i {
-				t.Fatalf("accepted=%d: snaps[%d] = [%d,%d), want [%d,%d)", accepted, i, ks.fromOffset, ks.toOffset, before+i-1, before+i)
+			// Commit rolls back via a live rewind (Restore(nil)) — the append-only
+			// buffer still holds [0, before+draft), so the edge captures go unused.
+			if accepted < draft {
+				if !c.Restore(nil, before+accepted) {
+					t.Fatalf("accepted=%d: live rewind failed", accepted)
+				}
+			}
+			for _, s := range snaps {
+				if s != nil {
+					s.Close()
+				}
+			}
+
+			want := before + draft
+			if accepted < draft {
+				want = before + accepted
+			}
+			if c.Offset() != want {
+				t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want)
+			}
+			if st := c.State(); len(st) == 2 && st[0].Dim(2) != want {
+				t.Fatalf("accepted=%d: state seq dim = %d, want %d", accepted, st[0].Dim(2), want)
 			}
 		}
-
-		// Commit rolls back via a live rewind (Restore(nil)) — the append-only
-		// buffer still holds [0, before+draft), so the edge captures go unused.
-		if accepted < draft {
-			if !c.Restore(nil, before+accepted) {
-				t.Fatalf("accepted=%d: live rewind failed", accepted)
-			}
-		}
-		for _, s := range snaps {
-			if s != nil {
-				s.Close()
-			}
-		}
-
-		want := before + draft
-		if accepted < draft {
-			want = before + accepted
-		}
-		if c.Offset() != want {
-			t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want)
-		}
-		if st := c.State(); len(st) == 2 && st[0].Dim(2) != want {
-			t.Fatalf("accepted=%d: state seq dim = %d, want %d", accepted, st[0].Dim(2), want)
-		}
-	}
+	})
 }
 
 // TestKVCaptureMergeSplit verifies that two adjacent edge-local captures merge
@@ -169,46 +170,46 @@ func TestKVCachePerTokenSnapshotRestore(t *testing.T) {
 // trie performs on stored snapshots (mergeWithChild / splitNode) — proving they
 // work on captured snapshots, not just freshly-taken ones.
 func TestKVCaptureMergeSplit(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const before = 6
+		c := NewKVCache()
+		fillKV(c, before)
 
-	const before = 6
-	c := NewKVCache()
-	fillKV(c, before)
+		// Schedule two interior offsets so the write captures [before, before+1) and
+		// [before+1, before+2).
+		c.PrepareSnapshots([]int{before + 1, before + 2})
+		k, v := batchKV(3)
+		c.Update(newKVBatch(before, 3), k, v)
 
-	// Schedule two interior offsets so the write captures [before, before+1) and
-	// [before+1, before+2).
-	c.PrepareSnapshots([]int{before + 1, before + 2})
-	k, v := batchKV(3)
-	c.Update(newKVBatch(before, 3), k, v)
+		snaps := c.TakeSnapshots()
+		a := snaps[0].(*kvSnapshot)
+		b := snaps[1].(*kvSnapshot)
+		if a.fromOffset != before || a.toOffset != before+1 {
+			t.Fatalf("snaps[0] = [%d,%d), want [%d,%d)", a.fromOffset, a.toOffset, before, before+1)
+		}
+		if b.fromOffset != before+1 || b.toOffset != before+2 {
+			t.Fatalf("snaps[1] = [%d,%d), want [%d,%d)", b.fromOffset, b.toOffset, before+1, before+2)
+		}
 
-	snaps := c.TakeSnapshots()
-	a := snaps[0].(*kvSnapshot)
-	b := snaps[1].(*kvSnapshot)
-	if a.fromOffset != before || a.toOffset != before+1 {
-		t.Fatalf("snaps[0] = [%d,%d), want [%d,%d)", a.fromOffset, a.toOffset, before, before+1)
-	}
-	if b.fromOffset != before+1 || b.toOffset != before+2 {
-		t.Fatalf("snaps[1] = [%d,%d), want [%d,%d)", b.fromOffset, b.toOffset, before+1, before+2)
-	}
+		// Merge the adjacent edges into [before, before+2).
+		merged := c.Merge(snaps[0], snaps[1]).(*kvSnapshot)
+		if merged.fromOffset != before || merged.toOffset != before+2 {
+			t.Fatalf("merged = [%d,%d), want [%d,%d)", merged.fromOffset, merged.toOffset, before, before+2)
+		}
 
-	// Merge the adjacent edges into [before, before+2).
-	merged := c.Merge(snaps[0], snaps[1]).(*kvSnapshot)
-	if merged.fromOffset != before || merged.toOffset != before+2 {
-		t.Fatalf("merged = [%d,%d), want [%d,%d)", merged.fromOffset, merged.toOffset, before, before+2)
-	}
-
-	// Split back at before+1 and confirm the halves match the originals.
-	p, ch := c.Split(merged, before+1)
-	ps := p.(*kvSnapshot)
-	cs := ch.(*kvSnapshot)
-	if ps.fromOffset != before || ps.toOffset != before+1 {
-		t.Fatalf("split parent = [%d,%d), want [%d,%d)", ps.fromOffset, ps.toOffset, before, before+1)
-	}
-	if cs.fromOffset != before+1 || cs.toOffset != before+2 {
-		t.Fatalf("split child = [%d,%d), want [%d,%d)", cs.fromOffset, cs.toOffset, before+1, before+2)
-	}
-	p.Close()
-	ch.Close()
+		// Split back at before+1 and confirm the halves match the originals.
+		p, ch := c.Split(merged, before+1)
+		ps := p.(*kvSnapshot)
+		cs := ch.(*kvSnapshot)
+		if ps.fromOffset != before || ps.toOffset != before+1 {
+			t.Fatalf("split parent = [%d,%d), want [%d,%d)", ps.fromOffset, ps.toOffset, before, before+1)
+		}
+		if cs.fromOffset != before+1 || cs.toOffset != before+2 {
+			t.Fatalf("split child = [%d,%d), want [%d,%d)", cs.fromOffset, cs.toOffset, before+1, before+2)
+		}
+		p.Close()
+		ch.Close()
+	})
 }
 
 // TestRotatingPerTokenSnapshotRestore exercises per-token capture on a
@@ -219,7 +220,7 @@ func TestKVCaptureMergeSplit(t *testing.T) {
 // observable as wrong positions, not just wrong shapes; the wrapped regime
 // forces concat's linearize branch to run on entry.
 func TestRotatingPerTokenSnapshotRestore(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	cases := []struct {
 		name   string
@@ -233,7 +234,7 @@ func TestRotatingPerTokenSnapshotRestore(t *testing.T) {
 
 	const draft = 4
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			for accepted := 0; accepted <= draft; accepted++ {
 				c := NewRotatingKVCache(tc.window)
 				fillTagged(c, tc.before)
@@ -294,7 +295,7 @@ func TestRotatingPerTokenSnapshotRestore(t *testing.T) {
 // same window data. Covers the restored content, a following decode write, and
 // the snapshot copying out correctly afterward, across wrap regimes.
 func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	cases := []struct {
 		name   string
@@ -308,7 +309,7 @@ func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) {
 
 	const draft, accepted = 4, 2
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			c := NewRotatingKVCache(tc.window)
 			fillTagged(c, tc.before)
 
@@ -386,61 +387,61 @@ func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) {
 // when a following write would destroy the window and copies the snapshot out —
 // the lazy mechanism paying for itself only when the data is about to be lost.
 func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window, before, draft, accepted = 4, 10, 4, 2
 
-	const window, before, draft, accepted = 4, 10, 4, 2
+		c := NewRotatingKVCache(window)
+		fillTagged(c, before)
 
-	c := NewRotatingKVCache(window)
-	fillTagged(c, before)
-
-	offsets := make([]int, draft)
-	for i := range offsets {
-		offsets[i] = before + i
-	}
-	c.PrepareSnapshots(offsets)
-
-	k, v := taggedKV(before, draft)
-	c.Update(newKVBatch(before, draft), k, v)
-
-	snaps := c.TakeSnapshots()
-	snap := snaps[accepted].(*rotatingSnapshot)
-	// Simulate trie ownership: a node sets a materialize hook on attach.
-	fired := 0
-	snap.SetMaterializeHook(func(int) { fired++ })
-
-	if !c.Restore(snap, before+accepted) {
-		t.Fatal("restore failed")
-	}
-	// Re-pointed, not copied out: still lazy, hook unfired.
-	if snap.keys != nil {
-		t.Fatal("hooked snapshot was copied out; expected the re-point fast path")
-	}
-	if fired != 0 {
-		t.Fatalf("materialize hook fired %d times on restore, want 0 (still lazy)", fired)
-	}
-
-	got := windowTags(t, c)
-	want := wantWindowTags(before+accepted, window)
-	if !slices.Equal(got, want) {
-		t.Fatalf("window tags = %v, want %v", got, want)
-	}
-
-	// A following decode write destroys the window's slots, so it copies the
-	// snapshot out — firing the hook exactly once.
-	wk, wv := taggedKV(c.Offset(), 1)
-	c.Update(newKVBatch(c.Offset(), 1), wk, wv)
-	if snap.keys == nil {
-		t.Fatal("following write did not copy out the snapshot")
-	}
-	if fired != 1 {
-		t.Fatalf("materialize hook fired %d times, want 1", fired)
-	}
-
-	for _, s := range snaps {
-		if s != nil {
-			s.Close()
+		offsets := make([]int, draft)
+		for i := range offsets {
+			offsets[i] = before + i
 		}
-	}
+		c.PrepareSnapshots(offsets)
+
+		k, v := taggedKV(before, draft)
+		c.Update(newKVBatch(before, draft), k, v)
+
+		snaps := c.TakeSnapshots()
+		snap := snaps[accepted].(*rotatingSnapshot)
+		// Simulate trie ownership: a node sets a materialize hook on attach.
+		fired := 0
+		snap.SetMaterializeHook(func(int) { fired++ })
+
+		if !c.Restore(snap, before+accepted) {
+			t.Fatal("restore failed")
+		}
+		// Re-pointed, not copied out: still lazy, hook unfired.
+		if snap.keys != nil {
+			t.Fatal("hooked snapshot was copied out; expected the re-point fast path")
+		}
+		if fired != 0 {
+			t.Fatalf("materialize hook fired %d times on restore, want 0 (still lazy)", fired)
+		}
+
+		got := windowTags(t, c)
+		want := wantWindowTags(before+accepted, window)
+		if !slices.Equal(got, want) {
+			t.Fatalf("window tags = %v, want %v", got, want)
+		}
+
+		// A following decode write destroys the window's slots, so it copies the
+		// snapshot out — firing the hook exactly once.
+		wk, wv := taggedKV(c.Offset(), 1)
+		c.Update(newKVBatch(c.Offset(), 1), wk, wv)
+		if snap.keys == nil {
+			t.Fatal("following write did not copy out the snapshot")
+		}
+		if fired != 1 {
+			t.Fatalf("materialize hook fired %d times, want 1", fired)
+		}
+
+		for _, s := range snaps {
+			if s != nil {
+				s.Close()
+			}
+		}
+	})
 }
 
 // TestRotatingSnapshotSingleTokenWrite mirrors the tail of chunked prefill: the
@@ -451,7 +452,7 @@ func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) {
 // not-yet-wrapped and wrapped regimes; the wrapped one exercises the ring-clone
 // fallback.
 func TestRotatingSnapshotSingleTokenWrite(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	cases := []struct {
 		name   string
@@ -463,7 +464,7 @@ func TestRotatingSnapshotSingleTokenWrite(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			c := NewRotatingKVCache(tc.window)
 			fillTagged(c, tc.before)
 
@@ -495,46 +496,46 @@ func TestRotatingSnapshotSingleTokenWrite(t *testing.T) {
 // batched write trims/rewrites the buffer the snapshot's slots lived in. This is
 // the case that forces a lazy snapshot to copy out before the later write destroys it.
 func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 6
+		c := NewRotatingKVCache(window)
 
-	const window = 6
-	c := NewRotatingKVCache(window)
+		// Schedule an offset in the first chunk and one in the second, then write
+		// both chunks before taking — the second chunk's concat trims past the first
+		// snapshot's window.
+		c.PrepareSnapshots([]int{4, 12})
 
-	// Schedule an offset in the first chunk and one in the second, then write
-	// both chunks before taking — the second chunk's concat trims past the first
-	// snapshot's window.
-	c.PrepareSnapshots([]int{4, 12})
+		k1, v1 := taggedKV(0, 8) // chunk 1: [0, 8)
+		c.Update(newKVBatch(0, 8), k1, v1)
+		k2, v2 := taggedKV(8, 8) // chunk 2: [8, 16)
+		c.Update(newKVBatch(8, 8), k2, v2)
 
-	k1, v1 := taggedKV(0, 8) // chunk 1: [0, 8)
-	c.Update(newKVBatch(0, 8), k1, v1)
-	k2, v2 := taggedKV(8, 8) // chunk 2: [8, 16)
-	c.Update(newKVBatch(8, 8), k2, v2)
+		snaps := c.TakeSnapshots()
+		if len(snaps) != 2 || snaps[0] == nil || snaps[1] == nil {
+			t.Fatalf("snapshots not captured: %v", snaps)
+		}
 
-	snaps := c.TakeSnapshots()
-	if len(snaps) != 2 || snaps[0] == nil || snaps[1] == nil {
-		t.Fatalf("snapshots not captured: %v", snaps)
-	}
+		// Restore the first-chunk snapshot (offset 4): its window predates the
+		// second chunk entirely, so the data must have survived the chunk-2 write.
+		if !c.Restore(snaps[0], 4) {
+			t.Fatal("restore to offset 4 failed")
+		}
+		if got, want := windowTags(t, c), wantWindowTags(4, window); !slices.Equal(got, want) {
+			t.Fatalf("offset 4 window tags = %v, want %v", got, want)
+		}
 
-	// Restore the first-chunk snapshot (offset 4): its window predates the
-	// second chunk entirely, so the data must have survived the chunk-2 write.
-	if !c.Restore(snaps[0], 4) {
-		t.Fatal("restore to offset 4 failed")
-	}
-	if got, want := windowTags(t, c), wantWindowTags(4, window); !slices.Equal(got, want) {
-		t.Fatalf("offset 4 window tags = %v, want %v", got, want)
-	}
-
-	// Restore the second-chunk snapshot (offset 12) on a fresh cache.
-	c2 := NewRotatingKVCache(window)
-	if !c2.Restore(snaps[1], 12) {
-		t.Fatal("restore to offset 12 failed")
-	}
-	if got, want := windowTags(t, c2), wantWindowTags(12, window); !slices.Equal(got, want) {
-		t.Fatalf("offset 12 window tags = %v, want %v", got, want)
-	}
-	for _, s := range snaps {
-		s.Close()
-	}
+		// Restore the second-chunk snapshot (offset 12) on a fresh cache.
+		c2 := NewRotatingKVCache(window)
+		if !c2.Restore(snaps[1], 12) {
+			t.Fatal("restore to offset 12 failed")
+		}
+		if got, want := windowTags(t, c2), wantWindowTags(12, window); !slices.Equal(got, want) {
+			t.Fatalf("offset 12 window tags = %v, want %v", got, want)
+		}
+		for _, s := range snaps {
+			s.Close()
+		}
+	})
 }
 
 // TestRotatingLazySnapshotSizeZeroUntilMaterialized verifies the speculation
@@ -544,211 +545,211 @@ func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) {
 // lazy snapshot reports Size() == 0 until a destructive write copies it out, at
 // which point the materialize hook fires with the newly-allocated bytes.
 func TestRotatingLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const window = 4
+		const before = 10 // wrapped
+		const draft = 4
+		c := NewRotatingKVCache(window)
+		fillTagged(c, before)
 
-	const window = 4
-	const before = 10 // wrapped
-	const draft = 4
-	c := NewRotatingKVCache(window)
-	fillTagged(c, before)
+		offsets := make([]int, draft)
+		for i := range offsets {
+			offsets[i] = before + i
+		}
+		c.PrepareSnapshots(offsets)
 
-	offsets := make([]int, draft)
-	for i := range offsets {
-		offsets[i] = before + i
-	}
-	c.PrepareSnapshots(offsets)
+		k, v := taggedKV(before, draft)
+		c.Update(newKVBatch(before, draft), k, v)
 
-	k, v := taggedKV(before, draft)
-	c.Update(newKVBatch(before, draft), k, v)
+		snaps := c.TakeSnapshots()
+		defer func() {
+			for _, s := range snaps {
+				if s != nil {
+					s.Close()
+				}
+			}
+		}()
 
-	snaps := c.TakeSnapshots()
-	defer func() {
-		for _, s := range snaps {
-			if s != nil {
-				s.Close()
+		// Interior captures (offsets after the start boundary) stay lazy: no write
+		// destroyed their slots, so keys is still nil and the issuing cache is live.
+		for i := 1; i < draft; i++ {
+			rs := snaps[i].(*rotatingSnapshot)
+			if rs.keys != nil {
+				t.Fatalf("snaps[%d] copied out (keys != nil); expected a live lazy snapshot", i)
+			}
+			if rs.cache == nil {
+				t.Fatalf("snaps[%d] has no issuing cache; expected a live lazy snapshot", i)
 			}
 		}
-	}()
 
-	// Interior captures (offsets after the start boundary) stay lazy: no write
-	// destroyed their slots, so keys is still nil and the issuing cache is live.
-	for i := 1; i < draft; i++ {
-		rs := snaps[i].(*rotatingSnapshot)
-		if rs.keys != nil {
-			t.Fatalf("snaps[%d] copied out (keys != nil); expected a live lazy snapshot", i)
+		lazy := snaps[1].(*rotatingSnapshot)
+		if lazy.Size() != 0 {
+			t.Fatalf("lazy rotating snapshot Size = %d, want 0", lazy.Size())
 		}
-		if rs.cache == nil {
-			t.Fatalf("snaps[%d] has no issuing cache; expected a live lazy snapshot", i)
+
+		var hookDelta int
+		lazy.SetMaterializeHook(func(delta int) { hookDelta = delta })
+
+		// Free the cache to force every outstanding lazy snapshot to copy out.
+		c.Free()
+
+		if lazy.keys == nil {
+			t.Fatal("Free did not materialize the lazy snapshot")
 		}
-	}
-
-	lazy := snaps[1].(*rotatingSnapshot)
-	if lazy.Size() != 0 {
-		t.Fatalf("lazy rotating snapshot Size = %d, want 0", lazy.Size())
-	}
-
-	var hookDelta int
-	lazy.SetMaterializeHook(func(delta int) { hookDelta = delta })
-
-	// Free the cache to force every outstanding lazy snapshot to copy out.
-	c.Free()
-
-	if lazy.keys == nil {
-		t.Fatal("Free did not materialize the lazy snapshot")
-	}
-	want := lazy.keys.NumBytes() + lazy.values.NumBytes()
-	if hookDelta != want {
-		t.Fatalf("hook fired with delta %d, want %d", hookDelta, want)
-	}
-	if lazy.Size() != want {
-		t.Fatalf("materialized Size = %d, want %d", lazy.Size(), want)
-	}
+		want := lazy.keys.NumBytes() + lazy.values.NumBytes()
+		if hookDelta != want {
+			t.Fatalf("hook fired with delta %d, want %d", hookDelta, want)
+		}
+		if lazy.Size() != want {
+			t.Fatalf("materialized Size = %d, want %d", lazy.Size(), want)
+		}
+	})
 }
 
 // TestPerTokenSnapshotPersistsAcrossWrites verifies that scheduled offsets
 // survive multiple writes until TakeSnapshots — the property prefill would rely
 // on to snapshot interior offsets without splitting its forward.
 func TestPerTokenSnapshotPersistsAcrossWrites(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
+		fillKV(c, 2)
 
-	c := NewKVCache()
-	fillKV(c, 2)
+		// Schedule offsets that span two separate writes. Offset 2 equals the
+		// schedule-time position, so it captures a zero-width range (nil); the rest
+		// are edge-local.
+		c.PrepareSnapshots([]int{2, 3, 5})
 
-	// Schedule offsets that span two separate writes. Offset 2 equals the
-	// schedule-time position, so it captures a zero-width range (nil); the rest
-	// are edge-local.
-	c.PrepareSnapshots([]int{2, 3, 5})
+		k1, v1 := batchKV(2) // reaches offsets 2,3
+		c.Update(newKVBatch(2, 2), k1, v1)
+		k2, v2 := batchKV(2) // reaches offsets 4,5
+		c.Update(newKVBatch(4, 2), k2, v2)
 
-	k1, v1 := batchKV(2) // reaches offsets 2,3
-	c.Update(newKVBatch(2, 2), k1, v1)
-	k2, v2 := batchKV(2) // reaches offsets 4,5
-	c.Update(newKVBatch(4, 2), k2, v2)
-
-	snaps := c.TakeSnapshots()
-	if len(snaps) != 3 {
-		t.Fatalf("got %d snapshots, want 3", len(snaps))
-	}
-	if snaps[0] != nil {
-		t.Fatalf("snaps[0] = %v, want nil (zero-width base)", snaps[0])
-	}
-	if s := snaps[1].(*kvSnapshot); s.fromOffset != 2 || s.toOffset != 3 {
-		t.Fatalf("snaps[1] = [%d,%d), want [2,3)", s.fromOffset, s.toOffset)
-	}
-	// Offset 5 was scheduled across the second write; its edge starts at the
-	// previous scheduled offset (3), confirming the base cursor only advances
-	// on capture so the snapshot range matches the trie edge between scheduled
-	// offsets — write boundaries between captures must not move it.
-	if s := snaps[2].(*kvSnapshot); s.fromOffset != 3 || s.toOffset != 5 {
-		t.Fatalf("snaps[2] = [%d,%d), want [3,5)", s.fromOffset, s.toOffset)
-	}
-
-	// Restore from the [2,3) edge snapshot to offset 3.
-	if !c.Restore(snaps[1], 3) {
-		t.Fatal("restore to offset 3 failed")
-	}
-	if c.Offset() != 3 {
-		t.Fatalf("offset = %d, want 3", c.Offset())
-	}
-	for _, s := range snaps {
-		if s != nil {
-			s.Close()
+		snaps := c.TakeSnapshots()
+		if len(snaps) != 3 {
+			t.Fatalf("got %d snapshots, want 3", len(snaps))
 		}
-	}
+		if snaps[0] != nil {
+			t.Fatalf("snaps[0] = %v, want nil (zero-width base)", snaps[0])
+		}
+		if s := snaps[1].(*kvSnapshot); s.fromOffset != 2 || s.toOffset != 3 {
+			t.Fatalf("snaps[1] = [%d,%d), want [2,3)", s.fromOffset, s.toOffset)
+		}
+		// Offset 5 was scheduled across the second write; its edge starts at the
+		// previous scheduled offset (3), confirming the base cursor only advances
+		// on capture so the snapshot range matches the trie edge between scheduled
+		// offsets — write boundaries between captures must not move it.
+		if s := snaps[2].(*kvSnapshot); s.fromOffset != 3 || s.toOffset != 5 {
+			t.Fatalf("snaps[2] = [%d,%d), want [3,5)", s.fromOffset, s.toOffset)
+		}
+
+		// Restore from the [2,3) edge snapshot to offset 3.
+		if !c.Restore(snaps[1], 3) {
+			t.Fatal("restore to offset 3 failed")
+		}
+		if c.Offset() != 3 {
+			t.Fatalf("offset = %d, want 3", c.Offset())
+		}
+		for _, s := range snaps {
+			if s != nil {
+				s.Close()
+			}
+		}
+	})
 }
 
 // TestRecurrentSnapshotSplitsAndSegmentedCapture verifies that SnapshotSplits
 // reports the interior scheduled offsets and PutSegmented captures them from the
 // per-boundary states, so each accepted count restores to a distinct state.
 func TestRecurrentSnapshotSplitsAndSegmentedCapture(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const convTail, convDim, nv, vd, kd = 3, 8, 2, 4, 4
+		c := NewRecurrentCache(convTail, convDim, nv, vd, kd)
+		c.Get(newKVBatch(0, 1), mlx.DTypeFloat16)
+		// Advance to offset 5 so the speculative forward starts there.
+		c.Put(newKVBatch(0, 5),
+			[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim)},
+			[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd)})
 
-	const convTail, convDim, nv, vd, kd = 3, 8, 2, 4, 4
-	c := NewRecurrentCache(convTail, convDim, nv, vd, kd)
-	c.Get(newKVBatch(0, 1), mlx.DTypeFloat16)
-	// Advance to offset 5 so the speculative forward starts there.
-	c.Put(newKVBatch(0, 5),
-		[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim)},
-		[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd)})
+		const before, draft = 5, 4
+		offsets := []int{before, before + 1, before + 2, before + 3}
+		c.PrepareSnapshots(offsets)
 
-	const before, draft = 5, 4
-	offsets := []int{before, before + 1, before + 2, before + 3}
-	c.PrepareSnapshots(offsets)
-
-	splits := c.SnapshotSplits(draft)
-	want := []int{1, 2, 3}
-	if len(splits) != len(want) {
-		t.Fatalf("SnapshotSplits = %v, want %v", splits, want)
-	}
-	for i := range want {
-		if splits[i] != want[i] {
+		splits := c.SnapshotSplits(draft)
+		want := []int{1, 2, 3}
+		if len(splits) != len(want) {
 			t.Fatalf("SnapshotSplits = %v, want %v", splits, want)
 		}
-	}
-
-	// Distinct per-boundary states (3 interior splits + the end) so restore
-	// targets are distinguishable.
-	mkConv := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim), s) }
-	mkDelta := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd), s) }
-	convStates := []*mlx.Array{mkConv(1), mkConv(2), mkConv(3), mkConv(4)}
-	deltaStates := []*mlx.Array{mkDelta(1), mkDelta(2), mkDelta(3), mkDelta(4)}
-
-	c.Put(newKVBatch(before, draft), convStates, deltaStates)
-	if c.Offset() != before+draft {
-		t.Fatalf("offset after segmented put = %d, want %d", c.Offset(), before+draft)
-	}
-
-	snaps := c.TakeSnapshots()
-	if len(snaps) != draft {
-		t.Fatalf("got %d snapshots, want %d", len(snaps), draft)
-	}
-	for i, s := range snaps {
-		if s == nil {
-			t.Fatalf("snapshot %d not captured", i)
+		for i := range want {
+			if splits[i] != want[i] {
+				t.Fatalf("SnapshotSplits = %v, want %v", splits, want)
+			}
 		}
-	}
 
-	// Full accept (no restore): the live state is the committed end boundary
-	// (value 4) at offset before+draft.
-	st := c.State()
-	mlx.Eval(st[1])
-	if got := st[1].Floats()[0]; got != 4 {
-		t.Fatalf("full-accept delta state = %v, want end boundary value 4", got)
-	}
+		// Distinct per-boundary states (3 interior splits + the end) so restore
+		// targets are distinguishable.
+		mkConv := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim), s) }
+		mkDelta := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd), s) }
+		convStates := []*mlx.Array{mkConv(1), mkConv(2), mkConv(3), mkConv(4)}
+		deltaStates := []*mlx.Array{mkDelta(1), mkDelta(2), mkDelta(3), mkDelta(4)}
 
-	// Each partial accept restores to offset before+accepted and must recover the
-	// distinct boundary state captured there: snaps[0] is the pre-forward state
-	// (value 0); snaps[i>=1] is the interior split boundary (value i). Recurrent
-	// snapshots are self-contained, so restores need not run in order.
-	for accepted := range draft {
-		if !c.Restore(snaps[accepted], before+accepted) {
-			t.Fatalf("accepted=%d: restore to before+%d failed", accepted, accepted)
+		c.Put(newKVBatch(before, draft), convStates, deltaStates)
+		if c.Offset() != before+draft {
+			t.Fatalf("offset after segmented put = %d, want %d", c.Offset(), before+draft)
 		}
-		if c.Offset() != before+accepted {
-			t.Fatalf("accepted=%d: offset after restore = %d, want %d", accepted, c.Offset(), before+accepted)
+
+		snaps := c.TakeSnapshots()
+		if len(snaps) != draft {
+			t.Fatalf("got %d snapshots, want %d", len(snaps), draft)
 		}
+		for i, s := range snaps {
+			if s == nil {
+				t.Fatalf("snapshot %d not captured", i)
+			}
+		}
+
+		// Full accept (no restore): the live state is the committed end boundary
+		// (value 4) at offset before+draft.
 		st := c.State()
 		mlx.Eval(st[1])
-		if got := st[1].Floats()[0]; got != float32(accepted) {
-			t.Fatalf("accepted=%d: restored delta state = %v, want boundary value %d", accepted, got, accepted)
+		if got := st[1].Floats()[0]; got != 4 {
+			t.Fatalf("full-accept delta state = %v, want end boundary value 4", got)
 		}
-	}
-	for _, s := range snaps {
-		s.Close()
-	}
+
+		// Each partial accept restores to offset before+accepted and must recover the
+		// distinct boundary state captured there: snaps[0] is the pre-forward state
+		// (value 0); snaps[i>=1] is the interior split boundary (value i). Recurrent
+		// snapshots are self-contained, so restores need not run in order.
+		for accepted := range draft {
+			if !c.Restore(snaps[accepted], before+accepted) {
+				t.Fatalf("accepted=%d: restore to before+%d failed", accepted, accepted)
+			}
+			if c.Offset() != before+accepted {
+				t.Fatalf("accepted=%d: offset after restore = %d, want %d", accepted, c.Offset(), before+accepted)
+			}
+			st := c.State()
+			mlx.Eval(st[1])
+			if got := st[1].Floats()[0]; got != float32(accepted) {
+				t.Fatalf("accepted=%d: restored delta state = %v, want boundary value %d", accepted, got, accepted)
+			}
+		}
+		for _, s := range snaps {
+			s.Close()
+		}
+	})
 }
 
 // TestPrepareSnapshotsPastOffsetPanics verifies scheduling an already-passed offset
 // is rejected.
 func TestPrepareSnapshotsPastOffsetPanics(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := NewKVCache()
+		fillKV(c, 5)
 
-	c := NewKVCache()
-	fillKV(c, 5)
-
-	defer func() {
-		if recover() == nil {
-			t.Fatal("expected panic for already-passed offset")
-		}
-	}()
-	c.PrepareSnapshots([]int{3})
+		defer func() {
+			if recover() == nil {
+				t.Fatal("expected panic for already-passed offset")
+			}
+		}()
+		c.PrepareSnapshots([]int{3})
+	})
 }

x/mlxrunner/dflash_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/cache"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
@@ -92,7 +93,7 @@ var _ base.BlockDraft = (*fakeBlockDraft)(nil)
 // newBlockTestSession wires a runner around a fakeBlockDraft and opens one
 // request's drafting session, returning the concrete session for
 // internal-state assertions.
-func newBlockTestSession(t *testing.T, predict map[int32]int32, blockSize int) (*Runner, *fakeBlockDraft, *dflashDraftSession, []cache.Cache) {
+func newBlockTestSession(t *mlxtest.T, predict map[int32]int32, blockSize int) (*Runner, *fakeBlockDraft, *dflashDraftSession, []cache.Cache) {
 	t.Helper()
 	r := mtpTestRunner(t, predict, []int32{7}, sampler.Options{})
 	caches, _ := newMTPTestCaches(2) // caches[0] target, caches[1] draft context
@@ -108,236 +109,243 @@ func draftTokensOf(caches []cache.Cache) []int32 {
 }
 
 func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) {
-	skipIfNoMLX(t)
-	_, draft, session, caches := newBlockTestSession(t, nil, 4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		_, draft, session, caches := newBlockTestSession(t, nil, 4)
 
-	// One prefill-sized run at the flush cap writes through immediately in a
-	// single context-only Draft call.
-	n := dflashPendingFlushTokens
-	ids := make([]int32, n)
-	for i := range ids {
-		ids[i] = int32(i % mtpTestVocab)
-	}
-	session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0, nil)
-	if got := len(draft.calls); got != 1 {
-		t.Fatalf("draft calls after cap-sized run = %d, want 1", got)
-	}
-	if got := caches[1].Offset(); got != n {
-		t.Fatalf("draft cache offset = %d, want %d", got, n)
-	}
+		// One prefill-sized run at the flush cap writes through immediately in a
+		// single context-only Draft call.
+		n := dflashPendingFlushTokens
+		ids := make([]int32, n)
+		for i := range ids {
+			ids[i] = int32(i % mtpTestVocab)
+		}
+		session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0, nil)
+		if got := len(draft.calls); got != 1 {
+			t.Fatalf("draft calls after cap-sized run = %d, want 1", got)
+		}
+		if got := caches[1].Offset(); got != n {
+			t.Fatalf("draft cache offset = %d, want %d", got, n)
+		}
 
-	// A run below the cap only buffers; settle writes it through, skipping
-	// the leading rows the flush already covered.
-	tail := []int32{1, 2, 3}
-	session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1, nil)
-	if got := len(draft.calls); got != 1 {
-		t.Fatalf("draft calls after buffered run = %d, want 1 (buffered)", got)
-	}
-	session.settle(nil)
-	want := blockCall{offset: int32(n), ctx: []int32{2, 3}}
-	if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
-		t.Fatalf("settle flush = %+v, want %+v", got, want)
-	}
-	if got := caches[1].Offset(); got != n+2 {
-		t.Fatalf("draft cache offset = %d, want %d (level with reports)", got, n+2)
-	}
+		// A run below the cap only buffers; settle writes it through, skipping
+		// the leading rows the flush already covered.
+		tail := []int32{1, 2, 3}
+		session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1, nil)
+		if got := len(draft.calls); got != 1 {
+			t.Fatalf("draft calls after buffered run = %d, want 1 (buffered)", got)
+		}
+		session.settle(nil)
+		want := blockCall{offset: int32(n), ctx: []int32{2, 3}}
+		if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
+			t.Fatalf("settle flush = %+v, want %+v", got, want)
+		}
+		if got := caches[1].Offset(); got != n+2 {
+			t.Fatalf("draft cache offset = %d, want %d (level with reports)", got, n+2)
+		}
+	})
 }
 
 func TestDFlashCommittedGapPanics(t *testing.T) {
-	skipIfNoMLX(t)
-	_, _, session, _ := newBlockTestSession(t, nil, 4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		_, _, session, _ := newBlockTestSession(t, nil, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
-	defer func() {
-		if recover() == nil {
-			t.Fatalf("committed run past the frontier did not panic")
-		}
-	}()
-	// The frontier is at slot 1; a run starting at 3 leaves slot 1..2 unfed.
-	session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3, nil)
+		session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
+		defer func() {
+			if recover() == nil {
+				t.Fatalf("committed run past the frontier did not panic")
+			}
+		}()
+		// The frontier is at slot 1; a run starting at 3 leaves slot 1..2 unfed.
+		session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3, nil)
+	})
 }
 
 func TestDFlashRestoredPrefixResumes(t *testing.T) {
-	skipIfNoMLX(t)
-	r := mtpTestRunner(t, nil, []int32{7}, sampler.Options{})
-	caches, _ := newMTPTestCaches(2)
-	draft := &fakeBlockDraft{blockSize: 4, maskToken: 6, draftCaches: caches[1:]}
-	r.cache.caches = caches
-	r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		r := mtpTestRunner(t, nil, []int32{7}, sampler.Options{})
+		caches, _ := newMTPTestCaches(2)
+		draft := &fakeBlockDraft{blockSize: 4, maskToken: 6, draftCaches: caches[1:]}
+		r.cache.caches = caches
+		r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
 
-	// A restored prefix arrives with the draft caches already written.
-	restored := []int32{1, 2, 3, 4, 5}
-	caches[1].(*fakeRewindableCache).feed(restored)
-	session := r.spec.drafter.open(nil).(*dflashDraftSession)
-	if session.ctxOffset != len(restored) {
-		t.Fatalf("ctxOffset = %d, want %d (synced to restored offset)", session.ctxOffset, len(restored))
-	}
+		// A restored prefix arrives with the draft caches already written.
+		restored := []int32{1, 2, 3, 4, 5}
+		caches[1].(*fakeRewindableCache).feed(restored)
+		session := r.spec.drafter.open(nil).(*dflashDraftSession)
+		if session.ctxOffset != len(restored) {
+			t.Fatalf("ctxOffset = %d, want %d (synced to restored offset)", session.ctxOffset, len(restored))
+		}
 
-	// The resumed prefill's run overlaps the restore point; only the rows
-	// past the frontier are buffered and written.
-	run := []int32{2, 3, 0, 1}
-	session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3, nil)
-	session.settle(nil)
-	want := blockCall{offset: 5, ctx: []int32{0, 1}}
-	if got := draft.calls[0]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
-		t.Fatalf("resume flush = %+v, want %+v", got, want)
-	}
-	if got, wantTok := draftTokensOf(caches), append(restored, 0, 1); !slices.Equal(got, wantTok) {
-		t.Fatalf("draft cache = %v, want %v", got, wantTok)
-	}
+		// The resumed prefill's run overlaps the restore point; only the rows
+		// past the frontier are buffered and written.
+		run := []int32{2, 3, 0, 1}
+		session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3, nil)
+		session.settle(nil)
+		want := blockCall{offset: 5, ctx: []int32{0, 1}}
+		if got := draft.calls[0]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
+			t.Fatalf("resume flush = %+v, want %+v", got, want)
+		}
+		if got, wantTok := draftTokensOf(caches), append(restored, 0, 1); !slices.Equal(got, wantTok) {
+			t.Fatalf("draft cache = %v, want %v", got, wantTok)
+		}
+	})
 }
 
 func TestDFlashProposeBounds(t *testing.T) {
-	skipIfNoMLX(t)
-	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5}
-	_, draft, session, _ := newBlockTestSession(t, predict, 4)
-	current := mlx.FromValues([]int32{1}, 1)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5}
+		_, draft, session, _ := newBlockTestSession(t, predict, 4)
+		current := mlx.FromValues([]int32{1}, 1)
 
-	// Nothing committed yet: no context to draft from.
-	if session.propose(current, 4) != nil {
-		t.Fatalf("propose with no context did not decline")
-	}
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
-	if session.propose(current, 0) != nil {
-		t.Fatalf("propose with no budget did not decline")
-	}
+		// Nothing committed yet: no context to draft from.
+		if session.propose(current, 4) != nil {
+			t.Fatalf("propose with no context did not decline")
+		}
+		session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
+		if session.propose(current, 0) != nil {
+			t.Fatalf("propose with no budget did not decline")
+		}
 
-	// The block caps the draft at blockSize-1 mask rows regardless of budget.
-	cand := session.propose(current, 10)
-	if cand == nil {
-		t.Fatalf("propose declined with context and budget")
-	}
-	mlx.Eval(cand.tokens)
-	if got := cand.tokens.Ints(); !slices.Equal(got, []int32{2, 3, 4}) {
-		t.Fatalf("draft tokens = %v, want [2 3 4]", got)
-	}
-	if got, want := draft.calls[0].block, []int32{1, 6, 6, 6}; !slices.Equal(got, want) {
-		t.Fatalf("block = %v, want %v (anchor plus blockSize-1 masks)", got, want)
-	}
+		// The block caps the draft at blockSize-1 mask rows regardless of budget.
+		cand := session.propose(current, 10)
+		if cand == nil {
+			t.Fatalf("propose declined with context and budget")
+		}
+		mlx.Eval(cand.tokens)
+		if got := cand.tokens.Ints(); !slices.Equal(got, []int32{2, 3, 4}) {
+			t.Fatalf("draft tokens = %v, want [2 3 4]", got)
+		}
+		if got, want := draft.calls[0].block, []int32{1, 6, 6, 6}; !slices.Equal(got, want) {
+			t.Fatalf("block = %v, want %v (anchor plus blockSize-1 masks)", got, want)
+		}
+	})
 }
 
 func TestDFlashBlockRewoundBeforeContextWrites(t *testing.T) {
-	skipIfNoMLX(t)
-	predict := map[int32]int32{2: 3, 3: 4, 4: 5}
-	_, draft, session, caches := newBlockTestSession(t, predict, 4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		predict := map[int32]int32{2: 3, 3: 4, 4: 5}
+		_, draft, session, caches := newBlockTestSession(t, predict, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
-	if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
-		t.Fatalf("propose declined")
-	}
-	// The proposal's block sits in the caches until the next write.
-	if got, want := draftTokensOf(caches), []int32{1, 2, 6, 6, 6}; !slices.Equal(got, want) {
-		t.Fatalf("draft cache after propose = %v, want %v", got, want)
-	}
+		session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
+		if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
+			t.Fatalf("propose declined")
+		}
+		// The proposal's block sits in the caches until the next write.
+		if got, want := draftTokensOf(caches), []int32{1, 2, 6, 6, 6}; !slices.Equal(got, want) {
+			t.Fatalf("draft cache after propose = %v, want %v", got, want)
+		}
 
-	// The next round's report rewinds the block before appending context, so
-	// the accepted tokens' rows land at their true slots.
-	run := []int32{2, 3, 4}
-	session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1, nil)
-	session.settle(nil)
-	if got, want := draftTokensOf(caches), []int32{1, 2, 3, 4}; !slices.Equal(got, want) {
-		t.Fatalf("draft cache after settle = %v, want %v (block rewound)", got, want)
-	}
-	if got := caches[1].Offset(); got != 4 {
-		t.Fatalf("draft cache offset = %d, want 4 (level with reports)", got)
-	}
-	want := blockCall{offset: 1, ctx: []int32{2, 3, 4}}
-	if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
-		t.Fatalf("context flush = %+v, want %+v", got, want)
-	}
+		// The next round's report rewinds the block before appending context, so
+		// the accepted tokens' rows land at their true slots.
+		run := []int32{2, 3, 4}
+		session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1, nil)
+		session.settle(nil)
+		if got, want := draftTokensOf(caches), []int32{1, 2, 3, 4}; !slices.Equal(got, want) {
+			t.Fatalf("draft cache after settle = %v, want %v (block rewound)", got, want)
+		}
+		if got := caches[1].Offset(); got != 4 {
+			t.Fatalf("draft cache offset = %d, want 4 (level with reports)", got)
+		}
+		want := blockCall{offset: 1, ctx: []int32{2, 3, 4}}
+		if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
+			t.Fatalf("context flush = %+v, want %+v", got, want)
+		}
+	})
 }
 
 func TestDFlashCloseDrainsOutstandingBlock(t *testing.T) {
-	skipIfNoMLX(t)
-	predict := map[int32]int32{2: 3, 3: 4, 4: 5}
-	_, _, session, caches := newBlockTestSession(t, predict, 4)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		predict := map[int32]int32{2: 3, 3: 4, 4: 5}
+		_, _, session, caches := newBlockTestSession(t, predict, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
-	if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
-		t.Fatalf("propose declined")
-	}
-	// A session that ends with a proposal in flight still leaves the caches
-	// level: close rewinds the block even with nothing pending to flush.
-	session.close()
-	if got, want := draftTokensOf(caches), []int32{1}; !slices.Equal(got, want) {
-		t.Fatalf("draft cache after close = %v, want %v", got, want)
-	}
+		session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
+		if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
+			t.Fatalf("propose declined")
+		}
+		// A session that ends with a proposal in flight still leaves the caches
+		// level: close rewinds the block even with nothing pending to flush.
+		session.close()
+		if got, want := draftTokensOf(caches), []int32{1}; !slices.Equal(got, want) {
+			t.Fatalf("draft cache after close = %v, want %v", got, want)
+		}
+	})
 }
 
 func TestDecodeBlockDraft(t *testing.T) {
-	skipIfNoMLX(t)
-	// The block draft mirrors the target chain, so one proposal round accepts
-	// every draft and the bonus token is the EOS.
-	const eos int32 = 7
-	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: eos, eos: 0}
-	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
-	caches, _ := newMTPTestCaches(2)
-	draft := &fakeBlockDraft{predict: predict, blockSize: 3, maskToken: 6, draftCaches: caches[1:]}
-	r.cache.caches = caches
-	r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
-	session, ch := newMTPTestSession(caches)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		// The block draft mirrors the target chain, so one proposal round accepts
+		// every draft and the bonus token is the EOS.
+		const eos int32 = 7
+		predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: eos, eos: 0}
+		r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
+		caches, _ := newMTPTestCaches(2)
+		draft := &fakeBlockDraft{predict: predict, blockSize: 3, maskToken: 6, draftCaches: caches[1:]}
+		r.cache.caches = caches
+		r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
+		session, ch := newMTPTestSession(caches)
 
-	req := Request{
-		Responses:         ch,
-		Tokens:            []int32{1},
-		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
-		SamplerOpts:       sampler.Options{},
-	}
-	spec := r.spec.open(req, nil)
-	if spec == nil || !spec.enabled {
-		t.Fatalf("open rejected a block-draft request")
-	}
-	pinDraftLimit(spec, 4)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil)
-	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
-		t.Fatalf("decode: %v", err)
-	}
-	d.close()
-	spec.close()
-
-	content, final := collectResponses(ch)
-	if content != "2345" {
-		t.Fatalf("content = %q, want %q", content, "2345")
-	}
-	if !final.Done || final.DoneReason != 0 {
-		t.Fatalf("final = %+v, want Done with EOS reason", final)
-	}
-	if want := []int32{2, 3, 4, 5, eos}; !slices.Equal(session.outputs, want) {
-		t.Fatalf("session outputs = %v, want %v", session.outputs, want)
-	}
-
-	// The unprimed drafter parks the first call, so two tokens decode as
-	// pipelined plain forwards; the resumed round then validates the current
-	// token and blockSize-1 drafts in one fused forward.
-	wantForwards := []forwardCall{{offset: 0, n: 1}, {offset: 1, n: 1}, {offset: 2, n: 3}}
-	model := r.Model.(*fakeMTPModel)
-	if !slices.Equal(model.forwards, wantForwards) {
-		t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards)
-	}
-
-	// Ending the parked stretch settles the buffered context through, so the
-	// proposal runs block-only; close's flush then writes the accepted rows
-	// after rewinding the block.
-	wantCalls := []blockCall{
-		{offset: 0, ctx: []int32{2, 3}},
-		{offset: 2, block: []int32{3, 6, 6}},
-		{offset: 2, ctx: []int32{4, 5, 7}},
-	}
-	if len(draft.calls) != len(wantCalls) {
-		t.Fatalf("draft calls = %+v, want %+v", draft.calls, wantCalls)
-	}
-	for i, want := range wantCalls {
-		got := draft.calls[i]
-		if got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || !slices.Equal(got.block, want.block) {
-			t.Fatalf("draft call %d = %+v, want %+v", i, got, want)
+		req := Request{
+			Responses:         ch,
+			Tokens:            []int32{1},
+			CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
+			SamplerOpts:       sampler.Options{},
 		}
-	}
+		spec := r.spec.open(req, nil)
+		if spec == nil || !spec.enabled {
+			t.Fatalf("open rejected a block-draft request")
+		}
+		pinDraftLimit(spec, 4)
+		d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil)
+		if err := r.decode(context.Background(), req, session, d, 0); err != nil {
+			t.Fatalf("decode: %v", err)
+		}
+		d.close()
+		spec.close()
 
-	// The draft caches end level with the target, holding only context rows.
-	if got, want := caches[1].Offset(), caches[0].Offset(); got != want {
-		t.Fatalf("draft cache offset = %d, want %d (level with target)", got, want)
-	}
-	if toks := draftTokensOf(caches); slices.Contains(toks, 6) {
-		t.Fatalf("draft cache retains block rows: %v", toks)
-	}
+		content, final := collectResponses(ch)
+		if content != "2345" {
+			t.Fatalf("content = %q, want %q", content, "2345")
+		}
+		if !final.Done || final.DoneReason != 0 {
+			t.Fatalf("final = %+v, want Done with EOS reason", final)
+		}
+		if want := []int32{2, 3, 4, 5, eos}; !slices.Equal(session.outputs, want) {
+			t.Fatalf("session outputs = %v, want %v", session.outputs, want)
+		}
+
+		// The unprimed drafter parks the first call, so two tokens decode as
+		// pipelined plain forwards; the resumed round then validates the current
+		// token and blockSize-1 drafts in one fused forward.
+		wantForwards := []forwardCall{{offset: 0, n: 1}, {offset: 1, n: 1}, {offset: 2, n: 3}}
+		model := r.Model.(*fakeMTPModel)
+		if !slices.Equal(model.forwards, wantForwards) {
+			t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards)
+		}
+
+		// Ending the parked stretch settles the buffered context through, so the
+		// proposal runs block-only; close's flush then writes the accepted rows
+		// after rewinding the block.
+		wantCalls := []blockCall{
+			{offset: 0, ctx: []int32{2, 3}},
+			{offset: 2, block: []int32{3, 6, 6}},
+			{offset: 2, ctx: []int32{4, 5, 7}},
+		}
+		if len(draft.calls) != len(wantCalls) {
+			t.Fatalf("draft calls = %+v, want %+v", draft.calls, wantCalls)
+		}
+		for i, want := range wantCalls {
+			got := draft.calls[i]
+			if got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || !slices.Equal(got.block, want.block) {
+				t.Fatalf("draft call %d = %+v, want %+v", i, got, want)
+			}
+		}
+
+		// The draft caches end level with the target, holding only context rows.
+		if got, want := caches[1].Offset(), caches[0].Offset(); got != want {
+			t.Fatalf("draft cache offset = %d, want %d (level with target)", got, want)
+		}
+		if toks := draftTokensOf(caches); slices.Contains(toks, 6) {
+			t.Fatalf("draft cache retains block rows: %v", toks)
+		}
+	})
 }

x/mlxrunner/grammar_mask_test.go

@@ -4,45 +4,47 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
 func TestApplyTokenMask(t *testing.T) {
-	skipIfNoMLX(t)
-	const (
-		bitsPerMaskWord       = 32
-		firstTokenID          = 0                   // Least-significant bit of the first mask word.
-		interiorTokenID       = 7                   // Last bit of the first mask word's low byte.
-		lastIDInFirstMaskWord = bitsPerMaskWord - 1 // Sign bit of the int32-backed first mask word.
-		lastVocabID           = 40                  // Final valid ID in a partially used second mask word.
-		vocabSize             = lastVocabID + 1
-	)
-	allowedIDs := []int{firstTokenID, interiorTokenID, lastIDInFirstMaskWord, lastVocabID}
-	word0 := uint32(1)<<firstTokenID |
-		uint32(1)<<interiorTokenID |
-		uint32(1)<<lastIDInFirstMaskWord
-	packed := []int32{
-		int32(word0),
-		int32(uint32(1) << (lastVocabID - bitsPerMaskWord)),
-	}
-	e := &grammarEngine{}
-	e.initMask(vocabSize)
-	logits := mlx.Zeros(mlx.DTypeFloat32, 1, vocabSize)
-	masked := e.apply(logits, mlx.FromValues(packed, 1, len(packed)))
-	mlx.Eval(masked)
-	got := masked.Floats()
-	for id := range vocabSize {
-		allowed := false
-		for _, a := range allowedIDs {
-			if id == a {
-				allowed = true
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const (
+			bitsPerMaskWord       = 32
+			firstTokenID          = 0                   // Least-significant bit of the first mask word.
+			interiorTokenID       = 7                   // Last bit of the first mask word's low byte.
+			lastIDInFirstMaskWord = bitsPerMaskWord - 1 // Sign bit of the int32-backed first mask word.
+			lastVocabID           = 40                  // Final valid ID in a partially used second mask word.
+			vocabSize             = lastVocabID + 1
+		)
+		allowedIDs := []int{firstTokenID, interiorTokenID, lastIDInFirstMaskWord, lastVocabID}
+		word0 := uint32(1)<<firstTokenID |
+			uint32(1)<<interiorTokenID |
+			uint32(1)<<lastIDInFirstMaskWord
+		packed := []int32{
+			int32(word0),
+			int32(uint32(1) << (lastVocabID - bitsPerMaskWord)),
+		}
+		e := &grammarEngine{}
+		e.initMask(vocabSize)
+		logits := mlx.Zeros(mlx.DTypeFloat32, 1, vocabSize)
+		masked := e.apply(logits, mlx.FromValues(packed, 1, len(packed)))
+		mlx.Eval(masked)
+		got := masked.Floats()
+		for id := range vocabSize {
+			allowed := false
+			for _, a := range allowedIDs {
+				if id == a {
+					allowed = true
+				}
+			}
+			if allowed && got[id] != 0 {
+				t.Fatalf("allowed token %d masked to %v", id, got[id])
+			}
+			if !allowed && !math.IsInf(float64(got[id]), -1) {
+				t.Fatalf("disallowed token %d = %v, want -Inf", id, got[id])
 			}
 		}
-		if allowed && got[id] != 0 {
-			t.Fatalf("allowed token %d masked to %v", id, got[id])
-		}
-		if !allowed && !math.IsInf(float64(got[id]), -1) {
-			t.Fatalf("disallowed token %d = %v, want -Inf", id, got[id])
-		}
-	}
+	})
 }

x/mlxrunner/media_test.go

@@ -4,6 +4,7 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/mlxrunner/model/base"
 )
@@ -72,52 +73,52 @@ func (m encodeCountingModel) EncodeMedia(item *base.PreparedItem, data *mlx.Arra
 }
 
 func TestBatchMediaLifecycle(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		calls := 0
+		prepared := &base.PreparedItem{
+			Range:     [2]int{2, 6},
+			MediaData: []float32{1, 2},
+			Dims:      []int{2},
+			Opaque:    7,
+		}
+		r := &Runner{Model: encodeCountingModel{calls: &calls}}
+		request := Request{
+			Tokens:     make([]int32, 8),
+			MediaItems: []mediaItem{{pos: 2, length: 4, item: prepared}},
+		}
 
-	calls := 0
-	prepared := &base.PreparedItem{
-		Range:     [2]int{2, 6},
-		MediaData: []float32{1, 2},
-		Dims:      []int{2},
-		Opaque:    7,
-	}
-	r := &Runner{Model: encodeCountingModel{calls: &calls}}
-	request := Request{
-		Tokens:     make([]int32, 8),
-		MediaItems: []mediaItem{{pos: 2, length: 4, item: prepared}},
-	}
+		m := r.openMedia(request)
+		if m == nil {
+			t.Fatal("openMedia returned nil for a media request")
+		}
+		if m.manifest[0].Pos != 2 || m.manifest[0].Opaque != 7 {
+			t.Fatalf("manifest = %+v", m.manifest[0])
+		}
 
-	m := r.openMedia(request)
-	if m == nil {
-		t.Fatal("openMedia returned nil for a media request")
-	}
-	if m.manifest[0].Pos != 2 || m.manifest[0].Opaque != 7 {
-		t.Fatalf("manifest = %+v", m.manifest[0])
-	}
+		if items := m.batchMedia(0, 2); items[0].Features != nil || calls != 0 {
+			t.Fatal("non-overlapping chunk encoded features")
+		}
+		if items := m.batchMedia(0, 4); items[0].Features == nil || calls != 1 {
+			t.Fatalf("overlap did not encode once (calls=%d)", calls)
+		}
+		if items := m.batchMedia(4, 2); items[0].Features == nil || calls != 1 {
+			t.Fatalf("second overlap re-encoded (calls=%d)", calls)
+		}
 
-	if items := m.batchMedia(0, 2); items[0].Features != nil || calls != 0 {
-		t.Fatal("non-overlapping chunk encoded features")
-	}
-	if items := m.batchMedia(0, 4); items[0].Features == nil || calls != 1 {
-		t.Fatalf("overlap did not encode once (calls=%d)", calls)
-	}
-	if items := m.batchMedia(4, 2); items[0].Features == nil || calls != 1 {
-		t.Fatalf("second overlap re-encoded (calls=%d)", calls)
-	}
+		m.release(4)
+		if m.manifest[0].Features == nil {
+			t.Fatal("release dropped features before the expansion was evaluated")
+		}
+		m.release(6)
+		if m.manifest[0].Features != nil {
+			t.Fatal("release kept features past the expansion end")
+		}
+		m.close()
 
-	m.release(4)
-	if m.manifest[0].Features == nil {
-		t.Fatal("release dropped features before the expansion was evaluated")
-	}
-	m.release(6)
-	if m.manifest[0].Features != nil {
-		t.Fatal("release kept features past the expansion end")
-	}
-	m.close()
-
-	if r.openMedia(Request{Tokens: make([]int32, 8)}) != nil {
-		t.Fatal("openMedia returned non-nil for a text-only request")
-	}
+		if r.openMedia(Request{Tokens: make([]int32, 8)}) != nil {
+			t.Fatal("openMedia returned non-nil for a text-only request")
+		}
+	})
 }
 
 // Two prompts that differ only in their image diverge at the expansion's

x/mlxrunner/mlx/act_test.go

@@ -4,7 +4,7 @@ import (
 	"math"
 	"testing"
 
-	"github.com/ollama/ollama/x/internal/mlxthread"
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 func TestGELUCompiledMatchesEager(t *testing.T) {
@@ -20,7 +20,7 @@ func TestGELUCompiledMatchesEager(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			withMLXThread(t, func() {
+			withMLXThread(t, func(t *mlxthreadtest.T) {
 				EnableCompile()
 				input := FromValues(values, len(values)).AsType(tt.dtype)
 				Pin(input)
@@ -53,34 +53,17 @@ func BenchmarkGELUCompiled(b *testing.B) {
 }
 
 func benchmarkGELU(b *testing.B, fn func(*Array) *Array) {
-	thread, err := mlxthread.Start("mlx-gelu-benchmark", func() error {
-		if err := CheckInit(); err != nil {
-			return err
-		}
-		if GPUIsAvailable() {
-			SetDefaultDeviceGPU()
-		}
-		EnableCompile()
-		return nil
-	})
-	if err != nil {
-		b.Skipf("MLX not available: %v", err)
-	}
-	defer func() {
-		if err := thread.Stop(b.Context(), func() {
-			Sweep()
-			ClearCache()
-			resetDefaultStreamCache()
-		}); err != nil {
-			b.Fatal(err)
-		}
-	}()
-
+	thread := mlxTestThread(b)
 	if err := thread.Do(b.Context(), func() error {
+		EnableCompile()
 		input := AddScalar(Zeros(DTypeBFloat16, 1, 4096, 8192), 1)
 		Eval(input)
 		Pin(input)
-		defer Unpin(input)
+		defer func() {
+			Unpin(input)
+			Sweep()
+			ClearCache()
+		}()
 
 		warmup := fn(input)
 		Eval(warmup)
@@ -100,7 +83,7 @@ func benchmarkGELU(b *testing.B, fn func(*Array) *Array) {
 
 func TestReLUSquared(t *testing.T) {
 	var got []float32
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		x := FromValues([]float32{-2, -0, 0.5, 2}, 4)
 		Pin(x)
 		defer Unpin(x)

x/mlxrunner/mlx/array_test.go

@@ -1,9 +1,13 @@
 package mlx
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
+)
 
 func TestFromValue(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		for got, want := range map[*Array]DType{
 			FromValue(true):              DTypeBool,
 			FromValue(false):             DTypeBool,
@@ -20,7 +24,7 @@ func TestFromValue(t *testing.T) {
 }
 
 func TestFromValues(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		for got, want := range map[*Array]DType{
 			FromValues([]bool{true, false, true}, 3):           DTypeBool,
 			FromValues([]uint8{1, 2, 3}, 3):                    DTypeUint8,
@@ -43,36 +47,40 @@ func TestFromValues(t *testing.T) {
 }
 
 func TestComparisonOpsAndBernoulli(t *testing.T) {
-	withMLXThread(t, func() {
-		testComparisonOpsAndBernoulli(t)
-	})
-}
-
-func testComparisonOpsAndBernoulli(t *testing.T) {
-	a := FromValues([]float32{1, 2, 3}, 3)
-	b := FromValues([]float32{1, 1, 4}, 3)
-	eq := a.Equal(b).AsType(DTypeInt32)
-	gt := a.Greater(b).AsType(DTypeInt32)
-	le := a.LessEqual(b).AsType(DTypeInt32)
-	bern := Bernoulli(FromValues([]float32{1, 0}, 2)).AsType(DTypeInt32)
-	Eval(eq, gt, le, bern)
-
-	for name, tc := range map[string]struct {
+	var tests []struct {
+		name string
 		got  []int32
 		want []int32
-	}{
-		"equal":     {eq.Ints(), []int32{1, 0, 0}},
-		"greater":   {gt.Ints(), []int32{0, 1, 0}},
-		"lessEqual": {le.Ints(), []int32{1, 0, 1}},
-		"bernoulli": {bern.Ints(), []int32{1, 0}},
-	} {
-		t.Run(name, func(t *testing.T) {
-			if len(tc.got) != len(tc.want) {
-				t.Fatalf("got %v, want %v", tc.got, tc.want)
+	}
+	withMLXThread(t, func(*mlxthreadtest.T) {
+		a := FromValues([]float32{1, 2, 3}, 3)
+		b := FromValues([]float32{1, 1, 4}, 3)
+		eq := a.Equal(b).AsType(DTypeInt32)
+		gt := a.Greater(b).AsType(DTypeInt32)
+		le := a.LessEqual(b).AsType(DTypeInt32)
+		bern := Bernoulli(FromValues([]float32{1, 0}, 2)).AsType(DTypeInt32)
+		Eval(eq, gt, le, bern)
+
+		tests = []struct {
+			name string
+			got  []int32
+			want []int32
+		}{
+			{name: "equal", got: eq.Ints(), want: []int32{1, 0, 0}},
+			{name: "greater", got: gt.Ints(), want: []int32{0, 1, 0}},
+			{name: "lessEqual", got: le.Ints(), want: []int32{1, 0, 1}},
+			{name: "bernoulli", got: bern.Ints(), want: []int32{1, 0}},
+		}
+	})
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if len(tt.got) != len(tt.want) {
+				t.Fatalf("got %v, want %v", tt.got, tt.want)
 			}
-			for i := range tc.want {
-				if tc.got[i] != tc.want[i] {
-					t.Fatalf("got %v, want %v", tc.got, tc.want)
+			for i := range tt.want {
+				if tt.got[i] != tt.want[i] {
+					t.Fatalf("got %v, want %v", tt.got, tt.want)
 				}
 			}
 		})

x/mlxrunner/mlx/compile_test.go

@@ -2,15 +2,17 @@ package mlx
 
 import (
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 func TestCompileFusion(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testCompileFusion(t)
 	})
 }
 
-func testCompileFusion(t *testing.T) {
+func testCompileFusion(t *mlxthreadtest.T) {
 	// Compile fuses the ops inside a function body into a single kernel,
 	// eliminating intermediate buffers. Use a diamond-shaped graph where
 	// two branches must be materialized simultaneously without fusion,
@@ -67,12 +69,12 @@ func testCompileFusion(t *testing.T) {
 }
 
 func TestCompileNested(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testCompileNested(t)
 	})
 }
 
-func testCompileNested(t *testing.T) {
+func testCompileNested(t *mlxthreadtest.T) {
 	// A compiled function that calls another compiled function should
 	// produce correct results. The inner function inlines via isTracing()
 	// during the outer's trace.
@@ -103,12 +105,12 @@ func testCompileNested(t *testing.T) {
 }
 
 func TestCompileCallbackPanicRecovers(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testCompileCallbackPanicRecovers(t)
 	})
 }
 
-func testCompileCallbackPanicRecovers(t *testing.T) {
+func testCompileCallbackPanicRecovers(t *mlxthreadtest.T) {
 	boom := Compile1("boom", func(a *Array) *Array {
 		panic("intentional test panic")
 	})
@@ -130,12 +132,12 @@ func testCompileCallbackPanicRecovers(t *testing.T) {
 }
 
 func TestCompileNoTrackingGrowth(t *testing.T) {
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testCompileNoTrackingGrowth(t)
 	})
 }
 
-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.

x/mlxrunner/mlx/depthwise_conv_test.go

@@ -3,21 +3,16 @@ package mlx
 import (
 	"fmt"
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 func TestDepthwiseConvSiLUMatchesGraph(t *testing.T) {
-	skipIfNoMLX(t)
-
-	// Collected, not reported inside the callback: a t.Fatal there calls
-	// runtime.Goexit, which the MLX worker cannot recover, so the job's result
-	// is never delivered and the test hangs until the binary timeout.
-	var mismatches []string
-	withMLXThread(t, func() {
-		mismatches = depthwiseConvSiLUMismatches()
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		for _, mismatch := range depthwiseConvSiLUMismatches() {
+			t.Error(mismatch)
+		}
 	})
-	for _, m := range mismatches {
-		t.Error(m)
-	}
 }
 
 func depthwiseConvSiLUMismatches() []string {

x/mlxrunner/mlx/gated_delta_test.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"math"
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 type gatedDeltaTestGeometry struct {
@@ -36,13 +38,12 @@ func gatedDeltaReference(in gatedDeltaTestInputs, captureAll bool) (y, nextState
 // metal::exp a float ulp apart, which after bf16 rounding leaves rare
 // differing inputs (beta sigmoid at -6.84375); the lattice here avoids them.
 func TestGatedDeltaMatchesGraph(t *testing.T) {
-	skipIfNoMLX(t)
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testGatedDeltaMatchesGraph(t)
 	})
 }
 
-func testGatedDeltaMatchesGraph(t *testing.T) {
+func testGatedDeltaMatchesGraph(t *mlxthreadtest.T) {
 	geometries := []gatedDeltaTestGeometry{
 		{Hk: 16, Dk: 128, Hv: 32, Dv: 128},
 		{Hk: 4, Dk: 64, Hv: 8, Dv: 32},
@@ -76,15 +77,14 @@ func testGatedDeltaMatchesGraph(t *testing.T) {
 }
 
 func TestGatedDeltaGraphRouting(t *testing.T) {
-	skipIfNoMLX(t)
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testGatedDeltaGraphRouting(t)
 	})
 }
 
 // Contract misses — T beyond the kernel cap and a float32 packed input —
 // run the same step as graph ops.
-func testGatedDeltaGraphRouting(t *testing.T) {
+func testGatedDeltaGraphRouting(t *mlxthreadtest.T) {
 	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
 
 	check := func(name string, in gatedDeltaTestInputs, captureAll bool) {
@@ -144,14 +144,13 @@ func scaledGatedDeltaRow(base gatedDeltaTestInputs, scale float32) gatedDeltaTes
 }
 
 func TestGatedDeltaBatchedRows(t *testing.T) {
-	skipIfNoMLX(t)
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testGatedDeltaBatchedRows(t)
 	})
 }
 
 // Each batched row must match its own single-row launch bit-for-bit.
-func testGatedDeltaBatchedRows(t *testing.T) {
+func testGatedDeltaBatchedRows(t *mlxthreadtest.T) {
 	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
 	for _, T := range []int{1, 5, 11} {
 		rows := []gatedDeltaTestInputs{gatedDeltaTestInputs36(g, T)}
@@ -181,8 +180,7 @@ func testGatedDeltaBatchedRows(t *testing.T) {
 }
 
 func TestGatedDeltaRaggedRows(t *testing.T) {
-	skipIfNoMLX(t)
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testGatedDeltaRaggedRows(t)
 	})
 }
@@ -191,7 +189,7 @@ func TestGatedDeltaRaggedRows(t *testing.T) {
 // ba tail poisoned to -inf — runs identity steps with zero output there: its
 // full output and final state must match a launch of only its real tokens
 // with zeros appended, while the full-length row is unaffected.
-func testGatedDeltaRaggedRows(t *testing.T) {
+func testGatedDeltaRaggedRows(t *mlxthreadtest.T) {
 	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
 	const T, realLen = 6, 4
 	row0 := gatedDeltaTestInputs36(g, T)

x/mlxrunner/mlx/mamba_kernels_test.go

@@ -4,12 +4,14 @@ import (
 	"fmt"
 	"math"
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 func TestMamba2ScanMatchesReference(t *testing.T) {
-	requireMamba2Metal(t)
-	var failures []error
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		requireMamba2Metal(t)
+		var failures []error
 		in := newMamba2TestInputs(1, 3, 2, 2, 2, 32)
 		gotY, gotState, interior := Mamba2Scan(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias, nil, false)
 		if len(interior) != 0 {
@@ -18,15 +20,15 @@ func TestMamba2ScanMatchesReference(t *testing.T) {
 		wantY, wantState := mamba2ScanReference(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias)
 		failures = appendArrayCloseError(failures, "mamba2 scan y", gotY, wantY, 1e-5)
 		failures = appendArrayCloseError(failures, "mamba2 scan state", gotState, wantState, 1e-5)
+		reportMamba2Failures(t, failures)
 	})
-	reportMamba2Failures(t, failures)
 }
 
 // Every interior state must match, not just one boundary.
 func TestMamba2ScanCaptureAllMatchesPerTokenReference(t *testing.T) {
-	requireMamba2Metal(t)
-	var failures []error
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		requireMamba2Metal(t)
+		var failures []error
 		const T = 4
 		in := newMamba2TestInputs(1, T, 2, 2, 2, 32)
 		gotY, gotEnd, gotInterior := Mamba2Scan(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias, nil, true)
@@ -50,14 +52,14 @@ func TestMamba2ScanCaptureAllMatchesPerTokenReference(t *testing.T) {
 			}
 		}
 		failures = appendArrayCloseError(failures, "captureAll end state", gotEnd, state, 1e-5)
+		reportMamba2Failures(t, failures)
 	})
-	reportMamba2Failures(t, failures)
 }
 
 func TestMamba2ScanGroupedStatesMatchRepeatedReference(t *testing.T) {
-	requireMamba2Metal(t)
-	var failures []error
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		requireMamba2Metal(t)
+		var failures []error
 		in := newMamba2TestInputs(1, 2, 4, 2, 2, 32)
 		gotY, gotState, _ := Mamba2Scan(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias, nil, false)
 		wantY, wantState := mamba2ScanReference(
@@ -68,15 +70,15 @@ func TestMamba2ScanGroupedStatesMatchRepeatedReference(t *testing.T) {
 		)
 		failures = appendArrayCloseError(failures, "grouped mamba2 y", gotY, wantY, 1e-5)
 		failures = appendArrayCloseError(failures, "grouped mamba2 state", gotState, wantState, 1e-5)
+		reportMamba2Failures(t, failures)
 	})
-	reportMamba2Failures(t, failures)
 }
 
 // A shape outside the kernel's contract must still compute the right answer
 // through the graph implementation rather than fail.
 func TestMamba2ScanUnsupportedShapeMatchesGraph(t *testing.T) {
-	var failures []error
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		var failures []error
 		in := newMamba2TestInputs(1, 2, 2, 2, 2, 31)
 		if _, ok := resolveMamba2ScanDims(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias); ok {
 			failures = append(failures, fmt.Errorf("resolveMamba2ScanDims ok=true for unsupported S=31 shape"))
@@ -85,16 +87,16 @@ func TestMamba2ScanUnsupportedShapeMatchesGraph(t *testing.T) {
 		wantY, wantState := mamba2ScanReference(in.hidden, in.bState, in.cState, in.dt, in.state, in.a, in.d, in.dtBias)
 		failures = appendArrayCloseError(failures, "unsupported-shape y", gotY, wantY, 1e-5)
 		failures = appendArrayCloseError(failures, "unsupported-shape state", gotState, wantState, 1e-5)
+		reportMamba2Failures(t, failures)
 	})
-	reportMamba2Failures(t, failures)
 }
 
 // A padded position must be an identity step. Without the mask it still decays
 // the state by exp(dt*a).
 func TestMamba2ScanPaddedRowIsIdentity(t *testing.T) {
-	requireMamba2Metal(t)
-	var failures []error
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		requireMamba2Metal(t)
+		var failures []error
 		const (
 			B = 2
 			L = 3
@@ -127,8 +129,8 @@ func TestMamba2ScanPaddedRowIsIdentity(t *testing.T) {
 				failures = append(failures, fmt.Errorf("padded output[%d] = %v, want 0", i, v))
 			}
 		}
+		reportMamba2Failures(t, failures)
 	})
-	reportMamba2Failures(t, failures)
 }
 
 type mamba2TestInputs struct {
@@ -187,9 +189,8 @@ func onesTest(dtype DType, shape ...int) *Array {
 	return AddScalar(Zeros(dtype, shape...), 1)
 }
 
-func requireMamba2Metal(t *testing.T) {
+func requireMamba2Metal(t *mlxthreadtest.T) {
 	t.Helper()
-	skipIfNoMLX(t)
 	if !MetalIsAvailable() {
 		t.Skip("MLX Metal not available")
 	}
@@ -210,7 +211,7 @@ func appendArrayCloseError(failures []error, name string, got, want *Array, tol
 	return failures
 }
 
-func reportMamba2Failures(t *testing.T, failures []error) {
+func reportMamba2Failures(t *mlxthreadtest.T, failures []error) {
 	t.Helper()
 	for _, err := range failures {
 		t.Error(err)

x/mlxrunner/mlx/memory_test.go

@@ -6,21 +6,19 @@ import (
 	"math"
 	"slices"
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 func TestSetWiredLimitRejectsOversizeWithoutChangingLimit(t *testing.T) {
-	skipIfNoMLX(t)
-	if !GPUIsAvailable() {
-		t.Skip("MLX GPU not available")
-	}
-
-	var testErr error
-	withMLXThread(t, func() {
-		testErr = checkWiredLimitRejectsOversize()
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		if !GPUIsAvailable() {
+			t.Skip("MLX GPU not available")
+		}
+		if err := checkWiredLimitRejectsOversize(); err != nil {
+			t.Fatal(err)
+		}
 	})
-	if testErr != nil {
-		t.Fatal(testErr)
-	}
 }
 
 func checkWiredLimitRejectsOversize() (err error) {

x/mlxrunner/mlx/ops_extra_test.go

@@ -3,14 +3,15 @@ package mlx
 import (
 	"math"
 	"testing"
+
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
 // fp4Values decodes an fp4 (E2M1) code to its value.
 var fp4Values = [16]float32{0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6}
 
 func TestDequantizeGlobalScale(t *testing.T) {
-	skipIfNoMLX(t)
-	withMLXThread(t, func() {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		testDequantizeGlobalScale(t)
 	})
 }
@@ -18,7 +19,7 @@ func TestDequantizeGlobalScale(t *testing.T) {
 // The quantized payload is built directly, the way an nvfp4 checkpoint ships
 // it: packed fp4 codes, e4m3 group-scale bytes, and a separate global scale.
 // Only the dequantize consumer path runs, so expectations are exact.
-func testDequantizeGlobalScale(t *testing.T) {
+func testDequantizeGlobalScale(t *mlxthreadtest.T) {
 	const rows, cols, group = 4, 64, 16
 	// Every group cycles through all 16 codes; group g of row r has scale
 	// 2^((r+g)%4-1), a power of two so every expected product is exact.

x/mlxrunner/mlx/thread_test.go

@@ -6,20 +6,11 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/ollama/ollama/x/internal/mlxthread"
+	"github.com/ollama/ollama/x/internal/mlxthreadtest"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}
-
-func startMLXThread(t *testing.T) *mlxthread.Thread {
-	t.Helper()
-
-	thread, err := mlxthread.Start("mlx-test", func() error {
+var testThread = sync.OnceValues(func() (*mlxthreadtest.Thread, error) {
+	return mlxthreadtest.Start("mlx-test", func() error {
 		if err := CheckInit(); err != nil {
 			return err
 		}
@@ -28,42 +19,26 @@ func startMLXThread(t *testing.T) *mlxthread.Thread {
 		}
 		return nil
 	})
+})
+
+func mlxTestThread(tb testing.TB) *mlxthreadtest.Thread {
+	tb.Helper()
+
+	thread, err := testThread()
 	if err != nil {
-		t.Skipf("MLX not available: %v", err)
+		tb.Skipf("MLX not available: %v", err)
 	}
 
 	return thread
 }
 
-func stopMLXThread(t *testing.T, thread *mlxthread.Thread) {
+func withMLXThread(t *testing.T, fn func(*mlxthreadtest.T)) {
 	t.Helper()
-
-	if err := thread.Stop(context.Background(), func() {
-		Sweep()
-		ClearCache()
-		resetDefaultStreamCache()
-	}); err != nil {
-		t.Fatal(err)
-	}
-}
-
-func withMLXThread(t *testing.T, fn func()) {
-	t.Helper()
-
-	thread := startMLXThread(t)
-	defer stopMLXThread(t, thread)
-
-	if err := thread.Do(context.Background(), func() error {
-		fn()
-		return nil
-	}); err != nil {
-		t.Fatal(err)
-	}
+	mlxthreadtest.Run(t, mlxTestThread(t), fn)
 }
 
 func TestThreadedMLXOperations(t *testing.T) {
-	thread := startMLXThread(t)
-	defer stopMLXThread(t, thread)
+	thread := mlxTestThread(t)
 
 	oldProcs := runtime.GOMAXPROCS(8)
 	defer runtime.GOMAXPROCS(oldProcs)

x/mlxrunner/model/embedding_test.go

@@ -3,78 +3,72 @@ package model
 import (
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/models/nn"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}
-
 func TestMakeEmbeddingLayerDense(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		weight := mlx.FromValues([]float32{
+			1, 2, 3, 4,
+			5, 6, 7, 8,
+		}, 2, 4).AsType(mlx.DTypeBFloat16)
 
-	weight := mlx.FromValues([]float32{
-		1, 2, 3, 4,
-		5, 6, 7, 8,
-	}, 2, 4).AsType(mlx.DTypeBFloat16)
+		emb := MakeEmbeddingLayer(map[string]*mlx.Array{
+			"model.embed_tokens.weight": weight,
+		}, "model.embed_tokens", 0, 0, "", nil)
 
-	emb := MakeEmbeddingLayer(map[string]*mlx.Array{
-		"model.embed_tokens.weight": weight,
-	}, "model.embed_tokens", 0, 0, "", nil)
-
-	dense, ok := emb.(*nn.Embedding)
-	if !ok {
-		t.Fatalf("embedding type = %T, want *nn.Embedding", emb)
-	}
-	if dense.Weight.DType() != mlx.DTypeBFloat16 {
-		t.Fatalf("embedding dtype = %v, want %v", dense.Weight.DType(), mlx.DTypeBFloat16)
-	}
-	if _, ok := emb.AsLinear().(*nn.Linear); !ok {
-		t.Fatalf("AsLinear type = %T, want *nn.Linear", emb.AsLinear())
-	}
+		dense, ok := emb.(*nn.Embedding)
+		if !ok {
+			t.Fatalf("embedding type = %T, want *nn.Embedding", emb)
+		}
+		if dense.Weight.DType() != mlx.DTypeBFloat16 {
+			t.Fatalf("embedding dtype = %v, want %v", dense.Weight.DType(), mlx.DTypeBFloat16)
+		}
+		if _, ok := emb.AsLinear().(*nn.Linear); !ok {
+			t.Fatalf("AsLinear type = %T, want *nn.Linear", emb.AsLinear())
+		}
+	})
 }
 
 func TestMakeEmbeddingLayerQuantized(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		denseWeight := mlx.FromValues(func() []float32 {
+			out := make([]float32, 2*64)
+			for i := range out {
+				out[i] = float32(i%17) / 8
+			}
+			return out
+		}(), 2, 64).AsType(mlx.DTypeBFloat16)
 
-	denseWeight := mlx.FromValues(func() []float32 {
-		out := make([]float32, 2*64)
-		for i := range out {
-			out[i] = float32(i%17) / 8
+		qw, scales, qbiases := mlx.Quantize(denseWeight, 64, 4, "affine")
+		mlx.Eval(qw, scales, qbiases)
+
+		emb := MakeEmbeddingLayer(map[string]*mlx.Array{
+			"model.embed_tokens.weight":       qw,
+			"model.embed_tokens.weight_scale": scales,
+			"model.embed_tokens.weight_qbias": qbiases,
+		}, "model.embed_tokens", 64, 4, "affine", nil)
+
+		qemb, ok := emb.(*nn.QuantizedEmbedding)
+		if !ok {
+			t.Fatalf("embedding type = %T, want *nn.QuantizedEmbedding", emb)
+		}
+		if qemb.GroupSize != 64 || qemb.Bits != 4 || qemb.Mode != "affine" {
+			t.Fatalf("quant params = (%d, %d, %q), want (64, 4, %q)", qemb.GroupSize, qemb.Bits, qemb.Mode, "affine")
 		}
-		return out
-	}(), 2, 64).AsType(mlx.DTypeBFloat16)
 
-	qw, scales, qbiases := mlx.Quantize(denseWeight, 64, 4, "affine")
-	mlx.Eval(qw, scales, qbiases)
-
-	emb := MakeEmbeddingLayer(map[string]*mlx.Array{
-		"model.embed_tokens.weight":       qw,
-		"model.embed_tokens.weight_scale": scales,
-		"model.embed_tokens.weight_qbias": qbiases,
-	}, "model.embed_tokens", 64, 4, "affine", nil)
-
-	qemb, ok := emb.(*nn.QuantizedEmbedding)
-	if !ok {
-		t.Fatalf("embedding type = %T, want *nn.QuantizedEmbedding", emb)
-	}
-	if qemb.GroupSize != 64 || qemb.Bits != 4 || qemb.Mode != "affine" {
-		t.Fatalf("quant params = (%d, %d, %q), want (64, 4, %q)", qemb.GroupSize, qemb.Bits, qemb.Mode, "affine")
-	}
-
-	indices := mlx.FromValues([]int32{1, 0}, 2)
-	out := emb.Forward(indices)
-	mlx.Eval(out)
-	if dims := out.Dims(); len(dims) != 2 || dims[0] != 2 || dims[1] != 64 {
-		t.Fatalf("embedding output dims = %v, want [2 64]", dims)
-	}
-	if _, ok := emb.AsLinear().(*nn.QuantizedLinear); !ok {
-		t.Fatalf("AsLinear type = %T, want *nn.QuantizedLinear", emb.AsLinear())
-	}
+		indices := mlx.FromValues([]int32{1, 0}, 2)
+		out := emb.Forward(indices)
+		mlx.Eval(out)
+		if dims := out.Dims(); len(dims) != 2 || dims[0] != 2 || dims[1] != 64 {
+			t.Fatalf("embedding output dims = %v, want [2 64]", dims)
+		}
+		if _, ok := emb.AsLinear().(*nn.QuantizedLinear); !ok {
+			t.Fatalf("AsLinear type = %T, want *nn.QuantizedLinear", emb.AsLinear())
+		}
+	})
 }
 
 func TestMakeEmbeddingLayerQuantizedGlobalScale(t *testing.T) {

x/mlxrunner/sample/logprob_test.go

@@ -1,5 +1,3 @@
-//go:build mlx
-
 package sample
 
 import (
@@ -7,6 +5,7 @@ import (
 	"sort"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
@@ -21,21 +20,21 @@ type logprobEntry struct {
 // and returns the greedily-sampled token id, its logprob, and the top-K
 // entries sorted descending by logprob. Logits must be a [vocab]-shaped
 // slice; the helper reshapes it to [1, vocab] before calling the sampler.
-func runSampleLogprobs(t *testing.T, logits []float32, topK int) (int32, float64, []logprobEntry) {
+func runSampleLogprobs(t *mlxtest.T, logits []float32, topK int) (int32, float64, []logprobEntry) {
 	t.Helper()
 
 	s := New(128)
-	defer func() {
+	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()...)
-	defer mlx.Unpin(res.Arrays()...)
+	t.Cleanup(func() { mlx.Unpin(res.Arrays()...) })
 	mlx.Sweep()
 	mlx.Eval(res.Arrays()...)
 
@@ -56,7 +55,7 @@ func runSampleLogprobs(t *testing.T, logits []float32, topK int) (int32, float64
 }
 
 func TestSampleLogprobsBasic(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	tests := []struct {
 		name           string
@@ -82,7 +81,7 @@ func TestSampleLogprobsBasic(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
 			selected, _, top := runSampleLogprobs(t, tt.logits, tt.topK)
 			if selected != tt.wantSelectedID {
 				t.Errorf("selected = %d, want %d", selected, tt.wantSelectedID)
@@ -95,28 +94,28 @@ func TestSampleLogprobsBasic(t *testing.T) {
 }
 
 func TestSampleLogprobsNumericalStability(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		logits := []float32{1000.0, 999.0, 998.0}
+		_, selLP, top := runSampleLogprobs(t, logits, 3)
 
-	logits := []float32{1000.0, 999.0, 998.0}
-	_, selLP, top := runSampleLogprobs(t, logits, 3)
-
-	if math.IsInf(selLP, 0) || math.IsNaN(selLP) {
-		t.Errorf("selected logprob is not finite: %f", selLP)
-	}
-	for i, e := range top {
-		if math.IsInf(e.logprob, 0) || math.IsNaN(e.logprob) {
-			t.Errorf("top[%d] logprob is not finite: %f", i, e.logprob)
+		if math.IsInf(selLP, 0) || math.IsNaN(selLP) {
+			t.Errorf("selected logprob is not finite: %f", selLP)
 		}
-	}
-	for i := 1; i < len(top); i++ {
-		if top[i].logprob > top[i-1].logprob {
-			t.Errorf("top logprobs not descending: %f > %f", top[i].logprob, top[i-1].logprob)
+		for i, e := range top {
+			if math.IsInf(e.logprob, 0) || math.IsNaN(e.logprob) {
+				t.Errorf("top[%d] logprob is not finite: %f", i, e.logprob)
+			}
 		}
-	}
+		for i := 1; i < len(top); i++ {
+			if top[i].logprob > top[i-1].logprob {
+				t.Errorf("top logprobs not descending: %f > %f", top[i].logprob, top[i-1].logprob)
+			}
+		}
+	})
 }
 
 func TestSampleLogprobsProbabilityCorrectness(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	tests := []struct {
 		name   string
@@ -129,7 +128,7 @@ func TestSampleLogprobsProbabilityCorrectness(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
 			selected, selLP, top := runSampleLogprobs(t, tt.logits, len(tt.logits))
 
 			if selLP > 0 {
@@ -174,7 +173,7 @@ func TestSampleLogprobsProbabilityCorrectness(t *testing.T) {
 }
 
 func TestSampleLogprobsSoftmaxCorrectness(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	tests := []struct {
 		name   string
@@ -188,7 +187,7 @@ func TestSampleLogprobsSoftmaxCorrectness(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
 			_, _, top := runSampleLogprobs(t, tt.logits, len(tt.logits))
 			if len(top) != len(tt.logits) {
 				t.Fatalf("top-K length = %d, want %d", len(top), len(tt.logits))
@@ -211,90 +210,90 @@ func TestSampleLogprobsSoftmaxCorrectness(t *testing.T) {
 }
 
 func TestSampleLogprobsSelectedTokenCorrectness(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		logits := []float32{3.0, 1.0, 2.0, 0.5}
 
-	logits := []float32{3.0, 1.0, 2.0, 0.5}
-
-	maxIdx := int32(0)
-	for i, v := range logits[1:] {
-		if v > logits[maxIdx] {
-			maxIdx = int32(i + 1)
+		maxIdx := int32(0)
+		for i, v := range logits[1:] {
+			if v > logits[maxIdx] {
+				maxIdx = int32(i + 1)
+			}
 		}
-	}
 
-	selected, selLP, top := runSampleLogprobs(t, logits, len(logits))
+		selected, selLP, top := runSampleLogprobs(t, logits, len(logits))
 
-	if selected != maxIdx {
-		t.Errorf("selected = %d, want argmax %d", selected, maxIdx)
-	}
+		if selected != maxIdx {
+			t.Errorf("selected = %d, want argmax %d", selected, maxIdx)
+		}
 
-	if top[0].id != maxIdx {
-		t.Errorf("top[0].id = %d, want argmax %d", top[0].id, maxIdx)
-	}
-	if math.Abs(top[0].logprob-selLP) > 1e-6 {
-		t.Errorf("top[0].logprob = %f, want selected %f", top[0].logprob, selLP)
-	}
+		if top[0].id != maxIdx {
+			t.Errorf("top[0].id = %d, want argmax %d", top[0].id, maxIdx)
+		}
+		if math.Abs(top[0].logprob-selLP) > 1e-6 {
+			t.Errorf("top[0].logprob = %f, want selected %f", top[0].logprob, selLP)
+		}
+	})
 }
 
 // TestBatchedLogprobsPerRow verifies that per-row logprobs in a batched
 // sample call match the per-slot reference. The numerically-stable softmax
 // must reduce along the last axis only, not over the whole batch.
 func TestBatchedLogprobsPerRow(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		rowA := []float32{2, 1, 0}
+		rowB := []float32{0, 5, 0}
 
-	rowA := []float32{2, 1, 0}
-	rowB := []float32{0, 5, 0}
+		_, wantA, _ := runSampleLogprobs(t, rowA, 0)
+		_, wantB, _ := runSampleLogprobs(t, rowB, 0)
 
-	_, wantA, _ := runSampleLogprobs(t, rowA, 0)
-	_, wantB, _ := runSampleLogprobs(t, rowB, 0)
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
+		s.Add(1, Options{Logprobs: true}, nil)
+		s.Add(2, Options{Logprobs: true}, nil)
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
+		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()
+		if len(got) != 2 {
+			t.Fatalf("Logprob length = %d, want 2", len(got))
+		}
+		if math.Abs(float64(got[0])-wantA) > 1e-5 {
+			t.Errorf("row 0 logprob = %f, want %f (per-slot reference)", got[0], wantA)
+		}
+		if math.Abs(float64(got[1])-wantB) > 1e-5 {
+			t.Errorf("row 1 logprob = %f, want %f (per-slot reference)", got[1], wantB)
+		}
 	})
-	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()
-	if len(got) != 2 {
-		t.Fatalf("Logprob length = %d, want 2", len(got))
-	}
-	if math.Abs(float64(got[0])-wantA) > 1e-5 {
-		t.Errorf("row 0 logprob = %f, want %f (per-slot reference)", got[0], wantA)
-	}
-	if math.Abs(float64(got[1])-wantB) > 1e-5 {
-		t.Errorf("row 1 logprob = %f, want %f (per-slot reference)", got[1], wantB)
-	}
 }
 
 func TestSampleLogprobsTopKOrdering(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		// Logits chosen so argmax order differs from index order.
+		logits := []float32{2.0, 5.0, 1.0, 4.0, 3.0}
+		wantOrder := []int32{1, 3, 4, 0, 2}
 
-	// Logits chosen so argmax order differs from index order.
-	logits := []float32{2.0, 5.0, 1.0, 4.0, 3.0}
-	wantOrder := []int32{1, 3, 4, 0, 2}
+		_, _, top := runSampleLogprobs(t, logits, len(logits))
 
-	_, _, top := runSampleLogprobs(t, logits, len(logits))
-
-	if len(top) != len(wantOrder) {
-		t.Fatalf("top-K length = %d, want %d", len(top), len(wantOrder))
-	}
-	for i, e := range top {
-		if e.id != wantOrder[i] {
-			t.Errorf("top[%d].id = %d, want %d", i, e.id, wantOrder[i])
+		if len(top) != len(wantOrder) {
+			t.Fatalf("top-K length = %d, want %d", len(top), len(wantOrder))
 		}
-	}
-	for i := 1; i < len(top); i++ {
-		if top[i].logprob > top[i-1].logprob {
-			t.Errorf("top[%d].logprob (%f) > top[%d].logprob (%f)",
-				i, top[i].logprob, i-1, top[i-1].logprob)
+		for i, e := range top {
+			if e.id != wantOrder[i] {
+				t.Errorf("top[%d].id = %d, want %d", i, e.id, wantOrder[i])
+			}
 		}
-	}
+		for i := 1; i < len(top); i++ {
+			if top[i].logprob > top[i-1].logprob {
+				t.Errorf("top[%d].logprob (%f) > top[%d].logprob (%f)",
+					i, top[i].logprob, i-1, top[i-1].logprob)
+			}
+		}
+	})
 }

x/mlxrunner/sample/sample_test.go

@@ -1,5 +1,3 @@
-//go:build mlx
-
 package sample
 
 import (
@@ -7,16 +5,10 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}
-
 // slotLogits builds a [1, V] logits tensor for a single-slot Sample call.
 func slotLogits(values []float32) *mlx.Array {
 	return mlx.FromValues(values, 1, len(values))
@@ -39,7 +31,7 @@ func batchLogits(rows ...[]float32) *mlx.Array {
 // sampleOne runs Sample on a freshly-added single slot and returns the
 // sampled token id. Used both for the single-slot options table and as the
 // reference oracle for the batched-equivalence test.
-func sampleOne(t *testing.T, opts Options, priorTokens []int32, values []float32) int32 {
+func sampleOne(t *mlxtest.T, opts Options, priorTokens []int32, values []float32) int32 {
 	t.Helper()
 	s := New(128)
 	t.Cleanup(func() {
@@ -62,7 +54,7 @@ func logOf(p float64) float32 { return float32(math.Log(p)) }
 // hand from the math of each transform, not from a second call into the
 // sampler — so a regression in any single transform shows up here.
 func TestSampleSingleSlotOptions(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	cases := []struct {
 		name   string
@@ -134,7 +126,7 @@ func TestSampleSingleSlotOptions(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			if got := sampleOne(t, tc.opts, tc.priors, tc.logits); got != tc.want {
 				t.Errorf("got %d, want %d", got, tc.want)
 			}
@@ -143,129 +135,129 @@ func TestSampleSingleSlotOptions(t *testing.T) {
 }
 
 func TestDistributionAppliesTopKBeforeTopP(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
+		s.Add(0, Options{Temperature: 1, TopK: 2, TopP: 0.7}, nil)
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
-	})
-	s.Add(0, Options{Temperature: 1, TopK: 2, TopP: 0.7}, nil)
+		dist := s.Distribution(0, slotLogits([]float32{logOf(0.6), logOf(0.2), logOf(0.2)}), nil)
+		mlx.Eval(dist.Arrays()...)
 
-	dist := s.Distribution(0, slotLogits([]float32{logOf(0.6), logOf(0.2), logOf(0.2)}), nil)
-	mlx.Eval(dist.Arrays()...)
+		ids := dist.IDs.Ints()
+		probs := dist.Probs.Floats()
+		if len(ids) != 2 || len(probs) != 2 {
+			t.Fatalf("support = ids %v probs %v, want 2 sparse entries", ids, probs)
+		}
 
-	ids := dist.IDs.Ints()
-	probs := dist.Probs.Floats()
-	if len(ids) != 2 || len(probs) != 2 {
-		t.Fatalf("support = ids %v probs %v, want 2 sparse entries", ids, probs)
-	}
-
-	foundTop := false
-	for i, id := range ids {
-		switch id {
-		case 0:
-			foundTop = true
-			if math.Abs(float64(probs[i]-1)) > 1e-5 {
-				t.Fatalf("top token prob = %v, want 1; ids=%v probs=%v", probs[i], ids, probs)
-			}
-		default:
-			if math.Abs(float64(probs[i])) > 1e-5 {
-				t.Fatalf("non-top token %d prob = %v, want 0; ids=%v probs=%v", id, probs[i], ids, probs)
+		foundTop := false
+		for i, id := range ids {
+			switch id {
+			case 0:
+				foundTop = true
+				if math.Abs(float64(probs[i]-1)) > 1e-5 {
+					t.Fatalf("top token prob = %v, want 1; ids=%v probs=%v", probs[i], ids, probs)
+				}
+			default:
+				if math.Abs(float64(probs[i])) > 1e-5 {
+					t.Fatalf("non-top token %d prob = %v, want 0; ids=%v probs=%v", id, probs[i], ids, probs)
+				}
 			}
 		}
-	}
-	if !foundTop {
-		t.Fatalf("top-k support %v did not include token 0", ids)
-	}
+		if !foundTop {
+			t.Fatalf("top-k support %v did not include token 0", ids)
+		}
+	})
 }
 
 func TestDistributionResidualUsesTargetSupport(t *testing.T) {
-	skipIfNoMLX(t)
-
-	target := Distribution{
-		IDs:   mlx.NewArrayInt32([]int32{2, 5}, []int32{1, 2}),
-		Probs: mlx.FromValues([]float32{0.7, 0.3}, 1, 2),
-	}
-	draft := Distribution{
-		IDs:   mlx.NewArrayInt32([]int32{2, 4}, []int32{1, 2}),
-		Probs: mlx.FromValues([]float32{0.2, 0.8}, 1, 2),
-	}
-
-	residual := target.ResidualAgainst(draft)
-	mlx.Eval(residual.Arrays()...)
-
-	ids := residual.IDs.Ints()
-	probs := residual.Probs.Floats()
-	want := map[int32]float64{2: 0.625, 5: 0.375}
-	if len(ids) != 2 || len(probs) != 2 {
-		t.Fatalf("residual = ids %v probs %v, want 2 sparse entries", ids, probs)
-	}
-	for i, id := range ids {
-		w, ok := want[id]
-		if !ok {
-			t.Fatalf("residual includes token %d outside target support: ids=%v probs=%v", id, ids, probs)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		target := Distribution{
+			IDs:   mlx.NewArrayInt32([]int32{2, 5}, []int32{1, 2}),
+			Probs: mlx.FromValues([]float32{0.7, 0.3}, 1, 2),
 		}
-		if math.Abs(float64(probs[i])-w) > 1e-5 {
-			t.Fatalf("residual token %d prob = %v, want %v; ids=%v probs=%v", id, probs[i], w, ids, probs)
+		draft := Distribution{
+			IDs:   mlx.NewArrayInt32([]int32{2, 4}, []int32{1, 2}),
+			Probs: mlx.FromValues([]float32{0.2, 0.8}, 1, 2),
 		}
-	}
+
+		residual := target.ResidualAgainst(draft)
+		mlx.Eval(residual.Arrays()...)
+
+		ids := residual.IDs.Ints()
+		probs := residual.Probs.Floats()
+		want := map[int32]float64{2: 0.625, 5: 0.375}
+		if len(ids) != 2 || len(probs) != 2 {
+			t.Fatalf("residual = ids %v probs %v, want 2 sparse entries", ids, probs)
+		}
+		for i, id := range ids {
+			w, ok := want[id]
+			if !ok {
+				t.Fatalf("residual includes token %d outside target support: ids=%v probs=%v", id, ids, probs)
+			}
+			if math.Abs(float64(probs[i])-w) > 1e-5 {
+				t.Fatalf("residual token %d prob = %v, want %v; ids=%v probs=%v", id, probs[i], w, ids, probs)
+			}
+		}
+	})
 }
 
 func TestSeededSamplingIsReproducible(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		seededSequence := func(seed int) []int32 {
+			s := New(128)
+			t.Cleanup(func() {
+				s.Free()
+				mlx.Sweep()
+			})
+			s.Add(0, Options{Temperature: 1, TopK: 4, Seed: seed, UseSeed: true}, nil)
 
-	seededSequence := func(seed int) []int32 {
-		s := New(128)
-		t.Cleanup(func() {
-			s.Free()
-			mlx.Sweep()
-		})
-		s.Add(0, Options{Temperature: 1, TopK: 4, Seed: seed, UseSeed: true}, nil)
-
-		logits := slotLogits([]float32{0, 0, 0, 0})
-		out := make([]int32, 32)
-		for i := range out {
-			token := s.Sample([]int{0}, logits).Token
-			mlx.Eval(token)
-			out[i] = token.Int()
+			logits := slotLogits([]float32{0, 0, 0, 0})
+			out := make([]int32, 32)
+			for i := range out {
+				token := s.Sample([]int{0}, logits).Token
+				mlx.Eval(token)
+				out[i] = token.Int()
+			}
+			return out
 		}
-		return out
-	}
 
-	a := seededSequence(1234)
-	b := seededSequence(1234)
-	if !slices.Equal(a, b) {
-		t.Fatalf("same seed produced different sequences:\n%v\n%v", a, b)
-	}
+		a := seededSequence(1234)
+		b := seededSequence(1234)
+		if !slices.Equal(a, b) {
+			t.Fatalf("same seed produced different sequences:\n%v\n%v", a, b)
+		}
 
-	c := seededSequence(5678)
-	if slices.Equal(a, c) {
-		t.Fatalf("different seeds produced the same sequence: %v", a)
-	}
+		c := seededSequence(5678)
+		if slices.Equal(a, c) {
+			t.Fatalf("different seeds produced the same sequence: %v", a)
+		}
+	})
 }
 
 func TestSeededBernoulliIsReproducible(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		seededMask := func() []int32 {
+			s := New(128)
+			t.Cleanup(func() {
+				s.Free()
+				mlx.Sweep()
+			})
+			s.Add(0, Options{Seed: 99, UseSeed: true}, nil)
 
-	seededMask := func() []int32 {
-		s := New(128)
-		t.Cleanup(func() {
-			s.Free()
-			mlx.Sweep()
-		})
-		s.Add(0, Options{Seed: 99, UseSeed: true}, nil)
+			mask := s.Bernoulli(0, mlx.FromValues([]float32{0.5, 0.5, 0.5, 0.5, 0.5, 0.5}, 6)).AsType(mlx.DTypeInt32)
+			mlx.Eval(mask)
+			return mask.Ints()
+		}
 
-		mask := s.Bernoulli(0, mlx.FromValues([]float32{0.5, 0.5, 0.5, 0.5, 0.5, 0.5}, 6)).AsType(mlx.DTypeInt32)
-		mlx.Eval(mask)
-		return mask.Ints()
-	}
-
-	a := seededMask()
-	b := seededMask()
-	if !slices.Equal(a, b) {
-		t.Fatalf("same seed produced different bernoulli masks:\n%v\n%v", a, b)
-	}
+		a := seededMask()
+		b := seededMask()
+		if !slices.Equal(a, b) {
+			t.Fatalf("same seed produced different bernoulli masks:\n%v\n%v", a, b)
+		}
+	})
 }
 
 // TestSampleHistoryWindow verifies that penalty history respects the
@@ -273,154 +265,154 @@ func TestSeededBernoulliIsReproducible(t *testing.T) {
 // and once the ring wraps, tokens that rotate out no longer contribute
 // to penalties.
 func TestSampleHistoryWindow(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
+		// RepeatLastN=2 with priors {1, 2, 3}: makeHistoryRow keeps only
+		// {2, 3}. Token 1 was trimmed — its penalty is NOT active.
+		s.Add(0, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{1, 2, 3})
+
+		// Step 1: logits favor token 1 (trimmed). If the trim were broken it
+		// would be penalized and the argmax would move.
+		step1 := s.Sample([]int{0}, slotLogits([]float32{0, 5, 0, 0, 0})).Token
+		mlx.Eval(step1)
+		if got := step1.Int(); got != 1 {
+			t.Fatalf("step 1 = %d, want 1 (token 1 trimmed from priors)", got)
+		}
+		// After step 1 the ring holds {1, 3}; token 2 has rotated out.
+
+		// Step 2: logits favor token 2 (rotated out). If the ring wrap were
+		// wrong, token 2 would still be penalized.
+		step2 := s.Sample([]int{0}, slotLogits([]float32{0, 0, 5, 0, 0})).Token
+		mlx.Eval(step2)
+		if got := step2.Int(); got != 2 {
+			t.Fatalf("step 2 = %d, want 2 (token 2 rotated out of ring)", got)
+		}
 	})
-
-	// RepeatLastN=2 with priors {1, 2, 3}: makeHistoryRow keeps only
-	// {2, 3}. Token 1 was trimmed — its penalty is NOT active.
-	s.Add(0, Options{RepeatLastN: 2, PresencePenalty: 10}, []int32{1, 2, 3})
-
-	// Step 1: logits favor token 1 (trimmed). If the trim were broken it
-	// would be penalized and the argmax would move.
-	step1 := s.Sample([]int{0}, slotLogits([]float32{0, 5, 0, 0, 0})).Token
-	mlx.Eval(step1)
-	if got := step1.Int(); got != 1 {
-		t.Fatalf("step 1 = %d, want 1 (token 1 trimmed from priors)", got)
-	}
-	// After step 1 the ring holds {1, 3}; token 2 has rotated out.
-
-	// Step 2: logits favor token 2 (rotated out). If the ring wrap were
-	// wrong, token 2 would still be penalized.
-	step2 := s.Sample([]int{0}, slotLogits([]float32{0, 0, 5, 0, 0})).Token
-	mlx.Eval(step2)
-	if got := step2.Int(); got != 2 {
-		t.Fatalf("step 2 = %d, want 2 (token 2 rotated out of ring)", got)
-	}
 }
 
 func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
-	})
+		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{1, 2})
+		draftTokens := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
+		scores := s.SpeculativeScores(0, batchLogits(
+			[]float32{0, 9, 9, 8, 0}, // history {1,2}; token 3 wins
+			[]float32{0, 0, 9, 9, 8}, // history {2,3}; token 4 wins
+			[]float32{0, 0, 9, 9, 8}, // history {3,4}; token 2 wins
+		), draftTokens)
+		tokens := scores.Argmax(-1, false).AsType(mlx.DTypeInt32)
+		mlx.Eval(tokens)
 
-	s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{1, 2})
-	draftTokens := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
-	scores := s.SpeculativeScores(0, batchLogits(
-		[]float32{0, 9, 9, 8, 0}, // history {1,2}; token 3 wins
-		[]float32{0, 0, 9, 9, 8}, // history {2,3}; token 4 wins
-		[]float32{0, 0, 9, 9, 8}, // history {3,4}; token 2 wins
-	), draftTokens)
-	tokens := scores.Argmax(-1, false).AsType(mlx.DTypeInt32)
-	mlx.Eval(tokens)
-
-	if got, want := tokens.Ints(), []int32{3, 4, 2}; len(got) != len(want) {
-		t.Fatalf("tokens = %v, want %v", got, want)
-	} else {
-		for i := range want {
-			if got[i] != want[i] {
-				t.Fatalf("tokens = %v, want %v", got, want)
+		if got, want := tokens.Ints(), []int32{3, 4, 2}; len(got) != len(want) {
+			t.Fatalf("tokens = %v, want %v", got, want)
+		} else {
+			for i := range want {
+				if got[i] != want[i] {
+					t.Fatalf("tokens = %v, want %v", got, want)
+				}
 			}
 		}
-	}
-	if s.byID[0].historyLen != 2 {
-		t.Fatalf("historyLen = %d, want 2", s.byID[0].historyLen)
-	}
+		if s.byID[0].historyLen != 2 {
+			t.Fatalf("historyLen = %d, want 2", s.byID[0].historyLen)
+		}
+	})
 }
 
 func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
-	})
+		// A proposal step passes one logits row with the chain's earlier drafts:
+		// the single row is the chain's final step, so every draft belongs to
+		// its history. Slot 0 exercises the batched history path (full ring),
+		// slot 1 the serial path (ring not yet full).
+		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{0, 1})
+		s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{0, 1})
+		prefix := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
 
-	// A proposal step passes one logits row with the chain's earlier drafts:
-	// the single row is the chain's final step, so every draft belongs to
-	// its history. Slot 0 exercises the batched history path (full ring),
-	// slot 1 the serial path (ring not yet full).
-	s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{0, 1})
-	s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{0, 1})
-	prefix := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
-
-	for _, seqID := range []int{0, 1} {
-		// Drafts 3 and 4 are penalized, so token 2 wins over the higher raw
-		// scores; with the drafts absent from the history, token 3 would.
-		dist := s.Distribution(seqID, batchLogits([]float32{0, 0, 9, 9, 8}), prefix)
-		mlx.Eval(dist.IDs)
-		if got := dist.IDs.Ints()[0]; got != 2 {
-			t.Fatalf("seq %d token = %d, want 2 (drafts 3 and 4 penalized)", seqID, got)
+		for _, seqID := range []int{0, 1} {
+			// Drafts 3 and 4 are penalized, so token 2 wins over the higher raw
+			// scores; with the drafts absent from the history, token 3 would.
+			dist := s.Distribution(seqID, batchLogits([]float32{0, 0, 9, 9, 8}), prefix)
+			mlx.Eval(dist.IDs)
+			if got := dist.IDs.Ints()[0]; got != 2 {
+				t.Fatalf("seq %d token = %d, want 2 (drafts 3 and 4 penalized)", seqID, got)
+			}
 		}
-	}
+	})
 }
 
 func TestDistributionMultiRowWithoutChain(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
-	})
+		// A block drafter's proposal batch samples every row from one call with
+		// no draft chain: each row sees the slot history unchanged. Slot 0
+		// exercises the batched history path (full ring), slot 1 the serial path
+		// (ring not yet full).
+		s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{3, 4})
+		s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{3, 4})
 
-	// A block drafter's proposal batch samples every row from one call with
-	// no draft chain: each row sees the slot history unchanged. Slot 0
-	// exercises the batched history path (full ring), slot 1 the serial path
-	// (ring not yet full).
-	s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{3, 4})
-	s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{3, 4})
-
-	for _, seqID := range []int{0, 1} {
-		// Tokens 3 and 4 are penalized in every row alike; rows 1 and 3
-		// share logits, so a chain alignment leaking between rows would
-		// split their winners.
-		dist := s.Distribution(seqID, batchLogits(
-			[]float32{0, 0, 8, 9, 9},
-			[]float32{0, 8, 0, 9, 9},
-			[]float32{0, 0, 8, 9, 9},
-		), nil)
-		top := dist.IDs.Slice(mlx.Slice(), mlx.Slice(0, 1))
-		mlx.Eval(top)
-		if got, want := top.Ints(), []int32{2, 1, 2}; !slices.Equal(got, want) {
-			t.Fatalf("seq %d top tokens = %v, want %v", seqID, got, want)
+		for _, seqID := range []int{0, 1} {
+			// Tokens 3 and 4 are penalized in every row alike; rows 1 and 3
+			// share logits, so a chain alignment leaking between rows would
+			// split their winners.
+			dist := s.Distribution(seqID, batchLogits(
+				[]float32{0, 0, 8, 9, 9},
+				[]float32{0, 8, 0, 9, 9},
+				[]float32{0, 0, 8, 9, 9},
+			), nil)
+			top := dist.IDs.Slice(mlx.Slice(), mlx.Slice(0, 1))
+			mlx.Eval(top)
+			if got, want := top.Ints(), []int32{2, 1, 2}; !slices.Equal(got, want) {
+				t.Fatalf("seq %d top tokens = %v, want %v", seqID, got, want)
+			}
 		}
-	}
+	})
 }
 
 func TestCommitBatchesRingWrites(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
 
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
-	})
+		s.Add(0, Options{RepeatLastN: 4, RepeatPenalty: 1.1}, []int32{10, 11, 12})
+		s.Commit(0, []int32{20, 21, 22})
+		s.Commit(0, []int32{30, 31, 32, 33, 34})
+		mlx.Eval(s.history)
 
-	s.Add(0, Options{RepeatLastN: 4, RepeatPenalty: 1.1}, []int32{10, 11, 12})
-	s.Commit(0, []int32{20, 21, 22})
-	s.Commit(0, []int32{30, 31, 32, 33, 34})
-	mlx.Eval(s.history)
-
-	got := s.history.Ints()
-	want := []int32{32, 33, 34, 31}
-	for i := range want {
-		if got[i] != want[i] {
-			t.Fatalf("history = %v, want %v", got, want)
+		got := s.history.Ints()
+		want := []int32{32, 33, 34, 31}
+		for i := range want {
+			if got[i] != want[i] {
+				t.Fatalf("history = %v, want %v", got, want)
+			}
 		}
-	}
-	if s.byID[0].historyLen != 11 {
-		t.Fatalf("historyLen = %d, want 11", s.byID[0].historyLen)
-	}
+		if s.byID[0].historyLen != 11 {
+			t.Fatalf("historyLen = %d, want 11", s.byID[0].historyLen)
+		}
+	})
 }
 
 // TestBatchSamplingPreservesPerSlotBehavior is the core equivalence test:
@@ -428,7 +420,7 @@ func TestCommitBatchesRingWrites(t *testing.T) {
 // serial on partial ring, subset/out-of-order), a batched Sample call must
 // produce the same token per row as running the same slot alone.
 func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	type slot struct {
 		id     int
@@ -482,7 +474,7 @@ func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) {
 			// Per-slot reference for each sampled seq.
 			want := make([]int32, len(tc.sample))
 			for i, id := range tc.sample {
@@ -522,32 +514,32 @@ func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) {
 // recycled row must start from its own priors only — no carryover from
 // the removed slot's history.
 func TestRemoveDoesNotLeakHistory(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		opts := Options{RepeatLastN: 1, PresencePenalty: 10}
+		s := New(128)
+		t.Cleanup(func() {
+			s.Free()
+			mlx.Sweep()
+		})
+		s.Add(1, opts, []int32{1})
+		s.Add(2, opts, []int32{2})
+		s.Remove(1)
+		s.Add(3, opts, []int32{0})
 
-	opts := Options{RepeatLastN: 1, PresencePenalty: 10}
-	s := New(128)
-	t.Cleanup(func() {
-		s.Free()
-		mlx.Sweep()
+		// Slot 2 retains history {2}; slot 3 retains history {0}. With
+		// equal logits and PresencePenalty=10 the argmax drops to the first
+		// unpenalized token.
+		res := s.Sample([]int{2, 3}, batchLogits(
+			[]float32{3, 3, 0},
+			[]float32{3, 3, 0},
+		))
+		mlx.Eval(res.Token)
+		tokens := res.Token.Ints()
+		if tokens[0] != 0 {
+			t.Errorf("slot 2 = %d, want 0 (token 2 penalized)", tokens[0])
+		}
+		if tokens[1] != 1 {
+			t.Errorf("slot 3 = %d, want 1 (token 0 penalized, no slot-1 carryover)", tokens[1])
+		}
 	})
-	s.Add(1, opts, []int32{1})
-	s.Add(2, opts, []int32{2})
-	s.Remove(1)
-	s.Add(3, opts, []int32{0})
-
-	// Slot 2 retains history {2}; slot 3 retains history {0}. With
-	// equal logits and PresencePenalty=10 the argmax drops to the first
-	// unpenalized token.
-	res := s.Sample([]int{2, 3}, batchLogits(
-		[]float32{3, 3, 0},
-		[]float32{3, 3, 0},
-	))
-	mlx.Eval(res.Token)
-	tokens := res.Token.Ints()
-	if tokens[0] != 0 {
-		t.Errorf("slot 2 = %d, want 0 (token 2 penalized)", tokens[0])
-	}
-	if tokens[1] != 1 {
-		t.Errorf("slot 3 = %d, want 1 (token 0 penalized, no slot-1 carryover)", tokens[1])
-	}
 }

x/models/gemma4/gemma4_moe_test.go

@@ -1,37 +1,12 @@
 package gemma4
 
 import (
-	"runtime"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func useMLXTestThread(t *testing.T) {
-	t.Helper()
-
-	runtime.LockOSThread()
-	initialized := false
-	t.Cleanup(func() {
-		if initialized {
-			mlx.Sweep()
-			mlx.ClearCache()
-			if mlx.GPUIsAvailable() {
-				mlx.SetDefaultDeviceGPU()
-			}
-		}
-		runtime.UnlockOSThread()
-	})
-
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-	initialized = true
-	if mlx.GPUIsAvailable() {
-		mlx.SetDefaultDeviceGPU()
-	}
-}
-
 // onesLike creates a tensor of the given shape filled with a small constant.
 func onesLike(shape ...int) *mlx.Array {
 	return mlx.AddScalar(mlx.Zeros(mlx.DTypeBFloat16, shape...), 0.01)
@@ -73,72 +48,72 @@ func newMoEBlock(cfg *TextConfig) *MoEBlock {
 }
 
 func TestMoERouterForward(t *testing.T) {
-	useMLXTestThread(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := tinyMoEConfig()
+		B, L := int32(1), int32(3)
+		x := onesLike(int(B), int(L), int(cfg.HiddenSize))
+		router := newRouter(cfg)
 
-	cfg := tinyMoEConfig()
-	B, L := int32(1), int32(3)
-	x := onesLike(int(B), int(L), int(cfg.HiddenSize))
-	router := newRouter(cfg)
+		scores, inds := router.Forward(x, cfg)
+		mlx.Eval(scores, inds)
 
-	scores, inds := router.Forward(x, cfg)
-	mlx.Eval(scores, inds)
+		sDims := scores.Dims()
+		iDims := inds.Dims()
+		t.Logf("scores shape: %v, inds shape: %v", sDims, iDims)
 
-	sDims := scores.Dims()
-	iDims := inds.Dims()
-	t.Logf("scores shape: %v, inds shape: %v", sDims, iDims)
-
-	if len(sDims) != 2 || sDims[0] != int(B*L) || sDims[1] != int(cfg.TopKExperts) {
-		t.Errorf("scores shape = %v, want [%d, %d]", sDims, B*L, cfg.TopKExperts)
-	}
-	if len(iDims) != 2 || iDims[0] != int(B*L) || iDims[1] != int(cfg.TopKExperts) {
-		t.Errorf("inds shape = %v, want [%d, %d]", iDims, B*L, cfg.TopKExperts)
-	}
+		if len(sDims) != 2 || sDims[0] != int(B*L) || sDims[1] != int(cfg.TopKExperts) {
+			t.Errorf("scores shape = %v, want [%d, %d]", sDims, B*L, cfg.TopKExperts)
+		}
+		if len(iDims) != 2 || iDims[0] != int(B*L) || iDims[1] != int(cfg.TopKExperts) {
+			t.Errorf("inds shape = %v, want [%d, %d]", iDims, B*L, cfg.TopKExperts)
+		}
+	})
 }
 
 func TestMoEBlockForward(t *testing.T) {
-	useMLXTestThread(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := tinyMoEConfig()
+		B, L := int32(1), int32(3)
+		x := onesLike(int(B), int(L), int(cfg.HiddenSize))
+		router := newRouter(cfg)
+		moe := newMoEBlock(cfg)
 
-	cfg := tinyMoEConfig()
-	B, L := int32(1), int32(3)
-	x := onesLike(int(B), int(L), int(cfg.HiddenSize))
-	router := newRouter(cfg)
-	moe := newMoEBlock(cfg)
+		scores, inds := router.Forward(x, cfg)
+		mlx.Eval(scores, inds)
 
-	scores, inds := router.Forward(x, cfg)
-	mlx.Eval(scores, inds)
+		out := moe.Forward(x, scores, inds, cfg)
+		mlx.Eval(out)
 
-	out := moe.Forward(x, scores, inds, cfg)
-	mlx.Eval(out)
+		outDims := out.Dims()
+		t.Logf("MoE output shape: %v", outDims)
 
-	outDims := out.Dims()
-	t.Logf("MoE output shape: %v", outDims)
-
-	if len(outDims) != 3 || outDims[0] != int(B) || outDims[1] != int(L) || outDims[2] != int(cfg.HiddenSize) {
-		t.Errorf("output shape = %v, want [%d, %d, %d]", outDims, B, L, cfg.HiddenSize)
-	}
+		if len(outDims) != 3 || outDims[0] != int(B) || outDims[1] != int(L) || outDims[2] != int(cfg.HiddenSize) {
+			t.Errorf("output shape = %v, want [%d, %d, %d]", outDims, B, L, cfg.HiddenSize)
+		}
+	})
 }
 
 func TestMoEBlockSortedForward(t *testing.T) {
-	useMLXTestThread(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := tinyMoEConfig()
+		B, L := int32(1), int32(128)
+		x := onesLike(int(B), int(L), int(cfg.HiddenSize))
+		router := newRouter(cfg)
+		moe := newMoEBlock(cfg)
 
-	cfg := tinyMoEConfig()
-	B, L := int32(1), int32(128)
-	x := onesLike(int(B), int(L), int(cfg.HiddenSize))
-	router := newRouter(cfg)
-	moe := newMoEBlock(cfg)
+		scores, inds := router.Forward(x, cfg)
+		mlx.Eval(scores, inds)
 
-	scores, inds := router.Forward(x, cfg)
-	mlx.Eval(scores, inds)
+		out := moe.Forward(x, scores, inds, cfg)
+		mlx.Eval(out)
 
-	out := moe.Forward(x, scores, inds, cfg)
-	mlx.Eval(out)
+		outDims := out.Dims()
+		t.Logf("MoE sorted output shape: %v", outDims)
 
-	outDims := out.Dims()
-	t.Logf("MoE sorted output shape: %v", outDims)
-
-	if len(outDims) != 3 || outDims[0] != int(B) || outDims[1] != int(L) || outDims[2] != int(cfg.HiddenSize) {
-		t.Errorf("output shape = %v, want [%d, %d, %d]", outDims, B, L, cfg.HiddenSize)
-	}
+		if len(outDims) != 3 || outDims[0] != int(B) || outDims[1] != int(L) || outDims[2] != int(cfg.HiddenSize) {
+			t.Errorf("output shape = %v, want [%d, %d, %d]", outDims, B, L, cfg.HiddenSize)
+		}
+	})
 }
 
 // TestLoadFusedExpertsQuantized verifies that a quantized, fused gate_up
@@ -147,7 +122,7 @@ func TestMoEBlockSortedForward(t *testing.T) {
 // through to the dense branch and was loaded unquantized (the memory bloat
 // bug this fix addresses).
 func TestLoadFusedExpertsQuantized(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.SkipIfUnavailable(t)
 
 	const E, I, H = 4, 8, 16
 	m := &Model{TextConfig: &TextConfig{QuantGroupSize: 16, QuantBits: 4, QuantMode: "nvfp4"}}
@@ -156,7 +131,7 @@ func TestLoadFusedExpertsQuantized(t *testing.T) {
 		"model.language_model.layers.0.experts",        // gemma HF (bare .experts.)
 		"model.language_model.layers.0.moe.switch_mlp", // create pipeline
 	} {
-		t.Run(prefix, func(t *testing.T) {
+		mlxtest.RunSubtest(t, prefix, func(t *mlxtest.T) {
 			gateUpKey := prefix + ".gate_up_proj"
 			downKey := prefix + ".down_proj"
 			tensors := map[string]*mlx.Array{
@@ -194,33 +169,33 @@ func TestLoadFusedExpertsQuantized(t *testing.T) {
 // TestLoadFusedExpertsDense verifies that a fused gate_up projection with no
 // scale companions is loaded onto the dense GatherMM path, kept fused.
 func TestLoadFusedExpertsDense(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		const E, I, H = 4, 8, 16
+		m := &Model{TextConfig: &TextConfig{}}
 
-	const E, I, H = 4, 8, 16
-	m := &Model{TextConfig: &TextConfig{}}
+		gateUpKey := "model.language_model.layers.0.experts.gate_up_proj"
+		downKey := "model.language_model.layers.0.experts.down_proj"
+		tensors := map[string]*mlx.Array{
+			gateUpKey: onesLike(E, 2*I, H),
+			downKey:   onesLike(E, I, H),
+		}
 
-	gateUpKey := "model.language_model.layers.0.experts.gate_up_proj"
-	downKey := "model.language_model.layers.0.experts.down_proj"
-	tensors := map[string]*mlx.Array{
-		gateUpKey: onesLike(E, 2*I, H),
-		downKey:   onesLike(E, I, H),
-	}
+		moe := &MoEBlock{}
+		m.loadFusedExperts(moe, tensors, gateUpKey, tensors[gateUpKey], downKey, tensors[downKey])
 
-	moe := &MoEBlock{}
-	m.loadFusedExperts(moe, tensors, gateUpKey, tensors[gateUpKey], downKey, tensors[downKey])
-
-	if moe.UseQuantized {
-		t.Error("UseQuantized = true, want false (no scales present)")
-	}
-	if !moe.UseFusedGateUp {
-		t.Error("UseFusedGateUp = false, want true")
-	}
-	if moe.GateUpWeight == nil || moe.DownWeight == nil {
-		t.Error("dense fused weights not set")
-	}
-	if moe.GateUpWeightQ != nil || moe.DownWeightQ != nil {
-		t.Error("quantized weights set on a dense block")
-	}
+		if moe.UseQuantized {
+			t.Error("UseQuantized = true, want false (no scales present)")
+		}
+		if !moe.UseFusedGateUp {
+			t.Error("UseFusedGateUp = false, want true")
+		}
+		if moe.GateUpWeight == nil || moe.DownWeight == nil {
+			t.Error("dense fused weights not set")
+		}
+		if moe.GateUpWeightQ != nil || moe.DownWeightQ != nil {
+			t.Error("quantized weights set on a dense block")
+		}
+	})
 }
 
 // TestRouterForwardMatchesLegacy verifies the optimized Router.Forward —
@@ -229,54 +204,54 @@ func TestLoadFusedExpertsDense(t *testing.T) {
 // normalized scores as the legacy path that softmaxes over every expert
 // first, gathers the top-k probabilities, then renormalizes.
 func TestRouterForwardMatchesLegacy(t *testing.T) {
-	useMLXTestThread(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &TextConfig{
+			HiddenSize:  8,
+			NumExperts:  4,
+			TopKExperts: 2,
+			RMSNormEps:  1e-6,
+			RouterScale: 0.5,
+		}
 
-	cfg := &TextConfig{
-		HiddenSize:  8,
-		NumExperts:  4,
-		TopKExperts: 2,
-		RMSNormEps:  1e-6,
-		RouterScale: 0.5,
-	}
+		// Distinct per-expert weight rows so top-k has a well-defined ordering
+		// (tied scores would let argpartition pick either tied expert and make
+		// the index comparison below flaky).
+		projWeight := mlx.FromValues([]float32{
+			0.10, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, // expert 0
+			0.30, 0.29, 0.28, 0.27, 0.26, 0.25, 0.24, 0.23, // expert 1
+			-0.05, -0.06, -0.07, -0.08, -0.09, -0.10, -0.11, -0.12, // expert 2
+			0.50, 0.48, 0.46, 0.44, 0.42, 0.40, 0.38, 0.36, // expert 3
+		}, int(cfg.NumExperts), int(cfg.HiddenSize))
 
-	// Distinct per-expert weight rows so top-k has a well-defined ordering
-	// (tied scores would let argpartition pick either tied expert and make
-	// the index comparison below flaky).
-	projWeight := mlx.FromValues([]float32{
-		0.10, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, // expert 0
-		0.30, 0.29, 0.28, 0.27, 0.26, 0.25, 0.24, 0.23, // expert 1
-		-0.05, -0.06, -0.07, -0.08, -0.09, -0.10, -0.11, -0.12, // expert 2
-		0.50, 0.48, 0.46, 0.44, 0.42, 0.40, 0.38, 0.36, // expert 3
-	}, int(cfg.NumExperts), int(cfg.HiddenSize))
+		scale := mlx.FromValues([]float32{
+			1.0, 0.9, 1.1, 1.0, 1.2, 0.8, 1.0, 1.05,
+		}, int(cfg.HiddenSize))
 
-	scale := mlx.FromValues([]float32{
-		1.0, 0.9, 1.1, 1.0, 1.2, 0.8, 1.0, 1.05,
-	}, int(cfg.HiddenSize))
+		r := &Router{
+			Proj:  linearFromWeight(projWeight),
+			Scale: scale,
+		}
 
-	r := &Router{
-		Proj:  linearFromWeight(projWeight),
-		Scale: scale,
-	}
+		// Varied x so different positions potentially hit different top-k.
+		x := mlx.FromValues([]float32{
+			0.2, -0.1, 0.3, 0.0, 0.4, -0.2, 0.1, 0.05,
+			-0.3, 0.2, -0.1, 0.4, -0.05, 0.3, 0.0, 0.2,
+			0.5, 0.4, -0.2, 0.1, -0.3, 0.0, 0.3, -0.1,
+		}, 1, 3, int(cfg.HiddenSize))
 
-	// Varied x so different positions potentially hit different top-k.
-	x := mlx.FromValues([]float32{
-		0.2, -0.1, 0.3, 0.0, 0.4, -0.2, 0.1, 0.05,
-		-0.3, 0.2, -0.1, 0.4, -0.05, 0.3, 0.0, 0.2,
-		0.5, 0.4, -0.2, 0.1, -0.3, 0.0, 0.3, -0.1,
-	}, 1, 3, int(cfg.HiddenSize))
+		gotScores, gotInds := r.Forward(x, cfg)
+		wantScores, wantInds := legacyRouterForward(r, x, cfg)
+		gotInds = gotInds.AsType(mlx.DTypeInt32)
+		wantInds = wantInds.AsType(mlx.DTypeInt32)
+		mlx.Eval(gotScores, gotInds, wantScores, wantInds)
 
-	gotScores, gotInds := r.Forward(x, cfg)
-	wantScores, wantInds := legacyRouterForward(r, x, cfg)
-	gotInds = gotInds.AsType(mlx.DTypeInt32)
-	wantInds = wantInds.AsType(mlx.DTypeInt32)
-	mlx.Eval(gotScores, gotInds, wantScores, wantInds)
-
-	if got, want := gotInds.Ints(), wantInds.Ints(); !intSlicesEqual(got, want) {
-		t.Fatalf("indices mismatch:\n  got  %v\n  want %v", got, want)
-	}
-	if got, want := gotScores.Floats(), wantScores.Floats(); !floatSlicesClose(got, want, 1e-5) {
-		t.Fatalf("scores mismatch:\n  got  %v\n  want %v", got, want)
-	}
+		if got, want := gotInds.Ints(), wantInds.Ints(); !intSlicesEqual(got, want) {
+			t.Fatalf("indices mismatch:\n  got  %v\n  want %v", got, want)
+		}
+		if got, want := gotScores.Floats(), wantScores.Floats(); !floatSlicesClose(got, want, 1e-5) {
+			t.Fatalf("scores mismatch:\n  got  %v\n  want %v", got, want)
+		}
+	})
 }
 
 // legacyRouterForward implements the pre-optimization router: full softmax

x/models/gemma4/gemma4_test.go

@@ -5,6 +5,7 @@ import (
 	"slices"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
@@ -25,8 +26,8 @@ func TestParseSuppressTokens(t *testing.T) {
 }
 
 func TestParseTextConfigE2B(t *testing.T) {
-	skipIfNoMLX(t)
-	data := []byte(`{
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
 			"hidden_size": 1536,
@@ -71,78 +72,79 @@ func TestParseTextConfigE2B(t *testing.T) {
 		}
 	}`)
 
-	cfg, err := parseTextConfig(data)
-	if err != nil {
-		t.Fatalf("parseTextConfig failed: %v", err)
-	}
+		cfg, err := parseTextConfig(data)
+		if err != nil {
+			t.Fatalf("parseTextConfig failed: %v", err)
+		}
 
-	// Basic fields.
-	if cfg.HiddenSize != 1536 {
-		t.Errorf("HiddenSize = %d, want 1536", cfg.HiddenSize)
-	}
-	if cfg.NumHiddenLayers != 35 {
-		t.Errorf("NumHiddenLayers = %d, want 35", cfg.NumHiddenLayers)
-	}
-	if cfg.GlobalHeadDim != 512 {
-		t.Errorf("GlobalHeadDim = %d, want 512", cfg.GlobalHeadDim)
-	}
-	if cfg.FinalLogitSoftcapping != 30.0 {
-		t.Errorf("FinalLogitSoftcapping = %f, want 30.0", cfg.FinalLogitSoftcapping)
-	}
-	if cfg.NumKVSharedLayers != 20 {
-		t.Errorf("NumKVSharedLayers = %d, want 20", cfg.NumKVSharedLayers)
-	}
-	if cfg.HiddenSizePerLayer != 256 {
-		t.Errorf("HiddenSizePerLayer = %d, want 256", cfg.HiddenSizePerLayer)
-	}
+		// Basic fields.
+		if cfg.HiddenSize != 1536 {
+			t.Errorf("HiddenSize = %d, want 1536", cfg.HiddenSize)
+		}
+		if cfg.NumHiddenLayers != 35 {
+			t.Errorf("NumHiddenLayers = %d, want 35", cfg.NumHiddenLayers)
+		}
+		if cfg.GlobalHeadDim != 512 {
+			t.Errorf("GlobalHeadDim = %d, want 512", cfg.GlobalHeadDim)
+		}
+		if cfg.FinalLogitSoftcapping != 30.0 {
+			t.Errorf("FinalLogitSoftcapping = %f, want 30.0", cfg.FinalLogitSoftcapping)
+		}
+		if cfg.NumKVSharedLayers != 20 {
+			t.Errorf("NumKVSharedLayers = %d, want 20", cfg.NumKVSharedLayers)
+		}
+		if cfg.HiddenSizePerLayer != 256 {
+			t.Errorf("HiddenSizePerLayer = %d, want 256", cfg.HiddenSizePerLayer)
+		}
 
-	// RoPE settings.
-	if cfg.SlidingRopeDims != 256 {
-		t.Errorf("SlidingRopeDims = %d, want 256", cfg.SlidingRopeDims)
-	}
-	if cfg.FullRopeDims != 512 {
-		t.Errorf("FullRopeDims = %d, want 512 (GlobalHeadDim, partial rotation handled via custom freqs)", cfg.FullRopeDims)
-	}
-	if cfg.SlidingRopeBase != 10000 {
-		t.Errorf("SlidingRopeBase = %f, want 10000", cfg.SlidingRopeBase)
-	}
-	if cfg.FullRopeBase != 1000000 {
-		t.Errorf("FullRopeBase = %f, want 1000000", cfg.FullRopeBase)
-	}
+		// RoPE settings.
+		if cfg.SlidingRopeDims != 256 {
+			t.Errorf("SlidingRopeDims = %d, want 256", cfg.SlidingRopeDims)
+		}
+		if cfg.FullRopeDims != 512 {
+			t.Errorf("FullRopeDims = %d, want 512 (GlobalHeadDim, partial rotation handled via custom freqs)", cfg.FullRopeDims)
+		}
+		if cfg.SlidingRopeBase != 10000 {
+			t.Errorf("SlidingRopeBase = %f, want 10000", cfg.SlidingRopeBase)
+		}
+		if cfg.FullRopeBase != 1000000 {
+			t.Errorf("FullRopeBase = %f, want 1000000", cfg.FullRopeBase)
+		}
 
-	// Attention scale.
-	if cfg.SlidingScale == 0 || cfg.FullScale == 0 {
-		t.Error("attention scales should be non-zero")
-	}
+		// Attention scale.
+		if cfg.SlidingScale == 0 || cfg.FullScale == 0 {
+			t.Error("attention scales should be non-zero")
+		}
 
-	// KV sharing map.
-	// First shared layer is 35 - 20 = 15.
-	if donor, ok := cfg.KVShareMap[15]; !ok || donor != 13 {
-		t.Errorf("KVShareMap[15] = %d, ok=%v; want 13, true", donor, ok)
-	}
-	if donor, ok := cfg.KVShareMap[19]; !ok || donor != 14 {
-		t.Errorf("KVShareMap[19] = %d, ok=%v; want 14, true (full attn donor)", donor, ok)
-	}
-	if donor, ok := cfg.KVShareMap[34]; !ok || donor != 14 {
-		t.Errorf("KVShareMap[34] = %d, ok=%v; want 14, true (full attn donor)", donor, ok)
-	}
-	// Layer 14 should not be shared.
-	if _, ok := cfg.KVShareMap[14]; ok {
-		t.Error("layer 14 should not be in KVShareMap (non-shared)")
-	}
+		// KV sharing map.
+		// First shared layer is 35 - 20 = 15.
+		if donor, ok := cfg.KVShareMap[15]; !ok || donor != 13 {
+			t.Errorf("KVShareMap[15] = %d, ok=%v; want 13, true", donor, ok)
+		}
+		if donor, ok := cfg.KVShareMap[19]; !ok || donor != 14 {
+			t.Errorf("KVShareMap[19] = %d, ok=%v; want 14, true (full attn donor)", donor, ok)
+		}
+		if donor, ok := cfg.KVShareMap[34]; !ok || donor != 14 {
+			t.Errorf("KVShareMap[34] = %d, ok=%v; want 14, true (full attn donor)", donor, ok)
+		}
+		// Layer 14 should not be shared.
+		if _, ok := cfg.KVShareMap[14]; ok {
+			t.Error("layer 14 should not be in KVShareMap (non-shared)")
+		}
 
-	// Donors.
-	if !cfg.KVDonors[13] {
-		t.Error("layer 13 should be a KV donor")
-	}
-	if !cfg.KVDonors[14] {
-		t.Error("layer 14 should be a KV donor")
-	}
+		// Donors.
+		if !cfg.KVDonors[13] {
+			t.Error("layer 13 should be a KV donor")
+		}
+		if !cfg.KVDonors[14] {
+			t.Error("layer 14 should be a KV donor")
+		}
+	})
 }
 
 func TestParseTextConfig26B(t *testing.T) {
-	skipIfNoMLX(t)
-	data := []byte(`{
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
 			"hidden_size": 2816,
@@ -188,40 +190,41 @@ func TestParseTextConfig26B(t *testing.T) {
 		}
 	}`)
 
-	cfg, err := parseTextConfig(data)
-	if err != nil {
-		t.Fatalf("parseTextConfig failed: %v", err)
-	}
+		cfg, err := parseTextConfig(data)
+		if err != nil {
+			t.Fatalf("parseTextConfig failed: %v", err)
+		}
 
-	if cfg.HiddenSize != 2816 {
-		t.Errorf("HiddenSize = %d, want 2816", cfg.HiddenSize)
-	}
-	if !cfg.AttentionKEqV {
-		t.Error("AttentionKEqV should be true")
-	}
-	if cfg.NumGlobalKeyValueHeads != 2 {
-		t.Errorf("NumGlobalKeyValueHeads = %d, want 2", cfg.NumGlobalKeyValueHeads)
-	}
-	if !cfg.EnableMoeBlock {
-		t.Error("EnableMoeBlock should be true")
-	}
-	if cfg.NumExperts != 128 {
-		t.Errorf("NumExperts = %d, want 128", cfg.NumExperts)
-	}
-	if cfg.TopKExperts != 8 {
-		t.Errorf("TopKExperts = %d, want 8", cfg.TopKExperts)
-	}
-	if cfg.ExpertIntermediateSize != 704 {
-		t.Errorf("ExpertIntermediateSize = %d, want 704", cfg.ExpertIntermediateSize)
-	}
-	if cfg.HiddenSizePerLayer != 0 {
-		t.Errorf("HiddenSizePerLayer = %d, want 0 (no PLE)", cfg.HiddenSizePerLayer)
-	}
+		if cfg.HiddenSize != 2816 {
+			t.Errorf("HiddenSize = %d, want 2816", cfg.HiddenSize)
+		}
+		if !cfg.AttentionKEqV {
+			t.Error("AttentionKEqV should be true")
+		}
+		if cfg.NumGlobalKeyValueHeads != 2 {
+			t.Errorf("NumGlobalKeyValueHeads = %d, want 2", cfg.NumGlobalKeyValueHeads)
+		}
+		if !cfg.EnableMoeBlock {
+			t.Error("EnableMoeBlock should be true")
+		}
+		if cfg.NumExperts != 128 {
+			t.Errorf("NumExperts = %d, want 128", cfg.NumExperts)
+		}
+		if cfg.TopKExperts != 8 {
+			t.Errorf("TopKExperts = %d, want 8", cfg.TopKExperts)
+		}
+		if cfg.ExpertIntermediateSize != 704 {
+			t.Errorf("ExpertIntermediateSize = %d, want 704", cfg.ExpertIntermediateSize)
+		}
+		if cfg.HiddenSizePerLayer != 0 {
+			t.Errorf("HiddenSizePerLayer = %d, want 0 (no PLE)", cfg.HiddenSizePerLayer)
+		}
+	})
 }
 
 func TestParseTextConfig31B(t *testing.T) {
-	skipIfNoMLX(t)
-	data := []byte(`{
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
 			"hidden_size": 5376,
@@ -268,168 +271,169 @@ func TestParseTextConfig31B(t *testing.T) {
 		}
 	}`)
 
-	cfg, err := parseTextConfig(data)
-	if err != nil {
-		t.Fatalf("parseTextConfig failed: %v", err)
-	}
+		cfg, err := parseTextConfig(data)
+		if err != nil {
+			t.Fatalf("parseTextConfig failed: %v", err)
+		}
 
-	if cfg.HiddenSize != 5376 {
-		t.Errorf("HiddenSize = %d, want 5376", cfg.HiddenSize)
-	}
-	if cfg.NumHiddenLayers != 60 {
-		t.Errorf("NumHiddenLayers = %d, want 60", cfg.NumHiddenLayers)
-	}
-	if !cfg.AttentionKEqV {
-		t.Error("AttentionKEqV should be true")
-	}
-	if cfg.NumGlobalKeyValueHeads != 4 {
-		t.Errorf("NumGlobalKeyValueHeads = %d, want 4", cfg.NumGlobalKeyValueHeads)
-	}
-	if cfg.NumKeyValueHeads != 16 {
-		t.Errorf("NumKeyValueHeads = %d, want 16", cfg.NumKeyValueHeads)
-	}
-	if cfg.NumKVSharedLayers != 0 {
-		t.Errorf("NumKVSharedLayers = %d, want 0", cfg.NumKVSharedLayers)
-	}
-	if cfg.HiddenSizePerLayer != 0 {
-		t.Errorf("HiddenSizePerLayer = %d, want 0 (no PLE)", cfg.HiddenSizePerLayer)
-	}
-	if cfg.SlidingWindow != 1024 {
-		t.Errorf("SlidingWindow = %d, want 1024", cfg.SlidingWindow)
-	}
+		if cfg.HiddenSize != 5376 {
+			t.Errorf("HiddenSize = %d, want 5376", cfg.HiddenSize)
+		}
+		if cfg.NumHiddenLayers != 60 {
+			t.Errorf("NumHiddenLayers = %d, want 60", cfg.NumHiddenLayers)
+		}
+		if !cfg.AttentionKEqV {
+			t.Error("AttentionKEqV should be true")
+		}
+		if cfg.NumGlobalKeyValueHeads != 4 {
+			t.Errorf("NumGlobalKeyValueHeads = %d, want 4", cfg.NumGlobalKeyValueHeads)
+		}
+		if cfg.NumKeyValueHeads != 16 {
+			t.Errorf("NumKeyValueHeads = %d, want 16", cfg.NumKeyValueHeads)
+		}
+		if cfg.NumKVSharedLayers != 0 {
+			t.Errorf("NumKVSharedLayers = %d, want 0", cfg.NumKVSharedLayers)
+		}
+		if cfg.HiddenSizePerLayer != 0 {
+			t.Errorf("HiddenSizePerLayer = %d, want 0 (no PLE)", cfg.HiddenSizePerLayer)
+		}
+		if cfg.SlidingWindow != 1024 {
+			t.Errorf("SlidingWindow = %d, want 1024", cfg.SlidingWindow)
+		}
 
-	// KV sharing should be empty (no shared layers).
-	if len(cfg.KVShareMap) != 0 {
-		t.Errorf("KVShareMap should be empty, got %d entries", len(cfg.KVShareMap))
-	}
+		// KV sharing should be empty (no shared layers).
+		if len(cfg.KVShareMap) != 0 {
+			t.Errorf("KVShareMap should be empty, got %d entries", len(cfg.KVShareMap))
+		}
 
-	// Layer types: pattern is 5 sliding + 1 full, repeating 10 times.
-	if !isLayerSliding(0, &cfg) {
-		t.Error("layer 0 should be sliding")
-	}
-	if isLayerSliding(5, &cfg) {
-		t.Error("layer 5 should be full attention")
-	}
-	if !isLayerSliding(6, &cfg) {
-		t.Error("layer 6 should be sliding")
-	}
-	if isLayerSliding(59, &cfg) {
-		t.Error("layer 59 should be full attention")
-	}
+		// Layer types: pattern is 5 sliding + 1 full, repeating 10 times.
+		if !isLayerSliding(0, &cfg) {
+			t.Error("layer 0 should be sliding")
+		}
+		if isLayerSliding(5, &cfg) {
+			t.Error("layer 5 should be full attention")
+		}
+		if !isLayerSliding(6, &cfg) {
+			t.Error("layer 6 should be sliding")
+		}
+		if isLayerSliding(59, &cfg) {
+			t.Error("layer 59 should be full attention")
+		}
+	})
 }
 
 func TestParseTextConfig12BUnified(t *testing.T) {
-	skipIfNoMLX(t)
-
-	layerTypes := make([]string, 0, 48)
-	for i := range 48 {
-		if i%6 == 5 {
-			layerTypes = append(layerTypes, "full_attention")
-		} else {
-			layerTypes = append(layerTypes, "sliding_attention")
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		layerTypes := make([]string, 0, 48)
+		for i := range 48 {
+			if i%6 == 5 {
+				layerTypes = append(layerTypes, "full_attention")
+			} else {
+				layerTypes = append(layerTypes, "sliding_attention")
+			}
 		}
-	}
 
-	data, err := json.Marshal(map[string]any{
-		"architectures": []string{"Gemma4UnifiedForConditionalGeneration"},
-		"model_type":    "gemma4_unified",
-		"text_config": map[string]any{
-			"hidden_size":                 3840,
-			"num_hidden_layers":           48,
-			"intermediate_size":           15360,
-			"num_attention_heads":         16,
-			"num_key_value_heads":         8,
-			"num_global_key_value_heads":  1,
-			"head_dim":                    256,
-			"global_head_dim":             512,
-			"vocab_size":                  262144,
-			"rms_norm_eps":                1e-6,
-			"max_position_embeddings":     131072,
-			"sliding_window":              1024,
-			"final_logit_softcapping":     30.0,
-			"use_double_wide_mlp":         false,
-			"num_kv_shared_layers":        0,
-			"hidden_size_per_layer_input": 0,
-			"vocab_size_per_layer_input":  262144,
-			"attention_k_eq_v":            true,
-			"enable_moe_block":            false,
-			"tie_word_embeddings":         true,
-			"layer_types":                 layerTypes,
-			"rope_parameters": map[string]any{
-				"full_attention": map[string]any{
-					"partial_rotary_factor": 0.25,
-					"rope_theta":            1000000.0,
-					"rope_type":             "proportional",
-				},
-				"sliding_attention": map[string]any{
-					"rope_theta": 10000.0,
-					"rope_type":  "default",
+		data, err := json.Marshal(map[string]any{
+			"architectures": []string{"Gemma4UnifiedForConditionalGeneration"},
+			"model_type":    "gemma4_unified",
+			"text_config": map[string]any{
+				"hidden_size":                 3840,
+				"num_hidden_layers":           48,
+				"intermediate_size":           15360,
+				"num_attention_heads":         16,
+				"num_key_value_heads":         8,
+				"num_global_key_value_heads":  1,
+				"head_dim":                    256,
+				"global_head_dim":             512,
+				"vocab_size":                  262144,
+				"rms_norm_eps":                1e-6,
+				"max_position_embeddings":     131072,
+				"sliding_window":              1024,
+				"final_logit_softcapping":     30.0,
+				"use_double_wide_mlp":         false,
+				"num_kv_shared_layers":        0,
+				"hidden_size_per_layer_input": 0,
+				"vocab_size_per_layer_input":  262144,
+				"attention_k_eq_v":            true,
+				"enable_moe_block":            false,
+				"tie_word_embeddings":         true,
+				"layer_types":                 layerTypes,
+				"rope_parameters": map[string]any{
+					"full_attention": map[string]any{
+						"partial_rotary_factor": 0.25,
+						"rope_theta":            1000000.0,
+						"rope_type":             "proportional",
+					},
+					"sliding_attention": map[string]any{
+						"rope_theta": 10000.0,
+						"rope_type":  "default",
+					},
 				},
 			},
-		},
+		})
+		if err != nil {
+			t.Fatalf("json.Marshal failed: %v", err)
+		}
+
+		cfg, err := parseTextConfig(data)
+		if err != nil {
+			t.Fatalf("parseTextConfig failed: %v", err)
+		}
+
+		if cfg.HiddenSize != 3840 {
+			t.Errorf("HiddenSize = %d, want 3840", cfg.HiddenSize)
+		}
+		if cfg.NumHiddenLayers != 48 {
+			t.Errorf("NumHiddenLayers = %d, want 48", cfg.NumHiddenLayers)
+		}
+		if cfg.IntermediateSize != 15360 {
+			t.Errorf("IntermediateSize = %d, want 15360", cfg.IntermediateSize)
+		}
+		if cfg.NumAttentionHeads != 16 {
+			t.Errorf("NumAttentionHeads = %d, want 16", cfg.NumAttentionHeads)
+		}
+		if cfg.NumKeyValueHeads != 8 {
+			t.Errorf("NumKeyValueHeads = %d, want 8", cfg.NumKeyValueHeads)
+		}
+		if cfg.NumGlobalKeyValueHeads != 1 {
+			t.Errorf("NumGlobalKeyValueHeads = %d, want 1", cfg.NumGlobalKeyValueHeads)
+		}
+		if !cfg.AttentionKEqV {
+			t.Error("AttentionKEqV should be true")
+		}
+		if cfg.EnableMoeBlock {
+			t.Error("EnableMoeBlock should be false")
+		}
+		if cfg.HiddenSizePerLayer != 0 {
+			t.Errorf("HiddenSizePerLayer = %d, want 0", cfg.HiddenSizePerLayer)
+		}
+		if cfg.NumKVSharedLayers != 0 {
+			t.Errorf("NumKVSharedLayers = %d, want 0", cfg.NumKVSharedLayers)
+		}
+		if len(cfg.KVShareMap) != 0 {
+			t.Errorf("KVShareMap should be empty, got %d entries", len(cfg.KVShareMap))
+		}
+		if cfg.FullRopeDims != 512 {
+			t.Errorf("FullRopeDims = %d, want 512", cfg.FullRopeDims)
+		}
+		if cfg.FullRopeFreqs == nil {
+			t.Error("FullRopeFreqs should be precomputed for proportional RoPE")
+		}
+		if isLayerSliding(5, &cfg) {
+			t.Error("layer 5 should be full attention")
+		}
+		if !isLayerSliding(6, &cfg) {
+			t.Error("layer 6 should be sliding")
+		}
+		if isLayerSliding(47, &cfg) {
+			t.Error("layer 47 should be full attention")
+		}
 	})
-	if err != nil {
-		t.Fatalf("json.Marshal failed: %v", err)
-	}
-
-	cfg, err := parseTextConfig(data)
-	if err != nil {
-		t.Fatalf("parseTextConfig failed: %v", err)
-	}
-
-	if cfg.HiddenSize != 3840 {
-		t.Errorf("HiddenSize = %d, want 3840", cfg.HiddenSize)
-	}
-	if cfg.NumHiddenLayers != 48 {
-		t.Errorf("NumHiddenLayers = %d, want 48", cfg.NumHiddenLayers)
-	}
-	if cfg.IntermediateSize != 15360 {
-		t.Errorf("IntermediateSize = %d, want 15360", cfg.IntermediateSize)
-	}
-	if cfg.NumAttentionHeads != 16 {
-		t.Errorf("NumAttentionHeads = %d, want 16", cfg.NumAttentionHeads)
-	}
-	if cfg.NumKeyValueHeads != 8 {
-		t.Errorf("NumKeyValueHeads = %d, want 8", cfg.NumKeyValueHeads)
-	}
-	if cfg.NumGlobalKeyValueHeads != 1 {
-		t.Errorf("NumGlobalKeyValueHeads = %d, want 1", cfg.NumGlobalKeyValueHeads)
-	}
-	if !cfg.AttentionKEqV {
-		t.Error("AttentionKEqV should be true")
-	}
-	if cfg.EnableMoeBlock {
-		t.Error("EnableMoeBlock should be false")
-	}
-	if cfg.HiddenSizePerLayer != 0 {
-		t.Errorf("HiddenSizePerLayer = %d, want 0", cfg.HiddenSizePerLayer)
-	}
-	if cfg.NumKVSharedLayers != 0 {
-		t.Errorf("NumKVSharedLayers = %d, want 0", cfg.NumKVSharedLayers)
-	}
-	if len(cfg.KVShareMap) != 0 {
-		t.Errorf("KVShareMap should be empty, got %d entries", len(cfg.KVShareMap))
-	}
-	if cfg.FullRopeDims != 512 {
-		t.Errorf("FullRopeDims = %d, want 512", cfg.FullRopeDims)
-	}
-	if cfg.FullRopeFreqs == nil {
-		t.Error("FullRopeFreqs should be precomputed for proportional RoPE")
-	}
-	if isLayerSliding(5, &cfg) {
-		t.Error("layer 5 should be full attention")
-	}
-	if !isLayerSliding(6, &cfg) {
-		t.Error("layer 6 should be sliding")
-	}
-	if isLayerSliding(47, &cfg) {
-		t.Error("layer 47 should be full attention")
-	}
 }
 
 func TestParseTextConfigE4B(t *testing.T) {
-	skipIfNoMLX(t)
-	data := []byte(`{
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
 			"hidden_size": 2560,
@@ -474,73 +478,74 @@ func TestParseTextConfigE4B(t *testing.T) {
 		}
 	}`)
 
-	cfg, err := parseTextConfig(data)
-	if err != nil {
-		t.Fatalf("parseTextConfig failed: %v", err)
-	}
+		cfg, err := parseTextConfig(data)
+		if err != nil {
+			t.Fatalf("parseTextConfig failed: %v", err)
+		}
 
-	if cfg.HiddenSize != 2560 {
-		t.Errorf("HiddenSize = %d, want 2560", cfg.HiddenSize)
-	}
-	if cfg.NumHiddenLayers != 42 {
-		t.Errorf("NumHiddenLayers = %d, want 42", cfg.NumHiddenLayers)
-	}
-	if cfg.IntermediateSize != 10240 {
-		t.Errorf("IntermediateSize = %d, want 10240", cfg.IntermediateSize)
-	}
-	if cfg.NumKeyValueHeads != 2 {
-		t.Errorf("NumKeyValueHeads = %d, want 2", cfg.NumKeyValueHeads)
-	}
-	if cfg.UseDoubleWideMLP {
-		t.Error("UseDoubleWideMLP should be false")
-	}
-	if cfg.NumKVSharedLayers != 18 {
-		t.Errorf("NumKVSharedLayers = %d, want 18", cfg.NumKVSharedLayers)
-	}
-	if cfg.HiddenSizePerLayer != 256 {
-		t.Errorf("HiddenSizePerLayer = %d, want 256 (has PLE)", cfg.HiddenSizePerLayer)
-	}
-	if cfg.AttentionKEqV {
-		t.Error("AttentionKEqV should be false")
-	}
-	if cfg.EnableMoeBlock {
-		t.Error("EnableMoeBlock should be false")
-	}
-	if cfg.SlidingWindow != 512 {
-		t.Errorf("SlidingWindow = %d, want 512", cfg.SlidingWindow)
-	}
+		if cfg.HiddenSize != 2560 {
+			t.Errorf("HiddenSize = %d, want 2560", cfg.HiddenSize)
+		}
+		if cfg.NumHiddenLayers != 42 {
+			t.Errorf("NumHiddenLayers = %d, want 42", cfg.NumHiddenLayers)
+		}
+		if cfg.IntermediateSize != 10240 {
+			t.Errorf("IntermediateSize = %d, want 10240", cfg.IntermediateSize)
+		}
+		if cfg.NumKeyValueHeads != 2 {
+			t.Errorf("NumKeyValueHeads = %d, want 2", cfg.NumKeyValueHeads)
+		}
+		if cfg.UseDoubleWideMLP {
+			t.Error("UseDoubleWideMLP should be false")
+		}
+		if cfg.NumKVSharedLayers != 18 {
+			t.Errorf("NumKVSharedLayers = %d, want 18", cfg.NumKVSharedLayers)
+		}
+		if cfg.HiddenSizePerLayer != 256 {
+			t.Errorf("HiddenSizePerLayer = %d, want 256 (has PLE)", cfg.HiddenSizePerLayer)
+		}
+		if cfg.AttentionKEqV {
+			t.Error("AttentionKEqV should be false")
+		}
+		if cfg.EnableMoeBlock {
+			t.Error("EnableMoeBlock should be false")
+		}
+		if cfg.SlidingWindow != 512 {
+			t.Errorf("SlidingWindow = %d, want 512", cfg.SlidingWindow)
+		}
 
-	// Layer types: pattern is 5 sliding + 1 full, repeating 7 times = 42 layers.
-	if !isLayerSliding(0, &cfg) {
-		t.Error("layer 0 should be sliding")
-	}
-	if isLayerSliding(5, &cfg) {
-		t.Error("layer 5 should be full attention")
-	}
-	if !isLayerSliding(6, &cfg) {
-		t.Error("layer 6 should be sliding")
-	}
-	if isLayerSliding(41, &cfg) {
-		t.Error("layer 41 should be full attention")
-	}
+		// Layer types: pattern is 5 sliding + 1 full, repeating 7 times = 42 layers.
+		if !isLayerSliding(0, &cfg) {
+			t.Error("layer 0 should be sliding")
+		}
+		if isLayerSliding(5, &cfg) {
+			t.Error("layer 5 should be full attention")
+		}
+		if !isLayerSliding(6, &cfg) {
+			t.Error("layer 6 should be sliding")
+		}
+		if isLayerSliding(41, &cfg) {
+			t.Error("layer 41 should be full attention")
+		}
 
-	// KV sharing: first shared = 42 - 18 = 24.
-	// Layer 24 is sliding, its donor should be the last non-shared sliding layer.
-	// Non-shared layers: 0-23. Last sliding in 0-23 is layer 22 (23=full).
-	if donor, ok := cfg.KVShareMap[24]; !ok {
-		t.Error("layer 24 should be in KVShareMap")
-	} else {
-		t.Logf("layer 24 donor = %d", donor)
-	}
-	// Layer 29 is full_attention (5th full), donor should be the last non-shared full layer.
-	// Non-shared full layers: 5, 11, 17, 23.
-	if donor, ok := cfg.KVShareMap[29]; !ok || donor != 23 {
-		t.Errorf("KVShareMap[29] = %d, ok=%v; want 23, true (full attn donor)", donor, ok)
-	}
-	// Layer 23 should NOT be shared (it's the last non-shared layer).
-	if _, ok := cfg.KVShareMap[23]; ok {
-		t.Error("layer 23 should not be in KVShareMap (non-shared)")
-	}
+		// KV sharing: first shared = 42 - 18 = 24.
+		// Layer 24 is sliding, its donor should be the last non-shared sliding layer.
+		// Non-shared layers: 0-23. Last sliding in 0-23 is layer 22 (23=full).
+		if donor, ok := cfg.KVShareMap[24]; !ok {
+			t.Error("layer 24 should be in KVShareMap")
+		} else {
+			t.Logf("layer 24 donor = %d", donor)
+		}
+		// Layer 29 is full_attention (5th full), donor should be the last non-shared full layer.
+		// Non-shared full layers: 5, 11, 17, 23.
+		if donor, ok := cfg.KVShareMap[29]; !ok || donor != 23 {
+			t.Errorf("KVShareMap[29] = %d, ok=%v; want 23, true (full attn donor)", donor, ok)
+		}
+		// Layer 23 should NOT be shared (it's the last non-shared layer).
+		if _, ok := cfg.KVShareMap[23]; ok {
+			t.Error("layer 23 should not be in KVShareMap (non-shared)")
+		}
+	})
 }
 
 func TestLayerTypeDetection(t *testing.T) {
@@ -644,9 +649,7 @@ func TestNewCachesAssistantSharedHistoryOrdering(t *testing.T) {
 }
 
 func TestResolveWeightPrefix(t *testing.T) {
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
+	mlxtest.SkipIfUnavailable(t)
 
 	tests := []struct {
 		name    string
@@ -659,7 +662,7 @@ func TestResolveWeightPrefix(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
 			dummy := mlx.FromValue(float32(1.0))
 			mlx.Eval(dummy)
 			tensors := map[string]*mlx.Array{tt.key: dummy}
@@ -670,10 +673,3 @@ func TestResolveWeightPrefix(t *testing.T) {
 		})
 	}
 }
-
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}

x/models/glimmer/glimmer_test.go

@@ -226,21 +226,21 @@ func TestNewCachesMatchesAttentionSchedule(t *testing.T) {
 }
 
 func TestUnembedReturnsFloat32Logits(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		input := mlx.FromValues([]float32{1}, 1, 1, 1).AsType(mlx.DTypeBFloat16)
+		weight := mlx.FromValues([]float32{1, 2}, 2, 1).AsType(mlx.DTypeBFloat16)
+		m := Model{
+			LMHead: nn.NewLinear(weight, nil),
+			Config: &Config{
+				OutputMultiplier:  0.19611613,
+				OutputSoftCapTemp: 20,
+			},
+		}
 
-	input := mlx.FromValues([]float32{1}, 1, 1, 1).AsType(mlx.DTypeBFloat16)
-	weight := mlx.FromValues([]float32{1, 2}, 2, 1).AsType(mlx.DTypeBFloat16)
-	m := Model{
-		LMHead: nn.NewLinear(weight, nil),
-		Config: &Config{
-			OutputMultiplier:  0.19611613,
-			OutputSoftCapTemp: 20,
-		},
-	}
-
-	if got := m.Unembed(input).DType(); got != mlx.DTypeFloat32 {
-		t.Fatalf("Unembed() dtype = %v, want %v", got, mlx.DTypeFloat32)
-	}
+		if got := m.Unembed(input).DType(); got != mlx.DTypeFloat32 {
+			t.Fatalf("Unembed() dtype = %v, want %v", got, mlx.DTypeFloat32)
+		}
+	})
 }
 
 func TestComputeImageSizeMatchesReference(t *testing.T) {
@@ -320,9 +320,7 @@ func TestApplyVisionRoPELayouts(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mlxtest.Setup(t)
-
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
 			x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 1, 1, 4)
 			cos := mlx.FromValues([]float32{0.5, 0.25}, 1, 2)
 			sin := mlx.FromValues([]float32{0.75, 0.125}, 1, 2)

x/models/nemotron_h/nemotron_h_test.go

@@ -323,52 +323,52 @@ func TestSupportsGatherQMM(t *testing.T) {
 }
 
 func TestApplyExpertWeightGlobalScale(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		weight := mlx.FromValues([]float32{
+			1, 2,
+			3, 4,
+			5, 6,
+			7, 8,
+		}, 2, 2, 2)
+		scale := mlx.FromValues([]float32{2, 3}, 2)
 
-	weight := mlx.FromValues([]float32{
-		1, 2,
-		3, 4,
-		5, 6,
-		7, 8,
-	}, 2, 2, 2)
-	scale := mlx.FromValues([]float32{2, 3}, 2)
+		got := applyExpertWeightGlobalScale(weight, scale)
+		mlx.Eval(got)
 
-	got := applyExpertWeightGlobalScale(weight, scale)
-	mlx.Eval(got)
-
-	assertAllClose(t, "scaled expert weight", got.Floats(), []float32{
-		2, 4,
-		6, 8,
-		15, 18,
-		21, 24,
-	}, 1e-5)
+		assertAllClose(t, "scaled expert weight", got.Floats(), []float32{
+			2, 4,
+			6, 8,
+			15, 18,
+			21, 24,
+		}, 1e-5)
+	})
 }
 
 // The production form leans on MLX's fused fast RMS norm, so check it against
 // a reference that spells out every step.
 func TestGatedGroupRMSNormMatchesElementwiseReference(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{MambaNumHeads: 4, MambaHeadDim: 2, NGroups: 2, LayerNormEpsilon: 1e-5}
+		inner := cfg.MambaNumHeads * cfg.MambaHeadDim
+		groupSize := inner / cfg.NGroups
+		B, L := int32(2), int32(3)
 
-	cfg := &Config{MambaNumHeads: 4, MambaHeadDim: 2, NGroups: 2, LayerNormEpsilon: 1e-5}
-	inner := cfg.MambaNumHeads * cfg.MambaHeadDim
-	groupSize := inner / cfg.NGroups
-	B, L := int32(2), int32(3)
+		y := testGatedValues(0.1, int(B), int(L), int(inner))
+		gate := testGatedValues(-0.2, int(B), int(L), int(inner))
+		weight := testGatedValues(0.9, int(inner))
 
-	y := testGatedValues(0.1, int(B), int(L), int(inner))
-	gate := testGatedValues(-0.2, int(B), int(L), int(inner))
-	weight := testGatedValues(0.9, int(inner))
+		got := gatedGroupRMSNorm(y, gate, weight, cfg, mlx.DTypeFloat32)
 
-	got := gatedGroupRMSNorm(y, gate, weight, cfg, mlx.DTypeFloat32)
+		ref := mlx.Mul(y, mlx.SiLU(gate))
+		ref = mlx.Reshape(ref, B, L, cfg.NGroups, groupSize)
+		variance := mlx.Mean(mlx.Mul(ref, ref), 3, true)
+		ref = mlx.Mul(ref, mlx.RSqrt(mlx.AddScalar(variance, cfg.LayerNormEpsilon)))
+		ref = mlx.Mul(ref, mlx.Reshape(weight, 1, 1, cfg.NGroups, groupSize))
+		ref = mlx.Reshape(ref, B, L, inner)
 
-	ref := mlx.Mul(y, mlx.SiLU(gate))
-	ref = mlx.Reshape(ref, B, L, cfg.NGroups, groupSize)
-	variance := mlx.Mean(mlx.Mul(ref, ref), 3, true)
-	ref = mlx.Mul(ref, mlx.RSqrt(mlx.AddScalar(variance, cfg.LayerNormEpsilon)))
-	ref = mlx.Mul(ref, mlx.Reshape(weight, 1, 1, cfg.NGroups, groupSize))
-	ref = mlx.Reshape(ref, B, L, inner)
-
-	mlx.Eval(got, ref)
-	assertAllClose(t, "gated group rmsnorm", got.Floats(), ref.Floats(), 1e-5)
+		mlx.Eval(got, ref)
+		assertAllClose(t, "gated group rmsnorm", got.Floats(), ref.Floats(), 1e-5)
+	})
 }
 
 func testGatedValues(seed float32, shape ...int) *mlx.Array {
@@ -383,7 +383,7 @@ func testGatedValues(seed float32, shape ...int) *mlx.Array {
 	return mlx.FromValues(vals, shape...)
 }
 
-func assertAllClose(t *testing.T, name string, got, want []float32, tol float64) {
+func assertAllClose(t *mlxtest.T, name string, got, want []float32, tol float64) {
 	t.Helper()
 	if len(got) != len(want) {
 		t.Fatalf("%s length = %d, want %d", name, len(got), len(want))

x/models/nn/nn_test.go

@@ -8,181 +8,176 @@ import (
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	mlxtest.Setup(t)
-}
-
 func approxEqual(a, b, tol float32) bool {
 	return float32(math.Abs(float64(a-b))) < tol
 }
 
 // TestLayerNormNoBias verifies LayerNorm without bias against manual computation.
 func TestLayerNormNoBias(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		// Input: [1, 4] — single row, 4 features
+		x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
+		weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4)
+		mlx.Eval(x, weight)
 
-	// Input: [1, 4] — single row, 4 features
-	x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
-	weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4)
-	mlx.Eval(x, weight)
+		ln := &LayerNorm{Weight: weight, Eps: 1e-5}
+		out := ln.Forward(x)
+		mlx.Eval(out)
 
-	ln := &LayerNorm{Weight: weight, Eps: 1e-5}
-	out := ln.Forward(x)
-	mlx.Eval(out)
-
-	data := out.Floats()
-	if len(data) != 4 {
-		t.Fatalf("expected 4 values, got %d", len(data))
-	}
-
-	// Manual LayerNorm: mean=2.5, var=1.25, std=sqrt(1.25+1e-5)
-	// normalized = (x - mean) / std
-	mean := float32(2.5)
-	variance := float32(1.25)
-	std := float32(math.Sqrt(float64(variance + 1e-5)))
-	for i, v := range []float32{1, 2, 3, 4} {
-		expected := (v - mean) / std
-		if !approxEqual(data[i], expected, 1e-4) {
-			t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i])
+		data := out.Floats()
+		if len(data) != 4 {
+			t.Fatalf("expected 4 values, got %d", len(data))
 		}
-	}
+
+		// Manual LayerNorm: mean=2.5, var=1.25, std=sqrt(1.25+1e-5)
+		// normalized = (x - mean) / std
+		mean := float32(2.5)
+		variance := float32(1.25)
+		std := float32(math.Sqrt(float64(variance + 1e-5)))
+		for i, v := range []float32{1, 2, 3, 4} {
+			expected := (v - mean) / std
+			if !approxEqual(data[i], expected, 1e-4) {
+				t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i])
+			}
+		}
+	})
 }
 
 // TestLayerNormWithBias verifies LayerNorm with weight and bias.
 func TestLayerNormWithBias(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
+		weight := mlx.FromValues([]float32{2, 2, 2, 2}, 4)
+		bias := mlx.FromValues([]float32{10, 20, 30, 40}, 4)
+		mlx.Eval(x, weight, bias)
 
-	x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
-	weight := mlx.FromValues([]float32{2, 2, 2, 2}, 4)
-	bias := mlx.FromValues([]float32{10, 20, 30, 40}, 4)
-	mlx.Eval(x, weight, bias)
+		ln := &LayerNorm{Weight: weight, Bias: bias, Eps: 1e-5}
+		out := ln.Forward(x)
+		mlx.Eval(out)
 
-	ln := &LayerNorm{Weight: weight, Bias: bias, Eps: 1e-5}
-	out := ln.Forward(x)
-	mlx.Eval(out)
-
-	data := out.Floats()
-	if len(data) != 4 {
-		t.Fatalf("expected 4 values, got %d", len(data))
-	}
-
-	mean := float32(2.5)
-	variance := float32(1.25)
-	std := float32(math.Sqrt(float64(variance + 1e-5)))
-	biases := []float32{10, 20, 30, 40}
-	for i, v := range []float32{1, 2, 3, 4} {
-		expected := ((v-mean)/std)*2 + biases[i]
-		if !approxEqual(data[i], expected, 1e-4) {
-			t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i])
+		data := out.Floats()
+		if len(data) != 4 {
+			t.Fatalf("expected 4 values, got %d", len(data))
 		}
-	}
+
+		mean := float32(2.5)
+		variance := float32(1.25)
+		std := float32(math.Sqrt(float64(variance + 1e-5)))
+		biases := []float32{10, 20, 30, 40}
+		for i, v := range []float32{1, 2, 3, 4} {
+			expected := ((v-mean)/std)*2 + biases[i]
+			if !approxEqual(data[i], expected, 1e-4) {
+				t.Errorf("index %d: expected %.6f, got %.6f", i, expected, data[i])
+			}
+		}
+	})
 }
 
 // TestLayerNormBatched verifies LayerNorm normalizes each row independently.
 func TestLayerNormBatched(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		// Input: [2, 3] — two rows
+		x := mlx.FromValues([]float32{
+			1, 2, 3,
+			10, 20, 30,
+		}, 2, 3)
+		weight := mlx.FromValues([]float32{1, 1, 1}, 3)
+		mlx.Eval(x, weight)
 
-	// Input: [2, 3] — two rows
-	x := mlx.FromValues([]float32{
-		1, 2, 3,
-		10, 20, 30,
-	}, 2, 3)
-	weight := mlx.FromValues([]float32{1, 1, 1}, 3)
-	mlx.Eval(x, weight)
+		ln := &LayerNorm{Weight: weight, Eps: 1e-5}
+		out := ln.Forward(x)
+		mlx.Eval(out)
 
-	ln := &LayerNorm{Weight: weight, Eps: 1e-5}
-	out := ln.Forward(x)
-	mlx.Eval(out)
-
-	data := out.Floats()
-	if len(data) != 6 {
-		t.Fatalf("expected 6 values, got %d", len(data))
-	}
-
-	// Each row should be independently normalized.
-	// Row 0: [1,2,3] mean=2, var=2/3
-	// Row 1: [10,20,30] mean=20, var=200/3
-	// After normalization both rows should have the same pattern
-	// since [10,20,30] = 10*[1,2,3], the normalized values are identical.
-	for i := range 3 {
-		if !approxEqual(data[i], data[i+3], 1e-4) {
-			t.Errorf("row 0 elem %d (%.6f) != row 1 elem %d (%.6f); expected identical normalized values",
-				i, data[i], i, data[i+3])
+		data := out.Floats()
+		if len(data) != 6 {
+			t.Fatalf("expected 6 values, got %d", len(data))
 		}
-	}
 
-	// Verify the normalized values sum to ~0 (mean-centered)
-	sum := data[0] + data[1] + data[2]
-	if !approxEqual(sum, 0, 1e-4) {
-		t.Errorf("normalized row sum should be ~0, got %.6f", sum)
-	}
+		// Each row should be independently normalized.
+		// Row 0: [1,2,3] mean=2, var=2/3
+		// Row 1: [10,20,30] mean=20, var=200/3
+		// After normalization both rows should have the same pattern
+		// since [10,20,30] = 10*[1,2,3], the normalized values are identical.
+		for i := range 3 {
+			if !approxEqual(data[i], data[i+3], 1e-4) {
+				t.Errorf("row 0 elem %d (%.6f) != row 1 elem %d (%.6f); expected identical normalized values",
+					i, data[i], i, data[i+3])
+			}
+		}
+
+		// Verify the normalized values sum to ~0 (mean-centered)
+		sum := data[0] + data[1] + data[2]
+		if !approxEqual(sum, 0, 1e-4) {
+			t.Errorf("normalized row sum should be ~0, got %.6f", sum)
+		}
+	})
 }
 
 // TestLayerNormDefaultEps verifies the default epsilon of 1e-5 is used when Eps is 0.
 func TestLayerNormDefaultEps(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
+		weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4)
+		mlx.Eval(x, weight)
 
-	x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 4)
-	weight := mlx.FromValues([]float32{1, 1, 1, 1}, 4)
-	mlx.Eval(x, weight)
+		// Eps=0 should use default 1e-5
+		ln0 := &LayerNorm{Weight: weight, Eps: 0}
+		out0 := ln0.Forward(x)
+		mlx.Eval(out0)
 
-	// Eps=0 should use default 1e-5
-	ln0 := &LayerNorm{Weight: weight, Eps: 0}
-	out0 := ln0.Forward(x)
-	mlx.Eval(out0)
+		lnExplicit := &LayerNorm{Weight: weight, Eps: 1e-5}
+		outExplicit := lnExplicit.Forward(x)
+		mlx.Eval(outExplicit)
 
-	lnExplicit := &LayerNorm{Weight: weight, Eps: 1e-5}
-	outExplicit := lnExplicit.Forward(x)
-	mlx.Eval(outExplicit)
-
-	d0 := out0.Floats()
-	dE := outExplicit.Floats()
-	for i := range d0 {
-		if !approxEqual(d0[i], dE[i], 1e-6) {
-			t.Errorf("index %d: Eps=0 gave %.6f, Eps=1e-5 gave %.6f", i, d0[i], dE[i])
+		d0 := out0.Floats()
+		dE := outExplicit.Floats()
+		for i := range d0 {
+			if !approxEqual(d0[i], dE[i], 1e-6) {
+				t.Errorf("index %d: Eps=0 gave %.6f, Eps=1e-5 gave %.6f", i, d0[i], dE[i])
+			}
 		}
-	}
+	})
 }
 
 func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) {
-	skipIfNoMLX(t)
-
-	weightVals := make([]float32, 3*32)
-	for i := range weightVals {
-		weightVals[i] = float32((i%11)-5) / 7
-	}
-	inputVals := make([]float32, 2*32)
-	for i := range inputVals {
-		inputVals[i] = float32((i%7)-3) / 5
-	}
-
-	weight := mlx.FromValues(weightVals, 3, 32).AsType(mlx.DTypeBFloat16)
-	input := mlx.FromValues(inputVals, 2, 32).AsType(mlx.DTypeBFloat16)
-	mlx.Eval(weight, input)
-
-	ql := NewQuantizedLinear(weight, nil, 32, 4, "mxfp4")
-	if ql.QBiases != nil {
-		t.Fatalf("mxfp4 qbiases = %v, want nil", ql.QBiases)
-	}
-
-	dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4", nil)
-	mlx.Eval(dequantizedWeight)
-
-	qOut := ql.Forward(input).AsType(mlx.DTypeFloat32)
-	dOut := NewLinear(dequantizedWeight, nil).Forward(input).AsType(mlx.DTypeFloat32)
-	mlx.Eval(qOut, dOut)
-
-	got := qOut.Floats()
-	want := dOut.Floats()
-	if len(got) != len(want) {
-		t.Fatalf("output length = %d, want %d", len(got), len(want))
-	}
-
-	for i := range got {
-		if !approxEqual(got[i], want[i], 1e-3) {
-			t.Fatalf("output[%d] = %.6f, want %.6f", i, got[i], want[i])
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		weightVals := make([]float32, 3*32)
+		for i := range weightVals {
+			weightVals[i] = float32((i%11)-5) / 7
 		}
-	}
+		inputVals := make([]float32, 2*32)
+		for i := range inputVals {
+			inputVals[i] = float32((i%7)-3) / 5
+		}
+
+		weight := mlx.FromValues(weightVals, 3, 32).AsType(mlx.DTypeBFloat16)
+		input := mlx.FromValues(inputVals, 2, 32).AsType(mlx.DTypeBFloat16)
+		mlx.Eval(weight, input)
+
+		ql := NewQuantizedLinear(weight, nil, 32, 4, "mxfp4")
+		if ql.QBiases != nil {
+			t.Fatalf("mxfp4 qbiases = %v, want nil", ql.QBiases)
+		}
+
+		dequantizedWeight := mlx.Dequantize(ql.Weight, ql.Scales, ql.QBiases, 32, 4, "mxfp4", nil)
+		mlx.Eval(dequantizedWeight)
+
+		qOut := ql.Forward(input).AsType(mlx.DTypeFloat32)
+		dOut := NewLinear(dequantizedWeight, nil).Forward(input).AsType(mlx.DTypeFloat32)
+		mlx.Eval(qOut, dOut)
+
+		got := qOut.Floats()
+		want := dOut.Floats()
+		if len(got) != len(want) {
+			t.Fatalf("output length = %d, want %d", len(got), len(want))
+		}
+
+		for i := range got {
+			if !approxEqual(got[i], want[i], 1e-3) {
+				t.Fatalf("output[%d] = %.6f, want %.6f", i, got[i], want[i])
+			}
+		}
+	})
 }
 
 func TestQuantizedEmbeddingAsLinearPreservesGlobalScale(t *testing.T) {

x/models/nn/recurrent_test.go

@@ -4,6 +4,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
@@ -36,30 +37,31 @@ func convFromKernel(w *mlx.Array) *Conv1d {
 // Guards a biased conv silently losing the fused kernel: depthwiseConvWeight
 // returning nil sends WithConvSiLU down separate graph ops.
 func TestCausalConv1DBiasTakesFusedPath(t *testing.T) {
-	skipIfNoMLX(t)
-	B, L, D, convTail := 2, 3, 4, 2
-	K := convTail + 1
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		B, L, D, convTail := 2, 3, 4, 2
+		K := convTail + 1
 
-	weight := fromValues(-0.3, D, K)
-	bias := fromValues(0.4, D)
-	conv := NewConv1d(mlx.ExpandDims(weight, 2), bias, 1, 0, 1, int32(D))
+		weight := fromValues(-0.3, D, K)
+		bias := fromValues(0.4, D)
+		conv := NewConv1d(mlx.ExpandDims(weight, 2), bias, 1, 0, 1, int32(D))
 
-	if depthwiseConvWeight(conv) == nil {
-		t.Fatal("depthwiseConvWeight = nil for a biased depthwise conv, so the fused path is skipped")
-	}
+		if depthwiseConvWeight(conv) == nil {
+			t.Fatal("depthwiseConvWeight = nil for a biased depthwise conv, so the fused path is skipped")
+		}
 
-	prior := fromValues(0.2, B, convTail, D)
-	input := fromValues(0.1, B, L, D)
-	b := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, B, L),
-		SeqOffsets:   []int32{0, 0},
-		SeqQueryLens: []int32{int32(L), int32(L)},
-	}
+		prior := fromValues(0.2, B, convTail, D)
+		input := fromValues(0.1, B, L, D)
+		b := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, B, L),
+			SeqOffsets:   []int32{0, 0},
+			SeqQueryLens: []int32{int32(L), int32(L)},
+		}
 
-	got, _ := CausalConv1D(b, input, conv, convTail, WithRecurrentState(prior, nil), WithConvSiLU())
-	want := mlx.SiLU(conv.Forward(mlx.Concatenate([]*mlx.Array{prior, input}, 1)))
-	mlx.Eval(got, want)
-	floatsClose(t, "biased fused conv+silu", got.Floats(), want.Floats(), 1e-5)
+		got, _ := CausalConv1D(b, input, conv, convTail, WithRecurrentState(prior, nil), WithConvSiLU())
+		want := mlx.SiLU(conv.Forward(mlx.Concatenate([]*mlx.Array{prior, input}, 1)))
+		mlx.Eval(got, want)
+		floatsClose(t, "biased fused conv+silu", got.Floats(), want.Floats(), 1e-5)
+	})
 }
 
 // TestCausalConv1DPaddedRowParity drives a B=2 batch with one short
@@ -68,103 +70,104 @@ func TestCausalConv1DBiasTakesFusedPath(t *testing.T) {
 // short row must be the row's last convTail real positions (not the
 // padded tail), (c) the full row must be unaffected.
 func TestCausalConv1DPaddedRowParity(t *testing.T) {
-	skipIfNoMLX(t)
-	L, D, convTail := 4, 3, 2
-	qLenShort := 2
-	K := convTail + 1
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		L, D, convTail := 4, 3, 2
+		qLenShort := 2
+		K := convTail + 1
 
-	weight := fromValues(0.2, D, K)
-	conv := convFromKernel(weight)
-	priorFull := fromValues(0.5, 2, convTail, D)
-	priorShort := mlx.SliceStartStop(priorFull,
-		[]int32{1, 0, 0},
-		[]int32{2, int32(convTail), int32(D)})
+		weight := fromValues(0.2, D, K)
+		conv := convFromKernel(weight)
+		priorFull := fromValues(0.5, 2, convTail, D)
+		priorShort := mlx.SliceStartStop(priorFull,
+			[]int32{1, 0, 0},
+			[]int32{2, int32(convTail), int32(D)})
 
-	// Pad row 1 with arbitrary values past qLenShort — the wrapper
-	// must zero them before convolving. Distinct values let us catch
-	// any leak.
-	inputFull := fromValues(1.0, 1, L, D)
-	inputShortReal := mlx.FromValues([]float32{
-		2.0, 2.1, 2.2,
-		2.3, 2.4, 2.5,
-	}, 1, qLenShort, D)
-	inputShortPad := mlx.FromValues([]float32{
-		99, 99, 99,
-		99, 99, 99,
-	}, 1, L-qLenShort, D)
-	inputShortFull := mlx.Concatenate([]*mlx.Array{inputShortReal, inputShortPad}, 1)
-	input := mlx.Concatenate([]*mlx.Array{inputFull, inputShortFull}, 0)
+		// Pad row 1 with arbitrary values past qLenShort — the wrapper
+		// must zero them before convolving. Distinct values let us catch
+		// any leak.
+		inputFull := fromValues(1.0, 1, L, D)
+		inputShortReal := mlx.FromValues([]float32{
+			2.0, 2.1, 2.2,
+			2.3, 2.4, 2.5,
+		}, 1, qLenShort, D)
+		inputShortPad := mlx.FromValues([]float32{
+			99, 99, 99,
+			99, 99, 99,
+		}, 1, L-qLenShort, D)
+		inputShortFull := mlx.Concatenate([]*mlx.Array{inputShortReal, inputShortPad}, 1)
+		input := mlx.Concatenate([]*mlx.Array{inputFull, inputShortFull}, 0)
 
-	b := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
-		SeqOffsets:   []int32{0, 0},
-		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
-	}
+		b := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
+			SeqOffsets:   []int32{0, 0},
+			SeqQueryLens: []int32{int32(L), int32(qLenShort)},
+		}
 
-	out, convStates := CausalConv1D(b, input, conv, convTail, WithRecurrentState(priorFull, nil))
-	nextConv := lastState(convStates)
-	mlx.Eval(out, nextConv)
+		out, convStates := CausalConv1D(b, input, conv, convTail, WithRecurrentState(priorFull, nil))
+		nextConv := lastState(convStates)
+		mlx.Eval(out, nextConv)
 
-	// Reference for row 0: B=1 unpadded length-L call.
-	refOut0, refConvStates0 := CausalConv1D(&batch.Batch{},
-		inputFull, conv, convTail,
-		WithRecurrentState(mlx.SliceStartStop(priorFull,
-			[]int32{0, 0, 0},
-			[]int32{1, int32(convTail), int32(D)}), nil))
-	refNextConv0 := lastState(refConvStates0)
-	// Reference for row 1: B=1 unpadded length-qLenShort call.
-	refOut1, refConvStates1 := CausalConv1D(&batch.Batch{},
-		inputShortReal, conv, convTail,
-		WithRecurrentState(priorShort, nil))
-	refNextConv1 := lastState(refConvStates1)
-	mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)
+		// Reference for row 0: B=1 unpadded length-L call.
+		refOut0, refConvStates0 := CausalConv1D(&batch.Batch{},
+			inputFull, conv, convTail,
+			WithRecurrentState(mlx.SliceStartStop(priorFull,
+				[]int32{0, 0, 0},
+				[]int32{1, int32(convTail), int32(D)}), nil))
+		refNextConv0 := lastState(refConvStates0)
+		// Reference for row 1: B=1 unpadded length-qLenShort call.
+		refOut1, refConvStates1 := CausalConv1D(&batch.Batch{},
+			inputShortReal, conv, convTail,
+			WithRecurrentState(priorShort, nil))
+		refNextConv1 := lastState(refConvStates1)
+		mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)
 
-	gotOut := out.Floats()
-	wantOut0 := refOut0.Floats()
-	wantOut1 := refOut1.Floats()
+		gotOut := out.Floats()
+		wantOut0 := refOut0.Floats()
+		wantOut1 := refOut1.Floats()
 
-	for q := range L {
-		for d := range D {
-			i := q*D + d
-			if gotOut[i] != wantOut0[i] {
-				t.Fatalf("row 0 out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[i], wantOut0[i])
+		for q := range L {
+			for d := range D {
+				i := q*D + d
+				if gotOut[i] != wantOut0[i] {
+					t.Fatalf("row 0 out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[i], wantOut0[i])
+				}
 			}
 		}
-	}
-	for q := range qLenShort {
-		for d := range D {
-			gotI := L*D + q*D + d
-			refI := q*D + d
-			if math.Abs(float64(gotOut[gotI]-wantOut1[refI])) > 1e-5 {
-				t.Fatalf("row 1 real out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[gotI], wantOut1[refI])
+		for q := range qLenShort {
+			for d := range D {
+				gotI := L*D + q*D + d
+				refI := q*D + d
+				if math.Abs(float64(gotOut[gotI]-wantOut1[refI])) > 1e-5 {
+					t.Fatalf("row 1 real out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[gotI], wantOut1[refI])
+				}
 			}
 		}
-	}
 
-	// nextConv: row 0 unaffected, row 1 must be the row's real tail
-	// (positions [qLenShort - convTail, qLenShort) of the per-row
-	// concat, i.e. the last two real input rows in this setup).
-	gotTail := nextConv.Floats()
-	wantTail0 := refNextConv0.Floats()
-	wantTail1 := refNextConv1.Floats()
-	for k := range convTail {
-		for d := range D {
-			i := k*D + d
-			if gotTail[i] != wantTail0[i] {
-				t.Fatalf("row 0 nextConv[k=%d,d=%d]: got %v, want %v", k, d, gotTail[i], wantTail0[i])
+		// nextConv: row 0 unaffected, row 1 must be the row's real tail
+		// (positions [qLenShort - convTail, qLenShort) of the per-row
+		// concat, i.e. the last two real input rows in this setup).
+		gotTail := nextConv.Floats()
+		wantTail0 := refNextConv0.Floats()
+		wantTail1 := refNextConv1.Floats()
+		for k := range convTail {
+			for d := range D {
+				i := k*D + d
+				if gotTail[i] != wantTail0[i] {
+					t.Fatalf("row 0 nextConv[k=%d,d=%d]: got %v, want %v", k, d, gotTail[i], wantTail0[i])
+				}
 			}
 		}
-	}
-	for k := range convTail {
-		for d := range D {
-			gotI := convTail*D + k*D + d
-			refI := k*D + d
-			if gotTail[gotI] != wantTail1[refI] {
-				t.Fatalf("row 1 nextConv[k=%d,d=%d]: got %v, want %v (must come from real positions, not the padded tail)",
-					k, d, gotTail[gotI], wantTail1[refI])
+		for k := range convTail {
+			for d := range D {
+				gotI := convTail*D + k*D + d
+				refI := k*D + d
+				if gotTail[gotI] != wantTail1[refI] {
+					t.Fatalf("row 1 nextConv[k=%d,d=%d]: got %v, want %v (must come from real positions, not the padded tail)",
+						k, d, gotTail[gotI], wantTail1[refI])
+				}
 			}
 		}
-	}
+	})
 }
 
 // gatedDeltaPackedInputs builds deterministic packed conv-output and
@@ -194,7 +197,7 @@ func slicePrefix(a *mlx.Array, lo, hi, n int32) *mlx.Array {
 }
 
 // floatsClose compares two flat float slices within tolerance.
-func floatsClose(t *testing.T, label string, got, want []float32, tol float64) {
+func floatsClose(t *mlxtest.T, label string, got, want []float32, tol float64) {
 	t.Helper()
 	if len(got) != len(want) {
 		t.Fatalf("%s: len %d, want %d", label, len(got), len(want))
@@ -212,82 +215,84 @@ func floatsClose(t *testing.T, label string, got, want []float32, tol float64) {
 // that each boundary state equals the single-shot state over the
 // corresponding prefix.
 func TestGatedDeltaSegmentEquivalence(t *testing.T) {
-	skipIfNoMLX(t)
-	B, T, Hk, Dk, Hv, Dv := 1, 5, 1, 32, 1, 32
-	packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
-	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
-	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(T)}}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		B, T, Hk, Dk, Hv, Dv := 1, 5, 1, 32, 1, 32
+		packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
+		prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
+		full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(T)}}
 
-	refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
-	if len(refStates) != 1 {
-		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
-	}
+		refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
+		if len(refStates) != 1 {
+			t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
+		}
 
-	cases := []struct {
-		name   string
-		splits []int
-	}{
-		{"perToken", []int{1, 2, 3, 4}},
-		{"sparse", []int{2}},
-	}
-	for _, tc := range cases {
-		segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
-			WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
-		mlx.Eval(refOut, segOut)
-		floatsClose(t, tc.name+" out", segOut.Floats(), refOut.Floats(), 1e-4)
-		if len(segStates) != len(tc.splits)+1 {
-			t.Fatalf("%s: got %d boundary states, want %d", tc.name, len(segStates), len(tc.splits)+1)
+		cases := []struct {
+			name   string
+			splits []int
+		}{
+			{"perToken", []int{1, 2, 3, 4}},
+			{"sparse", []int{2}},
 		}
-		boundaries := append(append([]int{}, tc.splits...), T)
-		for i, n := range boundaries {
-			_, want, _ := mlx.GatedDelta(
-				slicePrefix(packed, 0, 1, int32(n)), slicePrefix(ba, 0, 1, int32(n)),
-				dtBias, aExp, prior, nil, false)
-			mlx.Eval(segStates[i], want)
-			floatsClose(t, tc.name+" boundary delta", segStates[i].Floats(), want.Floats(), 1e-4)
+		for _, tc := range cases {
+			segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
+				WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
+			mlx.Eval(refOut, segOut)
+			floatsClose(t, tc.name+" out", segOut.Floats(), refOut.Floats(), 1e-4)
+			if len(segStates) != len(tc.splits)+1 {
+				t.Fatalf("%s: got %d boundary states, want %d", tc.name, len(segStates), len(tc.splits)+1)
+			}
+			boundaries := append(append([]int{}, tc.splits...), T)
+			for i, n := range boundaries {
+				_, want, _ := mlx.GatedDelta(
+					slicePrefix(packed, 0, 1, int32(n)), slicePrefix(ba, 0, 1, int32(n)),
+					dtBias, aExp, prior, nil, false)
+				mlx.Eval(segStates[i], want)
+				floatsClose(t, tc.name+" boundary delta", segStates[i].Floats(), want.Floats(), 1e-4)
+			}
 		}
-	}
+	})
 }
 
 // TestCausalConv1DSegmentEquivalence checks the conv segmented path matches the
 // single-shot conv for output, final conv tail, and each boundary conv state.
 func TestCausalConv1DSegmentEquivalence(t *testing.T) {
-	skipIfNoMLX(t)
-	B, L, D, convTail := 1, 4, 3, 2
-	K := convTail + 1
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		B, L, D, convTail := 1, 4, 3, 2
+		K := convTail + 1
 
-	input := fromValues(0.5, B, L, D)
-	prior := fromValues(-0.3, B, convTail, D)
-	weight := fromValues(0.2, D, K)
-	conv := convFromKernel(weight)
+		input := fromValues(0.5, B, L, D)
+		prior := fromValues(-0.3, B, convTail, D)
+		weight := fromValues(0.2, D, K)
+		conv := convFromKernel(weight)
 
-	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(L)}}
+		full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(L)}}
 
-	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
-	if len(refStates) != 1 {
-		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
-	}
+		refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
+		if len(refStates) != 1 {
+			t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
+		}
 
-	segOut, segStates := CausalConv1D(full, input, conv, convTail,
-		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
-	mlx.Eval(refOut, segOut)
+		segOut, segStates := CausalConv1D(full, input, conv, convTail,
+			WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
+		mlx.Eval(refOut, segOut)
 
-	floatsClose(t, "conv out", segOut.Floats(), refOut.Floats(), 1e-4)
-	if len(segStates) != 4 {
-		t.Fatalf("got %d boundary conv states, want 4", len(segStates))
-	}
-	mlx.Eval(lastState(segStates), lastState(refStates))
-	floatsClose(t, "conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
+		floatsClose(t, "conv out", segOut.Floats(), refOut.Floats(), 1e-4)
+		if len(segStates) != 4 {
+			t.Fatalf("got %d boundary conv states, want 4", len(segStates))
+		}
+		mlx.Eval(lastState(segStates), lastState(refStates))
+		floatsClose(t, "conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
 
-	for i := range segStates {
-		n := int32(i + 1)
-		pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
-		_, want := CausalConv1D(pb,
-			mlx.SliceStartStop(input, []int32{0, 0, 0}, []int32{int32(B), n, int32(D)}),
-			conv, convTail, WithRecurrentState(prior, nil))
-		mlx.Eval(segStates[i], lastState(want))
-		floatsClose(t, "boundary conv", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
-	}
+		for i := range segStates {
+			n := int32(i + 1)
+			pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
+			_, want := CausalConv1D(pb,
+				mlx.SliceStartStop(input, []int32{0, 0, 0}, []int32{int32(B), n, int32(D)}),
+				conv, convTail, WithRecurrentState(prior, nil))
+			mlx.Eval(segStates[i], lastState(want))
+			floatsClose(t, "boundary conv", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
+		}
+	})
 }
 
 // TestGatedDeltaSegmentEquivalenceBatched checks split forwards match the
@@ -296,57 +301,58 @@ func TestCausalConv1DSegmentEquivalence(t *testing.T) {
 // row's padded positions so a short row's boundary state freezes at its
 // real end.
 func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
-	skipIfNoMLX(t)
-	B, T, Hk, Dk, Hv, Dv := 2, 4, 1, 32, 1, 32
-	packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
-	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		B, T, Hk, Dk, Hv, Dv := 2, 4, 1, 32, 1, 32
+		packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
+		prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
 
-	// Row 0 full length T; row 1 ends at 3 (so segment [3,4) is all padding
-	// for row 1).
-	rowReal := []int32{int32(T), 3}
-	full := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, B, T),
-		SeqOffsets:   []int32{0, 0},
-		SeqQueryLens: rowReal,
-	}
-
-	refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
-
-	cases := []struct {
-		name   string
-		splits []int
-	}{
-		{"perToken", []int{1, 2, 3}},
-		{"sparse", []int{2}},
-	}
-	for _, tc := range cases {
-		segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
-			WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
-		mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
-		floatsClose(t, tc.name+" batched out", segOut.Floats(), refOut.Floats(), 1e-4)
-		floatsClose(t, tc.name+" batched final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
-		if len(segStates) != len(tc.splits)+1 {
-			t.Fatalf("%s: got %d boundary states, want %d", tc.name, len(segStates), len(tc.splits)+1)
+		// Row 0 full length T; row 1 ends at 3 (so segment [3,4) is all padding
+		// for row 1).
+		rowReal := []int32{int32(T), 3}
+		full := &batch.Batch{
+			InputIDs:     mlx.Zeros(mlx.DTypeInt32, B, T),
+			SeqOffsets:   []int32{0, 0},
+			SeqQueryLens: rowReal,
 		}
 
-		// Each row's boundary must equal a B=1 single-shot call over that
-		// row's real prefix: row 0 advances the full length, row 1 freezes
-		// once it reaches its real length.
-		boundaries := append(append([]int{}, tc.splits...), T)
-		for i, bound := range boundaries {
-			for r := range B {
-				n := min(int32(bound), rowReal[r])
-				lo, hi := int32(r), int32(r)+1
-				rowPrior := mlx.SliceStartStop(prior, []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
-				_, want, _ := mlx.GatedDelta(
-					slicePrefix(packed, lo, hi, n), slicePrefix(ba, lo, hi, n),
-					dtBias, aExp, rowPrior, nil, false)
-				gotRow := mlx.SliceStartStop(segStates[i], []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
-				mlx.Eval(gotRow, want)
-				floatsClose(t, tc.name+" batched boundary delta", gotRow.Floats(), want.Floats(), 1e-4)
+		refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
+
+		cases := []struct {
+			name   string
+			splits []int
+		}{
+			{"perToken", []int{1, 2, 3}},
+			{"sparse", []int{2}},
+		}
+		for _, tc := range cases {
+			segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
+				WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
+			mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
+			floatsClose(t, tc.name+" batched out", segOut.Floats(), refOut.Floats(), 1e-4)
+			floatsClose(t, tc.name+" batched final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
+			if len(segStates) != len(tc.splits)+1 {
+				t.Fatalf("%s: got %d boundary states, want %d", tc.name, len(segStates), len(tc.splits)+1)
+			}
+
+			// Each row's boundary must equal a B=1 single-shot call over that
+			// row's real prefix: row 0 advances the full length, row 1 freezes
+			// once it reaches its real length.
+			boundaries := append(append([]int{}, tc.splits...), T)
+			for i, bound := range boundaries {
+				for r := range B {
+					n := min(int32(bound), rowReal[r])
+					lo, hi := int32(r), int32(r)+1
+					rowPrior := mlx.SliceStartStop(prior, []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
+					_, want, _ := mlx.GatedDelta(
+						slicePrefix(packed, lo, hi, n), slicePrefix(ba, lo, hi, n),
+						dtBias, aExp, rowPrior, nil, false)
+					gotRow := mlx.SliceStartStop(segStates[i], []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
+					mlx.Eval(gotRow, want)
+					floatsClose(t, tc.name+" batched boundary delta", gotRow.Floats(), want.Floats(), 1e-4)
+				}
 			}
 		}
-	}
+	})
 }
 
 // TestCausalConv1DSegmentEquivalenceBatched is the conv analog of the gated-delta
@@ -354,46 +360,47 @@ func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
 // references for a ragged B>1 batch, where a short row must freeze its tail at
 // its real end rather than reach into padding.
 func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
-	skipIfNoMLX(t)
-	B, L, D, convTail := 2, 4, 3, 2
-	K := convTail + 1
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		B, L, D, convTail := 2, 4, 3, 2
+		K := convTail + 1
 
-	input := fromValues(0.5, B, L, D)
-	prior := fromValues(-0.3, B, convTail, D)
-	weight := fromValues(0.2, D, K)
-	conv := convFromKernel(weight)
+		input := fromValues(0.5, B, L, D)
+		prior := fromValues(-0.3, B, convTail, D)
+		weight := fromValues(0.2, D, K)
+		conv := convFromKernel(weight)
 
-	full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(L), 3}}
+		full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(L), 3}}
 
-	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
-	segOut, segStates := CausalConv1D(full, input, conv, convTail,
-		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
-	mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
+		refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
+		segOut, segStates := CausalConv1D(full, input, conv, convTail,
+			WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
+		mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
 
-	floatsClose(t, "batched conv out", segOut.Floats(), refOut.Floats(), 1e-4)
-	floatsClose(t, "batched conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
-	if len(segStates) != 4 {
-		t.Fatalf("got %d boundary conv states, want 4", len(segStates))
-	}
-
-	// Each row's boundary i (offset i+1) must equal a B=1 single-shot conv over
-	// that row's real prefix: row 0 advances the full length, row 1 freezes once
-	// it reaches its real length 3. Per-row B=1 references avoid the ambiguity of
-	// redeclaring a ragged length over a uniform input slice.
-	rowReal := []int32{int32(L), 3}
-	for i := range segStates {
-		for r := range B {
-			n := min(int32(i+1), rowReal[r])
-			rowPrior := mlx.SliceStartStop(prior,
-				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
-			rowInput := mlx.SliceStartStop(input,
-				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, n, int32(D)})
-			_, want := CausalConv1D(&batch.Batch{}, rowInput, conv, convTail,
-				WithRecurrentState(rowPrior, nil))
-			gotRow := mlx.SliceStartStop(segStates[i],
-				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
-			mlx.Eval(gotRow, lastState(want))
-			floatsClose(t, "batched boundary conv", gotRow.Floats(), lastState(want).Floats(), 1e-4)
+		floatsClose(t, "batched conv out", segOut.Floats(), refOut.Floats(), 1e-4)
+		floatsClose(t, "batched conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
+		if len(segStates) != 4 {
+			t.Fatalf("got %d boundary conv states, want 4", len(segStates))
 		}
-	}
+
+		// Each row's boundary i (offset i+1) must equal a B=1 single-shot conv over
+		// that row's real prefix: row 0 advances the full length, row 1 freezes once
+		// it reaches its real length 3. Per-row B=1 references avoid the ambiguity of
+		// redeclaring a ragged length over a uniform input slice.
+		rowReal := []int32{int32(L), 3}
+		for i := range segStates {
+			for r := range B {
+				n := min(int32(i+1), rowReal[r])
+				rowPrior := mlx.SliceStartStop(prior,
+					[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
+				rowInput := mlx.SliceStartStop(input,
+					[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, n, int32(D)})
+				_, want := CausalConv1D(&batch.Batch{}, rowInput, conv, convTail,
+					WithRecurrentState(rowPrior, nil))
+				gotRow := mlx.SliceStartStop(segStates[i],
+					[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
+				mlx.Eval(gotRow, lastState(want))
+				floatsClose(t, "batched boundary conv", gotRow.Floats(), lastState(want).Floats(), 1e-4)
+			}
+		}
+	})
 }

x/models/qwen3_5/gdn_projections_test.go

@@ -4,6 +4,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/models/nn"
 )
@@ -25,7 +26,7 @@ func patternArray(rows, cols int) *mlx.Array {
 	return mlx.FromValues(values, rows, cols).AsType(mlx.DTypeBFloat16)
 }
 
-func assertBitEqual(t *testing.T, label string, got, want *mlx.Array) {
+func assertBitEqual(t *mlxtest.T, label string, got, want *mlx.Array) {
 	t.Helper()
 	if got == nil || want == nil {
 		if got != want {
@@ -57,82 +58,85 @@ func scatterRows(packed *mlx.Array, perm []int32) *mlx.Array {
 }
 
 func TestPackGatedDeltaProjectionsNativeMatchesSplit(t *testing.T) {
-	skipIfNoMLX(t)
-	cfg := gdnTestConfig()
-	keyDim := int(cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim)
-	valueDim := int(cfg.LinearNumValueHeads * cfg.LinearValueHeadDim)
-	in := 8
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := gdnTestConfig()
+		keyDim := int(cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim)
+		valueDim := int(cfg.LinearNumValueHeads * cfg.LinearValueHeadDim)
+		in := 8
 
-	qkvW := patternArray(2*keyDim+valueDim, in)
-	zW := patternArray(valueDim, in)
-	bW := patternArray(int(cfg.LinearNumValueHeads), in)
-	aW := patternArray(int(cfg.LinearNumValueHeads), in)
+		qkvW := patternArray(2*keyDim+valueDim, in)
+		zW := patternArray(valueDim, in)
+		bW := patternArray(int(cfg.LinearNumValueHeads), in)
+		aW := patternArray(int(cfg.LinearNumValueHeads), in)
 
-	fromSplitQKVZ, fromSplitBA, err := packGatedDeltaProjections(
-		nn.NewLinear(qkvW, nil), nn.NewLinear(zW, nil),
-		nn.NewLinear(bW, nil), nn.NewLinear(aW, nil),
-		nil, nil, cfg,
-	)
-	if err != nil {
-		t.Fatal(err)
-	}
+		fromSplitQKVZ, fromSplitBA, err := packGatedDeltaProjections(
+			nn.NewLinear(qkvW, nil), nn.NewLinear(zW, nil),
+			nn.NewLinear(bW, nil), nn.NewLinear(aW, nil),
+			nil, nil, cfg,
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
 
-	packedQKVZ := mlx.Concatenate([]*mlx.Array{qkvW, zW}, 0)
-	packedBA := mlx.Concatenate([]*mlx.Array{bW, aW}, 0)
-	nativeQKVZ := scatterRows(packedQKVZ, nativeQKVZPerm(cfg))
-	nativeBA := scatterRows(packedBA, nativeBAPerm(cfg))
+		packedQKVZ := mlx.Concatenate([]*mlx.Array{qkvW, zW}, 0)
+		packedBA := mlx.Concatenate([]*mlx.Array{bW, aW}, 0)
+		nativeQKVZ := scatterRows(packedQKVZ, nativeQKVZPerm(cfg))
+		nativeBA := scatterRows(packedBA, nativeBAPerm(cfg))
 
-	fromNativeQKVZ, fromNativeBA, err := packGatedDeltaProjections(
-		nil, nil, nil, nil,
-		nn.NewLinear(nativeQKVZ, nil), nn.NewLinear(nativeBA, nil),
-		cfg,
-	)
-	if err != nil {
-		t.Fatal(err)
-	}
+		fromNativeQKVZ, fromNativeBA, err := packGatedDeltaProjections(
+			nil, nil, nil, nil,
+			nn.NewLinear(nativeQKVZ, nil), nn.NewLinear(nativeBA, nil),
+			cfg,
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
 
-	assertBitEqual(t, "qkvz", fromNativeQKVZ.(*nn.Linear).Weight, fromSplitQKVZ.(*nn.Linear).Weight)
-	assertBitEqual(t, "ba", fromNativeBA.(*nn.Linear).Weight, fromSplitBA.(*nn.Linear).Weight)
+		assertBitEqual(t, "qkvz", fromNativeQKVZ.(*nn.Linear).Weight, fromSplitQKVZ.(*nn.Linear).Weight)
+		assertBitEqual(t, "ba", fromNativeBA.(*nn.Linear).Weight, fromSplitBA.(*nn.Linear).Weight)
+	})
 }
 
 func TestConcatProjectionPairQuantized(t *testing.T) {
-	skipIfNoMLX(t)
-	in := 64
-	hi := nn.NewQuantizedLinear(patternArray(16, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
-	lo := nn.NewQuantizedLinear(patternArray(8, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		in := 64
+		hi := nn.NewQuantizedLinear(patternArray(16, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
+		lo := nn.NewQuantizedLinear(patternArray(8, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
 
-	packed, err := concatProjectionPair(hi, lo)
-	if err != nil {
-		t.Fatal(err)
-	}
-	q, ok := packed.(*nn.QuantizedLinear)
-	if !ok {
-		t.Fatalf("packed projection is %T, want *nn.QuantizedLinear", packed)
-	}
-	assertBitEqual(t, "weight", q.Weight, mlx.Concatenate([]*mlx.Array{hi.Weight, lo.Weight}, 0))
-	assertBitEqual(t, "scales", q.Scales, mlx.Concatenate([]*mlx.Array{hi.Scales, lo.Scales}, 0))
+		packed, err := concatProjectionPair(hi, lo)
+		if err != nil {
+			t.Fatal(err)
+		}
+		q, ok := packed.(*nn.QuantizedLinear)
+		if !ok {
+			t.Fatalf("packed projection is %T, want *nn.QuantizedLinear", packed)
+		}
+		assertBitEqual(t, "weight", q.Weight, mlx.Concatenate([]*mlx.Array{hi.Weight, lo.Weight}, 0))
+		assertBitEqual(t, "scales", q.Scales, mlx.Concatenate([]*mlx.Array{hi.Scales, lo.Scales}, 0))
+	})
 }
 
 func TestConcatProjectionPairMixedFallsBackToDense(t *testing.T) {
-	skipIfNoMLX(t)
-	in := 64
-	hi := nn.NewQuantizedLinear(patternArray(16, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
-	loW := patternArray(8, in)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		in := 64
+		hi := nn.NewQuantizedLinear(patternArray(16, in).AsType(mlx.DTypeFloat32), nil, 32, 4, "affine")
+		loW := patternArray(8, in)
 
-	packed, err := concatProjectionPair(hi, nn.NewLinear(loW, nil))
-	if err != nil {
-		t.Fatal(err)
-	}
-	dense, ok := packed.(*nn.Linear)
-	if !ok {
-		t.Fatalf("packed projection is %T, want *nn.Linear", packed)
-	}
-	want := mlx.Concatenate([]*mlx.Array{
-		mlx.Dequantize(hi.Weight, hi.Scales, hi.QBiases, hi.GroupSize, hi.Bits, hi.Mode, nil),
-		loW.AsType(mlx.DTypeFloat16),
-	}, 0)
-	if dense.Weight.Dim(0) != 24 {
-		t.Fatalf("packed rows = %d, want 24", dense.Weight.Dim(0))
-	}
-	assertBitEqual(t, "weight", dense.Weight, want.AsType(dense.Weight.DType()))
+		packed, err := concatProjectionPair(hi, nn.NewLinear(loW, nil))
+		if err != nil {
+			t.Fatal(err)
+		}
+		dense, ok := packed.(*nn.Linear)
+		if !ok {
+			t.Fatalf("packed projection is %T, want *nn.Linear", packed)
+		}
+		want := mlx.Concatenate([]*mlx.Array{
+			mlx.Dequantize(hi.Weight, hi.Scales, hi.QBiases, hi.GroupSize, hi.Bits, hi.Mode, nil),
+			loW.AsType(mlx.DTypeFloat16),
+		}, 0)
+		if dense.Weight.Dim(0) != 24 {
+			t.Fatalf("packed rows = %d, want 24", dense.Weight.Dim(0))
+		}
+		assertBitEqual(t, "weight", dense.Weight, want.AsType(dense.Weight.DType()))
+	})
 }

x/models/qwen3_5/qwen3_5_test.go

@@ -8,33 +8,26 @@ import (
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}
-
 func TestSanitizeConvWeight(t *testing.T) {
-	mlxtest.Setup(t)
-
-	tests := []struct {
-		name  string
-		shape []int
-		want  []int
-	}{
-		{name: "publisher layout", shape: []int{8, 1, 4}, want: []int{8, 4}},
-		{name: "imported layout", shape: []int{8, 4, 1}, want: []int{8, 4}},
-		{name: "already sanitized", shape: []int{8, 4}, want: []int{8, 4}},
-	}
-
-	for _, tt := range tests {
-		got := sanitizeConvWeight(mlx.Zeros(mlx.DTypeBFloat16, tt.shape...))
-		mlx.Eval(got)
-		if dims := got.Dims(); len(dims) != len(tt.want) || dims[0] != tt.want[0] || dims[1] != tt.want[1] {
-			t.Fatalf("%s: sanitizeConvWeight() shape = %v, want %v", tt.name, dims, tt.want)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		tests := []struct {
+			name  string
+			shape []int
+			want  []int
+		}{
+			{name: "publisher layout", shape: []int{8, 1, 4}, want: []int{8, 4}},
+			{name: "imported layout", shape: []int{8, 4, 1}, want: []int{8, 4}},
+			{name: "already sanitized", shape: []int{8, 4}, want: []int{8, 4}},
 		}
-	}
+
+		for _, tt := range tests {
+			got := sanitizeConvWeight(mlx.Zeros(mlx.DTypeBFloat16, tt.shape...))
+			mlx.Eval(got)
+			if dims := got.Dims(); len(dims) != len(tt.want) || dims[0] != tt.want[0] || dims[1] != tt.want[1] {
+				t.Fatalf("%s: sanitizeConvWeight() shape = %v, want %v", tt.name, dims, tt.want)
+			}
+		}
+	})
 }
 
 func TestParseConfigNestedDefaults(t *testing.T) {
@@ -209,182 +202,182 @@ func TestNewCachesLayout(t *testing.T) {
 }
 
 func TestLoadWeightsPreservesLinearAttentionNormWeightDType(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{
+			HiddenSize:            4,
+			IntermediateSize:      8,
+			NumHiddenLayers:       2,
+			NumAttentionHeads:     1,
+			NumKeyValueHeads:      1,
+			HeadDim:               4,
+			RMSNormEps:            1e-6,
+			TieWordEmbeddings:     true,
+			LayerTypes:            []string{"linear", "full"},
+			LinearNumValueHeads:   1,
+			LinearNumKeyHeads:     1,
+			LinearKeyHeadDim:      2,
+			LinearValueHeadDim:    2,
+			LinearConvKernelDim:   4,
+			FullAttentionInterval: 2,
+		}
 
-	cfg := &Config{
-		HiddenSize:            4,
-		IntermediateSize:      8,
-		NumHiddenLayers:       2,
-		NumAttentionHeads:     1,
-		NumKeyValueHeads:      1,
-		HeadDim:               4,
-		RMSNormEps:            1e-6,
-		TieWordEmbeddings:     true,
-		LayerTypes:            []string{"linear", "full"},
-		LinearNumValueHeads:   1,
-		LinearNumKeyHeads:     1,
-		LinearKeyHeadDim:      2,
-		LinearValueHeadDim:    2,
-		LinearConvKernelDim:   4,
-		FullAttentionInterval: 2,
-	}
+		m := &Model{
+			Config: cfg,
+			Layers: make([]*Layer, cfg.NumHiddenLayers),
+		}
 
-	m := &Model{
-		Config: cfg,
-		Layers: make([]*Layer, cfg.NumHiddenLayers),
-	}
+		bf16 := mlx.DTypeBFloat16
+		f32 := mlx.DTypeFloat32
+		tensors := map[string]*mlx.Array{
+			"model.embed_tokens.weight":                      mlx.FromValues([]float32{1, 2, 3, 4, 5, 6, 7, 8}, 2, 4).AsType(bf16),
+			"model.norm.weight":                              mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.0.input_layernorm.weight":          mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.0.post_attention_layernorm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.0.linear_attn.in_proj_qkv.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+			}, 6, 4),
+			"model.layers.0.linear_attn.in_proj_z.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+			}, 2, 4),
+			"model.layers.0.linear_attn.in_proj_b.weight": mlx.FromValues([]float32{1, 0, 0, 0}, 1, 4),
+			"model.layers.0.linear_attn.in_proj_a.weight": mlx.FromValues([]float32{0, 1, 0, 0}, 1, 4),
+			"model.layers.0.linear_attn.out_proj.weight": mlx.FromValues([]float32{
+				1, 0,
+				0, 1,
+				1, 1,
+				0, 0,
+			}, 4, 2),
+			"model.layers.0.linear_attn.conv1d.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+			}, 6, 4),
+			"model.layers.0.linear_attn.norm.weight": mlx.FromValues([]float32{1, 1}, 2),
+			"model.layers.0.linear_attn.dt_bias":     mlx.FromValues([]float32{0}, 1),
+			"model.layers.0.linear_attn.A_log":       mlx.FromValues([]float32{0}, 1),
+			"model.layers.0.mlp.gate_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+				0, 0, 1, 1,
+				1, 0, 0, 1,
+			}, 8, 4),
+			"model.layers.0.mlp.up_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+				0, 0, 1, 1,
+				1, 0, 0, 1,
+			}, 8, 4),
+			"model.layers.0.mlp.down_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0, 0, 0, 0, 0,
+				0, 1, 0, 0, 0, 0, 0, 0,
+				0, 0, 1, 0, 0, 0, 0, 0,
+				0, 0, 0, 1, 0, 0, 0, 0,
+			}, 4, 8),
+			"model.layers.1.input_layernorm.weight":          mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.1.post_attention_layernorm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.1.self_attn.q_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+				0, 0, 1, 1,
+				1, 0, 0, 1,
+			}, 8, 4),
+			"model.layers.1.self_attn.k_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+			}, 4, 4),
+			"model.layers.1.self_attn.v_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+			}, 4, 4),
+			"model.layers.1.self_attn.o_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+			}, 4, 4),
+			"model.layers.1.self_attn.q_norm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.1.self_attn.k_norm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
+			"model.layers.1.mlp.gate_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+				0, 0, 1, 1,
+				1, 0, 0, 1,
+			}, 8, 4),
+			"model.layers.1.mlp.up_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0,
+				0, 1, 0, 0,
+				0, 0, 1, 0,
+				0, 0, 0, 1,
+				1, 1, 0, 0,
+				0, 1, 1, 0,
+				0, 0, 1, 1,
+				1, 0, 0, 1,
+			}, 8, 4),
+			"model.layers.1.mlp.down_proj.weight": mlx.FromValues([]float32{
+				1, 0, 0, 0, 0, 0, 0, 0,
+				0, 1, 0, 0, 0, 0, 0, 0,
+				0, 0, 1, 0, 0, 0, 0, 0,
+				0, 0, 0, 1, 0, 0, 0, 0,
+			}, 4, 8),
+		}
 
-	bf16 := mlx.DTypeBFloat16
-	f32 := mlx.DTypeFloat32
-	tensors := map[string]*mlx.Array{
-		"model.embed_tokens.weight":                      mlx.FromValues([]float32{1, 2, 3, 4, 5, 6, 7, 8}, 2, 4).AsType(bf16),
-		"model.norm.weight":                              mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.0.input_layernorm.weight":          mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.0.post_attention_layernorm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.0.linear_attn.in_proj_qkv.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-		}, 6, 4),
-		"model.layers.0.linear_attn.in_proj_z.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-		}, 2, 4),
-		"model.layers.0.linear_attn.in_proj_b.weight": mlx.FromValues([]float32{1, 0, 0, 0}, 1, 4),
-		"model.layers.0.linear_attn.in_proj_a.weight": mlx.FromValues([]float32{0, 1, 0, 0}, 1, 4),
-		"model.layers.0.linear_attn.out_proj.weight": mlx.FromValues([]float32{
-			1, 0,
-			0, 1,
-			1, 1,
-			0, 0,
-		}, 4, 2),
-		"model.layers.0.linear_attn.conv1d.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-		}, 6, 4),
-		"model.layers.0.linear_attn.norm.weight": mlx.FromValues([]float32{1, 1}, 2),
-		"model.layers.0.linear_attn.dt_bias":     mlx.FromValues([]float32{0}, 1),
-		"model.layers.0.linear_attn.A_log":       mlx.FromValues([]float32{0}, 1),
-		"model.layers.0.mlp.gate_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-			0, 0, 1, 1,
-			1, 0, 0, 1,
-		}, 8, 4),
-		"model.layers.0.mlp.up_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-			0, 0, 1, 1,
-			1, 0, 0, 1,
-		}, 8, 4),
-		"model.layers.0.mlp.down_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0, 0, 0, 0, 0,
-			0, 1, 0, 0, 0, 0, 0, 0,
-			0, 0, 1, 0, 0, 0, 0, 0,
-			0, 0, 0, 1, 0, 0, 0, 0,
-		}, 4, 8),
-		"model.layers.1.input_layernorm.weight":          mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.1.post_attention_layernorm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.1.self_attn.q_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-			0, 0, 1, 1,
-			1, 0, 0, 1,
-		}, 8, 4),
-		"model.layers.1.self_attn.k_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-		}, 4, 4),
-		"model.layers.1.self_attn.v_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-		}, 4, 4),
-		"model.layers.1.self_attn.o_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-		}, 4, 4),
-		"model.layers.1.self_attn.q_norm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.1.self_attn.k_norm.weight": mlx.FromValues([]float32{1, 1, 1, 1}, 4),
-		"model.layers.1.mlp.gate_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-			0, 0, 1, 1,
-			1, 0, 0, 1,
-		}, 8, 4),
-		"model.layers.1.mlp.up_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0,
-			0, 1, 0, 0,
-			0, 0, 1, 0,
-			0, 0, 0, 1,
-			1, 1, 0, 0,
-			0, 1, 1, 0,
-			0, 0, 1, 1,
-			1, 0, 0, 1,
-		}, 8, 4),
-		"model.layers.1.mlp.down_proj.weight": mlx.FromValues([]float32{
-			1, 0, 0, 0, 0, 0, 0, 0,
-			0, 1, 0, 0, 0, 0, 0, 0,
-			0, 0, 1, 0, 0, 0, 0, 0,
-			0, 0, 0, 1, 0, 0, 0, 0,
-		}, 4, 8),
-	}
+		if err := m.LoadWeights(tensors); err != nil {
+			t.Fatalf("LoadWeights failed: %v", err)
+		}
 
-	if err := m.LoadWeights(tensors); err != nil {
-		t.Fatalf("LoadWeights failed: %v", err)
-	}
+		if got := m.Layers[0].InputNorm.Weight.DType(); got != f32 {
+			t.Fatalf("layer 0 input norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[0].PostAttentionNorm.Weight.DType(); got != f32 {
+			t.Fatalf("layer 0 post-attn norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[1].InputNorm.Weight.DType(); got != f32 {
+			t.Fatalf("layer 1 input norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[1].PostAttentionNorm.Weight.DType(); got != f32 {
+			t.Fatalf("layer 1 post-attn norm dtype = %v, want %v", got, f32)
+		}
 
-	if got := m.Layers[0].InputNorm.Weight.DType(); got != f32 {
-		t.Fatalf("layer 0 input norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[0].PostAttentionNorm.Weight.DType(); got != f32 {
-		t.Fatalf("layer 0 post-attn norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[1].InputNorm.Weight.DType(); got != f32 {
-		t.Fatalf("layer 1 input norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[1].PostAttentionNorm.Weight.DType(); got != f32 {
-		t.Fatalf("layer 1 post-attn norm dtype = %v, want %v", got, f32)
-	}
-
-	if got := m.Norm.Weight.DType(); got != f32 {
-		t.Fatalf("final norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[0].Linear.NormWeight.DType(); got != f32 {
-		t.Fatalf("linear-attn norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[1].FullAttn.QNorm.Weight.DType(); got != f32 {
-		t.Fatalf("q norm dtype = %v, want %v", got, f32)
-	}
-	if got := m.Layers[1].FullAttn.KNorm.Weight.DType(); got != f32 {
-		t.Fatalf("k norm dtype = %v, want %v", got, f32)
-	}
+		if got := m.Norm.Weight.DType(); got != f32 {
+			t.Fatalf("final norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[0].Linear.NormWeight.DType(); got != f32 {
+			t.Fatalf("linear-attn norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[1].FullAttn.QNorm.Weight.DType(); got != f32 {
+			t.Fatalf("q norm dtype = %v, want %v", got, f32)
+		}
+		if got := m.Layers[1].FullAttn.KNorm.Weight.DType(); got != f32 {
+			t.Fatalf("k norm dtype = %v, want %v", got, f32)
+		}
+	})
 }

x/models/qwen3_5/vision_test.go

@@ -13,12 +13,13 @@ import (
 )
 
 func TestVisionAdapterWeightsAreCollectable(t *testing.T) {
-	mlxtest.Setup(t)
-	weight := mlx.FromValue(float32(1))
-	adapter := &VisionAdapter{Model: &Model{VisionTower: &VisionTower{PosEmbed: weight}}}
-	if got := mlx.Collect(adapter); len(got) != 1 || got[0] != weight {
-		t.Fatalf("mlx.Collect(adapter) = %v, want the tower weight", got)
-	}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		weight := mlx.FromValue(float32(1))
+		adapter := &VisionAdapter{Model: &Model{VisionTower: &VisionTower{PosEmbed: weight}}}
+		if got := mlx.Collect(adapter); len(got) != 1 || got[0] != weight {
+			t.Fatalf("mlx.Collect(adapter) = %v, want the tower weight", got)
+		}
+	})
 }
 
 func TestSmartResize(t *testing.T) {

x/models/qwen4_exp/engram_test.go

@@ -10,96 +10,98 @@ import (
 )
 
 func TestEngramHashes(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		p := &PLE{
+			LayerMultipliers: mlx.FromValues([]int64{3, 5, 7}, 3),
+			HeadVocabSizes:   mlx.FromValues([]int64{11, 13, 17, 19}, 4),
+			HeadOffsets:      mlx.FromValues([]int64{0, 11, 24, 41}, 4),
+		}
+		cfg := &Config{NGramSize: 3, HeadsPerNGram: 2, EOSTokenID: 9}
+		b := &batch.Batch{InputIDs: mlx.FromValues([]int32{1, 9, 2, 3}, 1, 4)}
+		history := mlx.FromValues([]int64{7, 8}, 1, 2)
 
-	p := &PLE{
-		LayerMultipliers: mlx.FromValues([]int64{3, 5, 7}, 3),
-		HeadVocabSizes:   mlx.FromValues([]int64{11, 13, 17, 19}, 4),
-		HeadOffsets:      mlx.FromValues([]int64{0, 11, 24, 41}, 4),
-	}
-	cfg := &Config{NGramSize: 3, HeadsPerNGram: 2, EOSTokenID: 9}
-	b := &batch.Batch{InputIDs: mlx.FromValues([]int32{1, 9, 2, 3}, 1, 4)}
-	history := mlx.FromValues([]int64{7, 8}, 1, 2)
-
-	got, _ := p.hashes(b, history, cfg)
-	got = got.AsType(mlx.DTypeInt32)
-	mlx.Eval(got)
-	values := got.Ints()
-	want := []int32{10, 15, 33, 48, 8, 15, 28, 41, 10, 15, 27, 42, 3, 14, 33, 44}
-	if !slices.Equal(values, want) {
-		t.Fatalf("hashes = %v, want %v", values, want)
-	}
+		got, _ := p.hashes(b, history, cfg)
+		got = got.AsType(mlx.DTypeInt32)
+		mlx.Eval(got)
+		values := got.Ints()
+		want := []int32{10, 15, 33, 48, 8, 15, 28, 41, 10, 15, 27, 42, 3, 14, 33, 44}
+		if !slices.Equal(values, want) {
+			t.Fatalf("hashes = %v, want %v", values, want)
+		}
+	})
 }
 
 func TestEngramCacheCarriesChunkHistory(t *testing.T) {
-	mlxtest.Setup(t)
-	c := newEngramCache(2, 3, 1, 9)
-	t.Cleanup(c.Free)
-	b := &batch.Batch{
-		InputIDs:     mlx.FromValues([]int32{1, 2}, 1, 2),
-		SeqQueryLens: []int32{2},
-	}
-	c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1))
-	b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2)
-	c.put(b, mlx.FromValues([]int64{3, 4}, 1, 2), mlx.FromValues([]float32{30, 40}, 1, 2, 1))
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := newEngramCache(2, 3, 1, 9)
+		t.Cleanup(c.Free)
+		b := &batch.Batch{
+			InputIDs:     mlx.FromValues([]int32{1, 2}, 1, 2),
+			SeqQueryLens: []int32{2},
+		}
+		c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1))
+		b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2)
+		c.put(b, mlx.FromValues([]int64{3, 4}, 1, 2), mlx.FromValues([]float32{30, 40}, 1, 2, 1))
 
-	history := c.history.AsType(mlx.DTypeInt32)
-	mlx.Eval(history, c.convHistory)
-	if got, want := history.Ints(), []int32{3, 4}; !slices.Equal(got, want) {
-		t.Fatalf("token history = %v, want %v", got, want)
-	}
-	if got, want := c.convHistory.Floats(), []float32{20, 30, 40}; !slices.Equal(got, want) {
-		t.Fatalf("conv history = %v, want %v", got, want)
-	}
+		history := c.history.AsType(mlx.DTypeInt32)
+		mlx.Eval(history, c.convHistory)
+		if got, want := history.Ints(), []int32{3, 4}; !slices.Equal(got, want) {
+			t.Fatalf("token history = %v, want %v", got, want)
+		}
+		if got, want := c.convHistory.Floats(), []float32{20, 30, 40}; !slices.Equal(got, want) {
+			t.Fatalf("conv history = %v, want %v", got, want)
+		}
+	})
 }
 
 func TestEngramCacheRestoresScheduledSnapshot(t *testing.T) {
-	mlxtest.Setup(t)
-	c := newEngramCache(2, 3, 1, 9)
-	t.Cleanup(c.Free)
-	b := &batch.Batch{
-		InputIDs:     mlx.FromValues([]int32{1, 2}, 1, 2),
-		SeqQueryLens: []int32{2},
-	}
-	c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1))
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		c := newEngramCache(2, 3, 1, 9)
+		t.Cleanup(c.Free)
+		b := &batch.Batch{
+			InputIDs:     mlx.FromValues([]int32{1, 2}, 1, 2),
+			SeqQueryLens: []int32{2},
+		}
+		c.put(b, mlx.FromValues([]int64{1, 2}, 1, 2), mlx.FromValues([]float32{10, 20}, 1, 2, 1))
 
-	c.PrepareSnapshots([]int{3, 4})
-	b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2)
-	inputIDs := mlx.FromValues([]int64{3, 4}, 1, 2)
-	convInput := mlx.FromValues([]float32{30, 40}, 1, 2, 1)
-	c.put(b, inputIDs, convInput)
-	snapshots := c.TakeSnapshots()
-	if len(snapshots) != 2 || snapshots[0] == nil || snapshots[1] == nil {
-		t.Fatalf("TakeSnapshots() = %v, want two captured snapshots", snapshots)
-	}
-	for _, snapshot := range snapshots {
-		t.Cleanup(snapshot.Close)
-	}
-	first := snapshots[0].(*engramSnapshot)
-	mlx.Eval(first.history, first.convHistory)
-	if got, want := first.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) {
-		t.Fatalf("captured convolution history = %v, want %v", got, want)
-	}
+		c.PrepareSnapshots([]int{3, 4})
+		b.InputIDs = mlx.FromValues([]int32{3, 4}, 1, 2)
+		inputIDs := mlx.FromValues([]int64{3, 4}, 1, 2)
+		convInput := mlx.FromValues([]float32{30, 40}, 1, 2, 1)
+		c.put(b, inputIDs, convInput)
+		snapshots := c.TakeSnapshots()
+		if len(snapshots) != 2 || snapshots[0] == nil || snapshots[1] == nil {
+			t.Fatalf("TakeSnapshots() = %v, want two captured snapshots", snapshots)
+		}
+		for _, snapshot := range snapshots {
+			t.Cleanup(snapshot.Close)
+		}
+		first := snapshots[0].(*engramSnapshot)
+		mlx.Eval(first.history, first.convHistory)
+		if got, want := first.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) {
+			t.Fatalf("captured convolution history = %v, want %v", got, want)
+		}
 
-	b.InputIDs = mlx.FromValues([]int32{5, 6}, 1, 2)
-	c.put(b, mlx.FromValues([]int64{5, 6}, 1, 2), mlx.FromValues([]float32{50, 60}, 1, 2, 1))
-	if !c.Restore(snapshots[0], 3) {
-		t.Fatal("Restore(snapshot, 3) failed")
-	}
-	history := c.history.AsType(mlx.DTypeInt32)
-	mlx.Eval(history, c.convHistory)
-	if got, want := history.Ints(), []int32{2, 3}; !slices.Equal(got, want) {
-		t.Fatalf("restored token history = %v, want %v", got, want)
-	}
-	if got, want := c.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) {
-		t.Fatalf("restored convolution history = %v, want %v", got, want)
-	}
+		b.InputIDs = mlx.FromValues([]int32{5, 6}, 1, 2)
+		c.put(b, mlx.FromValues([]int64{5, 6}, 1, 2), mlx.FromValues([]float32{50, 60}, 1, 2, 1))
+		if !c.Restore(snapshots[0], 3) {
+			t.Fatal("Restore(snapshot, 3) failed")
+		}
+		history := c.history.AsType(mlx.DTypeInt32)
+		mlx.Eval(history, c.convHistory)
+		if got, want := history.Ints(), []int32{2, 3}; !slices.Equal(got, want) {
+			t.Fatalf("restored token history = %v, want %v", got, want)
+		}
+		if got, want := c.convHistory.Floats(), []float32{10, 20, 30}; !slices.Equal(got, want) {
+			t.Fatalf("restored convolution history = %v, want %v", got, want)
+		}
 
-	parent, child := c.Split(snapshots[1], 3)
-	if parent != nil || child != snapshots[1] {
-		t.Fatalf("Split(snapshot) = (%v, %v), want (nil, snapshot)", parent, child)
-	}
-	if merged := c.Merge(nil, child); merged != child {
-		t.Fatalf("Merge(nil, child) = %v, want child", merged)
-	}
+		parent, child := c.Split(snapshots[1], 3)
+		if parent != nil || child != snapshots[1] {
+			t.Fatalf("Split(snapshot) = (%v, %v), want (nil, snapshot)", parent, child)
+		}
+		if merged := c.Merge(nil, child); merged != child {
+			t.Fatalf("Merge(nil, child) = %v, want child", merged)
+		}
+	})
 }

x/models/qwen4_exp/hyper_connection_test.go

@@ -10,67 +10,67 @@ import (
 )
 
 func TestHyperConnectionMatchesReferenceFormula(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{HCCount: 2, HiddenSize: 2, HCLowRank: 2, RMSNormEps: 1e-6}
+		normWeight := []float32{1.1, 0.9, 1.2, 0.8}
+		downWeight := [][]float32{{0.2, -0.3, 0.4, 0.1}, {-0.1, 0.5, 0.2, -0.4}}
+		upWeight := [][]float32{{0.3, -0.2}, {0.1, 0.4}, {-0.5, 0.2}, {0.25, 0.15}}
+		injectWeight := [][]float32{{0.2, -0.1, 0.3, 0.4}, {-0.3, 0.2, 0.1, 0.5}}
+		input := []float32{1, 2, 3, 4}
 
-	cfg := &Config{HCCount: 2, HiddenSize: 2, HCLowRank: 2, RMSNormEps: 1e-6}
-	normWeight := []float32{1.1, 0.9, 1.2, 0.8}
-	downWeight := [][]float32{{0.2, -0.3, 0.4, 0.1}, {-0.1, 0.5, 0.2, -0.4}}
-	upWeight := [][]float32{{0.3, -0.2}, {0.1, 0.4}, {-0.5, 0.2}, {0.25, 0.15}}
-	injectWeight := [][]float32{{0.2, -0.1, 0.3, 0.4}, {-0.3, 0.2, 0.1, 0.5}}
-	input := []float32{1, 2, 3, 4}
-
-	h := &hyperConnection{
-		Norm:         &streamRMSNorm{Weight: mlx.FromValues(normWeight, 2, 2)},
-		InputMixDown: nn.NewLinear(matrix(downWeight), nil),
-		InputMixUp:   nn.NewLinear(matrix(upWeight), nil),
-		BlockInject:  nn.NewLinear(matrix(injectWeight), nil),
-	}
-	residual := mlx.FromValues(input, 1, 1, 4)
-	branch, state := h.Prepare(residual, cfg)
-	got := h.Inject(state, branch, cfg).AsType(mlx.DTypeFloat32)
-	reduced := h.Reduce(residual, cfg).AsType(mlx.DTypeFloat32)
-	mlx.Eval(got, reduced)
-
-	normed := append([]float32(nil), input...)
-	for stream := range int(cfg.HCCount) {
-		start := stream * int(cfg.HiddenSize)
-		var square float64
-		for i := range int(cfg.HiddenSize) {
-			x := float64(input[start+i])
-			square += x * x
+		h := &hyperConnection{
+			Norm:         &streamRMSNorm{Weight: mlx.FromValues(normWeight, 2, 2)},
+			InputMixDown: nn.NewLinear(matrix(downWeight), nil),
+			InputMixUp:   nn.NewLinear(matrix(upWeight), nil),
+			BlockInject:  nn.NewLinear(matrix(injectWeight), nil),
 		}
-		invRMS := 1 / math.Sqrt(square/float64(cfg.HiddenSize)+float64(cfg.RMSNormEps))
-		for i := range int(cfg.HiddenSize) {
-			normed[start+i] = float32(float64(input[start+i]) * invRMS * float64(normWeight[start+i]))
-		}
-	}
+		residual := mlx.FromValues(input, 1, 1, 4)
+		branch, state := h.Prepare(residual, cfg)
+		got := h.Inject(state, branch, cfg).AsType(mlx.DTypeFloat32)
+		reduced := h.Reduce(residual, cfg).AsType(mlx.DTypeFloat32)
+		mlx.Eval(got, reduced)
 
-	down := matvec(downWeight, normed)
-	for i := range down {
-		down[i] /= float32(cfg.HCCount)
-		down[i] *= 1 / (1 + float32(math.Exp(float64(-down[i]))))
-	}
-	mix := matvec(upWeight, down)
-	wantBranch := make([]float32, cfg.HiddenSize)
-	for stream := range int(cfg.HCCount) {
-		for i := range int(cfg.HiddenSize) {
-			j := stream*int(cfg.HiddenSize) + i
-			gate := 1 / (1 + float32(math.Exp(float64(-mix[j]))))
-			wantBranch[i] += gate * normed[j] / float32(cfg.HCCount)
+		normed := append([]float32(nil), input...)
+		for stream := range int(cfg.HCCount) {
+			start := stream * int(cfg.HiddenSize)
+			var square float64
+			for i := range int(cfg.HiddenSize) {
+				x := float64(input[start+i])
+				square += x * x
+			}
+			invRMS := 1 / math.Sqrt(square/float64(cfg.HiddenSize)+float64(cfg.RMSNormEps))
+			for i := range int(cfg.HiddenSize) {
+				normed[start+i] = float32(float64(input[start+i]) * invRMS * float64(normWeight[start+i]))
+			}
 		}
-	}
 
-	injection := matvec(injectWeight, normed)
-	want := append([]float32(nil), input...)
-	for stream := range int(cfg.HCCount) {
-		weight := 2 / (1 + float32(math.Exp(float64(-injection[stream]/float32(cfg.HCCount)))))
-		for i := range int(cfg.HiddenSize) {
-			want[stream*int(cfg.HiddenSize)+i] += weight * wantBranch[i]
+		down := matvec(downWeight, normed)
+		for i := range down {
+			down[i] /= float32(cfg.HCCount)
+			down[i] *= 1 / (1 + float32(math.Exp(float64(-down[i]))))
+		}
+		mix := matvec(upWeight, down)
+		wantBranch := make([]float32, cfg.HiddenSize)
+		for stream := range int(cfg.HCCount) {
+			for i := range int(cfg.HiddenSize) {
+				j := stream*int(cfg.HiddenSize) + i
+				gate := 1 / (1 + float32(math.Exp(float64(-mix[j]))))
+				wantBranch[i] += gate * normed[j] / float32(cfg.HCCount)
+			}
 		}
-	}
 
-	assertClose(t, "mixed branch", reduced.Floats(), wantBranch)
-	assertClose(t, "injected streams", got.Floats(), want)
+		injection := matvec(injectWeight, normed)
+		want := append([]float32(nil), input...)
+		for stream := range int(cfg.HCCount) {
+			weight := 2 / (1 + float32(math.Exp(float64(-injection[stream]/float32(cfg.HCCount)))))
+			for i := range int(cfg.HiddenSize) {
+				want[stream*int(cfg.HiddenSize)+i] += weight * wantBranch[i]
+			}
+		}
+
+		assertClose(t, "mixed branch", reduced.Floats(), wantBranch)
+		assertClose(t, "injected streams", got.Floats(), want)
+	})
 }
 
 func matrix(rows [][]float32) *mlx.Array {
@@ -91,7 +91,7 @@ func matvec(weight [][]float32, input []float32) []float32 {
 	return output
 }
 
-func assertClose(t *testing.T, name string, got, want []float32) {
+func assertClose(t *mlxtest.T, name string, got, want []float32) {
 	t.Helper()
 	if len(got) != len(want) {
 		t.Fatalf("%s length = %d, want %d", name, len(got), len(want))

x/models/qwen4_exp/qsa_test.go

@@ -12,114 +12,119 @@ import (
 )
 
 func TestQSASelectsCompressedBlocksAndCausalTail(t *testing.T) {
-	mlxtest.Setup(t)
-	cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4}
-	scores := mlx.FromValues([]float32{0.1, 4, 2, 3}, 1, 1, 4)
-	b := &batch.Batch{SeqOffsets: []int32{16}}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4}
+		scores := mlx.FromValues([]float32{0.1, 4, 2, 3}, 1, 1, 4)
+		b := &batch.Batch{SeqOffsets: []int32{16}}
 
-	indices, valid := qsaLogicalIndices(scores, b, 16, cfg)
-	indices = indices.AsType(mlx.DTypeInt32)
-	valid = valid.AsType(mlx.DTypeInt32)
-	mlx.Eval(indices, valid)
+		indices, valid := qsaLogicalIndices(scores, b, 16, cfg)
+		indices = indices.AsType(mlx.DTypeInt32)
+		valid = valid.AsType(mlx.DTypeInt32)
+		mlx.Eval(indices, valid)
 
-	values, mask := indices.Ints(), valid.Ints()
-	selected := make([]int32, 0, len(values))
-	for i, value := range values {
-		if mask[i] != 0 {
-			selected = append(selected, value)
+		values, mask := indices.Ints(), valid.Ints()
+		selected := make([]int32, 0, len(values))
+		for i, value := range values {
+			if mask[i] != 0 {
+				selected = append(selected, value)
+			}
 		}
-	}
-	slices.Sort(selected)
-	want := []int32{4, 5, 6, 7, 12, 13, 14, 15, 16}
-	if !slices.Equal(selected, want) {
-		t.Fatalf("selected indices = %v, want %v", selected, want)
-	}
+		slices.Sort(selected)
+		want := []int32{4, 5, 6, 7, 12, 13, 14, 15, 16}
+		if !slices.Equal(selected, want) {
+			t.Fatalf("selected indices = %v, want %v", selected, want)
+		}
+	})
 }
 
 func TestQSASelectionMasksFutureBlocks(t *testing.T) {
-	mlxtest.Setup(t)
-	cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4}
-	// Only block 0 and token 4 are visible. Give every future block a much
-	// larger score so the test fails if selection sees cached-but-causal-junk.
-	scores := mlx.FromValues([]float32{0.1, 100, 90, 80, 70}, 1, 1, 5)
-	b := &batch.Batch{SeqOffsets: []int32{4}}
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{IndexerBudget: 8, IndexerCompressRatio: 4}
+		// Only block 0 and token 4 are visible. Give every future block a much
+		// larger score so the test fails if selection sees cached-but-causal-junk.
+		scores := mlx.FromValues([]float32{0.1, 100, 90, 80, 70}, 1, 1, 5)
+		b := &batch.Batch{SeqOffsets: []int32{4}}
 
-	indices, valid := qsaLogicalIndices(scores, b, 20, cfg)
-	indices = indices.AsType(mlx.DTypeInt32)
-	valid = valid.AsType(mlx.DTypeInt32)
-	mlx.Eval(indices, valid)
+		indices, valid := qsaLogicalIndices(scores, b, 20, cfg)
+		indices = indices.AsType(mlx.DTypeInt32)
+		valid = valid.AsType(mlx.DTypeInt32)
+		mlx.Eval(indices, valid)
 
-	var selected []int32
-	for i, value := range indices.Ints() {
-		if valid.Ints()[i] != 0 {
-			selected = append(selected, value)
+		var selected []int32
+		for i, value := range indices.Ints() {
+			if valid.Ints()[i] != 0 {
+				selected = append(selected, value)
+			}
 		}
-	}
-	slices.Sort(selected)
-	if want := []int32{0, 1, 2, 3, 4}; !slices.Equal(selected, want) {
-		t.Fatalf("selected indices = %v, want %v", selected, want)
-	}
+		slices.Sort(selected)
+		if want := []int32{0, 1, 2, 3, 4}; !slices.Equal(selected, want) {
+			t.Fatalf("selected indices = %v, want %v", selected, want)
+		}
+	})
 }
 
 func TestQSASparseAttentionMatchesReference(t *testing.T) {
-	mlxtest.Setup(t)
-	cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
-	q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2)
-	k := mlx.FromValues([]float32{1, 0, 0, 1, 2, 0}, 1, 1, 3, 2)
-	v := mlx.FromValues([]float32{10, 1, 20, 2, 30, 3}, 1, 1, 3, 2)
-	indices := mlx.FromValues([]int32{2, 0}, 1, 1, 2)
-	valid := mlx.FromValues([]bool{true, true}, 1, 1, 2)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
+		q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2)
+		k := mlx.FromValues([]float32{1, 0, 0, 1, 2, 0}, 1, 1, 3, 2)
+		v := mlx.FromValues([]float32{10, 1, 20, 2, 30, 3}, 1, 1, 3, 2)
+		indices := mlx.FromValues([]int32{2, 0}, 1, 1, 2)
+		valid := mlx.FromValues([]bool{true, true}, 1, 1, 2)
 
-	out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
-	out = out.AsType(mlx.DTypeFloat32)
-	mlx.Eval(out)
-	got := out.Floats()
-	p2 := float32(math.Exp(2) / (math.Exp(2) + math.Exp(1)))
-	want := []float32{p2*30 + (1-p2)*10, p2*3 + (1-p2)*1}
-	for i := range want {
-		if math.Abs(float64(got[i]-want[i])) > 1e-4 {
-			t.Fatalf("sparse attention[%d] = %v, want %v", i, got[i], want[i])
+		out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
+		out = out.AsType(mlx.DTypeFloat32)
+		mlx.Eval(out)
+		got := out.Floats()
+		p2 := float32(math.Exp(2) / (math.Exp(2) + math.Exp(1)))
+		want := []float32{p2*30 + (1-p2)*10, p2*3 + (1-p2)*1}
+		for i := range want {
+			if math.Abs(float64(got[i]-want[i])) > 1e-4 {
+				t.Fatalf("sparse attention[%d] = %v, want %v", i, got[i], want[i])
+			}
 		}
-	}
+	})
 }
 
 func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) {
-	mlxtest.Setup(t)
-	cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
-	q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2)
-	// Row 0 is intentionally dominant junk. Only row 1 is logically valid.
-	k := mlx.FromValues([]float32{100, 0, 1, 0}, 1, 1, 2, 2)
-	v := mlx.FromValues([]float32{999, 999, 7, 3}, 1, 1, 2, 2)
-	indices := mlx.FromValues([]int32{0, 1}, 1, 1, 2)
-	valid := mlx.FromValues([]bool{false, true}, 1, 1, 2)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
+		q := mlx.FromValues([]float32{1, 0}, 1, 1, 1, 2)
+		// Row 0 is intentionally dominant junk. Only row 1 is logically valid.
+		k := mlx.FromValues([]float32{100, 0, 1, 0}, 1, 1, 2, 2)
+		v := mlx.FromValues([]float32{999, 999, 7, 3}, 1, 1, 2, 2)
+		indices := mlx.FromValues([]int32{0, 1}, 1, 1, 2)
+		valid := mlx.FromValues([]bool{false, true}, 1, 1, 2)
 
-	out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
-	out = out.AsType(mlx.DTypeFloat32)
-	mlx.Eval(out)
-	if got, want := out.Floats(), []float32{7, 3}; !slices.Equal(got, want) {
-		t.Fatalf("sparse attention = %v, want %v", got, want)
-	}
+		out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
+		out = out.AsType(mlx.DTypeFloat32)
+		mlx.Eval(out)
+		if got, want := out.Floats(), []float32{7, 3}; !slices.Equal(got, want) {
+			t.Fatalf("sparse attention = %v, want %v", got, want)
+		}
+	})
 }
 
 func TestQSASparseAttentionKeepsBatchRowsIndependent(t *testing.T) {
-	mlxtest.Setup(t)
-	cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
-	q := mlx.FromValues([]float32{1, 0, 1, 0}, 2, 1, 1, 2)
-	k := mlx.FromValues([]float32{
-		1, 0, 0, 1,
-		1, 0, 0, 1,
-	}, 2, 1, 2, 2)
-	v := mlx.FromValues([]float32{
-		10, 1, 20, 2,
-		30, 3, 40, 4,
-	}, 2, 1, 2, 2)
-	indices := mlx.FromValues([]int32{0, 1}, 2, 1, 1)
-	valid := mlx.FromValues([]bool{true, true}, 2, 1, 1)
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		cfg := &Config{NumKeyValueHeads: 1, Scale: 1}
+		q := mlx.FromValues([]float32{1, 0, 1, 0}, 2, 1, 1, 2)
+		k := mlx.FromValues([]float32{
+			1, 0, 0, 1,
+			1, 0, 0, 1,
+		}, 2, 1, 2, 2)
+		v := mlx.FromValues([]float32{
+			10, 1, 20, 2,
+			30, 3, 40, 4,
+		}, 2, 1, 2, 2)
+		indices := mlx.FromValues([]int32{0, 1}, 2, 1, 1)
+		valid := mlx.FromValues([]bool{true, true}, 2, 1, 1)
 
-	out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
-	out = out.AsType(mlx.DTypeFloat32)
-	mlx.Eval(out)
-	if got, want := out.Floats(), []float32{10, 1, 40, 4}; !slices.Equal(got, want) {
-		t.Fatalf("sparse attention = %v, want %v", got, want)
-	}
+		out := qsaSparseAttention(q, nn.NewKVHistory(k, v, nil), indices, valid, cfg)
+		out = out.AsType(mlx.DTypeFloat32)
+		mlx.Eval(out)
+		if got, want := out.Floats(), []float32{10, 1, 40, 4}; !slices.Equal(got, want) {
+			t.Fatalf("sparse attention = %v, want %v", got, want)
+		}
+	})
 }

x/tokenizer/tokenizer_benchmark_test.go

@@ -14,20 +14,6 @@ var (
 	benchmarkSinkTok *Tokenizer
 )
 
-const benchmarkWordPieceJSON = `{
-  "model": {
-    "type": "WordPiece",
-    "vocab": {
-      "[UNK]": 0,
-      "hello": 1,
-      "##world": 2,
-      "##ly": 3,
-      "##hello": 4
-    }
-  },
-  "added_tokens": []
-}`
-
 const benchmarkSentencePieceJSON = `{
   "model": {
     "type": "BPE",
@@ -194,33 +180,6 @@ func BenchmarkTokenizerLoadFromBytes(b *testing.B) {
 	})
 }
 
-func BenchmarkTokenizerEncodeWordPiece(b *testing.B) {
-	tok := benchmarkLoadFromBytes(b, []byte(benchmarkWordPieceJSON))
-	text := strings.Repeat("helloworldly", 16)
-
-	b.ReportAllocs()
-	b.SetBytes(int64(len(text)))
-	b.ResetTimer()
-
-	for range b.N {
-		benchmarkSinkIDs = tok.Encode(text, false)
-	}
-}
-
-func BenchmarkTokenizerDecodeWordPiece(b *testing.B) {
-	tok := benchmarkLoadFromBytes(b, []byte(benchmarkWordPieceJSON))
-	text := strings.Repeat("helloworldly", 16)
-	ids := tok.Encode(text, false)
-
-	b.ReportAllocs()
-	b.SetBytes(int64(len(text)))
-	b.ResetTimer()
-
-	for range b.N {
-		benchmarkSinkStr = tok.Decode(ids)
-	}
-}
-
 func BenchmarkTokenizerEncodeSentencePiece(b *testing.B) {
 	tok := benchmarkLoadFromBytes(b, []byte(benchmarkSentencePieceJSON))
 	text := strings.Repeat("hello world\n", 64)

x/transfer/transfer_test.go

@@ -1680,6 +1680,10 @@ func BenchmarkUploadThroughput(b *testing.B) {
 		case http.MethodPost:
 			w.Header().Set("Location", fmt.Sprintf("%s/v2/library/_/blobs/uploads/1", serverURL))
 			w.WriteHeader(http.StatusAccepted)
+		case http.MethodPatch:
+			io.Copy(io.Discard, r.Body)
+			w.Header().Set("Location", r.URL.Path)
+			w.WriteHeader(http.StatusAccepted)
 		case http.MethodPut:
 			io.Copy(io.Discard, r.Body)
 			w.WriteHeader(http.StatusCreated)