feat(cache): multi-slot LRU MRU partial block cache

[blightbow]

This PR proposes an in-memory cache for the trailing slice of prefill that cannot be committed to the disk cache. It captures the sub-block tail of a completed prefill (anything past the last full block_size boundary) so that a resubmitted prompt skips the avoidable partial recompute that block-aligned caching otherwise forces.

The cache is a bounded LRU dict keyed by parent_hash, with default capacity 4 entries (configurable via --mru-partial-max-entries). Entries are written at the end of store_cache for any prefill that produces a trailing partial, and consumed during cache-hit admission via apply_mru_partial. On a successful splice the entry is promoted to the LRU tail; on capacity overflow the oldest entry is evicted. On any apply-time mismatch (different tail tokens, layer-count mismatch, splice failure), only the mismatched entry is evicted. Sibling entries for other prefixes are preserved.

This is a memory-only structure, kept distinct from the disk and hot caches. No attempt is made to persist partials; the eviction discipline below ensures they cannot accumulate beyond capacity.

Scope

The optimization helps across the full range of single-prompt repeat and concurrent-prompt scenarios:

• Single-user iteration (swipes, regenerations, A/B testing with the same prompt) gets the "skip the partial recompute" benefit on every repeat.
• Multi-user or multi-conversation workloads up to mru_partial_max_entries distinct active prefixes get the same benefit per prefix. LRU eviction handles workloads beyond capacity gracefully: the most recently used prefixes stay warm, older ones fall back to a full recompute on next hit (which is the pre-cache baseline, so the worst case is "no benefit," not "regression").

The eviction gate on every mismatch protects correctness regardless of multiplicity; the operational footprint is bounded (see below).

Safety

The cache surface is small but the failure modes are sharp. The implementation:

• Refuses to stash on hybrid models. Any layer with a non-sliceable cache type (RotatingKVCache, ArraysCache, BatchRotatingKVCache, and so on) disables stashing entirely. Splicing into only the sliceable layers would create per-layer offset skew at decode and silently corrupt generation on Gemma 3, Mistral, and any future hybrid. Gated via an explicit KNOWN_SLICEABLE_CACHE_TYPES whitelist on omlx/cache/type_registry.py. Notably, the registry's supports_block_slicing flag is not trustworthy here. DefaultCacheHandler falls back to KVCache semantics and would report several unregistered batch/pool types as sliceable.
• Splices transactionally. Phase 1 builds replacement keys/values/offset for every layer with no mutation; phase 2 commits. A failure mid-loop rolls everything back and evicts the entry. No half-mutated cache state can reach a forward pass.
• Evicts on every miss kind. Token-prefix mismatch, length mismatch, layer-count mismatch, splice failure: each pops only the matching key, leaving siblings alone. BlockAwarePrefixCache.clear() wipes the whole dict (cache-corruption recovery path).
• Freed-paged-block guard. If the parent paged block has been freed between stash and apply, the apply path returns no-op rather than falling through to a None-keyed lookup, which would falsely match a short-prompt entry against an unrelated request whose parent is gone.
• Refuses on ambiguous cache layout. Multi-turn requests can produce cache buffers whose length doesn't unambiguously identify whether they're global-indexed or local-indexed; the stash path refuses rather than guess.
• Functions without disk writes. The gate keys on PagedSSDCacheManager instance presence, not on whether SSD writes are happening. Under hot_cache_only=True (set via settings or OMLX_HOT_CACHE_ONLY env), the disk writer thread is disabled but reconstruct still works via the hot tier's short-circuit; the MRU cache correctly remains active in that mode.

Test coverage in tests/test_prefix_cache.py::TestMRUPartialBlockCache and ::TestMRUPartialMultiSlot pins each of these, including a transactional-rollback test that mocks mx.concatenate to fail mid-loop and verifies no layer was mutated, a structural invariant test that pins kv_data to mx.array storage so future "optimize to CPU copy" regressions get caught, and parameterized LRU mechanics covering capacity eviction, apply-success promotion, and sibling preservation under mismatches.

Operational note: deferred Metal cache clear

