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,7 +20,7 @@ func newKVBatch(off, L int) *batch.Batch {
 }
 
 func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewKVCache()
 
 		for range 10 {
@@ -45,12 +39,13 @@ func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewKVCache()
 
 		for range 10 {
@@ -85,12 +80,13 @@ func TestKVCacheDataSurvivesSnapshotRestore(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewKVCache()
 
 		for range 10 {
@@ -129,12 +125,13 @@ func TestKVCacheSplitPreservesData(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewKVCache()
 
 		for range 10 {
@@ -165,10 +162,11 @@ func TestKVCacheSplitMergeRoundTripData(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewRotatingKVCache(4)
 
 		// Feed 10 tokens (window size 4, so positions 0-5 are evicted).
@@ -182,12 +180,13 @@ func TestRotatingKVCacheRestoreOutsideWindow(t *testing.T) {
 		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)
+	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.
@@ -226,13 +225,14 @@ func TestRotatingKVCacheSnapshotPreservesWindow(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewRotatingKVCache(4)
 
 		// Fill the window: 6 tokens into a size-4 window.
@@ -280,4 +280,5 @@ func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) {
 		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,8 +45,7 @@ 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
 
 		c := NewKVCache()
@@ -91,6 +91,7 @@ func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) {
 			t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing",
 				after-baseline, perToken)
 		}
+	})
 }
 
 // TestKVLazySnapshotSizeZeroUntilMaterialized verifies the accounting contract:
@@ -98,8 +99,7 @@ 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
 
 		c := NewKVCache()
@@ -132,6 +132,7 @@ func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
 		if snap.Size() != want {
 			t.Fatalf("materialized Size = %d, want %d", snap.Size(), want)
 		}
+	})
 }
 
 // TestKVLazySnapshotCopiedOutOnOverwrite verifies that after a rewind, a write
@@ -139,8 +140,7 @@ 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
 
 		c := NewKVCache()
@@ -179,13 +179,13 @@ func TestKVLazySnapshotCopiedOutOnOverwrite(t *testing.T) {
 			}
 		}
 		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
 
 		c := NewKVCache()
@@ -203,14 +203,14 @@ func TestKVLazySnapshotCopiedOutOnFree(t *testing.T) {
 			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
 
 		c := NewKVCache()
@@ -243,6 +243,7 @@ func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) {
 			t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base)
 		}
 		merged.Close()
+	})
 }
 
 // TestKVLazySnapshotSurvivesPathSwitch reproduces the switchToPath sequence that
@@ -251,8 +252,7 @@ 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
 
 		c := NewKVCache()
@@ -286,6 +286,7 @@ func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) {
 			}
 		}
 		leaf.Close()
+	})
 }
 
 // TestKVRestoreLiveLazySnapshotIsOffsetMove verifies the same-path rewind/rematch
@@ -294,8 +295,7 @@ 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
 
 		c := NewKVCache()
@@ -336,4 +336,5 @@ func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) {
 			}
 		}
 		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,7 +14,7 @@ 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)
+	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
@@ -50,10 +51,11 @@ func TestRecurrentCacheRestoreExactOffset(t *testing.T) {
 		if c.Offset() != 10 {
 			t.Fatalf("offset = %d, want 10", c.Offset())
 		}
+	})
 }
 
 func TestRecurrentCacheGetLazyInit(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewRecurrentCache(3, 4, 2, 4, 4)
 		b := &batch.Batch{
 			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
@@ -73,6 +75,7 @@ func TestRecurrentCacheGetLazyInit(t *testing.T) {
 		if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 {
 			t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32)
 		}
+	})
 }
 
 // TestRecurrentCachePaddedRoundTrip runs Get → CausalConv1D →
@@ -83,7 +86,7 @@ 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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		const convTail, convDim = 2, 6
 		const numVHeads, headVDim, headKDim = 1, 4, 6
 		const L = 4
@@ -205,10 +208,11 @@ func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
 				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)
+	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)
@@ -217,4 +221,5 @@ func TestRecurrentCachePutAdvances(t *testing.T) {
 		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,8 +39,11 @@ 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.
@@ -52,12 +55,12 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) {
 
 		finalPos := totalWrites - 1
 		kFinal, vFinal := perPosKV(finalPos)
-	b := &batch.Batch{
+		b = &batch.Batch{
 			InputIDs:     mlx.Zeros(mlx.DTypeInt32, 1, 1),
 			SeqOffsets:   []int32{int32(finalPos)},
 			SeqQueryLens: []int32{1},
 		}
-	history := c.Update(b, kFinal, vFinal)
+		history = c.Update(b, kFinal, vFinal)
 
 		// Reference: the in-window logical-position-ordered K and V are
 		// the last `window` per-position values (positions
@@ -70,15 +73,15 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) {
 			logicalKs = append(logicalKs, kp)
 			logicalVs = append(logicalVs, vp)
 		}
-	kLogical := mlx.Concatenate(logicalKs, 2)
-	vLogical := mlx.Concatenate(logicalVs, 2)
+		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)
+		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,18 +118,26 @@ 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)
+	var available bool
+	var q *mlx.Array
+	var b *batch.Batch
+	var cases []struct {
+		name string
+		h    *nn.KVHistory
+		mask nn.AttentionMask
+	}
+	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 {
@@ -142,8 +153,8 @@ func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) {
 			sliding.Update(newKVBatch(sliding.Offset(), 1), k, v)
 		}
 
-	b := newKVBatch(total-1, 1)
-	cases := []struct {
+		b = newKVBatch(total-1, 1)
+		cases = []struct {
 			name string
 			h    *nn.KVHistory
 			mask nn.AttentionMask
@@ -151,9 +162,13 @@ func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) {
 			{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,11 +188,13 @@ 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
 
+	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)
@@ -186,10 +203,11 @@ func TestRotatingKVCachePrefillParity(t *testing.T) {
 			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)
+		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,7 +275,7 @@ 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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		const H, D = 1, 4
 		const window = 4
 		const before = 5 // past wrap before the batched write
@@ -329,6 +347,7 @@ func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) {
 				t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i])
 			}
 		}
+	})
 }
 
 // TestRotatingKVCacheMLAParity drives a rotating cache with the MLA
@@ -337,7 +356,7 @@ 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)
+	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
@@ -368,4 +387,5 @@ func TestRotatingKVCacheMLAParity(t *testing.T) {
 				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,8 +76,7 @@ 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)
 
@@ -101,6 +101,7 @@ func TestRotatingKVCacheConcatMidRotationPreservesContext(t *testing.T) {
 		if c.Offset() != 11 {
 			t.Fatalf("offset=%d want 11", c.Offset())
 		}
+	})
 }
 
 // TestRotatingKVCacheConcatAlignedInvariant: with an aligned buffer
@@ -108,8 +109,7 @@ 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)
 
