diff --git a/.github/scripts/prepare_mlx_darwin.sh b/.github/scripts/prepare_mlx_darwin.sh new file mode 100755 index 000000000..6b2d083f1 --- /dev/null +++ b/.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}" diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 8a0fec6c1..733d39e64 100644 --- a/.github/workflows/test.yaml +++ b/.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 ./... diff --git a/app/server/server_test.go b/app/server/server_test.go index cb9674229..fb29fc169 100644 --- a/app/server/server_test.go +++ b/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 { diff --git a/app/ui/ui_test.go b/app/ui/ui_test.go index a63bd608f..95cc88a5a 100644 --- a/app/ui/ui_test.go +++ b/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, diff --git a/app/updater/updater.go b/app/updater/updater.go index e042807fe..53051c0e8 100644 --- a/app/updater/updater.go +++ b/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 } diff --git a/app/updater/updater_test.go b/app/updater/updater_test.go index 7a27599cd..f31ea20f1 100644 --- a/app/updater/updater_test.go +++ b/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 { diff --git a/x/internal/mlxtest/mlxtest.go b/x/internal/mlxtest/mlxtest.go index ef006abb2..d19fa1892 100644 --- a/x/internal/mlxtest/mlxtest.go +++ b/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) }) } diff --git a/x/internal/mlxthreadtest/mlxthreadtest.go b/x/internal/mlxthreadtest/mlxthreadtest.go new file mode 100644 index 000000000..5e43ea384 --- /dev/null +++ b/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) + } +} diff --git a/x/internal/mlxthreadtest/mlxthreadtest_test.go b/x/internal/mlxthreadtest/mlxthreadtest_test.go new file mode 100644 index 000000000..4009a02c1 --- /dev/null +++ b/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) + } +} diff --git a/x/internal/mlxthreadtest/threadid_darwin.go b/x/internal/mlxthreadtest/threadid_darwin.go new file mode 100644 index 000000000..48d3eb685 --- /dev/null +++ b/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) +} diff --git a/x/internal/mlxthreadtest/threadid_linux.go b/x/internal/mlxthreadtest/threadid_linux.go new file mode 100644 index 000000000..5c81aa74c --- /dev/null +++ b/x/internal/mlxthreadtest/threadid_linux.go @@ -0,0 +1,9 @@ +//go:build linux + +package mlxthreadtest + +import "syscall" + +func currentThreadID() uint64 { + return uint64(syscall.Gettid()) +} diff --git a/x/internal/mlxthreadtest/threadid_other.go b/x/internal/mlxthreadtest/threadid_other.go new file mode 100644 index 000000000..4f87279c3 --- /dev/null +++ b/x/internal/mlxthreadtest/threadid_other.go @@ -0,0 +1,7 @@ +//go:build !darwin && !linux && !windows + +package mlxthreadtest + +func currentThreadID() uint64 { + return 0 +} diff --git a/x/internal/mlxthreadtest/threadid_windows.go b/x/internal/mlxthreadtest/threadid_windows.go new file mode 100644 index 000000000..69454f1a1 --- /dev/null +++ b/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()) +} diff --git a/x/mlxrunner/cache/cache_test.go b/x/mlxrunner/cache/cache_test.go index 22ecbdbfa..408aa7da6 100644 --- a/x/mlxrunner/cache/cache_test.go +++ b/x/mlxrunner/cache/cache_test.go @@ -3,17 +3,11 @@ package cache import ( "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" ) -func skipIfNoMLX(t *testing.T) { - t.Helper() - if err := mlx.CheckInit(); err != nil { - t.Skipf("MLX not available: %v", err) - } -} - // newKVBatch builds a B=1 batch at SeqOffsets=off with all-real // queries (SeqQueryLens=L) — the standard single-sequence cache // test shape. @@ -26,258 +20,265 @@ func newKVBatch(off, L int) *batch.Batch { } func TestKVCacheSnapshotRestoreNeedBase(t *testing.T) { - skipIfNoMLX(t) - c := NewKVCache() + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - // Snapshot [5, 10). - snap := c.Snapshot(5) + // Snapshot [5, 10). + snap := c.Snapshot(5) - // Free the cache completely — offset is now 0. - c.Free() + // Free the cache completely — offset is now 0. + c.Free() - // Restore should fail because cache doesn't have data up to fromOffset=5. - if c.Restore(snap, 10) { - t.Fatal("expected Restore to fail with no base data") - } + // Restore should fail because cache doesn't have data up to fromOffset=5. + if c.Restore(snap, 10) { + t.Fatal("expected Restore to fail with no base data") + } + }) } // TestKVCacheDataSurvivesSnapshotRestore verifies that actual array data // is preserved through a snapshot→free→restore cycle. func TestKVCacheDataSurvivesSnapshotRestore(t *testing.T) { - skipIfNoMLX(t) - c := NewKVCache() + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - snap := c.Snapshot(0) - if snap == nil { - t.Fatal("Snapshot returned nil") - } + snap := c.Snapshot(0) + if snap == nil { + t.Fatal("Snapshot returned nil") + } - // Free and restore to a fresh cache. - c2 := NewKVCache() - if !c2.Restore(snap, 10) { - t.Fatal("Restore failed") - } - if c2.Offset() != 10 { - t.Fatalf("offset = %d, want 10", c2.Offset()) - } + // Free and restore to a fresh cache. + c2 := NewKVCache() + if !c2.Restore(snap, 10) { + t.Fatal("Restore failed") + } + if c2.Offset() != 10 { + t.Fatalf("offset = %d, want 10", c2.Offset()) + } - // Verify State() returns arrays with correct sequence dimension. - state := c2.State() - if len(state) != 2 { - t.Fatalf("State() returned %d arrays, want 2", len(state)) - } - // keys shape: [B, H, seqLen, Dk] - if state[0].Dim(2) != 10 { - t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2)) - } - if state[1].Dim(2) != 10 { - t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2)) - } + // Verify State() returns arrays with correct sequence dimension. + state := c2.State() + if len(state) != 2 { + t.Fatalf("State() returned %d arrays, want 2", len(state)) + } + // keys shape: [B, H, seqLen, Dk] + if state[0].Dim(2) != 10 { + t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2)) + } + if state[1].Dim(2) != 10 { + t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2)) + } + }) } // TestKVCacheSplitPreservesData verifies that split produces two snapshots // that can be sequentially restored to rebuild the original cache state. func TestKVCacheSplitPreservesData(t *testing.T) { - skipIfNoMLX(t) - c := NewKVCache() + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - snap := c.Snapshot(0) - parent, child := c.Split(snap, 5) - if parent == nil || child == nil { - t.Fatal("Split returned nil") - } + snap := c.Snapshot(0) + parent, child := c.Split(snap, 5) + if parent == nil || child == nil { + t.Fatal("Split returned nil") + } - // Restore parent → offset=5, seq dim=5. - c2 := NewKVCache() - if !c2.Restore(parent, 5) { - t.Fatal("Restore(parent) failed") - } - if c2.Offset() != 5 { - t.Fatalf("offset after parent = %d, want 5", c2.Offset()) - } - state := c2.State() - if state[0].Dim(2) != 5 { - t.Fatalf("keys seq dim after parent = %d, want 5", state[0].Dim(2)) - } + // Restore parent → offset=5, seq dim=5. + c2 := NewKVCache() + if !c2.Restore(parent, 5) { + t.Fatal("Restore(parent) failed") + } + if c2.Offset() != 5 { + t.Fatalf("offset after parent = %d, want 5", c2.Offset()) + } + state := c2.State() + if state[0].Dim(2) != 5 { + t.Fatalf("keys seq dim after parent = %d, want 5", state[0].Dim(2)) + } - // Restore child on top → offset=10, seq dim=10. - if !c2.Restore(child, 10) { - t.Fatal("Restore(child) failed") - } - if c2.Offset() != 10 { - t.Fatalf("offset after child = %d, want 10", c2.Offset()) - } - state = c2.State() - if state[0].Dim(2) != 10 { - t.Fatalf("keys seq dim after child = %d, want 10", state[0].Dim(2)) - } + // Restore child on top → offset=10, seq dim=10. + if !c2.Restore(child, 10) { + t.Fatal("Restore(child) failed") + } + if c2.Offset() != 10 { + t.Fatalf("offset after child = %d, want 10", c2.Offset()) + } + state = c2.State() + if state[0].Dim(2) != 10 { + t.Fatalf("keys seq dim after child = %d, want 10", state[0].Dim(2)) + } + }) } // TestKVCacheSplitMergeRoundTripData verifies that splitting and merging back // produces a snapshot equivalent to the original. func TestKVCacheSplitMergeRoundTripData(t *testing.T) { - skipIfNoMLX(t) - c := NewKVCache() + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - snap := c.Snapshot(0) - parent, child := c.Split(snap, 6) - merged := c.Merge(parent, child) - if merged == nil { - t.Fatal("Merge returned nil") - } + snap := c.Snapshot(0) + parent, child := c.Split(snap, 6) + merged := c.Merge(parent, child) + if merged == nil { + t.Fatal("Merge returned nil") + } - c2 := NewKVCache() - if !c2.Restore(merged, 10) { - t.Fatal("Restore(merged) failed") - } - if c2.Offset() != 10 { - t.Fatalf("offset = %d, want 10", c2.Offset()) - } + c2 := NewKVCache() + if !c2.Restore(merged, 10) { + t.Fatal("Restore(merged) failed") + } + if c2.Offset() != 10 { + t.Fatalf("offset = %d, want 10", c2.Offset()) + } - state := c2.State() - if state[0].Dim(2) != 10 { - t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2)) - } - if state[1].Dim(2) != 10 { - t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2)) - } + state := c2.State() + if state[0].Dim(2) != 10 { + t.Fatalf("keys seq dim = %d, want 10", state[0].Dim(2)) + } + if state[1].Dim(2) != 10 { + t.Fatalf("values seq dim = %d, want 10", state[1].Dim(2)) + } + }) } func TestRotatingKVCacheRestoreOutsideWindow(t *testing.T) { - skipIfNoMLX(t) - c := NewRotatingKVCache(4) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRotatingKVCache(4) - // Feed 10 tokens (window size 4, so positions 0-5 are evicted). - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + // Feed 10 tokens (window size 4, so positions 0-5 are evicted). + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - // Offset 3 is outside the window. - if c.Restore(nil, 3) { - t.Fatal("Restore(nil, 3) should fail when outside window") - } + // Offset 3 is outside the window. + if c.Restore(nil, 3) { + t.Fatal("Restore(nil, 3) should fail when outside window") + } + }) } // TestRotatingKVCacheSnapshotPreservesWindow verifies that after restoring // from a snapshot, the rotating cache has the correct window of data. func TestRotatingKVCacheSnapshotPreservesWindow(t *testing.T) { - skipIfNoMLX(t) - c := NewRotatingKVCache(4) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRotatingKVCache(4) - // Feed 10 tokens one at a time. Window size 4, so only last 4 are kept. - for range 10 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + // Feed 10 tokens one at a time. Window size 4, so only last 4 are kept. + for range 10 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - snap := c.Snapshot(0) - if snap == nil { - t.Fatal("Snapshot returned nil") - } + snap := c.Snapshot(0) + if snap == nil { + t.Fatal("Snapshot returned nil") + } - // Feed 5 more tokens. - for range 5 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + // Feed 5 more tokens. + for range 5 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - // Restore to offset 10. - if !c.Restore(snap, 10) { - t.Fatal("Restore failed") - } - if c.Offset() != 10 { - t.Fatalf("offset = %d, want 10", c.Offset()) - } + // Restore to offset 10. + if !c.Restore(snap, 10) { + t.Fatal("Restore failed") + } + if c.Offset() != 10 { + t.Fatalf("offset = %d, want 10", c.Offset()) + } - state := c.State() - if len(state) != 2 { - t.Fatalf("State() returned %d arrays, want 2", len(state)) - } - // Seq dim should be min(offset, maxSize) = min(10, 4) = 4. - seqDim := state[0].Dim(2) - if seqDim != 4 { - t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim) - } + state := c.State() + if len(state) != 2 { + t.Fatalf("State() returned %d arrays, want 2", len(state)) + } + // Seq dim should be min(offset, maxSize) = min(10, 4) = 4. + seqDim := state[0].Dim(2) + if seqDim != 4 { + t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim) + } + }) } // TestRotatingKVCacheRestoreFromSnapshot verifies that restoring from a // snapshot correctly preserves the write position (idx), so subsequent // single-token updates land in the right buffer slot. func TestRotatingKVCacheRestoreFromSnapshot(t *testing.T) { - skipIfNoMLX(t) - c := NewRotatingKVCache(4) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRotatingKVCache(4) - // Fill the window: 6 tokens into a size-4 window. - // After this, idx has wrapped and the buffer has rotated. - for range 6 { + // Fill the window: 6 tokens into a size-4 window. + // After this, idx has wrapped and the buffer has rotated. + for range 6 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } + if c.Offset() != 6 { + t.Fatalf("offset = %d, want 6", c.Offset()) + } + + snap := c.Snapshot(0) + + // Mutate the cache further so live state diverges from snapshot. + for range 3 { + k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } + + // Restore to snapshot state. + if !c.Restore(snap, 6) { + t.Fatal("Restore failed") + } + if c.Offset() != 6 { + t.Fatalf("offset after restore = %d, want 6", c.Offset()) + } + + // Feed one more token. If idx was restored correctly, this should + // produce a valid window of size 4 at offset 7. k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } - if c.Offset() != 6 { - t.Fatalf("offset = %d, want 6", c.Offset()) - } - snap := c.Snapshot(0) - - // Mutate the cache further so live state diverges from snapshot. - for range 3 { - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } - - // Restore to snapshot state. - if !c.Restore(snap, 6) { - t.Fatal("Restore failed") - } - if c.Offset() != 6 { - t.Fatalf("offset after restore = %d, want 6", c.Offset()) - } - - // Feed one more token. If idx was restored correctly, this should - // produce a valid window of size 4 at offset 7. - k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - - if c.Offset() != 7 { - t.Fatalf("offset after post-restore update = %d, want 7", c.Offset()) - } - state := c.State() - if len(state) != 2 { - t.Fatalf("State() returned %d arrays, want 2", len(state)) - } - seqDim := state[0].Dim(2) - if seqDim != 4 { - t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim) - } + if c.Offset() != 7 { + t.Fatalf("offset after post-restore update = %d, want 7", c.Offset()) + } + state := c.State() + if len(state) != 2 { + t.Fatalf("State() returned %d arrays, want 2", len(state)) + } + seqDim := state[0].Dim(2) + if seqDim != 4 { + t.Fatalf("keys seq dim = %d, want 4 (window size)", seqDim) + } + }) } diff --git a/x/mlxrunner/cache/lazy_test.go b/x/mlxrunner/cache/lazy_test.go index 22147e44f..4739e7e42 100644 --- a/x/mlxrunner/cache/lazy_test.go +++ b/x/mlxrunner/cache/lazy_test.go @@ -3,6 +3,7 @@ package cache import ( "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/mlx" ) @@ -44,53 +45,53 @@ func settledActiveMemory() int { // snapshot stays lazy and is discarded before any overwrite. Compare against the // bytes an eager per-token copy would cost. func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const before, draft, H, D = 16, 8, 4, 8 - const before, draft, H, D = 16, 8, 4, 8 + c := NewKVCache() + fillKV(c, before) - c := NewKVCache() - fillKV(c, before) - - offsets := make([]int, draft) - for i := range offsets { - offsets[i] = before + i - } - c.PrepareSnapshots(offsets) - - k, v := batchKV(draft) - c.Update(newKVBatch(before, draft), k, v) - - baseline := settledActiveMemory() - - snaps := c.TakeSnapshots() - // Every captured snapshot is a lazy snapshot (no owned buffer). - for i, s := range snaps { - if s == nil { - continue + offsets := make([]int, draft) + for i := range offsets { + offsets[i] = before + i } - if ks := s.(*kvSnapshot); ks.keys != nil { - t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i) - } - } + c.PrepareSnapshots(offsets) - // MTP commit: rewind to a partial accept, then discard all snapshots. - if !c.Restore(nil, before+draft/2) { - t.Fatal("live rewind failed") - } - for _, s := range snaps { - if s != nil { - s.Close() - } - } + k, v := batchKV(draft) + c.Update(newKVBatch(before, draft), k, v) - after := settledActiveMemory() - // Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but - // well under one per-token copy (draft tokens * keys+values). - perToken := (c.keys.NumBytes() + c.values.NumBytes()) / c.keys.Dim(2) - if after-baseline > perToken { - t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing", - after-baseline, perToken) - } + baseline := settledActiveMemory() + + snaps := c.TakeSnapshots() + // Every captured snapshot is a lazy snapshot (no owned buffer). + for i, s := range snaps { + if s == nil { + continue + } + if ks := s.(*kvSnapshot); ks.keys != nil { + t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i) + } + } + + // MTP commit: rewind to a partial accept, then discard all snapshots. + if !c.Restore(nil, before+draft/2) { + t.Fatal("live rewind failed") + } + for _, s := range snaps { + if s != nil { + s.Close() + } + } + + after := settledActiveMemory() + // Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but + // well under one per-token copy (draft tokens * keys+values). + perToken := (c.keys.NumBytes() + c.values.NumBytes()) / c.keys.Dim(2) + if after-baseline > perToken { + t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing", + after-baseline, perToken) + } + }) } // TestKVLazySnapshotSizeZeroUntilMaterialized verifies the accounting contract: @@ -98,40 +99,40 @@ func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) { // destructive write triggers copyOut the materialize hook fires with the // newly-allocated bytes and Size() reports the owned arrays. func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) + snap := c.Snapshot(5).(*kvSnapshot) + defer snap.Close() + if snap.Size() != 0 { + t.Fatalf("lazy snapshot Size = %d, want 0", snap.Size()) + } - snap := c.Snapshot(5).(*kvSnapshot) - defer snap.Close() - if snap.Size() != 0 { - t.Fatalf("lazy snapshot Size = %d, want 0", snap.Size()) - } + var hookDelta int + snap.SetMaterializeHook(func(delta int) { hookDelta = delta }) - var hookDelta int - snap.SetMaterializeHook(func(delta int) { hookDelta = delta }) + // Rewind and overwrite to force copyOut. + if !c.Restore(nil, 5) { + t.Fatal("rewind failed") + } + nk, nv := distinctKV(100, 3, H, D) + c.Update(newKVBatch(5, 3), nk, nv) - // Rewind and overwrite to force copyOut. - if !c.Restore(nil, 5) { - t.Fatal("rewind failed") - } - nk, nv := distinctKV(100, 3, H, D) - c.Update(newKVBatch(5, 3), nk, nv) - - if snap.keys == nil { - t.Fatal("snapshot was not materialized by overwriting write") - } - want := snap.keys.NumBytes() + snap.values.NumBytes() - if hookDelta != want { - t.Fatalf("hook fired with delta %d, want %d (owned bytes)", hookDelta, want) - } - if snap.Size() != want { - t.Fatalf("materialized Size = %d, want %d", snap.Size(), want) - } + if snap.keys == nil { + t.Fatal("snapshot was not materialized by overwriting write") + } + want := snap.keys.NumBytes() + snap.values.NumBytes() + if hookDelta != want { + t.Fatalf("hook fired with delta %d, want %d (owned bytes)", hookDelta, want) + } + if snap.Size() != want { + t.Fatalf("materialized Size = %d, want %d", snap.Size(), want) + } + }) } // TestKVLazySnapshotCopiedOutOnOverwrite verifies that after a rewind, a write @@ -139,110 +140,110 @@ func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) { // still reads the pre-overwrite data — both keys and values, across the whole // captured range (guarding the Slice+Contiguous copy-out representation). func TestKVLazySnapshotCopiedOutOnOverwrite(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + // Fill [0,10) with position-encoded values (keys p, values -p). + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - // Fill [0,10) with position-encoded values (keys p, values -p). - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) - - // Lazy snapshot [5,10). - snap := c.Snapshot(5).(*kvSnapshot) - if snap.keys != nil { - t.Fatal("Snapshot should return a lazy snapshot") - } - - // Rewind to 5 and overwrite [5,8) with different data. - if !c.Restore(nil, 5) { - t.Fatal("rewind failed") - } - nk, nv := distinctKV(100, 3, H, D) // positions 100,101,102 - c.Update(newKVBatch(5, 3), nk, nv) - - // The overwrite must have copied the lazy snapshot out beforehand, preserving - // the pre-overwrite keys and values across the whole [5,10) range. - if snap.keys == nil { - t.Fatal("lazy snapshot was not copied out before the overwriting write") - } - mlx.Eval(snap.keys, snap.values) - keys := snap.keys.Floats() - vals := snap.values.Floats() - for l := range snap.toOffset - snap.fromOffset { - wantK := float32(snap.fromOffset + l) - if got := keys[l*D]; got != wantK { - t.Fatalf("snapshot keys[%d] = %v, want %v (pre-overwrite data)", l, got, wantK) + // Lazy snapshot [5,10). + snap := c.Snapshot(5).(*kvSnapshot) + if snap.keys != nil { + t.Fatal("Snapshot should return a lazy snapshot") } - if got := vals[l*D]; got != -wantK { - t.Fatalf("snapshot values[%d] = %v, want %v (pre-overwrite data)", l, got, -wantK) + + // Rewind to 5 and overwrite [5,8) with different data. + if !c.Restore(nil, 5) { + t.Fatal("rewind failed") } - } - snap.Close() + nk, nv := distinctKV(100, 3, H, D) // positions 100,101,102 + c.Update(newKVBatch(5, 3), nk, nv) + + // The overwrite must have copied the lazy snapshot out beforehand, preserving + // the pre-overwrite keys and values across the whole [5,10) range. + if snap.keys == nil { + t.Fatal("lazy snapshot was not copied out before the overwriting write") + } + mlx.Eval(snap.keys, snap.values) + keys := snap.keys.Floats() + vals := snap.values.Floats() + for l := range snap.toOffset - snap.fromOffset { + wantK := float32(snap.fromOffset + l) + if got := keys[l*D]; got != wantK { + t.Fatalf("snapshot keys[%d] = %v, want %v (pre-overwrite data)", l, got, wantK) + } + if got := vals[l*D]; got != -wantK { + t.Fatalf("snapshot values[%d] = %v, want %v (pre-overwrite data)", l, got, -wantK) + } + } + snap.Close() + }) } // TestKVLazySnapshotCopiedOutOnFree verifies Free copies out outstanding lazy snapshots // so they survive after the cache buffer is gone. func TestKVLazySnapshotCopiedOutOnFree(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) + snap := c.Snapshot(5).(*kvSnapshot) + c.Free() - snap := c.Snapshot(5).(*kvSnapshot) - c.Free() - - if snap.keys == nil { - t.Fatal("lazy snapshot was not copied out on Free") - } - mlx.Eval(snap.keys) - if got := firstKeyAt(snap.keys, 0, D); got != 5 { - t.Fatalf("snapshot[0] key = %v, want 5 (data preserved through Free)", got) - } - snap.Close() + if snap.keys == nil { + t.Fatal("lazy snapshot was not copied out on Free") + } + mlx.Eval(snap.keys) + if got := firstKeyAt(snap.keys, 0, D); got != 5 { + t.Fatalf("snapshot[0] key = %v, want 5 (data preserved through Free)", got) + } + snap.Close() + }) } // TestKVLazySnapshotSplitMergeNoCopy verifies Split of a lazy snapshot and Merge of two // adjacent lazy snapshots are pure arithmetic — they produce lazy snapshots and allocate // nothing — while still tracking the correct offsets and data. func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) + base := settledActiveMemory() - base := settledActiveMemory() + // Lazy snapshot [2,10), split at 5. + snap := c.Snapshot(2) + p, ch := c.Split(snap, 5) + ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot) + if ps.keys != nil || cs.keys != nil { + t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)") + } + if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 { + t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset) + } - // Lazy snapshot [2,10), split at 5. - snap := c.Snapshot(2) - p, ch := c.Split(snap, 5) - ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot) - if ps.keys != nil || cs.keys != nil { - t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)") - } - if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 { - t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset) - } + // Merge them back into [2,10). + merged := c.Merge(p, ch).(*kvSnapshot) + if merged.keys != nil { + t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)") + } + if merged.fromOffset != 2 || merged.toOffset != 10 { + t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset) + } - // Merge them back into [2,10). - merged := c.Merge(p, ch).(*kvSnapshot) - if merged.keys != nil { - t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)") - } - if merged.fromOffset != 2 || merged.toOffset != 10 { - t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset) - } - - if after := settledActiveMemory(); after > base { - t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base) - } - merged.Close() + if after := settledActiveMemory(); after > base { + t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base) + } + merged.Close() + }) } // TestKVLazySnapshotSurvivesPathSwitch reproduces the switchToPath sequence that @@ -251,41 +252,41 @@ func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) { // different path (Restore feeds new data via appendKV, overwriting the old // leaf's slots). The paged-out snapshot must still hold the original tokens. func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + // Active path tokens [0,10): positions 0..9. + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - // Active path tokens [0,10): positions 0..9. - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) + // Page out the diverging leaf [4,10) as a lazy snapshot (what switchToPath does + // before rewinding). + leaf := c.Snapshot(4).(*kvSnapshot) - // Page out the diverging leaf [4,10) as a lazy snapshot (what switchToPath does - // before rewinding). - leaf := c.Snapshot(4).(*kvSnapshot) - - // Rewind to the common ancestor at 4 (Restore(nil): offset move only). - if !c.Restore(nil, 4) { - t.Fatal("rewind to ancestor failed") - } - - // Page in the new path [4,9): positions 200..204, overwriting the old leaf's - // slots. appendKV must copy the leaf lazy snapshot out first. - nk, nv := distinctKV(200, 5, H, D) - c.Update(newKVBatch(4, 5), nk, nv) - - if leaf.keys == nil { - t.Fatal("leaf lazy snapshot was not copied out during page-in") - } - mlx.Eval(leaf.keys) - // leaf covers [4,10): its keys are the original positions 4..9. - keys := leaf.keys.Floats() - for l := range leaf.toOffset - leaf.fromOffset { - if got, want := keys[l*D], float32(leaf.fromOffset+l); got != want { - t.Fatalf("paged-out leaf key[%d] = %v, want %v (original path data)", l, got, want) + // Rewind to the common ancestor at 4 (Restore(nil): offset move only). + if !c.Restore(nil, 4) { + t.Fatal("rewind to ancestor failed") } - } - leaf.Close() + + // Page in the new path [4,9): positions 200..204, overwriting the old leaf's + // slots. appendKV must copy the leaf lazy snapshot out first. + nk, nv := distinctKV(200, 5, H, D) + c.Update(newKVBatch(4, 5), nk, nv) + + if leaf.keys == nil { + t.Fatal("leaf lazy snapshot was not copied out during page-in") + } + mlx.Eval(leaf.keys) + // leaf covers [4,10): its keys are the original positions 4..9. + keys := leaf.keys.Floats() + for l := range leaf.toOffset - leaf.fromOffset { + if got, want := keys[l*D], float32(leaf.fromOffset+l); got != want { + t.Fatalf("paged-out leaf key[%d] = %v, want %v (original path data)", l, got, want) + } + } + leaf.Close() + }) } // TestKVRestoreLiveLazySnapshotIsOffsetMove verifies the same-path rewind/rematch @@ -294,46 +295,46 @@ func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) { // the offset without cloning or replaying. This is the switchToPath sequence // where a paged-out leaf is restored before any write displaced it. func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 4, 8 - const H, D = 4, 8 + c := NewKVCache() + k, v := distinctKV(0, 10, H, D) + c.Update(newKVBatch(0, 10), k, v) - c := NewKVCache() - k, v := distinctKV(0, 10, H, D) - c.Update(newKVBatch(0, 10), k, v) - - // Page out the leaf [5,10) as a lazy snapshot, then rewind (offset move only). - snap := c.Snapshot(5).(*kvSnapshot) - if !c.Restore(nil, 5) { - t.Fatal("rewind failed") - } - - base := settledActiveMemory() - - // Restore the snapshot back to 10. Its slots [5,10) were never overwritten, - // so it is still lazy and the data is already in the buffer — a pure offset - // move, no allocation. - if !c.Restore(snap, 10) { - t.Fatal("restore failed") - } - if snap.keys != nil { - t.Fatal("snapshot was copied out; expected the offset-move fast path") - } - if c.Offset() != 10 { - t.Fatalf("offset after restore = %d, want 10", c.Offset()) - } - if after := settledActiveMemory(); after > base { - t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base) - } - - // The buffer still holds the original positions 0..9. - st := c.State() - mlx.Eval(st[0]) - keys := st[0].Floats() - for l := range 10 { - if got := keys[l*D]; got != float32(l) { - t.Fatalf("restored key[%d] = %v, want %v", l, got, float32(l)) + // Page out the leaf [5,10) as a lazy snapshot, then rewind (offset move only). + snap := c.Snapshot(5).(*kvSnapshot) + if !c.Restore(nil, 5) { + t.Fatal("rewind failed") } - } - snap.Close() + + base := settledActiveMemory() + + // Restore the snapshot back to 10. Its slots [5,10) were never overwritten, + // so it is still lazy and the data is already in the buffer — a pure offset + // move, no allocation. + if !c.Restore(snap, 10) { + t.Fatal("restore failed") + } + if snap.keys != nil { + t.Fatal("snapshot was copied out; expected the offset-move fast path") + } + if c.Offset() != 10 { + t.Fatalf("offset after restore = %d, want 10", c.Offset()) + } + if after := settledActiveMemory(); after > base { + t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base) + } + + // The buffer still holds the original positions 0..9. + st := c.State() + mlx.Eval(st[0]) + keys := st[0].Floats() + for l := range 10 { + if got := keys[l*D]; got != float32(l) { + t.Fatalf("restored key[%d] = %v, want %v", l, got, float32(l)) + } + } + snap.Close() + }) } diff --git a/x/mlxrunner/cache/recurrent_test.go b/x/mlxrunner/cache/recurrent_test.go index ce1e12c4f..4a4624d51 100644 --- a/x/mlxrunner/cache/recurrent_test.go +++ b/x/mlxrunner/cache/recurrent_test.go @@ -4,6 +4,7 @@ import ( "math" "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/models/nn" @@ -13,66 +14,68 @@ import ( // only succeeds when target exactly matches the snapshot's offset. Recurrent // state is cumulative, so it can't be rewound or fast-forwarded. func TestRecurrentCacheRestoreExactOffset(t *testing.T) { - skipIfNoMLX(t) - c := NewRecurrentCache(3, 12, 4, 8, 8) - b1 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1)} - c.Get(b1, mlx.DTypeFloat16) // lazy-init + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRecurrentCache(3, 12, 4, 8, 8) + b1 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1)} + c.Get(b1, mlx.DTypeFloat16) // lazy-init - keep := func() ([]*mlx.Array, []*mlx.Array) { - s := c.State() - return []*mlx.Array{s[0]}, []*mlx.Array{s[1]} - } + keep := func() ([]*mlx.Array, []*mlx.Array) { + s := c.State() + return []*mlx.Array{s[0]}, []*mlx.Array{s[1]} + } - b10 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 10), SeqQueryLens: []int32{10}} - cs, ds := keep() - c.Put(b10, cs, ds) // advance to 10 + b10 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 10), SeqQueryLens: []int32{10}} + cs, ds := keep() + c.Put(b10, cs, ds) // advance to 10 - snap := c.Snapshot(0) // snap.offset == 10 + snap := c.Snapshot(0) // snap.offset == 10 - b5 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 5), SeqQueryLens: []int32{5}} - cs, ds = keep() - c.Put(b5, cs, ds) // cache now at 15 + b5 := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 5), SeqQueryLens: []int32{5}} + cs, ds = keep() + c.Put(b5, cs, ds) // cache now at 15 - // target < snap.offset: fails (can't rewind past snapshot) - if c.Restore(snap, 5) { - t.Fatal("Restore(snap, 5) should fail — target != snap.offset") - } + // target < snap.offset: fails (can't rewind past snapshot) + if c.Restore(snap, 5) { + t.Fatal("Restore(snap, 5) should fail — target != snap.offset") + } - // target > snap.offset: fails (can't advance without feeding tokens) - if c.Restore(snap, 15) { - t.Fatal("Restore(snap, 15) should fail — target != snap.offset") - } + // target > snap.offset: fails (can't advance without feeding tokens) + if c.Restore(snap, 15) { + t.Fatal("Restore(snap, 15) should fail — target != snap.offset") + } - // target == snap.offset: succeeds - if !c.Restore(snap, 10) { - t.Fatal("Restore(snap, 10) should succeed — target == snap.offset") - } - if c.Offset() != 10 { - t.Fatalf("offset = %d, want 10", c.Offset()) - } + // target == snap.offset: succeeds + if !c.Restore(snap, 10) { + t.Fatal("Restore(snap, 10) should succeed — target == snap.offset") + } + if c.Offset() != 10 { + t.Fatalf("offset = %d, want 10", c.Offset()) + } + }) } func TestRecurrentCacheGetLazyInit(t *testing.T) { - skipIfNoMLX(t) - c := NewRecurrentCache(3, 4, 2, 4, 4) - b := &batch.Batch{ - InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), - SeqOffsets: []int32{0}, - SeqQueryLens: []int32{1}, - } - h := c.Get(b, mlx.DTypeBFloat16) - if c.Offset() != 0 { - t.Fatalf("Get should not advance; got offset %d", c.Offset()) - } - if h.ConvState() == nil || h.DeltaState() == nil { - t.Fatal("history should expose conv/delta tensors") - } - if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 { - t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16) - } - if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 { - t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32) - } + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRecurrentCache(3, 4, 2, 4, 4) + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{1}, + } + h := c.Get(b, mlx.DTypeBFloat16) + if c.Offset() != 0 { + t.Fatalf("Get should not advance; got offset %d", c.Offset()) + } + if h.ConvState() == nil || h.DeltaState() == nil { + t.Fatal("history should expose conv/delta tensors") + } + if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16) + } + if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 { + t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32) + } + }) } // TestRecurrentCachePaddedRoundTrip runs Get → CausalConv1D → @@ -83,138 +86,140 @@ func TestRecurrentCacheGetLazyInit(t *testing.T) { // Pins the recurrent contract: a forward with padding produces the // same end-state as a forward with the real-prefix-only input. func TestRecurrentCachePaddedRoundTrip(t *testing.T) { - skipIfNoMLX(t) - const convTail, convDim = 2, 6 - const numVHeads, headVDim, headKDim = 1, 4, 6 - const L = 4 - const qLen = 2 + mlxtest.Run(t, func(t *mlxtest.T) { + const convTail, convDim = 2, 6 + const numVHeads, headVDim, headKDim = 1, 4, 6 + const L = 4 + const qLen = 2 - // Distinct values for the real prefix and large junk in the padded - // tail so any leak from padded positions is visible. - const packedDim = 2*headKDim + numVHeads*headVDim - mkPacked := func(seed float32, T int) (packed, ba *mlx.Array) { - pv := make([]float32, T*packedDim) - bv := make([]float32, T*2*numVHeads) - for i := range pv { - pv[i] = seed + 0.05*float32(i) + // Distinct values for the real prefix and large junk in the padded + // tail so any leak from padded positions is visible. + const packedDim = 2*headKDim + numVHeads*headVDim + mkPacked := func(seed float32, T int) (packed, ba *mlx.Array) { + pv := make([]float32, T*packedDim) + bv := make([]float32, T*2*numVHeads) + for i := range pv { + pv[i] = seed + 0.05*float32(i) + } + for i := range bv { + bv[i] = seed - 0.02*float32(i) + } + return mlx.FromValues(pv, 1, T, packedDim), mlx.FromValues(bv, 1, T, 2*numVHeads) } - for i := range bv { - bv[i] = seed - 0.02*float32(i) + mkPackedPadded := func() (packed, ba *mlx.Array) { + pReal, baReal := mkPacked(0.3, qLen) + pPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, packedDim), 99) + baPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 2*numVHeads), 99) + return mlx.Concatenate([]*mlx.Array{pReal, pPad}, 1), mlx.Concatenate([]*mlx.Array{baReal, baPad}, 1) } - return mlx.FromValues(pv, 1, T, packedDim), mlx.FromValues(bv, 1, T, 2*numVHeads) - } - mkPackedPadded := func() (packed, ba *mlx.Array) { - pReal, baReal := mkPacked(0.3, qLen) - pPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, packedDim), 99) - baPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 2*numVHeads), 99) - return mlx.Concatenate([]*mlx.Array{pReal, pPad}, 1), mlx.Concatenate([]*mlx.Array{baReal, baPad}, 1) - } - dtBias := mlx.FromValues([]float32{0.3}, numVHeads) - aExp := mlx.FromValues([]float32{0.12}, numVHeads) + dtBias := mlx.FromValues([]float32{0.3}, numVHeads) + aExp := mlx.FromValues([]float32{0.12}, numVHeads) - // The conv input dimension must match the cache's convDim. - mkConvInput := func(seed float32, T int) *mlx.Array { - vals := make([]float32, 1*T*convDim) - for i := range vals { - vals[i] = seed + 0.05*float32(i) + // The conv input dimension must match the cache's convDim. + mkConvInput := func(seed float32, T int) *mlx.Array { + vals := make([]float32, 1*T*convDim) + for i := range vals { + vals[i] = seed + 0.05*float32(i) + } + return mlx.FromValues(vals, 1, T, convDim) } - return mlx.FromValues(vals, 1, T, convDim) - } - mkWeight := func(seed float32) *mlx.Array { - vals := make([]float32, convDim*(convTail+1)) - for i := range vals { - vals[i] = seed + 0.1*float32(i) + mkWeight := func(seed float32) *mlx.Array { + vals := make([]float32, convDim*(convTail+1)) + for i := range vals { + vals[i] = seed + 0.1*float32(i) + } + return mlx.FromValues(vals, convDim, convTail+1) } - return mlx.FromValues(vals, convDim, convTail+1) - } - weight := mkWeight(0.2) - // Build the depthwise causal Conv1d as the model does at load time: the - // [C, K] kernel becomes [C, K, 1] and the conv is grouped per channel. - conv := nn.NewConv1d(mlx.ExpandDims(weight, 2), nil, 1, 0, 1, convDim) + weight := mkWeight(0.2) + // Build the depthwise causal Conv1d as the model does at load time: the + // [C, K] kernel becomes [C, K, 1] and the conv is grouped per channel. + conv := nn.NewConv1d(mlx.ExpandDims(weight, 2), nil, 1, 0, 1, convDim) - runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) { - var convInput *mlx.Array - if T == L { - realPart := mkConvInput(0.4, qLen) - padPart := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, T-qLen, convDim), 99) - convInput = mlx.Concatenate([]*mlx.Array{realPart, padPart}, 1) - } else { - convInput = mkConvInput(0.4, T) + runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) { + var convInput *mlx.Array + if T == L { + realPart := mkConvInput(0.4, qLen) + padPart := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, T-qLen, convDim), 99) + convInput = mlx.Concatenate([]*mlx.Array{realPart, padPart}, 1) + } else { + convInput = mkConvInput(0.4, T) + } + + history := c.Get(b, mlx.DTypeFloat32) + _, convStates := nn.CausalConv1D(b, convInput, conv, convTail, + nn.WithRecurrentHistory(history)) + + var packed, ba *mlx.Array + if T == L { + packed, ba = mkPackedPadded() + } else { + packed, ba = mkPacked(0.3, T) + } + _, deltaStates := nn.GatedDelta(b, packed, ba, dtBias, aExp, nn.WithRecurrentHistory(history)) + + c.Put(b, convStates, deltaStates) + return convStates[len(convStates)-1], deltaStates[len(deltaStates)-1] } - history := c.Get(b, mlx.DTypeFloat32) - _, convStates := nn.CausalConv1D(b, convInput, conv, convTail, - nn.WithRecurrentHistory(history)) - - var packed, ba *mlx.Array - if T == L { - packed, ba = mkPackedPadded() - } else { - packed, ba = mkPacked(0.3, T) + // Padded forward. + cPad := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim) + bPad := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, L), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{int32(qLen)}, } - _, deltaStates := nn.GatedDelta(b, packed, ba, dtBias, aExp, nn.WithRecurrentHistory(history)) - - c.Put(b, convStates, deltaStates) - return convStates[len(convStates)-1], deltaStates[len(deltaStates)-1] - } - - // Padded forward. - cPad := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim) - bPad := &batch.Batch{ - InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, L), - SeqOffsets: []int32{0}, - SeqQueryLens: []int32{int32(qLen)}, - } - nextConvPad, deltaPad := runForward(cPad, bPad, L) - mlx.Eval(nextConvPad, deltaPad) - if got := cPad.Offset(); got != qLen { - t.Fatalf("padded forward: Offset() = %d, want %d (must advance by SeqQueryLens, not L)", got, qLen) - } - - // Unpadded reference. - cRef := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim) - bRef := &batch.Batch{ - InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, qLen), - SeqOffsets: []int32{0}, - SeqQueryLens: []int32{int32(qLen)}, - } - nextConvRef, deltaRef := runForward(cRef, bRef, qLen) - mlx.Eval(nextConvRef, deltaRef) - if got := cRef.Offset(); got != qLen { - t.Fatalf("unpadded forward: Offset() = %d, want %d", got, qLen) - } - - gp := nextConvPad.Floats() - gr := nextConvRef.Floats() - if len(gp) != len(gr) { - t.Fatalf("nextConv shape mismatch: padded %d vs unpadded %d", len(gp), len(gr)) - } - for i := range gp { - if math.Abs(float64(gp[i]-gr[i])) > 1e-4 { - t.Fatalf("nextConv[%d]: padded=%v unpadded=%v (padding leaked into conv state)", i, gp[i], gr[i]) + nextConvPad, deltaPad := runForward(cPad, bPad, L) + mlx.Eval(nextConvPad, deltaPad) + if got := cPad.Offset(); got != qLen { + t.Fatalf("padded forward: Offset() = %d, want %d (must advance by SeqQueryLens, not L)", got, qLen) } - } - dp := deltaPad.Floats() - dr := deltaRef.Floats() - if len(dp) != len(dr) { - t.Fatalf("delta state shape mismatch: padded %d vs unpadded %d", len(dp), len(dr)) - } - for i := range dp { - if math.Abs(float64(dp[i]-dr[i])) > 1e-3 { - t.Fatalf("delta state[%d]: padded=%v unpadded=%v (padding leaked into recurrent state)", i, dp[i], dr[i]) + // Unpadded reference. + cRef := NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim) + bRef := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, qLen), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{int32(qLen)}, } - } + nextConvRef, deltaRef := runForward(cRef, bRef, qLen) + mlx.Eval(nextConvRef, deltaRef) + if got := cRef.Offset(); got != qLen { + t.Fatalf("unpadded forward: Offset() = %d, want %d", got, qLen) + } + + gp := nextConvPad.Floats() + gr := nextConvRef.Floats() + if len(gp) != len(gr) { + t.Fatalf("nextConv shape mismatch: padded %d vs unpadded %d", len(gp), len(gr)) + } + for i := range gp { + if math.Abs(float64(gp[i]-gr[i])) > 1e-4 { + t.Fatalf("nextConv[%d]: padded=%v unpadded=%v (padding leaked into conv state)", i, gp[i], gr[i]) + } + } + + dp := deltaPad.Floats() + dr := deltaRef.Floats() + if len(dp) != len(dr) { + t.Fatalf("delta state shape mismatch: padded %d vs unpadded %d", len(dp), len(dr)) + } + for i := range dp { + if math.Abs(float64(dp[i]-dr[i])) > 1e-3 { + t.Fatalf("delta state[%d]: padded=%v unpadded=%v (padding leaked into recurrent state)", i, dp[i], dr[i]) + } + } + }) } func TestRecurrentCachePutAdvances(t *testing.T) { - skipIfNoMLX(t) - c := NewRecurrentCache(3, 4, 2, 4, 4) - b := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 2), SeqQueryLens: []int32{2}} - newConv := mlx.Zeros(mlx.DTypeFloat16, 1, 3, 4) - newDelta := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 4, 4) - c.Put(b, []*mlx.Array{newConv}, []*mlx.Array{newDelta}) - if c.Offset() != 2 { - t.Fatalf("cache offset not advanced: %d", c.Offset()) - } + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRecurrentCache(3, 4, 2, 4, 4) + b := &batch.Batch{InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 2), SeqQueryLens: []int32{2}} + newConv := mlx.Zeros(mlx.DTypeFloat16, 1, 3, 4) + newDelta := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 4, 4) + c.Put(b, []*mlx.Array{newConv}, []*mlx.Array{newDelta}) + if c.Offset() != 2 { + t.Fatalf("cache offset not advanced: %d", c.Offset()) + } + }) } diff --git a/x/mlxrunner/cache/rotating_attention_test.go b/x/mlxrunner/cache/rotating_attention_test.go index a7fbdae74..6d2d6d318 100644 --- a/x/mlxrunner/cache/rotating_attention_test.go +++ b/x/mlxrunner/cache/rotating_attention_test.go @@ -4,6 +4,7 @@ import ( "math" "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/models/nn" @@ -19,7 +20,6 @@ import ( // applier's gather composing the caller's logical mask back into // storage order) must equal the logical-order reference. func TestRotatingKVCacheDecodeParity(t *testing.T) { - skipIfNoMLX(t) const H, D = 1, 4 const window = 4 const totalWrites = 7 // past wrap (window=4); last write is the L=1 decode @@ -39,46 +39,49 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) { return } - q := mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D) - mlx.Eval(q) + var q, kLogical, vLogical, logicalMask *mlx.Array + var b *batch.Batch + var history *nn.KVHistory + mlxtest.Run(t, func(t *mlxtest.T) { + q = mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D) - // Drive the cache: write positions 0..totalWrites-2 as a "history", - // then position totalWrites-1 is the actual L=1 decode under test. - c := NewRotatingKVCache(window) - for pos := range totalWrites - 1 { - k, v := perPosKV(pos) - c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) - } + // Drive the cache: write positions 0..totalWrites-2 as a "history", + // then position totalWrites-1 is the actual L=1 decode under test. + c := NewRotatingKVCache(window) + for pos := range totalWrites - 1 { + k, v := perPosKV(pos) + c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v) + } - finalPos := totalWrites - 1 - kFinal, vFinal := perPosKV(finalPos) - b := &batch.Batch{ - InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), - SeqOffsets: []int32{int32(finalPos)}, - SeqQueryLens: []int32{1}, - } - history := c.Update(b, kFinal, vFinal) + finalPos := totalWrites - 1 + kFinal, vFinal := perPosKV(finalPos) + b = &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{int32(finalPos)}, + SeqQueryLens: []int32{1}, + } + history = c.Update(b, kFinal, vFinal) - // Reference: the in-window logical-position-ordered K and V are - // the last `window` per-position values (positions - // [finalPos-window+1, finalPos]). Build them in that order. - startPos := max(finalPos-window+1, 0) - logicalKs := make([]*mlx.Array, 0, window) - logicalVs := make([]*mlx.Array, 0, window) - for pos := startPos; pos <= finalPos; pos++ { - kp, vp := perPosKV(pos) - logicalKs = append(logicalKs, kp) - logicalVs = append(logicalVs, vp) - } - kLogical := mlx.Concatenate(logicalKs, 2) - vLogical := mlx.Concatenate(logicalVs, 2) + // Reference: the in-window logical-position-ordered K and V are + // the last `window` per-position values (positions + // [finalPos-window+1, finalPos]). Build them in that order. + startPos := max(finalPos-window+1, 0) + logicalKs := make([]*mlx.Array, 0, window) + logicalVs := make([]*mlx.Array, 0, window) + for pos := startPos; pos <= finalPos; pos++ { + kp, vp := perPosKV(pos) + logicalKs = append(logicalKs, kp) + logicalVs = append(logicalVs, vp) + } + kLogical = mlx.Concatenate(logicalKs, 2) + vLogical = mlx.Concatenate(logicalVs, 2) - // A logical-order ArrayMask with distinct, non-trivial values per - // key column. Picked so each column's contribution to softmax is - // distinct — the test fails if the cache's gather permutes the - // columns wrong before the kernel sees them. - maskVals := []float32{0.1, -0.3, 0.7, -0.2} - logicalMask := mlx.FromValues(maskVals, 1, 1, 1, window) + // A logical-order ArrayMask with distinct, non-trivial values per + // key column. Picked so each column's contribution to softmax is + // distinct — the test fails if the cache's gather permutes the + // columns wrong before the kernel sees them. + logicalMask = mlx.FromValues([]float32{0.1, -0.3, 0.7, -0.2}, 1, 1, 1, window) + }) cases := []struct { name string @@ -95,7 +98,7 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) { } for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { got := nn.ScaledDotProductAttention(b, q, scale, nn.WithKVHistory(history), nn.WithMask(tc.model)) @@ -115,45 +118,57 @@ func TestRotatingKVCacheDecodeParity(t *testing.T) { } func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) { - skipIfNoMLX(t) - if !mlx.MetalIsAvailable() { - t.Skip("MLX Metal not available") - } const H, D = 1, 4 const window = 4 const total = 7 const scale = 1.0 - q := mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D) - mlx.Eval(q) - - full := NewKVCache() - sliding := NewRotatingKVCache(window) - for pos := range total { - kVals := make([]float32, H*D) - vVals := make([]float32, H*D) - for i := range kVals { - kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i) - vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i) - } - k := mlx.FromValues(kVals, 1, H, 1, D) - v := mlx.FromValues(vVals, 1, H, 1, D) - full.Update(newKVBatch(full.Offset(), 1), k, v) - sliding.Update(newKVBatch(sliding.Offset(), 1), k, v) - } - - b := newKVBatch(total-1, 1) - cases := []struct { + var available bool + var q *mlx.Array + var b *batch.Batch + var cases []struct { name string h *nn.KVHistory mask nn.AttentionMask - }{ - {name: "full", h: full.View(b), mask: nn.CausalMask()}, - {name: "sliding", h: sliding.View(b), mask: nn.CausalMask()}, + } + mlxtest.Run(t, func(t *mlxtest.T) { + available = mlx.MetalIsAvailable() + if !available { + return + } + + q = mlx.FromValues([]float32{0.7, -0.4, 0.2, 0.9}, 1, H, 1, D) + full := NewKVCache() + sliding := NewRotatingKVCache(window) + for pos := range total { + kVals := make([]float32, H*D) + vVals := make([]float32, H*D) + for i := range kVals { + kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i) + vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i) + } + k := mlx.FromValues(kVals, 1, H, 1, D) + v := mlx.FromValues(vVals, 1, H, 1, D) + full.Update(newKVBatch(full.Offset(), 1), k, v) + sliding.Update(newKVBatch(sliding.Offset(), 1), k, v) + } + + b = newKVBatch(total-1, 1) + cases = []struct { + name string + h *nn.KVHistory + mask nn.AttentionMask + }{ + {name: "full", h: full.View(b), mask: nn.CausalMask()}, + {name: "sliding", h: sliding.View(b), mask: nn.CausalMask()}, + } + }) + if !available { + t.Skip("MLX Metal not available") } for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { got := nn.ScaledDotProductAttention(b, q, scale, nn.WithKVHistory(tc.h), nn.WithMask(tc.mask)) want := mlx.FastScaledDotProductAttention(q, tc.h.K(), tc.h.V(), scale, "", nil) @@ -173,23 +188,26 @@ func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) { // matches a reference computed from the same K/V with the model mask // and window restriction composed manually. func TestRotatingKVCachePrefillParity(t *testing.T) { - skipIfNoMLX(t) const H, L, D = 1, 6, 4 const window = 4 const scale = 1.0 - qVals := make([]float32, 1*H*L*D) - kVals := make([]float32, 1*H*L*D) - vVals := make([]float32, 1*H*L*D) - for i := range qVals { - qVals[i] = 0.5 + 0.05*float32(i) - kVals[i] = -0.3 + 0.07*float32(i) - vVals[i] = 0.3 + 0.03*float32(i) - } - q := mlx.FromValues(qVals, 1, H, L, D) - k := mlx.FromValues(kVals, 1, H, L, D) - v := mlx.FromValues(vVals, 1, H, L, D) - b := newKVBatch(0, L) + var q, k, v *mlx.Array + var b *batch.Batch + mlxtest.Run(t, func(t *mlxtest.T) { + qVals := make([]float32, 1*H*L*D) + kVals := make([]float32, 1*H*L*D) + vVals := make([]float32, 1*H*L*D) + for i := range qVals { + qVals[i] = 0.5 + 0.05*float32(i) + kVals[i] = -0.3 + 0.07*float32(i) + vVals[i] = 0.3 + 0.03*float32(i) + } + q = mlx.FromValues(qVals, 1, H, L, D) + k = mlx.FromValues(kVals, 1, H, L, D) + v = mlx.FromValues(vVals, 1, H, L, D) + b = newKVBatch(0, L) + }) cases := []struct { name string @@ -205,7 +223,7 @@ func TestRotatingKVCachePrefillParity(t *testing.T) { negInf := float32(math.Inf(-1)) for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { c := NewRotatingKVCache(window) history := c.Update(b, k, v) @@ -257,78 +275,79 @@ func TestRotatingKVCachePrefillParity(t *testing.T) { // an identical write with no snapshots scheduled. Capture happens after the // write via lazy snapshots, so it must not perturb the write itself. func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) { - skipIfNoMLX(t) - const H, D = 1, 4 - const window = 4 - const before = 5 // past wrap before the batched write - const draft = 4 - const scale = 1.0 + mlxtest.Run(t, func(t *mlxtest.T) { + const H, D = 1, 4 + const window = 4 + const before = 5 // past wrap before the batched write + const draft = 4 + const scale = 1.0 - perPosKV := func(pos int) (k, v *mlx.Array) { - kVals := make([]float32, H*D) - vVals := make([]float32, H*D) - for i := range kVals { - kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i) - vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i) - } - return mlx.FromValues(kVals, 1, H, 1, D), mlx.FromValues(vVals, 1, H, 1, D) - } - - // Build the batched K/V for offsets [before, before+draft). - batchK := make([]*mlx.Array, draft) - batchV := make([]*mlx.Array, draft) - for i := range draft { - batchK[i], batchV[i] = perPosKV(before + i) - } - kBatch := mlx.Concatenate(batchK, 2) - vBatch := mlx.Concatenate(batchV, 2) - - qVals := make([]float32, H*draft*D) - for i := range qVals { - qVals[i] = 0.5 + 0.05*float32(i) - } - q := mlx.FromValues(qVals, 1, H, draft, D) - b := newKVBatch(before, draft) - - // Run the same write twice: once with snapshots scheduled, once without. - run := func(schedule bool) []float32 { - c := NewRotatingKVCache(window) - for pos := range before { - k, v := perPosKV(pos) - c.Update(newKVBatch(c.Offset(), 1), k, v) - } - if schedule { - offsets := make([]int, draft) - for i := range offsets { - offsets[i] = before + i + perPosKV := func(pos int) (k, v *mlx.Array) { + kVals := make([]float32, H*D) + vVals := make([]float32, H*D) + for i := range kVals { + kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i) + vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i) } - c.PrepareSnapshots(offsets) + return mlx.FromValues(kVals, 1, H, 1, D), mlx.FromValues(vVals, 1, H, 1, D) } - history := c.Update(b, kBatch, vBatch) - out := nn.ScaledDotProductAttention(b, q, scale, - nn.WithKVHistory(history), - nn.WithMask(nn.CausalMask())) - if schedule { - for _, s := range c.TakeSnapshots() { - if s != nil { - s.Close() + + // Build the batched K/V for offsets [before, before+draft). + batchK := make([]*mlx.Array, draft) + batchV := make([]*mlx.Array, draft) + for i := range draft { + batchK[i], batchV[i] = perPosKV(before + i) + } + kBatch := mlx.Concatenate(batchK, 2) + vBatch := mlx.Concatenate(batchV, 2) + + qVals := make([]float32, H*draft*D) + for i := range qVals { + qVals[i] = 0.5 + 0.05*float32(i) + } + q := mlx.FromValues(qVals, 1, H, draft, D) + b := newKVBatch(before, draft) + + // Run the same write twice: once with snapshots scheduled, once without. + run := func(schedule bool) []float32 { + c := NewRotatingKVCache(window) + for pos := range before { + k, v := perPosKV(pos) + c.Update(newKVBatch(c.Offset(), 1), k, v) + } + if schedule { + offsets := make([]int, draft) + for i := range offsets { + offsets[i] = before + i + } + c.PrepareSnapshots(offsets) + } + history := c.Update(b, kBatch, vBatch) + out := nn.ScaledDotProductAttention(b, q, scale, + nn.WithKVHistory(history), + nn.WithMask(nn.CausalMask())) + if schedule { + for _, s := range c.TakeSnapshots() { + if s != nil { + s.Close() + } } } + mlx.Eval(out) + return out.Floats() } - mlx.Eval(out) - return out.Floats() - } - withSnap := run(true) - noSnap := run(false) - if len(withSnap) != len(noSnap) { - t.Fatalf("output length %d vs %d", len(withSnap), len(noSnap)) - } - for i := range noSnap { - if math.Abs(float64(withSnap[i]-noSnap[i])) > 1e-5 { - t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i]) + withSnap := run(true) + noSnap := run(false) + if len(withSnap) != len(noSnap) { + t.Fatalf("output length %d vs %d", len(withSnap), len(noSnap)) } - } + for i := range noSnap { + if math.Abs(float64(withSnap[i]-noSnap[i])) > 1e-5 { + t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i]) + } + } + }) } // TestRotatingKVCacheMLAParity drives a rotating cache with the MLA @@ -337,35 +356,36 @@ func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) { // a manual reference. Pins the cache+MLA integration that // glm4_moe_lite uses in production. func TestRotatingKVCacheMLAParity(t *testing.T) { - skipIfNoMLX(t) - const H, L, D, valueDim = 1, 3, 6, 4 - const scale = 1.0 - const window = 8 // window >= L so no window restriction + mlxtest.Run(t, func(t *mlxtest.T) { + const H, L, D, valueDim = 1, 3, 6, 4 + const scale = 1.0 + const window = 8 // window >= L so no window restriction - kVals := make([]float32, 1*H*L*D) - for i := range kVals { - kVals[i] = 0.1 * float32(i+1) - } - k := mlx.FromValues(kVals, 1, H, L, D) - v := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, 0) - - q := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, D) - b := newKVBatch(0, L) - - c := NewRotatingKVCache(window) - history := c.Update(b, k, v) - got := nn.ScaledDotProductAttention(b, q, scale, - nn.WithMLAHistory(history, valueDim), - nn.WithMask(nn.CausalMask())) - - vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim)) - want := mlx.FastScaledDotProductAttention(q, k, vRef, scale, "causal", nil) - - mlx.Eval(got, want) - gs, ws := got.Floats(), want.Floats() - for i := range ws { - if math.Abs(float64(gs[i]-ws[i])) > 1e-5 { - t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i]) + kVals := make([]float32, 1*H*L*D) + for i := range kVals { + kVals[i] = 0.1 * float32(i+1) } - } + k := mlx.FromValues(kVals, 1, H, L, D) + v := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, 0) + + q := mlx.Zeros(mlx.DTypeFloat32, 1, H, L, D) + b := newKVBatch(0, L) + + c := NewRotatingKVCache(window) + history := c.Update(b, k, v) + got := nn.ScaledDotProductAttention(b, q, scale, + nn.WithMLAHistory(history, valueDim), + nn.WithMask(nn.CausalMask())) + + vRef := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, valueDim)) + want := mlx.FastScaledDotProductAttention(q, k, vRef, scale, "causal", nil) + + mlx.Eval(got, want) + gs, ws := got.Floats(), want.Floats() + for i := range ws { + if math.Abs(float64(gs[i]-ws[i])) > 1e-5 { + t.Fatalf("index %d: got %v, want %v", i, gs[i], ws[i]) + } + } + }) } diff --git a/x/mlxrunner/cache/rotating_multiturn_test.go b/x/mlxrunner/cache/rotating_multiturn_test.go index 03991a1b9..2b6ce30c0 100644 --- a/x/mlxrunner/cache/rotating_multiturn_test.go +++ b/x/mlxrunner/cache/rotating_multiturn_test.go @@ -3,6 +3,7 @@ package cache import ( "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/mlx" ) @@ -27,7 +28,7 @@ func multiTokenKV(ids []float32) (*mlx.Array, *mlx.Array) { // stateIDs returns the ids currently in the cache in slot order (logical // after a concat, physical/rotated after a single-token update). -func stateIDs(t *testing.T, c *RotatingKVCache) []float32 { +func stateIDs(t *mlxtest.T, c *RotatingKVCache) []float32 { t.Helper() state := c.State() if state == nil { @@ -75,32 +76,32 @@ func feedSingle(c *RotatingKVCache, id float32) { // pre-existing tokens in logical order so the first Q of the new batch // has a full sliding window. func TestRotatingKVCacheConcatMidRotationPreservesContext(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + c := NewRotatingKVCache(window) - const window = 4 - c := NewRotatingKVCache(window) + nextID := feedMulti(c, 1, 3) + for range 6 { + feedSingle(c, nextID) + nextID++ + } + if c.Offset() != 9 { + t.Fatalf("setup: offset=%d want 9", c.Offset()) + } + if c.idx >= c.maxSize { + t.Fatalf("setup: expected mid-rotation idx (<%d), got %d", c.maxSize, c.idx) + } - nextID := feedMulti(c, 1, 3) - for range 6 { - feedSingle(c, nextID) - nextID++ - } - if c.Offset() != 9 { - t.Fatalf("setup: offset=%d want 9", c.Offset()) - } - if c.idx >= c.maxSize { - t.Fatalf("setup: expected mid-rotation idx (<%d), got %d", c.maxSize, c.idx) - } - - feedMulti(c, 10, 2) - got := stateIDs(t, c) - want := []float32{7, 8, 9, 10, 11} - if !equalSlice(got, want) { - t.Fatalf("post-concat window=%v want %v", got, want) - } - if c.Offset() != 11 { - t.Fatalf("offset=%d want 11", c.Offset()) - } + feedMulti(c, 10, 2) + got := stateIDs(t, c) + want := []float32{7, 8, 9, 10, 11} + if !equalSlice(got, want) { + t.Fatalf("post-concat window=%v want %v", got, want) + } + if c.Offset() != 11 { + t.Fatalf("offset=%d want 11", c.Offset()) + } + }) } // TestRotatingKVCacheConcatAlignedInvariant: with an aligned buffer @@ -108,38 +109,38 @@ func TestRotatingKVCacheConcatMidRotationPreservesContext(t *testing.T) { // tokens plus the full new batch. This is the chunked-prefill contract // x/mlxrunner/pipeline.go relies on. func TestRotatingKVCacheConcatAlignedInvariant(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + c := NewRotatingKVCache(window) - const window = 4 - c := NewRotatingKVCache(window) - - // Chunk 1 fills past maxSize, leaving Dim == maxSize aligned. - feedMulti(c, 1, 6) - // Chunk 2: the buffer is intentionally oversized to (maxSize-1) + L - // so the first new Q has its full window in scope for this forward. - feedMulti(c, 7, 3) - got := stateIDs(t, c) - want := []float32{4, 5, 6, 7, 8, 9} - if !equalSlice(got, want) { - t.Fatalf("post-chunk-2 buffer=%v want %v", got, want) - } - - // The next decode trims oversize back to maxSize; order may be - // physical (rotated), so check as a set. - feedSingle(c, 10) - got = stateIDs(t, c) - if len(got) != window { - t.Fatalf("post-decode Dim=%d want %d", len(got), window) - } - seen := map[float32]bool{} - for _, v := range got { - seen[v] = true - } - for _, w := range []float32{7, 8, 9, 10} { - if !seen[w] { - t.Fatalf("post-decode window missing %v (got %v)", w, got) + // Chunk 1 fills past maxSize, leaving Dim == maxSize aligned. + feedMulti(c, 1, 6) + // Chunk 2: the buffer is intentionally oversized to (maxSize-1) + L + // so the first new Q has its full window in scope for this forward. + feedMulti(c, 7, 3) + got := stateIDs(t, c) + want := []float32{4, 5, 6, 7, 8, 9} + if !equalSlice(got, want) { + t.Fatalf("post-chunk-2 buffer=%v want %v", got, want) } - } + + // The next decode trims oversize back to maxSize; order may be + // physical (rotated), so check as a set. + feedSingle(c, 10) + got = stateIDs(t, c) + if len(got) != window { + t.Fatalf("post-decode Dim=%d want %d", len(got), window) + } + seen := map[float32]bool{} + for _, v := range got { + seen[v] = true + } + for _, w := range []float32{7, 8, 9, 10} { + if !seen[w] { + t.Fatalf("post-decode window missing %v (got %v)", w, got) + } + } + }) } // TestRotatingKVCacheConcatAfterDecodeGrowsBuffer: update() grows the @@ -148,20 +149,20 @@ func TestRotatingKVCacheConcatAlignedInvariant(t *testing.T) { // has not wrapped. Those trailing slots are zero padding and must not // be pulled back into the live window on the next concat. func TestRotatingKVCacheConcatAfterDecodeGrowsBuffer(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 512 + c := NewRotatingKVCache(window) - const window = 512 - c := NewRotatingKVCache(window) + feedMulti(c, 1, 3) + feedSingle(c, 4) + feedMulti(c, 5, 3) - feedMulti(c, 1, 3) - feedSingle(c, 4) - feedMulti(c, 5, 3) - - got := stateIDs(t, c) - want := []float32{1, 2, 3, 4, 5, 6, 7} - if !equalSlice(got, want) { - t.Fatalf("growing-buffer concat=%v want %v", got, want) - } + got := stateIDs(t, c) + want := []float32{1, 2, 3, 4, 5, 6, 7} + if !equalSlice(got, want) { + t.Fatalf("growing-buffer concat=%v want %v", got, want) + } + }) } // TestRotatingKVCacheConcatAfterLiveRewind: x/mlxrunner/cache.go calls @@ -171,53 +172,53 @@ func TestRotatingKVCacheConcatAfterDecodeGrowsBuffer(t *testing.T) { // tokens. A subsequent concat must drop those, not treat them as wrapped // window content. func TestRotatingKVCacheConcatAfterLiveRewind(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 8 + c := NewRotatingKVCache(window) - const window = 8 - c := NewRotatingKVCache(window) + // Grow the buffer to exactly maxSize without wrapping. + feedMulti(c, 1, 2) + for id := float32(3); id <= 8; id++ { + feedSingle(c, id) + } + if c.Offset() != window { + t.Fatalf("setup: offset=%d want %d", c.Offset(), window) + } - // Grow the buffer to exactly maxSize without wrapping. - feedMulti(c, 1, 2) - for id := float32(3); id <= 8; id++ { - feedSingle(c, id) - } - if c.Offset() != window { - t.Fatalf("setup: offset=%d want %d", c.Offset(), window) - } + if !c.Restore(nil, 2) { + t.Fatalf("live rewind to 2 failed") + } + if c.Offset() != 2 { + t.Fatalf("post-rewind offset=%d want 2", c.Offset()) + } - if !c.Restore(nil, 2) { - t.Fatalf("live rewind to 2 failed") - } - if c.Offset() != 2 { - t.Fatalf("post-rewind offset=%d want 2", c.Offset()) - } - - feedMulti(c, 9, 3) - got := stateIDs(t, c) - want := []float32{1, 2, 9, 10, 11} - if !equalSlice(got, want) { - t.Fatalf("post-rewind concat=%v want %v", got, want) - } - if c.Offset() != 5 { - t.Fatalf("offset=%d want 5", c.Offset()) - } + feedMulti(c, 9, 3) + got := stateIDs(t, c) + want := []float32{1, 2, 9, 10, 11} + if !equalSlice(got, want) { + t.Fatalf("post-rewind concat=%v want %v", got, want) + } + if c.Offset() != 5 { + t.Fatalf("offset=%d want 5", c.Offset()) + } + }) } // TestRotatingKVCacheConcatGrowingBuffer: when oldLen < maxSize the trim // formula drops to non-positive and all pre-existing tokens are kept. func TestRotatingKVCacheConcatGrowingBuffer(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + c := NewRotatingKVCache(window) - const window = 4 - c := NewRotatingKVCache(window) - - feedMulti(c, 1, 2) - feedMulti(c, 3, 2) - got := stateIDs(t, c) - want := []float32{1, 2, 3, 4} - if !equalSlice(got, want) { - t.Fatalf("growing buffer=%v want %v", got, want) - } + feedMulti(c, 1, 2) + feedMulti(c, 3, 2) + got := stateIDs(t, c) + want := []float32{1, 2, 3, 4} + if !equalSlice(got, want) { + t.Fatalf("growing buffer=%v want %v", got, want) + } + }) } // TestRotatingKVCacheRunnerChunkedPrefill mirrors the @@ -225,114 +226,114 @@ func TestRotatingKVCacheConcatGrowingBuffer(t *testing.T) { // repeated L>1 Update() calls on a single cache. Scaled-down proxy for // the Gemma 4 26B case (sliding_window=1024, prefillChunkSize=2048). func TestRotatingKVCacheRunnerChunkedPrefill(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + c := NewRotatingKVCache(window) - const window = 4 - c := NewRotatingKVCache(window) - - feedMulti(c, 1, 8) - if c.Offset() != 8 { - t.Fatalf("chunk 1: offset=%d want 8", c.Offset()) - } - - feedMulti(c, 9, 8) - got := stateIDs(t, c) - want := []float32{6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} - if !equalSlice(got, want) { - t.Fatalf("chunk 2: buffer=%v want %v", got, want) - } - - feedMulti(c, 17, 4) - got = stateIDs(t, c) - want = []float32{14, 15, 16, 17, 18, 19, 20} - if !equalSlice(got, want) { - t.Fatalf("chunk 3: buffer=%v want %v", got, want) - } - - // Decode trims oversize back to maxSize; order may be physical. - feedSingle(c, 21) - got = stateIDs(t, c) - if len(got) != window { - t.Fatalf("post-decode Dim=%d want %d", len(got), window) - } - seen := map[float32]bool{} - for _, v := range got { - seen[v] = true - } - for _, w := range []float32{18, 19, 20, 21} { - if !seen[w] { - t.Fatalf("post-decode window missing %v (got %v)", w, got) + feedMulti(c, 1, 8) + if c.Offset() != 8 { + t.Fatalf("chunk 1: offset=%d want 8", c.Offset()) } - } + + feedMulti(c, 9, 8) + got := stateIDs(t, c) + want := []float32{6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + if !equalSlice(got, want) { + t.Fatalf("chunk 2: buffer=%v want %v", got, want) + } + + feedMulti(c, 17, 4) + got = stateIDs(t, c) + want = []float32{14, 15, 16, 17, 18, 19, 20} + if !equalSlice(got, want) { + t.Fatalf("chunk 3: buffer=%v want %v", got, want) + } + + // Decode trims oversize back to maxSize; order may be physical. + feedSingle(c, 21) + got = stateIDs(t, c) + if len(got) != window { + t.Fatalf("post-decode Dim=%d want %d", len(got), window) + } + seen := map[float32]bool{} + for _, v := range got { + seen[v] = true + } + for _, w := range []float32{18, 19, 20, 21} { + if !seen[w] { + t.Fatalf("post-decode window missing %v (got %v)", w, got) + } + } + }) } // TestRotatingKVCacheMultiTurnChatSimulation walks a prefill → decode → // prefill sequence and checks that each new prefill retains the last // (maxSize-1) pre-existing tokens in logical order. func TestRotatingKVCacheMultiTurnChatSimulation(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + c := NewRotatingKVCache(window) - const window = 4 - c := NewRotatingKVCache(window) + nextID := feedMulti(c, 1, 2) + for range 5 { + feedSingle(c, nextID) + nextID++ + } + if c.Offset() != 7 { + t.Fatalf("turn 1: offset=%d want 7", c.Offset()) + } - nextID := feedMulti(c, 1, 2) - for range 5 { - feedSingle(c, nextID) - nextID++ - } - if c.Offset() != 7 { - t.Fatalf("turn 1: offset=%d want 7", c.Offset()) - } + feedMulti(c, nextID, 3) + nextID += 3 + got := stateIDs(t, c) + want := []float32{5, 6, 7, 8, 9, 10} + if !equalSlice(got, want) { + t.Fatalf("turn 2 prefill buffer=%v want %v", got, want) + } - feedMulti(c, nextID, 3) - nextID += 3 - got := stateIDs(t, c) - want := []float32{5, 6, 7, 8, 9, 10} - if !equalSlice(got, want) { - t.Fatalf("turn 2 prefill buffer=%v want %v", got, want) - } + for range 4 { + feedSingle(c, nextID) + nextID++ + } + if c.Offset() != 14 { + t.Fatalf("turn 2 decode: offset=%d want 14", c.Offset()) + } - for range 4 { - feedSingle(c, nextID) - nextID++ - } - if c.Offset() != 14 { - t.Fatalf("turn 2 decode: offset=%d want 14", c.Offset()) - } - - feedMulti(c, nextID, 2) - got = stateIDs(t, c) - want = []float32{12, 13, 14, 15, 16} - if !equalSlice(got, want) { - t.Fatalf("turn 3 prefill buffer=%v want %v", got, want) - } + feedMulti(c, nextID, 2) + got = stateIDs(t, c) + want = []float32{12, 13, 14, 15, 16} + if !equalSlice(got, want) { + t.Fatalf("turn 3 prefill buffer=%v want %v", got, want) + } + }) } // TestRotatingKVCacheOffsetTracking: Offset() is the monotonic logical // token count through any mix of Update() calls — Gemma 4 uses // donorEntry.Offset - L for the consumer's RoPE offset. func TestRotatingKVCacheOffsetTracking(t *testing.T) { - skipIfNoMLX(t) - - c := NewRotatingKVCache(4) - nextID := feedMulti(c, 1, 3) - if c.Offset() != 3 { - t.Fatalf("after prefill 3: offset=%d want 3", c.Offset()) - } - for i := range 5 { - feedSingle(c, nextID) - nextID++ - if c.Offset() != 3+i+1 { - t.Fatalf("after decode %d: offset=%d want %d", i, c.Offset(), 3+i+1) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewRotatingKVCache(4) + nextID := feedMulti(c, 1, 3) + if c.Offset() != 3 { + t.Fatalf("after prefill 3: offset=%d want 3", c.Offset()) } - } - nextID = feedMulti(c, nextID, 2) - if c.Offset() != 10 { - t.Fatalf("after turn-2 prefill: offset=%d want 10", c.Offset()) - } - // L > maxSize concat. - feedMulti(c, nextID, 7) - if c.Offset() != 17 { - t.Fatalf("after large prefill: offset=%d want 17", c.Offset()) - } + for i := range 5 { + feedSingle(c, nextID) + nextID++ + if c.Offset() != 3+i+1 { + t.Fatalf("after decode %d: offset=%d want %d", i, c.Offset(), 3+i+1) + } + } + nextID = feedMulti(c, nextID, 2) + if c.Offset() != 10 { + t.Fatalf("after turn-2 prefill: offset=%d want 10", c.Offset()) + } + // L > maxSize concat. + feedMulti(c, nextID, 7) + if c.Offset() != 17 { + t.Fatalf("after large prefill: offset=%d want 17", c.Offset()) + } + }) } diff --git a/x/mlxrunner/cache/snapshot_capture_test.go b/x/mlxrunner/cache/snapshot_capture_test.go index eaa47835e..eb81df1a3 100644 --- a/x/mlxrunner/cache/snapshot_capture_test.go +++ b/x/mlxrunner/cache/snapshot_capture_test.go @@ -4,6 +4,7 @@ import ( "slices" "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/mlx" ) @@ -55,7 +56,7 @@ func fillTagged(c Attention, n int) { // windowTags reads c's logical window and returns the per-position tag (the // absolute offset each slot holds, recovered as value-1). It uses element 0 of // each token, which taggedKV set uniformly. Returns nil if the window is empty. -func windowTags(t *testing.T, c *RotatingKVCache) []int { +func windowTags(t *mlxtest.T, c *RotatingKVCache) []int { t.Helper() state := c.State() if len(state) == 0 { @@ -100,68 +101,68 @@ func wantWindowTags(offset, window int) []int { // speculation commit path — a live rewind, since KV is append-only — restores // the cache to each accepted offset with the prefix intact. func TestKVCachePerTokenSnapshotRestore(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const before = 6 + const draft = 4 - const before = 6 - const draft = 4 + for accepted := 0; accepted <= draft; accepted++ { + c := NewKVCache() + fillKV(c, before) - for accepted := 0; accepted <= draft; accepted++ { - c := NewKVCache() - fillKV(c, before) + offsets := make([]int, draft) + for i := range offsets { + offsets[i] = before + i + } + c.PrepareSnapshots(offsets) - offsets := make([]int, draft) - for i := range offsets { - offsets[i] = before + i - } - c.PrepareSnapshots(offsets) + k, v := batchKV(draft) + c.Update(newKVBatch(before, draft), k, v) + if c.Offset() != before+draft { + t.Fatalf("accepted=%d: offset after write = %d, want %d", accepted, c.Offset(), before+draft) + } - k, v := batchKV(draft) - c.Update(newKVBatch(before, draft), k, v) - if c.Offset() != before+draft { - t.Fatalf("accepted=%d: offset after write = %d, want %d", accepted, c.Offset(), before+draft) - } + snaps := c.TakeSnapshots() + if len(snaps) != draft { + t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft) + } - snaps := c.TakeSnapshots() - if len(snaps) != draft { - t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft) - } + // Captures are edge-local: offset before is zero-width (nil), and each + // later offset holds exactly the single token [before+i-1, before+i). + if snaps[0] != nil { + t.Fatalf("accepted=%d: snaps[0] = %v, want nil (zero-width base)", accepted, snaps[0]) + } + for i := 1; i < draft; i++ { + ks := snaps[i].(*kvSnapshot) + if ks.fromOffset != before+i-1 || ks.toOffset != before+i { + t.Fatalf("accepted=%d: snaps[%d] = [%d,%d), want [%d,%d)", accepted, i, ks.fromOffset, ks.toOffset, before+i-1, before+i) + } + } - // Captures are edge-local: offset before is zero-width (nil), and each - // later offset holds exactly the single token [before+i-1, before+i). - if snaps[0] != nil { - t.Fatalf("accepted=%d: snaps[0] = %v, want nil (zero-width base)", accepted, snaps[0]) - } - for i := 1; i < draft; i++ { - ks := snaps[i].(*kvSnapshot) - if ks.fromOffset != before+i-1 || ks.toOffset != before+i { - t.Fatalf("accepted=%d: snaps[%d] = [%d,%d), want [%d,%d)", accepted, i, ks.fromOffset, ks.toOffset, before+i-1, before+i) + // Commit rolls back via a live rewind (Restore(nil)) — the append-only + // buffer still holds [0, before+draft), so the edge captures go unused. + if accepted < draft { + if !c.Restore(nil, before+accepted) { + t.Fatalf("accepted=%d: live rewind failed", accepted) + } + } + for _, s := range snaps { + if s != nil { + s.Close() + } + } + + want := before + draft + if accepted < draft { + want = before + accepted + } + if c.Offset() != want { + t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want) + } + if st := c.State(); len(st) == 2 && st[0].Dim(2) != want { + t.Fatalf("accepted=%d: state seq dim = %d, want %d", accepted, st[0].Dim(2), want) } } - - // Commit rolls back via a live rewind (Restore(nil)) — the append-only - // buffer still holds [0, before+draft), so the edge captures go unused. - if accepted < draft { - if !c.Restore(nil, before+accepted) { - t.Fatalf("accepted=%d: live rewind failed", accepted) - } - } - for _, s := range snaps { - if s != nil { - s.Close() - } - } - - want := before + draft - if accepted < draft { - want = before + accepted - } - if c.Offset() != want { - t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want) - } - if st := c.State(); len(st) == 2 && st[0].Dim(2) != want { - t.Fatalf("accepted=%d: state seq dim = %d, want %d", accepted, st[0].Dim(2), want) - } - } + }) } // TestKVCaptureMergeSplit verifies that two adjacent edge-local captures merge @@ -169,46 +170,46 @@ func TestKVCachePerTokenSnapshotRestore(t *testing.T) { // trie performs on stored snapshots (mergeWithChild / splitNode) — proving they // work on captured snapshots, not just freshly-taken ones. func TestKVCaptureMergeSplit(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const before = 6 + c := NewKVCache() + fillKV(c, before) - const before = 6 - c := NewKVCache() - fillKV(c, before) + // Schedule two interior offsets so the write captures [before, before+1) and + // [before+1, before+2). + c.PrepareSnapshots([]int{before + 1, before + 2}) + k, v := batchKV(3) + c.Update(newKVBatch(before, 3), k, v) - // Schedule two interior offsets so the write captures [before, before+1) and - // [before+1, before+2). - c.PrepareSnapshots([]int{before + 1, before + 2}) - k, v := batchKV(3) - c.Update(newKVBatch(before, 3), k, v) + snaps := c.TakeSnapshots() + a := snaps[0].(*kvSnapshot) + b := snaps[1].(*kvSnapshot) + if a.fromOffset != before || a.toOffset != before+1 { + t.Fatalf("snaps[0] = [%d,%d), want [%d,%d)", a.fromOffset, a.toOffset, before, before+1) + } + if b.fromOffset != before+1 || b.toOffset != before+2 { + t.Fatalf("snaps[1] = [%d,%d), want [%d,%d)", b.fromOffset, b.toOffset, before+1, before+2) + } - snaps := c.TakeSnapshots() - a := snaps[0].(*kvSnapshot) - b := snaps[1].(*kvSnapshot) - if a.fromOffset != before || a.toOffset != before+1 { - t.Fatalf("snaps[0] = [%d,%d), want [%d,%d)", a.fromOffset, a.toOffset, before, before+1) - } - if b.fromOffset != before+1 || b.toOffset != before+2 { - t.Fatalf("snaps[1] = [%d,%d), want [%d,%d)", b.fromOffset, b.toOffset, before+1, before+2) - } + // Merge the adjacent edges into [before, before+2). + merged := c.Merge(snaps[0], snaps[1]).(*kvSnapshot) + if merged.fromOffset != before || merged.toOffset != before+2 { + t.Fatalf("merged = [%d,%d), want [%d,%d)", merged.fromOffset, merged.toOffset, before, before+2) + } - // Merge the adjacent edges into [before, before+2). - merged := c.Merge(snaps[0], snaps[1]).(*kvSnapshot) - if merged.fromOffset != before || merged.toOffset != before+2 { - t.Fatalf("merged = [%d,%d), want [%d,%d)", merged.fromOffset, merged.toOffset, before, before+2) - } - - // Split back at before+1 and confirm the halves match the originals. - p, ch := c.Split(merged, before+1) - ps := p.(*kvSnapshot) - cs := ch.(*kvSnapshot) - if ps.fromOffset != before || ps.toOffset != before+1 { - t.Fatalf("split parent = [%d,%d), want [%d,%d)", ps.fromOffset, ps.toOffset, before, before+1) - } - if cs.fromOffset != before+1 || cs.toOffset != before+2 { - t.Fatalf("split child = [%d,%d), want [%d,%d)", cs.fromOffset, cs.toOffset, before+1, before+2) - } - p.Close() - ch.Close() + // Split back at before+1 and confirm the halves match the originals. + p, ch := c.Split(merged, before+1) + ps := p.(*kvSnapshot) + cs := ch.(*kvSnapshot) + if ps.fromOffset != before || ps.toOffset != before+1 { + t.Fatalf("split parent = [%d,%d), want [%d,%d)", ps.fromOffset, ps.toOffset, before, before+1) + } + if cs.fromOffset != before+1 || cs.toOffset != before+2 { + t.Fatalf("split child = [%d,%d), want [%d,%d)", cs.fromOffset, cs.toOffset, before+1, before+2) + } + p.Close() + ch.Close() + }) } // TestRotatingPerTokenSnapshotRestore exercises per-token capture on a @@ -219,7 +220,7 @@ func TestKVCaptureMergeSplit(t *testing.T) { // observable as wrong positions, not just wrong shapes; the wrapped regime // forces concat's linearize branch to run on entry. func TestRotatingPerTokenSnapshotRestore(t *testing.T) { - skipIfNoMLX(t) + mlxtest.SkipIfUnavailable(t) cases := []struct { name string @@ -233,7 +234,7 @@ func TestRotatingPerTokenSnapshotRestore(t *testing.T) { const draft = 4 for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { for accepted := 0; accepted <= draft; accepted++ { c := NewRotatingKVCache(tc.window) fillTagged(c, tc.before) @@ -294,7 +295,7 @@ func TestRotatingPerTokenSnapshotRestore(t *testing.T) { // same window data. Covers the restored content, a following decode write, and // the snapshot copying out correctly afterward, across wrap regimes. func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) { - skipIfNoMLX(t) + mlxtest.SkipIfUnavailable(t) cases := []struct { name string @@ -308,7 +309,7 @@ func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) { const draft, accepted = 4, 2 for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { c := NewRotatingKVCache(tc.window) fillTagged(c, tc.before) @@ -386,61 +387,61 @@ func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) { // when a following write would destroy the window and copies the snapshot out — // the lazy mechanism paying for itself only when the data is about to be lost. func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window, before, draft, accepted = 4, 10, 4, 2 - const window, before, draft, accepted = 4, 10, 4, 2 + c := NewRotatingKVCache(window) + fillTagged(c, before) - c := NewRotatingKVCache(window) - fillTagged(c, before) - - offsets := make([]int, draft) - for i := range offsets { - offsets[i] = before + i - } - c.PrepareSnapshots(offsets) - - k, v := taggedKV(before, draft) - c.Update(newKVBatch(before, draft), k, v) - - snaps := c.TakeSnapshots() - snap := snaps[accepted].(*rotatingSnapshot) - // Simulate trie ownership: a node sets a materialize hook on attach. - fired := 0 - snap.SetMaterializeHook(func(int) { fired++ }) - - if !c.Restore(snap, before+accepted) { - t.Fatal("restore failed") - } - // Re-pointed, not copied out: still lazy, hook unfired. - if snap.keys != nil { - t.Fatal("hooked snapshot was copied out; expected the re-point fast path") - } - if fired != 0 { - t.Fatalf("materialize hook fired %d times on restore, want 0 (still lazy)", fired) - } - - got := windowTags(t, c) - want := wantWindowTags(before+accepted, window) - if !slices.Equal(got, want) { - t.Fatalf("window tags = %v, want %v", got, want) - } - - // A following decode write destroys the window's slots, so it copies the - // snapshot out — firing the hook exactly once. - wk, wv := taggedKV(c.Offset(), 1) - c.Update(newKVBatch(c.Offset(), 1), wk, wv) - if snap.keys == nil { - t.Fatal("following write did not copy out the snapshot") - } - if fired != 1 { - t.Fatalf("materialize hook fired %d times, want 1", fired) - } - - for _, s := range snaps { - if s != nil { - s.Close() + offsets := make([]int, draft) + for i := range offsets { + offsets[i] = before + i } - } + c.PrepareSnapshots(offsets) + + k, v := taggedKV(before, draft) + c.Update(newKVBatch(before, draft), k, v) + + snaps := c.TakeSnapshots() + snap := snaps[accepted].(*rotatingSnapshot) + // Simulate trie ownership: a node sets a materialize hook on attach. + fired := 0 + snap.SetMaterializeHook(func(int) { fired++ }) + + if !c.Restore(snap, before+accepted) { + t.Fatal("restore failed") + } + // Re-pointed, not copied out: still lazy, hook unfired. + if snap.keys != nil { + t.Fatal("hooked snapshot was copied out; expected the re-point fast path") + } + if fired != 0 { + t.Fatalf("materialize hook fired %d times on restore, want 0 (still lazy)", fired) + } + + got := windowTags(t, c) + want := wantWindowTags(before+accepted, window) + if !slices.Equal(got, want) { + t.Fatalf("window tags = %v, want %v", got, want) + } + + // A following decode write destroys the window's slots, so it copies the + // snapshot out — firing the hook exactly once. + wk, wv := taggedKV(c.Offset(), 1) + c.Update(newKVBatch(c.Offset(), 1), wk, wv) + if snap.keys == nil { + t.Fatal("following write did not copy out the snapshot") + } + if fired != 1 { + t.Fatalf("materialize hook fired %d times, want 1", fired) + } + + for _, s := range snaps { + if s != nil { + s.Close() + } + } + }) } // TestRotatingSnapshotSingleTokenWrite mirrors the tail of chunked prefill: the @@ -451,7 +452,7 @@ func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) { // not-yet-wrapped and wrapped regimes; the wrapped one exercises the ring-clone // fallback. func TestRotatingSnapshotSingleTokenWrite(t *testing.T) { - skipIfNoMLX(t) + mlxtest.SkipIfUnavailable(t) cases := []struct { name string @@ -463,7 +464,7 @@ func TestRotatingSnapshotSingleTokenWrite(t *testing.T) { } for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { + mlxtest.RunSubtest(t, tc.name, func(t *mlxtest.T) { c := NewRotatingKVCache(tc.window) fillTagged(c, tc.before) @@ -495,46 +496,46 @@ func TestRotatingSnapshotSingleTokenWrite(t *testing.T) { // batched write trims/rewrites the buffer the snapshot's slots lived in. This is // the case that forces a lazy snapshot to copy out before the later write destroys it. func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 6 + c := NewRotatingKVCache(window) - const window = 6 - c := NewRotatingKVCache(window) + // Schedule an offset in the first chunk and one in the second, then write + // both chunks before taking — the second chunk's concat trims past the first + // snapshot's window. + c.PrepareSnapshots([]int{4, 12}) - // Schedule an offset in the first chunk and one in the second, then write - // both chunks before taking — the second chunk's concat trims past the first - // snapshot's window. - c.PrepareSnapshots([]int{4, 12}) + k1, v1 := taggedKV(0, 8) // chunk 1: [0, 8) + c.Update(newKVBatch(0, 8), k1, v1) + k2, v2 := taggedKV(8, 8) // chunk 2: [8, 16) + c.Update(newKVBatch(8, 8), k2, v2) - k1, v1 := taggedKV(0, 8) // chunk 1: [0, 8) - c.Update(newKVBatch(0, 8), k1, v1) - k2, v2 := taggedKV(8, 8) // chunk 2: [8, 16) - c.Update(newKVBatch(8, 8), k2, v2) + snaps := c.TakeSnapshots() + if len(snaps) != 2 || snaps[0] == nil || snaps[1] == nil { + t.Fatalf("snapshots not captured: %v", snaps) + } - snaps := c.TakeSnapshots() - if len(snaps) != 2 || snaps[0] == nil || snaps[1] == nil { - t.Fatalf("snapshots not captured: %v", snaps) - } + // Restore the first-chunk snapshot (offset 4): its window predates the + // second chunk entirely, so the data must have survived the chunk-2 write. + if !c.Restore(snaps[0], 4) { + t.Fatal("restore to offset 4 failed") + } + if got, want := windowTags(t, c), wantWindowTags(4, window); !slices.Equal(got, want) { + t.Fatalf("offset 4 window tags = %v, want %v", got, want) + } - // Restore the first-chunk snapshot (offset 4): its window predates the - // second chunk entirely, so the data must have survived the chunk-2 write. - if !c.Restore(snaps[0], 4) { - t.Fatal("restore to offset 4 failed") - } - if got, want := windowTags(t, c), wantWindowTags(4, window); !slices.Equal(got, want) { - t.Fatalf("offset 4 window tags = %v, want %v", got, want) - } - - // Restore the second-chunk snapshot (offset 12) on a fresh cache. - c2 := NewRotatingKVCache(window) - if !c2.Restore(snaps[1], 12) { - t.Fatal("restore to offset 12 failed") - } - if got, want := windowTags(t, c2), wantWindowTags(12, window); !slices.Equal(got, want) { - t.Fatalf("offset 12 window tags = %v, want %v", got, want) - } - for _, s := range snaps { - s.Close() - } + // Restore the second-chunk snapshot (offset 12) on a fresh cache. + c2 := NewRotatingKVCache(window) + if !c2.Restore(snaps[1], 12) { + t.Fatal("restore to offset 12 failed") + } + if got, want := windowTags(t, c2), wantWindowTags(12, window); !slices.Equal(got, want) { + t.Fatalf("offset 12 window tags = %v, want %v", got, want) + } + for _, s := range snaps { + s.Close() + } + }) } // TestRotatingLazySnapshotSizeZeroUntilMaterialized verifies the speculation @@ -544,211 +545,211 @@ func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) { // lazy snapshot reports Size() == 0 until a destructive write copies it out, at // which point the materialize hook fires with the newly-allocated bytes. func TestRotatingLazySnapshotSizeZeroUntilMaterialized(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const window = 4 + const before = 10 // wrapped + const draft = 4 + c := NewRotatingKVCache(window) + fillTagged(c, before) - const window = 4 - const before = 10 // wrapped - const draft = 4 - c := NewRotatingKVCache(window) - fillTagged(c, before) + offsets := make([]int, draft) + for i := range offsets { + offsets[i] = before + i + } + c.PrepareSnapshots(offsets) - offsets := make([]int, draft) - for i := range offsets { - offsets[i] = before + i - } - c.PrepareSnapshots(offsets) + k, v := taggedKV(before, draft) + c.Update(newKVBatch(before, draft), k, v) - k, v := taggedKV(before, draft) - c.Update(newKVBatch(before, draft), k, v) + snaps := c.TakeSnapshots() + defer func() { + for _, s := range snaps { + if s != nil { + s.Close() + } + } + }() - snaps := c.TakeSnapshots() - defer func() { - for _, s := range snaps { - if s != nil { - s.Close() + // Interior captures (offsets after the start boundary) stay lazy: no write + // destroyed their slots, so keys is still nil and the issuing cache is live. + for i := 1; i < draft; i++ { + rs := snaps[i].(*rotatingSnapshot) + if rs.keys != nil { + t.Fatalf("snaps[%d] copied out (keys != nil); expected a live lazy snapshot", i) + } + if rs.cache == nil { + t.Fatalf("snaps[%d] has no issuing cache; expected a live lazy snapshot", i) } } - }() - // Interior captures (offsets after the start boundary) stay lazy: no write - // destroyed their slots, so keys is still nil and the issuing cache is live. - for i := 1; i < draft; i++ { - rs := snaps[i].(*rotatingSnapshot) - if rs.keys != nil { - t.Fatalf("snaps[%d] copied out (keys != nil); expected a live lazy snapshot", i) + lazy := snaps[1].(*rotatingSnapshot) + if lazy.Size() != 0 { + t.Fatalf("lazy rotating snapshot Size = %d, want 0", lazy.Size()) } - if rs.cache == nil { - t.Fatalf("snaps[%d] has no issuing cache; expected a live lazy snapshot", i) + + var hookDelta int + lazy.SetMaterializeHook(func(delta int) { hookDelta = delta }) + + // Free the cache to force every outstanding lazy snapshot to copy out. + c.Free() + + if lazy.keys == nil { + t.Fatal("Free did not materialize the lazy snapshot") } - } - - lazy := snaps[1].(*rotatingSnapshot) - if lazy.Size() != 0 { - t.Fatalf("lazy rotating snapshot Size = %d, want 0", lazy.Size()) - } - - var hookDelta int - lazy.SetMaterializeHook(func(delta int) { hookDelta = delta }) - - // Free the cache to force every outstanding lazy snapshot to copy out. - c.Free() - - if lazy.keys == nil { - t.Fatal("Free did not materialize the lazy snapshot") - } - want := lazy.keys.NumBytes() + lazy.values.NumBytes() - if hookDelta != want { - t.Fatalf("hook fired with delta %d, want %d", hookDelta, want) - } - if lazy.Size() != want { - t.Fatalf("materialized Size = %d, want %d", lazy.Size(), want) - } + want := lazy.keys.NumBytes() + lazy.values.NumBytes() + if hookDelta != want { + t.Fatalf("hook fired with delta %d, want %d", hookDelta, want) + } + if lazy.Size() != want { + t.Fatalf("materialized Size = %d, want %d", lazy.Size(), want) + } + }) } // TestPerTokenSnapshotPersistsAcrossWrites verifies that scheduled offsets // survive multiple writes until TakeSnapshots — the property prefill would rely // on to snapshot interior offsets without splitting its forward. func TestPerTokenSnapshotPersistsAcrossWrites(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() + fillKV(c, 2) - c := NewKVCache() - fillKV(c, 2) + // Schedule offsets that span two separate writes. Offset 2 equals the + // schedule-time position, so it captures a zero-width range (nil); the rest + // are edge-local. + c.PrepareSnapshots([]int{2, 3, 5}) - // Schedule offsets that span two separate writes. Offset 2 equals the - // schedule-time position, so it captures a zero-width range (nil); the rest - // are edge-local. - c.PrepareSnapshots([]int{2, 3, 5}) + k1, v1 := batchKV(2) // reaches offsets 2,3 + c.Update(newKVBatch(2, 2), k1, v1) + k2, v2 := batchKV(2) // reaches offsets 4,5 + c.Update(newKVBatch(4, 2), k2, v2) - k1, v1 := batchKV(2) // reaches offsets 2,3 - c.Update(newKVBatch(2, 2), k1, v1) - k2, v2 := batchKV(2) // reaches offsets 4,5 - c.Update(newKVBatch(4, 2), k2, v2) - - snaps := c.TakeSnapshots() - if len(snaps) != 3 { - t.Fatalf("got %d snapshots, want 3", len(snaps)) - } - if snaps[0] != nil { - t.Fatalf("snaps[0] = %v, want nil (zero-width base)", snaps[0]) - } - if s := snaps[1].(*kvSnapshot); s.fromOffset != 2 || s.toOffset != 3 { - t.Fatalf("snaps[1] = [%d,%d), want [2,3)", s.fromOffset, s.toOffset) - } - // Offset 5 was scheduled across the second write; its edge starts at the - // previous scheduled offset (3), confirming the base cursor only advances - // on capture so the snapshot range matches the trie edge between scheduled - // offsets — write boundaries between captures must not move it. - if s := snaps[2].(*kvSnapshot); s.fromOffset != 3 || s.toOffset != 5 { - t.Fatalf("snaps[2] = [%d,%d), want [3,5)", s.fromOffset, s.toOffset) - } - - // Restore from the [2,3) edge snapshot to offset 3. - if !c.Restore(snaps[1], 3) { - t.Fatal("restore to offset 3 failed") - } - if c.Offset() != 3 { - t.Fatalf("offset = %d, want 3", c.Offset()) - } - for _, s := range snaps { - if s != nil { - s.Close() + snaps := c.TakeSnapshots() + if len(snaps) != 3 { + t.Fatalf("got %d snapshots, want 3", len(snaps)) } - } + if snaps[0] != nil { + t.Fatalf("snaps[0] = %v, want nil (zero-width base)", snaps[0]) + } + if s := snaps[1].(*kvSnapshot); s.fromOffset != 2 || s.toOffset != 3 { + t.Fatalf("snaps[1] = [%d,%d), want [2,3)", s.fromOffset, s.toOffset) + } + // Offset 5 was scheduled across the second write; its edge starts at the + // previous scheduled offset (3), confirming the base cursor only advances + // on capture so the snapshot range matches the trie edge between scheduled + // offsets — write boundaries between captures must not move it. + if s := snaps[2].(*kvSnapshot); s.fromOffset != 3 || s.toOffset != 5 { + t.Fatalf("snaps[2] = [%d,%d), want [3,5)", s.fromOffset, s.toOffset) + } + + // Restore from the [2,3) edge snapshot to offset 3. + if !c.Restore(snaps[1], 3) { + t.Fatal("restore to offset 3 failed") + } + if c.Offset() != 3 { + t.Fatalf("offset = %d, want 3", c.Offset()) + } + for _, s := range snaps { + if s != nil { + s.Close() + } + } + }) } // TestRecurrentSnapshotSplitsAndSegmentedCapture verifies that SnapshotSplits // reports the interior scheduled offsets and PutSegmented captures them from the // per-boundary states, so each accepted count restores to a distinct state. func TestRecurrentSnapshotSplitsAndSegmentedCapture(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + const convTail, convDim, nv, vd, kd = 3, 8, 2, 4, 4 + c := NewRecurrentCache(convTail, convDim, nv, vd, kd) + c.Get(newKVBatch(0, 1), mlx.DTypeFloat16) + // Advance to offset 5 so the speculative forward starts there. + c.Put(newKVBatch(0, 5), + []*mlx.Array{mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim)}, + []*mlx.Array{mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd)}) - const convTail, convDim, nv, vd, kd = 3, 8, 2, 4, 4 - c := NewRecurrentCache(convTail, convDim, nv, vd, kd) - c.Get(newKVBatch(0, 1), mlx.DTypeFloat16) - // Advance to offset 5 so the speculative forward starts there. - c.Put(newKVBatch(0, 5), - []*mlx.Array{mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim)}, - []*mlx.Array{mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd)}) + const before, draft = 5, 4 + offsets := []int{before, before + 1, before + 2, before + 3} + c.PrepareSnapshots(offsets) - const before, draft = 5, 4 - offsets := []int{before, before + 1, before + 2, before + 3} - c.PrepareSnapshots(offsets) - - splits := c.SnapshotSplits(draft) - want := []int{1, 2, 3} - if len(splits) != len(want) { - t.Fatalf("SnapshotSplits = %v, want %v", splits, want) - } - for i := range want { - if splits[i] != want[i] { + splits := c.SnapshotSplits(draft) + want := []int{1, 2, 3} + if len(splits) != len(want) { t.Fatalf("SnapshotSplits = %v, want %v", splits, want) } - } - - // Distinct per-boundary states (3 interior splits + the end) so restore - // targets are distinguishable. - mkConv := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim), s) } - mkDelta := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd), s) } - convStates := []*mlx.Array{mkConv(1), mkConv(2), mkConv(3), mkConv(4)} - deltaStates := []*mlx.Array{mkDelta(1), mkDelta(2), mkDelta(3), mkDelta(4)} - - c.Put(newKVBatch(before, draft), convStates, deltaStates) - if c.Offset() != before+draft { - t.Fatalf("offset after segmented put = %d, want %d", c.Offset(), before+draft) - } - - snaps := c.TakeSnapshots() - if len(snaps) != draft { - t.Fatalf("got %d snapshots, want %d", len(snaps), draft) - } - for i, s := range snaps { - if s == nil { - t.Fatalf("snapshot %d not captured", i) + for i := range want { + if splits[i] != want[i] { + t.Fatalf("SnapshotSplits = %v, want %v", splits, want) + } } - } - // Full accept (no restore): the live state is the committed end boundary - // (value 4) at offset before+draft. - st := c.State() - mlx.Eval(st[1]) - if got := st[1].Floats()[0]; got != 4 { - t.Fatalf("full-accept delta state = %v, want end boundary value 4", got) - } + // Distinct per-boundary states (3 interior splits + the end) so restore + // targets are distinguishable. + mkConv := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim), s) } + mkDelta := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd), s) } + convStates := []*mlx.Array{mkConv(1), mkConv(2), mkConv(3), mkConv(4)} + deltaStates := []*mlx.Array{mkDelta(1), mkDelta(2), mkDelta(3), mkDelta(4)} - // Each partial accept restores to offset before+accepted and must recover the - // distinct boundary state captured there: snaps[0] is the pre-forward state - // (value 0); snaps[i>=1] is the interior split boundary (value i). Recurrent - // snapshots are self-contained, so restores need not run in order. - for accepted := range draft { - if !c.Restore(snaps[accepted], before+accepted) { - t.Fatalf("accepted=%d: restore to before+%d failed", accepted, accepted) + c.Put(newKVBatch(before, draft), convStates, deltaStates) + if c.Offset() != before+draft { + t.Fatalf("offset after segmented put = %d, want %d", c.Offset(), before+draft) } - if c.Offset() != before+accepted { - t.Fatalf("accepted=%d: offset after restore = %d, want %d", accepted, c.Offset(), before+accepted) + + snaps := c.TakeSnapshots() + if len(snaps) != draft { + t.Fatalf("got %d snapshots, want %d", len(snaps), draft) } + for i, s := range snaps { + if s == nil { + t.Fatalf("snapshot %d not captured", i) + } + } + + // Full accept (no restore): the live state is the committed end boundary + // (value 4) at offset before+draft. st := c.State() mlx.Eval(st[1]) - if got := st[1].Floats()[0]; got != float32(accepted) { - t.Fatalf("accepted=%d: restored delta state = %v, want boundary value %d", accepted, got, accepted) + if got := st[1].Floats()[0]; got != 4 { + t.Fatalf("full-accept delta state = %v, want end boundary value 4", got) } - } - for _, s := range snaps { - s.Close() - } + + // Each partial accept restores to offset before+accepted and must recover the + // distinct boundary state captured there: snaps[0] is the pre-forward state + // (value 0); snaps[i>=1] is the interior split boundary (value i). Recurrent + // snapshots are self-contained, so restores need not run in order. + for accepted := range draft { + if !c.Restore(snaps[accepted], before+accepted) { + t.Fatalf("accepted=%d: restore to before+%d failed", accepted, accepted) + } + if c.Offset() != before+accepted { + t.Fatalf("accepted=%d: offset after restore = %d, want %d", accepted, c.Offset(), before+accepted) + } + st := c.State() + mlx.Eval(st[1]) + if got := st[1].Floats()[0]; got != float32(accepted) { + t.Fatalf("accepted=%d: restored delta state = %v, want boundary value %d", accepted, got, accepted) + } + } + for _, s := range snaps { + s.Close() + } + }) } // TestPrepareSnapshotsPastOffsetPanics verifies scheduling an already-passed offset // is rejected. func TestPrepareSnapshotsPastOffsetPanics(t *testing.T) { - skipIfNoMLX(t) + mlxtest.Run(t, func(t *mlxtest.T) { + c := NewKVCache() + fillKV(c, 5) - c := NewKVCache() - fillKV(c, 5) - - defer func() { - if recover() == nil { - t.Fatal("expected panic for already-passed offset") - } - }() - c.PrepareSnapshots([]int{3}) + defer func() { + if recover() == nil { + t.Fatal("expected panic for already-passed offset") + } + }() + c.PrepareSnapshots([]int{3}) + }) } diff --git a/x/mlxrunner/dflash_test.go b/x/mlxrunner/dflash_test.go index f9c515809..3f4a89f1c 100644 --- a/x/mlxrunner/dflash_test.go +++ b/x/mlxrunner/dflash_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/ollama/ollama/api" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/cache" "github.com/ollama/ollama/x/mlxrunner/mlx" @@ -92,7 +93,7 @@ var _ base.BlockDraft = (*fakeBlockDraft)(nil) // newBlockTestSession wires a runner around a fakeBlockDraft and opens one // request's drafting session, returning the concrete session for // internal-state assertions. -func newBlockTestSession(t *testing.T, predict map[int32]int32, blockSize int) (*Runner, *fakeBlockDraft, *dflashDraftSession, []cache.Cache) { +func newBlockTestSession(t *mlxtest.T, predict map[int32]int32, blockSize int) (*Runner, *fakeBlockDraft, *dflashDraftSession, []cache.Cache) { t.Helper() r := mtpTestRunner(t, predict, []int32{7}, sampler.Options{}) caches, _ := newMTPTestCaches(2) // caches[0] target, caches[1] draft context @@ -108,236 +109,243 @@ func draftTokensOf(caches []cache.Cache) []int32 { } func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) { - skipIfNoMLX(t) - _, draft, session, caches := newBlockTestSession(t, nil, 4) + mlxtest.Run(t, func(t *mlxtest.T) { + _, draft, session, caches := newBlockTestSession(t, nil, 4) - // One prefill-sized run at the flush cap writes through immediately in a - // single context-only Draft call. - n := dflashPendingFlushTokens - ids := make([]int32, n) - for i := range ids { - ids[i] = int32(i % mtpTestVocab) - } - session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0, nil) - if got := len(draft.calls); got != 1 { - t.Fatalf("draft calls after cap-sized run = %d, want 1", got) - } - if got := caches[1].Offset(); got != n { - t.Fatalf("draft cache offset = %d, want %d", got, n) - } + // One prefill-sized run at the flush cap writes through immediately in a + // single context-only Draft call. + n := dflashPendingFlushTokens + ids := make([]int32, n) + for i := range ids { + ids[i] = int32(i % mtpTestVocab) + } + session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0, nil) + if got := len(draft.calls); got != 1 { + t.Fatalf("draft calls after cap-sized run = %d, want 1", got) + } + if got := caches[1].Offset(); got != n { + t.Fatalf("draft cache offset = %d, want %d", got, n) + } - // A run below the cap only buffers; settle writes it through, skipping - // the leading rows the flush already covered. - tail := []int32{1, 2, 3} - session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1, nil) - if got := len(draft.calls); got != 1 { - t.Fatalf("draft calls after buffered run = %d, want 1 (buffered)", got) - } - session.settle(nil) - want := blockCall{offset: int32(n), ctx: []int32{2, 3}} - if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { - t.Fatalf("settle flush = %+v, want %+v", got, want) - } - if got := caches[1].Offset(); got != n+2 { - t.Fatalf("draft cache offset = %d, want %d (level with reports)", got, n+2) - } + // A run below the cap only buffers; settle writes it through, skipping + // the leading rows the flush already covered. + tail := []int32{1, 2, 3} + session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1, nil) + if got := len(draft.calls); got != 1 { + t.Fatalf("draft calls after buffered run = %d, want 1 (buffered)", got) + } + session.settle(nil) + want := blockCall{offset: int32(n), ctx: []int32{2, 3}} + if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { + t.Fatalf("settle flush = %+v, want %+v", got, want) + } + if got := caches[1].Offset(); got != n+2 { + t.Fatalf("draft cache offset = %d, want %d (level with reports)", got, n+2) + } + }) } func TestDFlashCommittedGapPanics(t *testing.T) { - skipIfNoMLX(t) - _, _, session, _ := newBlockTestSession(t, nil, 4) + mlxtest.Run(t, func(t *mlxtest.T) { + _, _, session, _ := newBlockTestSession(t, nil, 4) - session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) - defer func() { - if recover() == nil { - t.Fatalf("committed run past the frontier did not panic") - } - }() - // The frontier is at slot 1; a run starting at 3 leaves slot 1..2 unfed. - session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3, nil) + session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) + defer func() { + if recover() == nil { + t.Fatalf("committed run past the frontier did not panic") + } + }() + // The frontier is at slot 1; a run starting at 3 leaves slot 1..2 unfed. + session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3, nil) + }) } func TestDFlashRestoredPrefixResumes(t *testing.T) { - skipIfNoMLX(t) - r := mtpTestRunner(t, nil, []int32{7}, sampler.Options{}) - caches, _ := newMTPTestCaches(2) - draft := &fakeBlockDraft{blockSize: 4, maskToken: 6, draftCaches: caches[1:]} - r.cache.caches = caches - r.spec = newSpeculation(r, draft, caches[:1], caches[1:]) + mlxtest.Run(t, func(t *mlxtest.T) { + r := mtpTestRunner(t, nil, []int32{7}, sampler.Options{}) + caches, _ := newMTPTestCaches(2) + draft := &fakeBlockDraft{blockSize: 4, maskToken: 6, draftCaches: caches[1:]} + r.cache.caches = caches + r.spec = newSpeculation(r, draft, caches[:1], caches[1:]) - // A restored prefix arrives with the draft caches already written. - restored := []int32{1, 2, 3, 4, 5} - caches[1].(*fakeRewindableCache).feed(restored) - session := r.spec.drafter.open(nil).(*dflashDraftSession) - if session.ctxOffset != len(restored) { - t.Fatalf("ctxOffset = %d, want %d (synced to restored offset)", session.ctxOffset, len(restored)) - } + // A restored prefix arrives with the draft caches already written. + restored := []int32{1, 2, 3, 4, 5} + caches[1].(*fakeRewindableCache).feed(restored) + session := r.spec.drafter.open(nil).(*dflashDraftSession) + if session.ctxOffset != len(restored) { + t.Fatalf("ctxOffset = %d, want %d (synced to restored offset)", session.ctxOffset, len(restored)) + } - // The resumed prefill's run overlaps the restore point; only the rows - // past the frontier are buffered and written. - run := []int32{2, 3, 0, 1} - session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3, nil) - session.settle(nil) - want := blockCall{offset: 5, ctx: []int32{0, 1}} - if got := draft.calls[0]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { - t.Fatalf("resume flush = %+v, want %+v", got, want) - } - if got, wantTok := draftTokensOf(caches), append(restored, 0, 1); !slices.Equal(got, wantTok) { - t.Fatalf("draft cache = %v, want %v", got, wantTok) - } + // The resumed prefill's run overlaps the restore point; only the rows + // past the frontier are buffered and written. + run := []int32{2, 3, 0, 1} + session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3, nil) + session.settle(nil) + want := blockCall{offset: 5, ctx: []int32{0, 1}} + if got := draft.calls[0]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { + t.Fatalf("resume flush = %+v, want %+v", got, want) + } + if got, wantTok := draftTokensOf(caches), append(restored, 0, 1); !slices.Equal(got, wantTok) { + t.Fatalf("draft cache = %v, want %v", got, wantTok) + } + }) } func TestDFlashProposeBounds(t *testing.T) { - skipIfNoMLX(t) - predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5} - _, draft, session, _ := newBlockTestSession(t, predict, 4) - current := mlx.FromValues([]int32{1}, 1) + mlxtest.Run(t, func(t *mlxtest.T) { + predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5} + _, draft, session, _ := newBlockTestSession(t, predict, 4) + current := mlx.FromValues([]int32{1}, 1) - // Nothing committed yet: no context to draft from. - if session.propose(current, 4) != nil { - t.Fatalf("propose with no context did not decline") - } - session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) - if session.propose(current, 0) != nil { - t.Fatalf("propose with no budget did not decline") - } + // Nothing committed yet: no context to draft from. + if session.propose(current, 4) != nil { + t.Fatalf("propose with no context did not decline") + } + session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) + if session.propose(current, 0) != nil { + t.Fatalf("propose with no budget did not decline") + } - // The block caps the draft at blockSize-1 mask rows regardless of budget. - cand := session.propose(current, 10) - if cand == nil { - t.Fatalf("propose declined with context and budget") - } - mlx.Eval(cand.tokens) - if got := cand.tokens.Ints(); !slices.Equal(got, []int32{2, 3, 4}) { - t.Fatalf("draft tokens = %v, want [2 3 4]", got) - } - if got, want := draft.calls[0].block, []int32{1, 6, 6, 6}; !slices.Equal(got, want) { - t.Fatalf("block = %v, want %v (anchor plus blockSize-1 masks)", got, want) - } + // The block caps the draft at blockSize-1 mask rows regardless of budget. + cand := session.propose(current, 10) + if cand == nil { + t.Fatalf("propose declined with context and budget") + } + mlx.Eval(cand.tokens) + if got := cand.tokens.Ints(); !slices.Equal(got, []int32{2, 3, 4}) { + t.Fatalf("draft tokens = %v, want [2 3 4]", got) + } + if got, want := draft.calls[0].block, []int32{1, 6, 6, 6}; !slices.Equal(got, want) { + t.Fatalf("block = %v, want %v (anchor plus blockSize-1 masks)", got, want) + } + }) } func TestDFlashBlockRewoundBeforeContextWrites(t *testing.T) { - skipIfNoMLX(t) - predict := map[int32]int32{2: 3, 3: 4, 4: 5} - _, draft, session, caches := newBlockTestSession(t, predict, 4) + mlxtest.Run(t, func(t *mlxtest.T) { + predict := map[int32]int32{2: 3, 3: 4, 4: 5} + _, draft, session, caches := newBlockTestSession(t, predict, 4) - session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) - if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil { - t.Fatalf("propose declined") - } - // The proposal's block sits in the caches until the next write. - if got, want := draftTokensOf(caches), []int32{1, 2, 6, 6, 6}; !slices.Equal(got, want) { - t.Fatalf("draft cache after propose = %v, want %v", got, want) - } + session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) + if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil { + t.Fatalf("propose declined") + } + // The proposal's block sits in the caches until the next write. + if got, want := draftTokensOf(caches), []int32{1, 2, 6, 6, 6}; !slices.Equal(got, want) { + t.Fatalf("draft cache after propose = %v, want %v", got, want) + } - // The next round's report rewinds the block before appending context, so - // the accepted tokens' rows land at their true slots. - run := []int32{2, 3, 4} - session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1, nil) - session.settle(nil) - if got, want := draftTokensOf(caches), []int32{1, 2, 3, 4}; !slices.Equal(got, want) { - t.Fatalf("draft cache after settle = %v, want %v (block rewound)", got, want) - } - if got := caches[1].Offset(); got != 4 { - t.Fatalf("draft cache offset = %d, want 4 (level with reports)", got) - } - want := blockCall{offset: 1, ctx: []int32{2, 3, 4}} - if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { - t.Fatalf("context flush = %+v, want %+v", got, want) - } + // The next round's report rewinds the block before appending context, so + // the accepted tokens' rows land at their true slots. + run := []int32{2, 3, 4} + session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1, nil) + session.settle(nil) + if got, want := draftTokensOf(caches), []int32{1, 2, 3, 4}; !slices.Equal(got, want) { + t.Fatalf("draft cache after settle = %v, want %v (block rewound)", got, want) + } + if got := caches[1].Offset(); got != 4 { + t.Fatalf("draft cache offset = %d, want 4 (level with reports)", got) + } + want := blockCall{offset: 1, ctx: []int32{2, 3, 4}} + if got := draft.calls[1]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil { + t.Fatalf("context flush = %+v, want %+v", got, want) + } + }) } func TestDFlashCloseDrainsOutstandingBlock(t *testing.T) { - skipIfNoMLX(t) - predict := map[int32]int32{2: 3, 3: 4, 4: 5} - _, _, session, caches := newBlockTestSession(t, predict, 4) + mlxtest.Run(t, func(t *mlxtest.T) { + predict := map[int32]int32{2: 3, 3: 4, 4: 5} + _, _, session, caches := newBlockTestSession(t, predict, 4) - session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) - if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil { - t.Fatalf("propose declined") - } - // A session that ends with a proposal in flight still leaves the caches - // level: close rewinds the block even with nothing pending to flush. - session.close() - if got, want := draftTokensOf(caches), []int32{1}; !slices.Equal(got, want) { - t.Fatalf("draft cache after close = %v, want %v", got, want) - } + session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil) + if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil { + t.Fatalf("propose declined") + } + // A session that ends with a proposal in flight still leaves the caches + // level: close rewinds the block even with nothing pending to flush. + session.close() + if got, want := draftTokensOf(caches), []int32{1}; !slices.Equal(got, want) { + t.Fatalf("draft cache after close = %v, want %v", got, want) + } + }) } func TestDecodeBlockDraft(t *testing.T) { - skipIfNoMLX(t) - // The block draft mirrors the target chain, so one proposal round accepts - // every draft and the bonus token is the EOS. - const eos int32 = 7 - predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: eos, eos: 0} - r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) - caches, _ := newMTPTestCaches(2) - draft := &fakeBlockDraft{predict: predict, blockSize: 3, maskToken: 6, draftCaches: caches[1:]} - r.cache.caches = caches - r.spec = newSpeculation(r, draft, caches[:1], caches[1:]) - session, ch := newMTPTestSession(caches) + mlxtest.Run(t, func(t *mlxtest.T) { + // The block draft mirrors the target chain, so one proposal round accepts + // every draft and the bonus token is the EOS. + const eos int32 = 7 + predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: eos, eos: 0} + r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) + caches, _ := newMTPTestCaches(2) + draft := &fakeBlockDraft{predict: predict, blockSize: 3, maskToken: 6, draftCaches: caches[1:]} + r.cache.caches = caches + r.spec = newSpeculation(r, draft, caches[:1], caches[1:]) + session, ch := newMTPTestSession(caches) - req := Request{ - Responses: ch, - Tokens: []int32{1}, - CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, - SamplerOpts: sampler.Options{}, - } - spec := r.spec.open(req, nil) - if spec == nil || !spec.enabled { - t.Fatalf("open rejected a block-draft request") - } - pinDraftLimit(spec, 4) - d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil) - if err := r.decode(context.Background(), req, session, d, 0); err != nil { - t.Fatalf("decode: %v", err) - } - d.close() - spec.close() - - content, final := collectResponses(ch) - if content != "2345" { - t.Fatalf("content = %q, want %q", content, "2345") - } - if !final.Done || final.DoneReason != 0 { - t.Fatalf("final = %+v, want Done with EOS reason", final) - } - if want := []int32{2, 3, 4, 5, eos}; !slices.Equal(session.outputs, want) { - t.Fatalf("session outputs = %v, want %v", session.outputs, want) - } - - // The unprimed drafter parks the first call, so two tokens decode as - // pipelined plain forwards; the resumed round then validates the current - // token and blockSize-1 drafts in one fused forward. - wantForwards := []forwardCall{{offset: 0, n: 1}, {offset: 1, n: 1}, {offset: 2, n: 3}} - model := r.Model.(*fakeMTPModel) - if !slices.Equal(model.forwards, wantForwards) { - t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards) - } - - // Ending the parked stretch settles the buffered context through, so the - // proposal runs block-only; close's flush then writes the accepted rows - // after rewinding the block. - wantCalls := []blockCall{ - {offset: 0, ctx: []int32{2, 3}}, - {offset: 2, block: []int32{3, 6, 6}}, - {offset: 2, ctx: []int32{4, 5, 7}}, - } - if len(draft.calls) != len(wantCalls) { - t.Fatalf("draft calls = %+v, want %+v", draft.calls, wantCalls) - } - for i, want := range wantCalls { - got := draft.calls[i] - if got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || !slices.Equal(got.block, want.block) { - t.Fatalf("draft call %d = %+v, want %+v", i, got, want) + req := Request{ + Responses: ch, + Tokens: []int32{1}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, } - } + spec := r.spec.open(req, nil) + if spec == nil || !spec.enabled { + t.Fatalf("open rejected a block-draft request") + } + pinDraftLimit(spec, 4) + d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil) + if err := r.decode(context.Background(), req, session, d, 0); err != nil { + t.Fatalf("decode: %v", err) + } + d.close() + spec.close() - // The draft caches end level with the target, holding only context rows. - if got, want := caches[1].Offset(), caches[0].Offset(); got != want { - t.Fatalf("draft cache offset = %d, want %d (level with target)", got, want) - } - if toks := draftTokensOf(caches); slices.Contains(toks, 6) { - t.Fatalf("draft cache retains block rows: %v", toks) - } + content, final := collectResponses(ch) + if content != "2345" { + t.Fatalf("content = %q, want %q", content, "2345") + } + if !final.Done || final.DoneReason != 0 { + t.Fatalf("final = %+v, want Done with EOS reason", final) + } + if want := []int32{2, 3, 4, 5, eos}; !slices.Equal(session.outputs, want) { + t.Fatalf("session outputs = %v, want %v", session.outputs, want) + } + + // The unprimed drafter parks the first call, so two tokens decode as + // pipelined plain forwards; the resumed round then validates the current + // token and blockSize-1 drafts in one fused forward. + wantForwards := []forwardCall{{offset: 0, n: 1}, {offset: 1, n: 1}, {offset: 2, n: 3}} + model := r.Model.(*fakeMTPModel) + if !slices.Equal(model.forwards, wantForwards) { + t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards) + } + + // Ending the parked stretch settles the buffered context through, so the + // proposal runs block-only; close's flush then writes the accepted rows + // after rewinding the block. + wantCalls := []blockCall{ + {offset: 0, ctx: []int32{2, 3}}, + {offset: 2, block: []int32{3, 6, 6}}, + {offset: 2, ctx: []int32{4, 5, 7}}, + } + if len(draft.calls) != len(wantCalls) { + t.Fatalf("draft calls = %+v, want %+v", draft.calls, wantCalls) + } + for i, want := range wantCalls { + got := draft.calls[i] + if got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || !slices.Equal(got.block, want.block) { + t.Fatalf("draft call %d = %+v, want %+v", i, got, want) + } + } + + // The draft caches end level with the target, holding only context rows. + if got, want := caches[1].Offset(), caches[0].Offset(); got != want { + t.Fatalf("draft cache offset = %d, want %d (level with target)", got, want) + } + if toks := draftTokensOf(caches); slices.Contains(toks, 6) { + t.Fatalf("draft cache retains block rows: %v", toks) + } + }) } diff --git a/x/mlxrunner/grammar_mask_test.go b/x/mlxrunner/grammar_mask_test.go index a87ca6387..638d38cba 100644 --- a/x/mlxrunner/grammar_mask_test.go +++ b/x/mlxrunner/grammar_mask_test.go @@ -4,45 +4,47 @@ import ( "math" "testing" + "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/mlx" ) func TestApplyTokenMask(t *testing.T) { - skipIfNoMLX(t) - const ( - bitsPerMaskWord = 32 - firstTokenID = 0 // Least-significant bit of the first mask word. - interiorTokenID = 7 // Last bit of the first mask word's low byte. - lastIDInFirstMaskWord = bitsPerMaskWord - 1 // Sign bit of the int32-backed first mask word. - lastVocabID = 40 // Final valid ID in a partially used second mask word. - vocabSize = lastVocabID + 1 - ) - allowedIDs := []int{firstTokenID, interiorTokenID, lastIDInFirstMaskWord, lastVocabID} - word0 := uint32(1)<