The post-completion deferred clear (#435, #557) is suppressed for one extra _DEFERRED_CLEAR_DELAY window when any MRU entry is warm at the deadline. Warm partials are a strong predictor that the same prompt will return immediately and would benefit from the still-resident lazy KV tensors. The suppression is bounded at one suppression per deferral epoch (a fresh budget arms only on transition from _deferred_clear_at is None); the next deadline fires regardless of MRU state. This avoids the failure mode where hot-prompt repeats refresh the budget faster than _DEFERRED_CLEAR_DELAY and the pool-bloat mitigation (#411) is silently defeated.

Memory accounting

Each entry holds real mx.array allocations (via mx.copy) and counts automatically against mx.get_active_memory(). Every runtime memory enforcement and telemetry path in this codebase (process enforcer, scheduler limit checks, periodic-clear threshold) reads from there. Worst case is one block_size of KV per entry, held alive between completions and admissions.

Under hot_cache_only=True, the hot cache and the MRU dict share the same in-memory KV headroom envelope. Operators running that mode at high --mru-partial-max-entries should size --hot-cache-max-size accordingly: the two settings are co-tenants of the same budget, not independent dials. Default 4 is conservative enough that this rarely matters in practice.

Configuration

--mru-partial-max-entries N (default 4, matching the dflash max_entries precedent in #1120) sets the maximum simultaneous entries. 0 disables the feature entirely (silent fallback to "no MRU" behavior, mirroring the --hot-cache-max-size 0 convention). Also available via mru_partial_max_entries in CacheSettings. Operators with high-concurrency workloads can opt up; operators with memory pressure can opt down or off.

Admin endpoint symmetry

The /api/ssd-cache/clear admin endpoint now also wipes MRU partials for each loaded scheduler, via a new BlockAwarePrefixCache.clear_mru_partials() method. Without this, partials would chain from paged-block hashes whose KV bytes were just flushed by the endpoint, violating the operator's "drop all warm caches" intent. The same clear_mru_partials() hook is the seam intended for the future /api/hot-cache/clear endpoint introduced in #1183: a one-line addition at the same loop once that PR lands.

Observability

The MRU partial cache plugs into the same observability surface #1183 established for the prefix and hot/disk tiers, so operators tune --mru-partial-max-entries from the same dashboard they use for memory and disk hit rates.

New counters on PrefixCacheStats: mru_partial_stashes (entry writes, including same-key replacements), mru_partial_hits (successful splices), mru_partial_evictions (capacity-overflow + apply-miss + admin-clear wipes; the cache-corruption clear() path intentionally zeros all counters via reset_stats() rather than incrementing here), mru_partial_tokens_saved (the direct compute-saved measure: prefill tokens that did not have to be re-run). New gauges: mru_partial_entries and mru_partial_max_entries. Derived mru_partial_hit_rate = hits / stashes — the "stash payoff" ratio — surfaces both windowed and cumulative in CacheRateTracker.

The dashboard mirrors the hot-cache pattern: a header gauge ("MRU tails N/M entries"), rate-strip cells for the hit rate and tokens-saved counter, and a per-model column. All gated on mru_partial_max_entries > 0 so default-off configurations don't see the surface.

Base

This branch is based on #1183 to use the cache-tier observability and hot-cache architecture as the integration substrate. Targeting merge after #1183.

[blightbow]

Converting to a draft PR. I was attempting to rebase this on unmerged PR #1183 to implement MRU in the warm memory cache. During the rebase I identified test regressions in #1183 that need addressing. (TestRuntimeCacheObservability) This PR is now a draft pending resolution in #1183.

[blightbow]

The PR has been rescoped to function alongside the hot cache in addition to the SSD cache. To avoid commit churn, I rebased on ivaniguarans's #1183 and used it to surface dashboard metrics for the LRU cache.

Doing some additional local testing, then marking this as ready for review. (gated by merge of #1183)

[jundot]

Thanks for the careful work, the transactional splice and hybrid carve-out shows you thought through the failure modes.

My take is that avoiding at most ~2K tokens of partial-tail re-prefill (block_size upper bound) feels like a narrow win for the amount of new code this brings. Default 4 slots, hybrid models excluded, plus the permanent surface of transactional rollback, deferred-clear interaction, and hot-cache memory cotenancy is a lot for an optimization that prefix cache already covers at block boundaries.

This is just my read though, curious how you see the tradeoff. If partial-tail recompute is hurting a specific workload you're targeting, that context would help me look at this differently.

[blightbow]

@jundot Prefill is the main performance weakness on Apple silicon: if there is any edge to be shaven off of prefill, no matter how marginal, I attack it on principle. 😀

My goal was to solidify the approach, measure performance difference with and without, then make a judgement on worth. I know it's a lot of code that might not justify the effort, but worst case scenario others will know not to try.

The code approach has solidified. Just need to sit down and do the profiling...been distracted the last few days. I'll mark as ready to review when I have evidence.