@@ -140,6 +140,7 @@ func TestRotatingKVCacheConcatAlignedInvariant(t *testing.T) {
 				t.Fatalf("post-decode window missing %v (got %v)", w, got)
 			}
 		}
+	})
 }
 
 // TestRotatingKVCacheConcatAfterDecodeGrowsBuffer: update() grows the
@@ -148,8 +149,7 @@ 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)
 
@@ -162,6 +162,7 @@ func TestRotatingKVCacheConcatAfterDecodeGrowsBuffer(t *testing.T) {
 		if !equalSlice(got, want) {
 			t.Fatalf("growing-buffer concat=%v want %v", got, want)
 		}
+	})
 }
 
 // TestRotatingKVCacheConcatAfterLiveRewind: x/mlxrunner/cache.go calls
@@ -171,8 +172,7 @@ 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)
 
@@ -201,13 +201,13 @@ func TestRotatingKVCacheConcatAfterLiveRewind(t *testing.T) {
 		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)
 
@@ -218,6 +218,7 @@ func TestRotatingKVCacheConcatGrowingBuffer(t *testing.T) {
 		if !equalSlice(got, want) {
 			t.Fatalf("growing buffer=%v want %v", got, want)
 		}
+	})
 }
 
 // TestRotatingKVCacheRunnerChunkedPrefill mirrors the
@@ -225,8 +226,7 @@ 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)
 
@@ -264,14 +264,14 @@ func TestRotatingKVCacheRunnerChunkedPrefill(t *testing.T) {
 				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)
 
@@ -306,14 +306,14 @@ func TestRotatingKVCacheMultiTurnChatSimulation(t *testing.T) {
 		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)
-
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := NewRotatingKVCache(4)
 		nextID := feedMulti(c, 1, 3)
 		if c.Offset() != 3 {
@@ -335,4 +335,5 @@ func TestRotatingKVCacheOffsetTracking(t *testing.T) {
 		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,8 +101,7 @@ 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
 
@@ -162,6 +162,7 @@ func TestKVCachePerTokenSnapshotRestore(t *testing.T) {
 				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,8 +170,7 @@ 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)
@@ -209,6 +209,7 @@ func TestKVCaptureMergeSplit(t *testing.T) {
 		}
 		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,8 +387,7 @@ 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
 
 		c := NewRotatingKVCache(window)
@@ -441,6 +441,7 @@ func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) {
 				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,8 +496,7 @@ 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)
 
@@ -535,6 +535,7 @@ func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) {
 		for _, s := range snaps {
 			s.Close()
 		}
+	})
 }
 
 // TestRotatingLazySnapshotSizeZeroUntilMaterialized verifies the speculation
@@ -544,8 +545,7 @@ 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
@@ -603,14 +603,14 @@ func TestRotatingLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
 		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)
 
@@ -654,14 +654,14 @@ func TestPerTokenSnapshotPersistsAcrossWrites(t *testing.T) {
 				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)
@@ -735,13 +735,13 @@ func TestRecurrentSnapshotSplitsAndSegmentedCapture(t *testing.T) {
 		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)
 
@@ -751,4 +751,5 @@ func TestPrepareSnapshotsPastOffsetPanics(t *testing.T) {
 			}
 		}()
 		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,7 +109,7 @@ func draftTokensOf(caches []cache.Cache) []int32 {
 }
 
 func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) {
-	skipIfNoMLX(t)
+	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
@@ -141,10 +142,11 @@ func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) {
 		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)
+	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)
@@ -155,10 +157,11 @@ func TestDFlashCommittedGapPanics(t *testing.T) {
 		}()
 		// 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)
+	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:]}
@@ -185,10 +188,11 @@ func TestDFlashRestoredPrefixResumes(t *testing.T) {
 		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)
+	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)
@@ -214,10 +218,11 @@ func TestDFlashProposeBounds(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		predict := map[int32]int32{2: 3, 3: 4, 4: 5}
 		_, draft, session, caches := newBlockTestSession(t, predict, 4)
 
@@ -245,10 +250,11 @@ func TestDFlashBlockRewoundBeforeContextWrites(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		predict := map[int32]int32{2: 3, 3: 4, 4: 5}
 		_, _, session, caches := newBlockTestSession(t, predict, 4)
 
@@ -262,10 +268,11 @@ func TestDFlashCloseDrainsOutstandingBlock(t *testing.T) {
 		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)
+	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
@@ -340,4 +347,5 @@ func TestDecodeBlockDraft(t *testing.T) {
 		if toks := draftTokensOf(caches); slices.Contains(toks, 6) {
 			t.Fatalf("draft cache retains block rows: %v", toks)
 		}
+	})
 }

x/mlxrunner/grammar_mask_test.go

@@ -4,11 +4,12 @@ import (
 	"math"
 	"testing"
 
+	"github.com/ollama/ollama/x/internal/mlxtest"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
 func TestApplyTokenMask(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		const (
 			bitsPerMaskWord       = 32
 			firstTokenID          = 0                   // Least-significant bit of the first mask word.
@@ -45,4 +46,5 @@ func TestApplyTokenMask(t *testing.T) {
 				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,8 +73,7 @@ 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},
@@ -118,6 +118,7 @@ func TestBatchMediaLifecycle(t *testing.T) {
 		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,12 +47,12 @@ func TestFromValues(t *testing.T) {
 }
 
 func TestComparisonOpsAndBernoulli(t *testing.T) {
-	withMLXThread(t, func() {
-		testComparisonOpsAndBernoulli(t)
-	})
-}
-
-func testComparisonOpsAndBernoulli(t *testing.T) {
+	var tests []struct {
+		name string
+		got  []int32
+		want []int32
+	}
+	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)
@@ -57,22 +61,26 @@ func testComparisonOpsAndBernoulli(t *testing.T) {
 		bern := Bernoulli(FromValues([]float32{1, 0}, 2)).AsType(DTypeInt32)
 		Eval(eq, gt, le, bern)
 
-	for name, tc := range map[string]struct {
+		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)
+			{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 i := range tc.want {
-				if tc.got[i] != tc.want[i] {
-					t.Fatalf("got %v, want %v", tc.got, tc.want)
+	})
+
+	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 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()
-	})
-	for _, m := range mismatches {
-		t.Error(m)
+	withMLXThread(t, func(t *mlxthreadtest.T) {
+		for _, mismatch := range depthwiseConvSiLUMismatches() {
+			t.Error(mismatch)
 		}
+	})
 }
 
 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) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		requireMamba2Metal(t)
 		var failures []error
-	withMLXThread(t, func() {
 		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)
+	})
 }
 
 // Every interior state must match, not just one boundary.
 func TestMamba2ScanCaptureAllMatchesPerTokenReference(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		requireMamba2Metal(t)
 		var failures []error
-	withMLXThread(t, func() {
 		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)
+	})
 }
 
 func TestMamba2ScanGroupedStatesMatchRepeatedReference(t *testing.T) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		requireMamba2Metal(t)
 		var failures []error
-	withMLXThread(t, func() {
 		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)
+	})
 }
 
 // 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) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		var failures []error
