fix(ssd-cache): snapshot KV on producer thread + handle bf16↔numpy

[CBribiescas]

Summary

Fixes two distinct bugs that together made --ssd-cache-dir crash the server on any text-only model under --continuous-batching --enable-prefix-cache whenever the prefix cache filled and triggered eviction. Affects every top model I could test (qwen3-coder-30b, gpt-oss-20b, gpt-oss-120b).

Bug 1 — std::runtime_error: There is no Stream(gpu, N) in current thread.

When an eviction triggered ssd_tier.enqueue_spill(), the live mx.array references for layer.keys / layer.values were placed on a queue and consumed by a background writer thread. The writer then called np.array(mx_array) inside serialize_layer, forcing an MLX materialization. MLX ≥ 0.31.2 ties each request's compute to a thread-local Stream(gpu, N) (N == request uid); the writer thread had no such stream registered, so the conversion aborted the entire process.

Same architectural class as issue #496 / PR #479, but at a new site — the SSD spill writer is its own background daemon thread, not the BatchedEngine MLLM inference thread #479 fixes.

Fix

• Add LayerSerializer.snapshot_layer(layer) -> dict, called on the producer (request handler) thread inside enqueue_spill before the queue handoff. The snapshot materializes each mx.array to numpy while the request's stream is still valid.
• serialize_layer now takes the snapshot dict instead of the live layer — the writer thread can't accidentally touch MLX because the API doesn't give it any MLX objects.
• Both KV/RotatingKVCache and ArraysCache (Mamba) serializers implement the new method.

Bug 2 — RuntimeError: Item size 2 for PEP 3118 buffer format string B does not match the dtype B item size 1.

np.array(mx.array<bfloat16>) raises because numpy has no native bfloat16 and the mlx ≥ 0.31 buffer-protocol export confuses it. qwen3-coder-30b and many modern MLX models use bf16 KV cache, so this was masked behind Bug 1 on every model that actually evicted.

Fix

• New helper ssd_cache._mx_to_numpy_safe(arr) try/excepts the bf16 case and falls back to arr.astype(mx.float32) + mx.eval + np.array.
• The original dtype is recorded in the snapshot/manifest as a string (e.g. "bfloat16").
• scheduler._reconstruct_ssd_layers casts the restored array back to the original dtype via mx.array(numpy_fp32).astype(dt) so the model sees its native KV layout (preserves memory + avoids implicit-cast surprises in attention kernels).
• fp32 on disk is 2× bf16, accepted as a transient trade-off for the spill path.

Plus: 4 stale unit tests fixed

test_evict_lru_calls_ssd_spill, test_check_ssd_returns_none_on_ram_hit, test_full_round_trip, test_metrics_accuracy were all failing on clean v0.4.0rc1 base. The MemoryCacheConfig.min_prefix_tokens default is 128; the tests used 10-token sequences, so every store() was silently rejected as "short prefix" and the cache stayed empty. Pass min_prefix_tokens=1 explicitly. (Bug 1 was masking these too — even if store() worked, the writer would crash before any test could assert on spill_count.)

Plus: new regression test

test_snapshot_runs_on_caller_thread uses a ThreadBoundArray fake that raises if __array__() is accessed off the creator thread, proving enqueue_spill materializes layers on the calling thread (not the writer). This is the architectural invariant Bug 1 violated.

Verification

Host: Apple M5 Pro Max, 128 GB unified memory, macOS Tahoe.

Unit tests

tests/test_ssd_cache.py — 59 passed (was 54 passed + 4 failed; +1 new test)
tests/test_memory_cache.py — 47 passed (non-regression)

Lint clean: ruff check vllm_mlx/ssd_cache.py vllm_mlx/scheduler.py tests/test_ssd_cache.py --select E,F,W --ignore E402,E501,E731,F811,F841 passes. black --check passes.

Runtime verification

Same harness for all three: start vllm-mlx with --continuous-batching --enable-prefix-cache --cache-memory-mb 50 --ssd-cache-dir /tmp/test --ssd-cache-max-gb 5, send 9 distinct ≥200-token prompts to force eviction.

