feat(cache): persist GDN checkpoints as SSD sidecars - #2569
Conversation
|
Hi — I’m proposing this as an opt-in approach to keep long-context GDN prefix caching memory-bounded while preserving durable rewind and restart hits. The implementation, failure behavior, and measured 173K/256K results are documented in the PR body. The final lifecycle/security review is included, and the unit suite passes on Python 3.11, 3.12, and 3.13. I’d especially appreciate feedback on the split cache format, shared LRU/accounting model, and whether this is a suitable direction for upstream. I’m happy to narrow or restructure the change if maintainers would prefer a different integration shape. |
|
Thanks for the thorough work here, the validation depth made this review much easier. I went through the split store/restore paths, the sidecar lifecycle, and the flag-off surface against the #1423/#2465 history, and the design holds up well, so I merged this as is. The inline write on queue saturation closes the counter pinning case from #1423 and bounds the staging RAM growth behind #2465, which I consider the right call. One gap surfaced in my local smoke (Qwen3.5-0.8B hybrid, block size 2048). The sidecar signature hashes raw layer class names, so a request resumed from a warm hit extracts SizedArraysCache and commits its checkpoints under a different digest directory than the ArraysCache one restore derives from block metadata, which left a repaired sidecar invisible after restart and would quietly shrink multi-turn coverage toward the last cold-stored boundary. The fallback stayed safe and outputs matched throughout, so this is a coverage gap rather than a correctness bug. I fixed it in ded2bbe by routing the sidecar signature inputs through the existing |
2 participants
feat(cache): persist GDN recurrent state as bounded SSD sidecars
Summary
This PR separates non-sliceable GDN recurrent state from the sliceable prefix-cache payload and persists the recurrent state as durable, per-boundary SSD sidecars.
The goal is to make long-context Qwen3.6 GDN serving practical on unified-memory Macs without retaining one full recurrent state in RAM for every historical boundary. On restore, oMLX materializes only the newest compatible GDN checkpoint needed for the selected KV prefix. Older checkpoints remain on SSD and can be selected if a request rewinds.
The implementation is opt-in and backward-compatible. With the feature disabled, the existing cache layout and restore path remain available.
Background
This is a follow-up to #1887, which identified the cross-turn
ArraysCache/GDN tip chain as a distinct source of hot-cache growth. Increasing the GDN block size reduces the number of intra-request boundaries, but it does not by itself prevent historical recurrent-state tips from accumulating across turns.It is also related to #2465, where a long DeepSeek V4 request showed severe memory growth while many boundary snapshots were being persisted after the response had already completed. The common issue is that an SSD cache write can still require a large in-memory staging lifetime when the full recurrent payload is retained with every boundary.
Qwen3.6-27B is currently one of the most capable local models for this workload, and the Qwen team has officially announced Qwen3.8. That announcement is a practical reason to make the long-context cache path more memory-efficient before the next model generation becomes a common local-serving target.
This PR does not claim Qwen3.8 validation, KDA validation, or general validation of every future Qwen model. The measured model is
Qwen3.6-27B-8bit-mtpon Apple Silicon.Design
The cache is split into two coordinated layers:
ArraysCache/SizedArraysCacherecurrent state is extracted into an opaque sidecar for the same full block boundary.Each sidecar is addressed by the source block identity and a cache-layout signature. Sidecars live in a separate
_gdn_sidecarsnamespace, are written through a temporary staging file, and are atomically promoted into the durable index. The index is rebuilt from file metadata at startup without eagerly loading every recurrent tensor.Restore is deliberately lazy:
Historical GDN checkpoints are not all hydrated into the hot cache. If the newest sidecar is absent, corrupt, incompatible, or deliberately removed, restore walks back to the nearest valid full-block checkpoint. If no valid checkpoint is available, the request follows the existing safe miss/re-prefill path.
Main cache blocks and GDN sidecars participate in the same SSD budget and LRU accounting. This prevents the sidecar namespace from becoming an unbounded second cache. The scheduler also bounds request-local sidecar staging by bytes and applies backpressure before an unbounded in-memory fallback can form.
Internal cache statistics track sidecar count and bytes, checkpoint loads, the selected restore endpoint, and walkbacks. Boundary staging also records its byte high-water mark and backpressure time for benchmark diagnostics.
Safety and fallback behavior
Configuration used for the benchmark
The candidate profile used the following effective settings:
{ "cache.gdn_ssd_split_enabled": true, "cache.gdn_ssd_pending_max_size": "512MB", "cache.ssd_cache_max_size": "128GB", "cache.hot_cache_max_size": "24GB", "model.max_context_window": 262144, "model.arrays_cache_tip_retention": 0, "model.turboquant_kv_enabled": false, "model.mtp_enabled": true }The legacy comparison profile used the same model, context limit, KV precision, MTP setting, SSD limit, hot-cache limit, and request payload, with
gdn_ssd_split_enabled=falseand the previousarrays_cache_tip_retention=16behavior.The measurements below were collected on an Apple M5 Max with 128 GiB unified memory, running macOS 26.5.2 (25F84), arm64, and oMLX 0.5.7. Each repeated cold profile used a fresh isolated base path and cache, one request at a time, on port 18150. The operational server on port 8080 was not used for the benchmark.
Validation completed
Unit and integration coverage
Review hardening after the initial draft
The final review found and fixed several lifecycle cases that were not covered by the first implementation: the scheduler's latest boundary is now committed even though it is intentionally absent from the intermediate-snapshot map; a readable but malformed newest sidecar walks back instead of aborting reconstruction; and a deduplicated KV block recreates an independently evicted sidecar before reuse.
The storage path now coalesces deterministic duplicate boundaries, retains pending-byte reservations until the owning raw buffer is actually released, detaches promotion files outside request/session cleanup, protects an existing checkpoint during failed replacement, and restores conservative byte accounting whenever safe deletion is refused. Sidecar deletion uses no-follow directory descriptors, commit timestamps survive restart LRU reconstruction, and reset/abort cleanup waits for asynchronous cache-store ownership to end. Regression tests cover each of these paths.
173K exact replay
The same 173,008-token Pi coding-session payload was replayed with
temperature=0, streaming enabled, andmax_tokens=32.The cold split run created 84 GDN sidecars totaling 12,932,773,056 bytes. The warm replay returned the same response fingerprint as the cold run, proving that the exact prefix was restored rather than recomputed. The selected GDN sidecar load was approximately 0.30 ms in the warm replay.
Rewind probes also selected the expected full-block endpoints without a full re-prefill:
For a fault-injection run in which the newest full-block sidecar was made unavailable, the request completed with endpoint 169,984, 3,024 uncached tokens, and one GDN walkback. This demonstrates the intended safe degradation behavior.
256K cold run
The 256K cold run used a 262,112-token prompt, which leaves room below the 262,144-token model limit for the request protocol and generated output.
The cold request created 127 GDN sidecars totaling 19,553,121,168 bytes. The immediate warm replay restored endpoint 260,096 with 2,016 uncached tokens, one GDN checkpoint load, zero walkbacks, and a 0.532 ms checkpoint-load latency. Warm effective PP was 18,286.93 tok/s and the response fingerprint matched the cold request. The reported RSS is the process resident-set peak during the cold-prefill interval and should not be confused with SSD sidecar capacity.
A fresh server process then indexed the existing cache and replayed the same 262,112-token prompt. It restored endpoint 260,096 from the durable cache, left only 2,016 tokens to prefill, loaded one GDN checkpoint with zero walkbacks, and completed with effective PP 15,406.30 tok/s, TG 34.58 tok/s, TTFT 21.36 s, and wall time 21.59 s. The checkpoint load took 3.045 ms. All 127 sidecars were indexed after restart, and the response fingerprint matched the original cold and warm responses. This closes the post-restart durability check at the full tested context length.
Repeated A/B speed and memory table
The following controlled A/B alternated three fresh Split profiles and three fresh Legacy profiles. Every run used the identical 173,008-token payload,
temperature=0, streaming, and a 512-token decode limit. Values are medians with the observed min-max range in parentheses.Split passes the predefined performance criterion: median PP and TG are not worse than Legacy by more than 5%. The memory effect is stable across all three repetitions: Split stayed within a 0.09 GiB RSS range, Legacy within 0.04 GiB, and the median reduction was 12.006 GiB. Every run reported zero cache errors and zero SSD write drops. Split sidecar staging peaked at 153,944,064 bytes, with no more than 0.156 ms recorded backpressure.
Four of the six runs produced the same response fingerprint, including cross-profile adjacent runs. The first Split and first Legacy observations produced different fingerprints despite
temperature=0. This does not correlate with the cache layout and is consistent with native-MTP execution not being bitwise deterministic, but a dedicated MTP-off control would be required to prove that cause. The image and 256K cold/warm/restart checks did produce matching fingerprints.These measurements do not establish that SSD sidecar persistence is universally faster than recomputing GDN state. They establish the narrower claim that bounded sidecar persistence caused no material PP/TG regression in this repeated workload while substantially reducing peak RSS.
Actual-image VLM cold/warm smoke
A deterministic one-pixel PNG was attached to the first six messages of the same Pi session, producing a 14,117-token multimodal prompt. The cold request completed without image-processing or cache errors and created six GDN sidecars.
The warm checkpoint load took 0.317 ms, both requests returned the same response fingerprint, and the profile reported zero cache errors and zero SSD write drops. This confirms that the split checkpoint path remains usable when the serving request includes actual image input.
Limitations and remaining work
Scope, changed files, and ownership
The implementation scope is the GDN split persistence/restore path, its bounded staging and accounting, configuration propagation, and focused regression coverage. The changed production files are:
omlx/cache/paged_ssd_cache.py: durable GDN sidecar metadata/index, atomic commit, startup scan, shared budget/LRU, and sidecar statistics.omlx/cache/prefix_cache.py: split-layout detection, lazy compatible-sidecar restore, endpoint selection, walkback, and restore diagnostics.omlx/cache/boundary_snapshot_store.pyandomlx/scheduler.py: boundary extraction, bounded staging/backpressure, persistence dispatch, and safe fallback wiring.omlx/cache/type_registry.py,omlx/config.py, andomlx/settings.py: cache-type recognition and opt-in configuration propagation.tests/test_prefix_cache_gdn_split.py,tests/test_gdn_sidecar_index.py,tests/test_boundary_snapshot_store.py,tests/test_settings.py, and focused scheduler coverage: regression coverage.Merge checklist