-	withMLXThread(t, func() {
 		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)
+	})
 }
 
 // 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) {
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		requireMamba2Metal(t)
 		var failures []error
-	withMLXThread(t, func() {
 		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)
+	})
 }
 
 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)
+	withMLXThread(t, func(t *mlxthreadtest.T) {
 		if !GPUIsAvailable() {
 			t.Skip("MLX GPU not available")
 		}
-
-	var testErr error
-	withMLXThread(t, func() {
-		testErr = checkWiredLimitRejectsOversize()
-	})
-	if testErr != nil {
-		t.Fatal(testErr)
+		if err := checkWiredLimitRejectsOversize(); err != nil {
+			t.Fatal(err)
 		}
+	})
 }
 
 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,20 +3,13 @@ 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,
@@ -36,11 +29,11 @@ func TestMakeEmbeddingLayerDense(t *testing.T) {
 		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 {
@@ -75,6 +68,7 @@ func TestMakeEmbeddingLayerQuantized(t *testing.T) {
 		if _, ok := emb.AsLinear().(*nn.QuantizedLinear); !ok {
 			t.Fatalf("AsLinear type = %T, want *nn.QuantizedLinear", emb.AsLinear())
 		}
+	})
 }
 
 func TestMakeEmbeddingLayerQuantizedGlobalScale(t *testing.T) {

x/mlxrunner/mtp_test.go

@@ -20,15 +20,6 @@ import (
 	"github.com/ollama/ollama/x/tokenizer"
 )
 
-// skipIfNoMLX skips when MLX is unavailable and pins the test to its OS
-// thread. The pin is load-bearing: MLX caches its default stream per thread,
-// so a goroutine that migrates mid-run panics with "There is no Stream(gpu, 0)
-// in current thread".
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	mlxtest.Setup(t)
-}
-
 // The MTP fakes make hidden state and logits the same tensor (Forward returns
 // one-hot logits, Unembed is the identity), so tests fully script target and
 // draft predictions.
@@ -191,10 +182,15 @@ func (d *fakeKVDraft) Unembed(x *mlx.Array) *mlx.Array { return x }
 
 var _ base.DraftModel = (*fakeKVDraft)(nil)
 
+type fatalTester interface {
+	Helper()
+	Fatalf(string, ...any)
+}
+
 // newTestTokenizer builds a byte-level BPE tokenizer over single-character
 // tokens "0".."7" with the given EOS ids, so Decode(id) yields that digit and
 // IsEOS reports membership.
-func newTestTokenizer(t *testing.T, eos []int32) *tokenizer.Tokenizer {
+func newTestTokenizer(t fatalTester, eos []int32) *tokenizer.Tokenizer {
 	t.Helper()
 	vocab := make(map[string]int32, mtpTestVocab)
 	for i := range mtpTestVocab {
@@ -222,7 +218,7 @@ func newTestTokenizer(t *testing.T, eos []int32) *tokenizer.Tokenizer {
 
 // mtpTestRunner wires a Runner with the MTP fakes and a real sampler
 // registered with opts.
-func mtpTestRunner(t *testing.T, predict map[int32]int32, eos []int32, opts sampler.Options) *Runner {
+func mtpTestRunner(t *mlxtest.T, predict map[int32]int32, eos []int32, opts sampler.Options) *Runner {
 	t.Helper()
 	tok := newTestTokenizer(t, eos)
 	r := &Runner{
@@ -263,7 +259,7 @@ func resultIDs(results []sampler.Result) []int32 {
 }
 
 func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Target predicts 1->2->3->4 along the accepted chain; the draft proposed
 		// exactly that, so every draft token is accepted and the bonus token is the
 		// target's prediction after the last accepted token.
@@ -296,10 +292,11 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) {
 		if got := caches[0].Offset(); got != 4 {
 			t.Fatalf("cache offset = %d, want 4 (current + all drafts kept)", got)
 		}
+	})
 }
 
 func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Target predicts 1->2->9 but the draft proposed 2 then 7: the second draft
 		// token mismatches, so only the first is accepted and the bonus is the
 		// target's own prediction (3) at the rejection point.
@@ -333,10 +330,11 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) {
 		if got := caches[0].Offset(); got != 2 {
 			t.Fatalf("cache offset = %d, want 2 (rolled back to accepted)", got)
 		}
+	})
 }
 
 func TestAcceptMTPDraftsGreedyEOS(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The second accepted draft token is EOS: it is recorded but stops
 		// generation and no bonus token is produced. The EOS's own KV is rolled
 		// back so the caches rest one token behind the recorded outputs.
@@ -374,10 +372,11 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) {
 		if got := caches[0].Offset(); got != 2 {
 			t.Fatalf("cache offset = %d, want 2 (one behind the recorded outputs, EOS dropped)", got)
 		}
+	})
 }
 
 func TestRunMTPDecodeGreedy(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The seed token 1 is the last prefill token; its prediction (2) is the
 		// first generated token. The decode then walks 2->3->4->EOS. The draft
 		// proposes the correct chain so steps are accepted in a single forward.
@@ -436,10 +435,11 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
 		if !slices.Equal(draft.calls, wantDraft) {
 			t.Fatalf("draft calls = %v, want %v", draft.calls, wantDraft)
 		}
+	})
 }
 
 func TestRunMTPDecodeSampled(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The same chain at temperature 1: because oneHotLogits uses a large gap,
 		// the proposal and target distributions are effectively point masses, so the
 		// rejection-sampling accept path that the sampled and greedy paths now share
@@ -480,10 +480,11 @@ func TestRunMTPDecodeSampled(t *testing.T) {
 		if got := []int32{2, 3, 4, eos}; !slices.Equal(session.outputs, got) {
 			t.Fatalf("session outputs = %v, want %v", session.outputs, got)
 		}
+	})
 }
 
 func TestRunMTPDecodeWarmDrafter(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// A drafter warmed by a prefill report proposes in the very first round:
 		// the last prompt token seeds the decode loop and is forwarded fused
 		// with the drafts, so generation runs no plain forward at all.
@@ -540,10 +541,11 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 		if !slices.Equal(draft.calls, wantDraft) {
 			t.Fatalf("draft calls = %v, want %v", draft.calls, wantDraft)
 		}
+	})
 }
 
 func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// An accepted EOS inside the draft ends the round at a terminator, not a
 		// target rejection. The persisted acceptance model records outcomes only
 		// up to the EOS: the positions past it lie beyond where the round stopped,
@@ -599,10 +601,11 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
 				t.Fatalf("position %d: seen = %d, rate = %v, want 1 and 1", i, acc.seen[i], acc.rate[i])
 			}
 		}