Model	Pre-fix	Post-fix
Qwen3-Coder-30B-A3B-Instruct-MLX-4bit (bf16 KV)	crashes Stream(gpu, 2) after request 6	✅ 400 spill files / 462 MB, 8 evictions, no crash
mlx-community/gpt-oss-20b-MXFP4-Q8	crashes Stream(gpu, 3) after request 4	✅ 130 spill files / 115 MB, 5 evictions, 8/8 hits on replay
mlx-community/gpt-oss-120b-MXFP4-Q8	crashes Stream(gpu, 3) after request 3	✅ 266 spill files / 242 MB, 7 evictions, 8/8 hits on replay

Related

• Issue Vision requests fail with RuntimeError: There is no Stream(gpu, 0) in current thread #496 (closed) — the original Stream(gpu, N) thread-local invariant write-up.
• Issue Define stronger MLX generation stream ownership invariant #468 (open) — "Define stronger MLX generation stream ownership invariant" — this PR is one site of the broader work.
• PR fix: run BatchedEngine MLLM on dedicated MLXWorkerThread to prevent cross-thread stream errors #479 (open) — MLXWorkerThread for BatchedEngine MLLM. Orthogonal; this PR doesn't depend on it and doesn't conflict (different thread, different site).
• PR fix: harden SSD cache review follow-ups #481 (merged) — addressed SQLite/locking concerns from issue Track unresolved SSD cache tiering review concerns from PR #309 #480 but did not touch the writer-thread stream crash.

Test plan

• CI: tests/test_ssd_cache.py + tests/test_memory_cache.py green on all matrix Pythons.
• Manual: vllm-mlx serve <bf16-KV model> --continuous-batching --enable-prefix-cache --cache-memory-mb 50 --ssd-cache-dir /tmp/test, send N>cache-capacity prompts, confirm /v1/cache/stats shows evictions > 0, files appear under /tmp/test/data/, no process abort.
• Manual: replay the first prompt after eviction, confirm SSD reload path runs without dtype errors.

🤖 Generated with Claude Code

[CBribiescas]

Polite refresh — this has been quiet since 05-22. The fix follows the same shape as PR #479 / issue #468 (Stream(gpu, N) thread ownership), just at the SSD spill writer site rather than the BatchedEngine MLLM thread. CI is 9/9 green and the fix is pre/post-verified on qwen3-coder-30b and gpt-oss-20b/120b. Happy to address concerns or split the bf16↔numpy fallback off if it helps narrow scope.

[Thump604]

Reviewed the SSD spill/threading path and the dtype restore path. This looks correct to me.

What I checked:

• SSDCacheTier.enqueue_spill() now snapshots each layer on the caller/request thread before handing work to the writer queue.
• _writer_loop() / _write_entry() now receive only serializer + numpy snapshot data, so the SSD writer thread no longer needs to touch live MLX arrays.
• KVCacheSerializer and ArraysCacheSerializer both keep the old manifest shape plus original dtype hints when a bf16 upcast was needed.
• Scheduler._reconstruct_ssd_layers() casts restored arrays back when those dtype hints are present.
• The new test_snapshot_runs_on_caller_thread directly guards the producer-thread snapshot invariant.

Local verification on the PR worktree:

python -m pytest tests/test_ssd_cache.py tests/test_memory_cache.py -q
# 106 passed
ruff check vllm_mlx/ssd_cache.py vllm_mlx/scheduler.py tests/test_ssd_cache.py --select E,F,W --ignore E402,E501,E731,F811,F841
black --check vllm_mlx/ssd_cache.py vllm_mlx/scheduler.py tests/test_ssd_cache.py
git diff --check upstream/main...HEAD

All passed. I did not rerun the large-model manual SSD spill matrix locally; this approval is based on code review, focused tests, CI, and the runtime evidence in the PR body.

[CBribiescas]

Friendly merge nudge — approved 05-30, CI still 9/9 green and branch is mergeable. Happy to rebase if anything's drifted.

[waybarrios]

Carefully checked - it is well-done