+	})
 }
 
 func TestDecodePlain(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The same chain with no speculationSession: decode's pipelined loop runs,
 		// dispatching the forward that produces the next token before the
 		// current one is emitted.
@@ -646,10 +649,11 @@ func TestDecodePlain(t *testing.T) {
 		if !slices.Equal(model.forwards, wantForwards) {
 			t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards)
 		}
+	})
 }
 
 func TestDecodeCancelledMidStream(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Cancelling while accepted drafts stream must leave the session
 		// consistent: every token committed to the caches is recorded in
 		// session.outputs, no speculation snapshot schedule is left pending on
@@ -700,10 +704,11 @@ func TestDecodeCancelledMidStream(t *testing.T) {
 				t.Fatalf("snapshot #%d [%d,%d) leaked: never closed", i, s.from, s.to)
 			}
 		}
+	})
 }
 
 func TestLayoutRidesEveryForward(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The request's opaque layout state must reach every forward: parked
 		// pipelined dispatches, the fused verification forward, and the draft
 		// model's own forwards alike.
@@ -748,6 +753,7 @@ func TestLayoutRidesEveryForward(t *testing.T) {
 				t.Fatalf("draft forward %d layout = %v", i, l)
 			}
 		}
+	})
 }
 
 // pinDraftLimit fixes an engine's draft length for the whole run: decode
@@ -761,7 +767,7 @@ func pinDraftLimit(spec *speculationSession, limit int) {
 // testDecoder builds the decoder TextGenerationPipeline would construct for
 // this request, with the draft length pinned to a fixed width; tests close it
 // explicitly so close-time effects are visible to assertions.
-func testDecoder(t *testing.T, r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
+func testDecoder(t *mlxtest.T, r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
 	if spec := r.spec.open(req, nil); spec != nil {
 		if spec.enabled {
 			pinDraftLimit(spec, 4)
@@ -772,7 +778,7 @@ func testDecoder(t *testing.T, r *Runner, req Request, caches []cache.Cache, see
 }
 
 func TestDecodeKVDraft(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// A draft with its own KV cache mirroring the target chain
 		// 1->2->3->4->5->6->EOS.
 		// The unprimed drafter parks the first call, whose pipelined tokens pair
@@ -852,10 +858,11 @@ func TestDecodeKVDraft(t *testing.T) {
 		if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, eos}; !slices.Equal(got, want) {
 			t.Fatalf("draft cache = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The draft mispredicts mid-chain: it proposes 6 where the target's own
 		// next token is 4, so the round is accepted only up to the rejection and the
 		// loop re-proposes from the target's correction. The speculative draft KV
@@ -924,10 +931,11 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 		if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, eos}; !slices.Equal(got, want) {
 			t.Fatalf("draft cache = %v, want %v (rejected proposals rolled back)", got, want)
 		}
+	})
 }
 
 func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// A request that cannot speculate (logprobs) on a model whose draft has
 		// its own KV cache still maintains it: the speculationSession permanently parks,
 		// so the inner pipelined decoder reports each forwarded token, the pairs
@@ -990,10 +998,11 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 		if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, eos, 0}; !slices.Equal(got, want) {
 			t.Fatalf("draft cache = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Prefill attaches its scheduled snapshots only at offsets every cache
 		// has crossed, so the pipeline settles the drafter with the seed first:
 		// the completed frontier pair brings the draft cache level with the
@@ -1027,10 +1036,11 @@ func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
 		if !reflect.DeepEqual(draft.extends, wantExtends) {
 			t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends)
 		}
+	})
 }
 
 func TestFlushMediaHeldUntilEmbedded(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// The deferred flush embeds prompt tokens after prefill has released the
 		// media features, so the session holds delivered feature rows itself:
 		// each flush carries the held rows, a row is dropped once the flush's
@@ -1067,10 +1077,11 @@ func TestFlushMediaHeldUntilEmbedded(t *testing.T) {
 		if want := [][]int{{1}, nil, {1}}; !reflect.DeepEqual(got, want) {
 			t.Fatalf("flush media = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// A committed run longer than the pending-flush cap still writes the draft
 		// caches in a single head forward: the run's completed pairs coalesce into
 		// one batched extend at the run's start, rather than splitting at the cap.
@@ -1111,10 +1122,11 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
 		if got := caches[1].Offset(); got != n-1 {
 			t.Fatalf("draft cache offset = %d, want %d (all completed pairs)", got, n-1)
 		}
+	})
 }
 
 func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// A finished generation levels the draft with the target, its boundary
 		// pair naming the never-committed EOS. The next request restores one
 		// token below the match (the draft look-ahead), so the re-evaluated
@@ -1167,10 +1179,11 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
 		if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, 1}; !slices.Equal(got, want) {
 			t.Fatalf("draft cache = %v, want %v (stale EOS pair rewritten)", got, want)
 		}
+	})
 }
 
 func TestDecodeParkedDraftResume(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Leaving a parked stretch: the inner decoder's in-flight sample is
 		// emitted without any new forward, becomes the round's current, and the
 		// next call drafts from it — with the draft pairs contiguous across the
@@ -1256,6 +1269,7 @@ func TestDecodeParkedDraftResume(t *testing.T) {
 		if got, want := caches[1].Offset(), 6; got != want {
 			t.Fatalf("draft cache offset = %d, want %d (lockstep with target)", got, want)
 		}
+	})
 }
 
 // newMTPTestCaches returns n rewindable fake caches sharing one snapshot

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,8 +94,7 @@ 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)
 
@@ -113,10 +111,11 @@ func TestSampleLogprobsNumericalStability(t *testing.T) {
 				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,8 +210,7 @@ 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}
 
 		maxIdx := int32(0)
@@ -234,14 +232,14 @@ func TestSampleLogprobsSelectedTokenCorrectness(t *testing.T) {
 		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}
 
@@ -272,11 +270,11 @@ func TestBatchedLogprobsPerRow(t *testing.T) {
 		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}
@@ -297,4 +295,5 @@ func TestSampleLogprobsTopKOrdering(t *testing.T) {
 					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,8 +135,7 @@ 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()
@@ -178,11 +169,11 @@ func TestDistributionAppliesTopKBeforeTopP(t *testing.T) {
 		if !foundTop {
 			t.Fatalf("top-k support %v did not include token 0", ids)
 		}
+	})
 }
 
 func TestDistributionResidualUsesTargetSupport(t *testing.T) {
-	skipIfNoMLX(t)
-
+	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),
@@ -210,11 +201,11 @@ func TestDistributionResidualUsesTargetSupport(t *testing.T) {
 				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() {
@@ -243,11 +234,11 @@ func TestSeededSamplingIsReproducible(t *testing.T) {
 		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() {
@@ -266,6 +257,7 @@ func TestSeededBernoulliIsReproducible(t *testing.T) {
 		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,8 +265,7 @@ 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()
@@ -301,11 +292,11 @@ func TestSampleHistoryWindow(t *testing.T) {
 		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()
@@ -334,11 +325,11 @@ func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) {
 		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()
@@ -362,11 +353,11 @@ func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
 				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()
@@ -395,11 +386,11 @@ func TestDistributionMultiRowWithoutChain(t *testing.T) {
 				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()
@@ -421,6 +412,7 @@ func TestCommitBatchesRingWrites(t *testing.T) {
 		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,8 +514,7 @@ 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() {
@@ -550,4 +541,5 @@ func TestRemoveDoesNotLeakHistory(t *testing.T) {
 		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,8 +48,7 @@ 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))
@@ -93,11 +67,11 @@ func TestMoERouterForward(t *testing.T) {
 		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))
@@ -116,11 +90,11 @@ func TestMoEBlockForward(t *testing.T) {
 		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))
@@ -139,6 +113,7 @@ func TestMoEBlockSortedForward(t *testing.T) {
 		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,8 +169,7 @@ 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{}}
 
@@ -221,6 +195,7 @@ func TestLoadFusedExpertsDense(t *testing.T) {
 		if moe.GateUpWeightQ != nil || moe.DownWeightQ != nil {
 			t.Error("quantized weights set on a dense block")
 		}
+	})
 }
 
 // TestRouterForwardMatchesLegacy verifies the optimized Router.Forward —
@@ -229,8 +204,7 @@ 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,
@@ -277,6 +251,7 @@ func TestRouterForwardMatchesLegacy(t *testing.T) {
 		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,7 +26,7 @@ func TestParseSuppressTokens(t *testing.T) {
 }
 
 func TestParseTextConfigE2B(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
@@ -138,10 +139,11 @@ func TestParseTextConfigE2B(t *testing.T) {
 		if !cfg.KVDonors[14] {
 			t.Error("layer 14 should be a KV donor")
 		}
+	})
 }
 
 func TestParseTextConfig26B(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
@@ -217,10 +219,11 @@ func TestParseTextConfig26B(t *testing.T) {
 		if cfg.HiddenSizePerLayer != 0 {
 			t.Errorf("HiddenSizePerLayer = %d, want 0 (no PLE)", cfg.HiddenSizePerLayer)
 		}
+	})
 }
 
 func TestParseTextConfig31B(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
@@ -316,11 +319,11 @@ func TestParseTextConfig31B(t *testing.T) {
 		if isLayerSliding(59, &cfg) {
 			t.Error("layer 59 should be full attention")
 		}
+	})
 }
 
 func TestParseTextConfig12BUnified(t *testing.T) {
-	skipIfNoMLX(t)
-
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		layerTypes := make([]string, 0, 48)
 		for i := range 48 {
 			if i%6 == 5 {
@@ -425,10 +428,11 @@ func TestParseTextConfig12BUnified(t *testing.T) {
 		if isLayerSliding(47, &cfg) {
 			t.Error("layer 47 should be full attention")
 		}
+	})
 }
 
 func TestParseTextConfigE4B(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		data := []byte(`{
 		"architectures": ["Gemma4ForConditionalGeneration"],
 		"text_config": {
@@ -541,6 +545,7 @@ func TestParseTextConfigE4B(t *testing.T) {
 		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,8 +226,7 @@ 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{
@@ -241,6 +240,7 @@ func TestUnembedReturnsFloat32Logits(t *testing.T) {
 		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/laguna/laguna_test.go

@@ -6,13 +6,14 @@ import (
 	"strings"
 	"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"
 )
 
 func TestParseConfigLagunaXS(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
 		"hidden_size": 2048,
@@ -86,10 +87,11 @@ func TestParseConfigLagunaXS(t *testing.T) {
 		if got := numHeadsForLayer(&cfg, 1); got != 64 {
 			t.Fatalf("numHeadsForLayer(1) = %d, want 64", got)
 		}
+	})
 }
 
 func TestParseConfigLagunaFP8RopeScaling(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"hidden_size": 2048,
 		"intermediate_size": 8192,
@@ -115,10 +117,11 @@ func TestParseConfigLagunaFP8RopeScaling(t *testing.T) {
 		if cfg.FullRopeDim != 64 {
 			t.Fatalf("FullRopeDim = %d, want 64", cfg.FullRopeDim)
 		}
+	})
 }
 
 func TestParseConfigLagunaGASchema(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
 		"hidden_size": 2048,
@@ -187,10 +190,11 @@ func TestParseConfigLagunaGASchema(t *testing.T) {
 		if !layerUsesMoE(&cfg, 1) {
 			t.Fatal("layer 1 should use MoE")
 		}
+	})
 }
 
 func TestTinyLagunaLoadAndForward(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
 		"hidden_size": 8,
@@ -274,10 +278,11 @@ func TestTinyLagunaLoadAndForward(t *testing.T) {
 				t.Fatalf("logits[%d] is not finite: %v", i, v)
 			}
 		}
+	})
 }
 
 func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
 		"hidden_size": 8,
@@ -327,10 +332,11 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
 		if got, want := moe.SwitchMLP.GateUpWeight.Dims(), []int{2, 8, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
 			t.Fatalf("GateUpWeight dims = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
 		"hidden_size": 8,
@@ -389,10 +395,11 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
 		if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
 			t.Fatalf("GateWeight dims = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := &Config{
 			HiddenSize:                8,
 			IntermediateSize:          12,
@@ -457,10 +464,11 @@ func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) {
 		if moe.SwitchMLP.DownWeight == nil || moe.SwitchMLP.DownWeightQ != nil || !moe.SwitchMLP.DownWeightSourceLayout {
 			t.Fatal("expected BF16 down expert weights to retain source layout")
 		}
+	})
 }
 
 func TestSparseMoERouteBiasAffectsSelectionNotRoutingWeights(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := &Config{
 			HiddenSize:       1,
 			NumExperts:       2,
@@ -498,10 +506,11 @@ func TestSparseMoERouteBiasAffectsSelectionNotRoutingWeights(t *testing.T) {
 		if got := scores.Floats(); len(got) != 1 || math.Abs(float64(got[0]-probVals[0])) > 1e-6 {
 			t.Fatalf("routing weights = %v, want [%v] using unbiased sigmoid scores", got, probVals[0])
 		}
+	})
 }
 
 func TestLagunaSigmoidTopK8CompiledMatchesEager(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		gates := make([]float32, 2*16)
 		for i := range gates {
 			gates[i] = float32((i%13)-6) * 0.2
@@ -538,10 +547,11 @@ func TestLagunaSigmoidTopK8CompiledMatchesEager(t *testing.T) {
 		if got, want := gotIndices.Ints(), wantIndices.Ints(); !slices.Equal(got, want) {
 			t.Fatalf("indices = %v, want %v", got, want)
 		}
+	})
 }
 
 func TestLagunaSwiGLUGatheredGateScaleCompiledMatchesEager(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		gateValues := make([]float32, 2*8*4)
 		upValues := make([]float32, len(gateValues))
 		for i := range gateValues {
@@ -566,10 +576,11 @@ func TestLagunaSwiGLUGatheredGateScaleCompiledMatchesEager(t *testing.T) {
 		want = want.AsType(mlx.DTypeFloat32)
 		mlx.Eval(got, want)
 		assertFloatSlicesClose(t, got.Floats(), want.Floats(), 1e-6)
+	})
 }
 
 func TestLagunaMoEWeightedSumCompiledMatchesEager(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		expertValues := make([]float32, 1*2*8*4)
 		scoreValues := make([]float32, 1*2*8)
 		for i := range expertValues {
@@ -600,10 +611,11 @@ func TestLagunaMoEWeightedSumCompiledMatchesEager(t *testing.T) {
 		mlx.Eval(gotAdd, gotAdd2, wantAdd, wantAdd2)
 		assertFloatSlicesClose(t, gotAdd.Floats(), wantAdd.Floats(), 1e-6)
 		assertFloatSlicesClose(t, gotAdd2.Floats(), wantAdd2.Floats(), 1e-6)
+	})
 }
 
 func TestSwitchMLPFusedGateUpMatchesSeparate(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := &Config{HiddenSize: 4, NumExpertsPerTok: 2}
 		B, L := int32(2), int32(3)
 		xVals := make([]float32, int(B*L*cfg.HiddenSize))
@@ -637,10 +649,11 @@ func TestSwitchMLPFusedGateUpMatchesSeparate(t *testing.T) {
 		gotSeparateF32 := gotSeparate.AsType(mlx.DTypeFloat32)
 		mlx.Eval(gotFusedF32, gotSeparateF32)
 		assertFloatSlicesClose(t, gotFusedF32.Floats(), gotSeparateF32.Floats(), 1e-5)
+	})
 }
 
 func TestSwitchMLPMixedQuantizedGateUpDenseDownMatchesDense(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := &Config{HiddenSize: 32, NumExpertsPerTok: 2}
 		x := makePatternExpertWeight(1, 2, int(cfg.HiddenSize), 0.013)
 		indices := mlx.FromValues([]int32{0, 1, 1, 0}, 2, int(cfg.NumExpertsPerTok))
@@ -674,10 +687,11 @@ func TestSwitchMLPMixedQuantizedGateUpDenseDownMatchesDense(t *testing.T) {
 		want := dense.Forward(x, indices, cfg).AsType(mlx.DTypeFloat32)
 		mlx.Eval(got, want)
 		assertFloatSlicesClose(t, got.Floats(), want.Floats(), 0.02)
+	})
 }
 
 func TestDenseExpertWeightForGatherMMDequantizesQuantizedWeight(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		weight := makePatternExpertWeight(2, 4, 32, 0.011)
 		qweight, scales, qbiases := mlx.Quantize(weight, 32, 8, "mxfp8")
 		mlx.Eval(qweight, scales)
@@ -701,10 +715,11 @@ func TestDenseExpertWeightForGatherMMDequantizesQuantizedWeight(t *testing.T) {
 		if got.DType() == mlx.DTypeUint32 {
 			t.Fatal("dense expert fallback kept packed U32 weight")
 		}
+	})
 }
 
 func TestCombinedTensorGlobalScaleIgnoresInputGlobalScale(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		tensors := map[string]*mlx.Array{
 			"proj.weight.global_scale":       mlx.FromValues([]float32{0.25}, 1),
 			"proj.weight.input_global_scale": mlx.FromValues([]float32{8}, 1),
@@ -719,6 +734,7 @@ func TestCombinedTensorGlobalScaleIgnoresInputGlobalScale(t *testing.T) {
 		if len(vals) != 1 || vals[0] != 0.25 {
 			t.Fatalf("combinedTensorGlobalScale = %v, want [0.25]", vals)
 		}
+	})
 }
 
 func tinyLagunaTensors() map[string]*mlx.Array {
@@ -770,7 +786,7 @@ func makePatternExpertWeight(numExperts, rows, cols int, scale float32) *mlx.Arr
 	return makeExpertWeight(vals, numExperts, rows, cols)
 }
 
-func assertFloatSlicesClose(t *testing.T, got, want []float32, tol float64) {
+func assertFloatSlicesClose(t *mlxtest.T, got, want []float32, tol float64) {
 	t.Helper()
 	if len(got) != len(want) {
 		t.Fatalf("length mismatch: got %d want %d", len(got), len(want))
@@ -797,10 +813,3 @@ func ones(n int) *mlx.Array {
 	}
 	return mlx.FromValues(vals, n)
 }
-
-func skipIfNoMLX(t *testing.T) {
-	t.Helper()
-	if err := mlx.CheckInit(); err != nil {
-		t.Skipf("MLX not available: %v", err)
-	}
-}

x/models/nemotron_h/nemotron_h_test.go

@@ -323,8 +323,7 @@ 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,
@@ -342,13 +341,13 @@ func TestApplyExpertWeightGlobalScale(t *testing.T) {
 			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
@@ -369,6 +368,7 @@ func TestGatedGroupRMSNormMatchesElementwiseReference(t *testing.T) {
 
 		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,19 +8,13 @@ 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)
@@ -46,12 +40,12 @@ func TestLayerNormNoBias(t *testing.T) {
 				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)
@@ -76,12 +70,12 @@ func TestLayerNormWithBias(t *testing.T) {
 				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,
@@ -116,12 +110,12 @@ func TestLayerNormBatched(t *testing.T) {
 		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)
@@ -142,11 +136,11 @@ func TestLayerNormDefaultEps(t *testing.T) {
 				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)
-
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		weightVals := make([]float32, 3*32)
 		for i := range weightVals {
 			weightVals[i] = float32((i%11)-5) / 7
@@ -183,6 +177,7 @@ func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) {
 				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,7 +37,7 @@ 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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		B, L, D, convTail := 2, 3, 4, 2
 		K := convTail + 1
 
@@ -60,6 +61,7 @@ func TestCausalConv1DBiasTakesFusedPath(t *testing.T) {
 		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,7 +70,7 @@ 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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, D, convTail := 4, 3, 2
 		qLenShort := 2
 		K := convTail + 1
@@ -165,6 +167,7 @@ func TestCausalConv1DPaddedRowParity(t *testing.T) {
 				}
 			}
 		}
+	})
 }
 
 // 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,7 +215,7 @@ 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)
+	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)
@@ -247,12 +250,13 @@ func TestGatedDeltaSegmentEquivalence(t *testing.T) {
 				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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		B, L, D, convTail := 1, 4, 3, 2
 		K := convTail + 1
 
@@ -288,6 +292,7 @@ func TestCausalConv1DSegmentEquivalence(t *testing.T) {
 			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,7 +301,7 @@ 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)
+	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)
@@ -347,6 +352,7 @@ func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
 				}
 			}
 		}
+	})
 }
 
 // TestCausalConv1DSegmentEquivalenceBatched is the conv analog of the gated-delta
@@ -354,7 +360,7 @@ 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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		B, L, D, convTail := 2, 4, 3, 2
 		K := convTail + 1
 
@@ -396,4 +402,5 @@ func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
 				floatsClose(t, "batched boundary conv", gotRow.Floats(), lastState(want).Floats(), 1e-4)
 			}
 		}
+	})
 }

x/models/nn/sdpa_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"
 )
@@ -30,7 +31,7 @@ func newBatch(seqOffsets []int32, L int, qLens []int32) *batch.Batch {
 }
 
 func TestAttentionMaskZero(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		var m AttentionMask
 		if !m.IsZero() {
 			t.Fatal("zero value should report IsZero")
@@ -50,10 +51,11 @@ func TestAttentionMaskZero(t *testing.T) {
 				t.Fatalf("zero mask should materialize all zeros; got[%d] = %v", i, v)
 			}
 		}
+	})
 }
 
 func TestAttentionMaskAsArrayCausal(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 6
 		b := newBatch([]int32{2}, L, nil)
 		arr := CausalMask().AsArray(b, K, mlx.DTypeFloat32)
@@ -81,10 +83,11 @@ func TestAttentionMaskAsArrayCausal(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestAttentionMaskRelaxLazy(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		// Relax must not materialize a tensor — the perf invariant the
 		// causal-flag fast path relies on. Everything else (predicates,
 		// AsArray contents) is exercised by the materialization tests.
@@ -94,6 +97,7 @@ func TestAttentionMaskRelaxLazy(t *testing.T) {
 		if m.array != nil {
 			t.Fatal("Relax should not materialize a tensor")
 		}
+	})
 }
 
 // TestAttentionMaskRelaxNoopRectsMatchCausal pins the contract that
@@ -101,7 +105,7 @@ func TestAttentionMaskRelaxLazy(t *testing.T) {
 // inside the causal triangle — must produce the same materialized
 // tensor as plain causal.
 func TestAttentionMaskRelaxNoopRectsMatchCausal(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 6
 		b := newBatch([]int32{0}, L, nil)
 		want := CausalMask().AsArray(b, K, mlx.DTypeFloat32)
@@ -127,10 +131,11 @@ func TestAttentionMaskRelaxNoopRectsMatchCausal(t *testing.T) {
 				}
 			}
 		}
+	})
 }
 
 func TestAttentionMaskAsArrayWithRelax(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 6
 		b := newBatch([]int32{0}, L, nil)
 		arr := CausalMask().Relax(0, 1, 3, 2, 5).AsArray(b, K, mlx.DTypeFloat32)
@@ -158,10 +163,11 @@ func TestAttentionMaskAsArrayWithRelax(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestAttentionMaskAsArrayPerRow(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 3, 5
 		b := newBatch([]int32{0, 2}, L, nil)
 		m := CausalMask().
@@ -205,10 +211,11 @@ func TestAttentionMaskAsArrayPerRow(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestQPaddingMask(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L := 4
 		// Row 0 fully real; row 1 has 2 real queries.
 		b := newBatch([]int32{0, 0}, L, []int32{int32(L), 2})
@@ -229,10 +236,11 @@ func TestQPaddingMask(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestKPaddingMask(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		K := 5
 		// Row 0 full keys; row 1 has 3 real keys.
 		b := newBatch([]int32{0, 0}, 4, nil)
@@ -253,19 +261,21 @@ func TestKPaddingMask(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestQPaddingMaskZeroWhenFull(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		b := newBatch([]int32{0}, 4, nil)
 		m := QPaddingMask(b, mlx.DTypeFloat32)
 		if !m.IsZero() {
 			t.Fatal("QPaddingMask at full queries should be zero")
 		}
+	})
 }
 
 func TestKPaddingMaskZeroWhenFull(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		K := 4
 		b := newBatch([]int32{0}, 4, nil)
 		kLens := []int32{int32(K)}
@@ -273,10 +283,11 @@ func TestKPaddingMaskZeroWhenFull(t *testing.T) {
 		if !m.IsZero() {
 			t.Fatal("KPaddingMask at full keys should be zero")
 		}
+	})
 }
 
 func TestAttentionMaskCombineCausal(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		var z AttentionMask
 		got := z.Intersect(CausalMask())
 		if !got.IsCausal() {
@@ -290,10 +301,11 @@ func TestAttentionMaskCombineCausal(t *testing.T) {
 		if !got.IsCausal() {
 			t.Fatal("causal + causal should stay pure causal")
 		}
+	})
 }
 
 func TestAttentionMaskCombineRelaxDroppedAgainstCausal(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		relaxed := CausalMask().Relax(0, 1, 3, 2, 5)
 		got := relaxed.Intersect(CausalMask())
 		if !got.IsCausal() {
@@ -310,10 +322,11 @@ func TestAttentionMaskCombineRelaxDroppedAgainstCausal(t *testing.T) {
 		if !got.IsCausal() {
 			t.Fatal("disjoint relaxations on two causals should drop and stay pure causal")
 		}
+	})
 }
 
 func TestAttentionMaskCombineRelaxIntersect(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 6, 6
 		b := newBatch([]int32{0}, L, nil)
 
@@ -350,10 +363,11 @@ func TestAttentionMaskCombineRelaxIntersect(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], vals[i])
 			}
 		}
+	})
 }
 
 func TestAttentionMaskCombineRelaxKeptAgainstNonCausal(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 6
 		b := newBatch([]int32{0}, L, nil)
 
@@ -395,10 +409,11 @@ func TestAttentionMaskCombineRelaxKeptAgainstNonCausal(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], vals[i])
 			}
 		}
+	})
 }
 
 func TestAttentionMaskCombineArrays(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		a := mlx.FromValues([]float32{0, 0, 0, 0}, 1, 1, 2, 2)
 		bb := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 1, 2, 2)
 		sum := ArrayMask(a).Intersect(ArrayMask(bb))
@@ -413,10 +428,11 @@ func TestAttentionMaskCombineArrays(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, want[i], got[i])
 			}
 		}
+	})
 }
 
 func TestAttentionMaskRelaxPanicOnArray(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		a := mlx.FromValues([]float32{0}, 1, 1, 1, 1)
 		defer func() {
 			if r := recover(); r == nil {
@@ -424,10 +440,11 @@ func TestAttentionMaskRelaxPanicOnArray(t *testing.T) {
 			}
 		}()
 		ArrayMask(a).Relax(0, 0, 1, 0, 1)
+	})
 }
 
 func TestAttentionMaskRelaxPanicOnZero(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		defer func() {
 			if r := recover(); r == nil {
 				t.Fatal("Relax on zero mask should panic")
@@ -435,6 +452,7 @@ func TestAttentionMaskRelaxPanicOnZero(t *testing.T) {
 		}()
 		var z AttentionMask
 		z.Relax(0, 0, 1, 0, 1)
+	})
 }
 
 func sameF(a, b float32) bool {
@@ -468,7 +486,7 @@ func sdpaInputs(L, K int) (q, k, v *mlx.Array) {
 }
 
 func TestSDPACausalParity(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 4
 		q, k, v := sdpaInputs(L, K)
 		b := newBatch([]int32{int32(K - L)}, L, nil)
@@ -484,10 +502,11 @@ func TestSDPACausalParity(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i])
 			}
 		}
+	})
 }
 
 func TestSDPAZeroMaskParity(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 4, 4
 		q, k, v := sdpaInputs(L, K)
 		b := newBatch([]int32{0}, L, nil)
@@ -500,10 +519,11 @@ func TestSDPAZeroMaskParity(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i])
 			}
 		}
+	})
 }
 
 func TestSDPAArrayMaskParity(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 3, 3
 		q, k, v := sdpaInputs(L, K)
 		b := newBatch([]int32{0}, L, nil)
@@ -524,10 +544,11 @@ func TestSDPAArrayMaskParity(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i])
 			}
 		}
+	})
 }
 
 func TestSDPARelaxMaskMaterializes(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, K := 3, 5
 		q, k, v := sdpaInputs(L, K)
 		b := newBatch([]int32{int32(K - L)}, L, nil)
@@ -544,10 +565,11 @@ func TestSDPARelaxMaskMaterializes(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i])
 			}
 		}
+	})
 }
 
 func TestSDPAPanicsWithBothKVAndHistory(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L := 3
 		q, k, v := sdpaInputs(L, L)
 		b := newBatch([]int32{0}, L, nil)
@@ -558,10 +580,11 @@ func TestSDPAPanicsWithBothKVAndHistory(t *testing.T) {
 			}
 		}()
 		ScaledDotProductAttention(b, q, 1.0, WithKV(k, v, []int32{int32(L)}), WithKVHistory(history))
+	})
 }
 
 func TestSDPAMLAHistorySlicesVFromK(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		L, D, valueDim := 2, 5, 3
 		kBuf := make([]float32, 1*1*L*D)
 		for i := range kBuf {
@@ -585,10 +608,11 @@ func TestSDPAMLAHistorySlicesVFromK(t *testing.T) {
 				t.Fatalf("index %d: want %v, got %v", i, ws[i], gs[i])
 			}
 		}
+	})
 }
 
 func TestSDPAPanicsWithoutKV(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		q := mlx.FromValues(make([]float32, 4), 1, 1, 1, 4)
 		b := newBatch([]int32{0}, 1, nil)
 		defer func() {
@@ -597,6 +621,7 @@ func TestSDPAPanicsWithoutKV(t *testing.T) {
 			}
 		}()
 		ScaledDotProductAttention(b, q, 1.0)
+	})
 }
 
 // fillTensor builds a [B, H, T, D] float32 tensor whose entries are
@@ -617,7 +642,7 @@ func fillTensor(seed float32, B, H, T, D int) *mlx.Array {
 // central multi-sequence contract: right-padded rows must produce
 // per-row outputs that don't depend on the padded tails.
 func TestSDPAMultiSequenceParity(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		const H, D = 1, 4
 		const L, K = 4, 6
 		const qShort, kShort = 2, 2
@@ -675,4 +700,5 @@ func TestSDPAMultiSequenceParity(t *testing.T) {
 				}
 			}
 		}
+	})
 }

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,7 +58,7 @@ func scatterRows(packed *mlx.Array, perm []int32) *mlx.Array {
 }
 
 func TestPackGatedDeltaProjectionsNativeMatchesSplit(t *testing.T) {
-	skipIfNoMLX(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := gdnTestConfig()
 		keyDim := int(cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim)
 		valueDim := int(cfg.LinearNumValueHeads * cfg.LinearValueHeadDim)
@@ -93,10 +94,11 @@ func TestPackGatedDeltaProjectionsNativeMatchesSplit(t *testing.T) {
 
 		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)
+	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")
@@ -111,10 +113,11 @@ func TestConcatProjectionPairQuantized(t *testing.T) {
 		}
 		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)
+	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)
@@ -135,4 +138,5 @@ func TestConcatProjectionPairMixedFallsBackToDense(t *testing.T) {
 			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,16 +8,8 @@ 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)
-
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		tests := []struct {
 			name  string
 			shape []int
@@ -35,6 +27,7 @@ func TestSanitizeConvWeight(t *testing.T) {
 				t.Fatalf("%s: sanitizeConvWeight() shape = %v, want %v", tt.name, dims, tt.want)
 			}
 		}
+	})
 }
 
 func TestParseConfigNestedDefaults(t *testing.T) {
@@ -209,8 +202,7 @@ 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,
@@ -387,4 +379,5 @@ func TestLoadWeightsPreservesLinearAttentionNormWeightDType(t *testing.T) {
 		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)
+	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,8 +10,7 @@ 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),
@@ -29,10 +28,11 @@ func TestEngramHashes(t *testing.T) {
 		if !slices.Equal(values, want) {
 			t.Fatalf("hashes = %v, want %v", values, want)
 		}
+	})
 }
 
 func TestEngramCacheCarriesChunkHistory(t *testing.T) {
-	mlxtest.Setup(t)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := newEngramCache(2, 3, 1, 9)
 		t.Cleanup(c.Free)
 		b := &batch.Batch{
@@ -51,10 +51,11 @@ func TestEngramCacheCarriesChunkHistory(t *testing.T) {
 		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)
+	mlxtest.Run(t, func(t *mlxtest.T) {
 		c := newEngramCache(2, 3, 1, 9)
 		t.Cleanup(c.Free)
 		b := &batch.Batch{
@@ -102,4 +103,5 @@ func TestEngramCacheRestoresScheduledSnapshot(t *testing.T) {
 		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,8 +10,7 @@ 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}}
@@ -71,6 +70,7 @@ func TestHyperConnectionMatchesReferenceFormula(t *testing.T) {
 
 		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,7 +12,7 @@ import (
 )
 
 func TestQSASelectsCompressedBlocksAndCausalTail(t *testing.T) {
-	mlxtest.Setup(t)
+	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}}
@@ -34,10 +34,11 @@ func TestQSASelectsCompressedBlocksAndCausalTail(t *testing.T) {
 		if !slices.Equal(selected, want) {
 			t.Fatalf("selected indices = %v, want %v", selected, want)
 		}
+	})
 }
 
 func TestQSASelectionMasksFutureBlocks(t *testing.T) {
-	mlxtest.Setup(t)
+	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.
@@ -59,10 +60,11 @@ func TestQSASelectionMasksFutureBlocks(t *testing.T) {
 		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)
+	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)
@@ -81,10 +83,11 @@ func TestQSASparseAttentionMatchesReference(t *testing.T) {
 				t.Fatalf("sparse attention[%d] = %v, want %v", i, got[i], want[i])
 			}
 		}
+	})
 }
 
 func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) {
-	mlxtest.Setup(t)
+	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.
@@ -99,10 +102,11 @@ func TestQSASparseAttentionIgnoresInvalidRows(t *testing.T) {
 		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)
+	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{
@@ -122,4 +126,5 @@ func TestQSASparseAttentionKeepsBatchRowsIndependent(t *testing.T) {
 		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)