[oniisanlabs]

Summary

Addresses four unaddressed bottlenecks in rec_mlx prefill, prefix caching, and unified-memory residency on 48 GB Apple Silicon Macs while preserving 100% bit-exact (0 ULP / mx.array_equal) parity across all existing rollback, shared-KV, and grammar tests:

1. 1-Token-Exact Sub-Chunk Prefix Rewind (rec_mlx/cache.py, rec_mlx/prefix.py):

   • In QATSession.prefill, PromptCache.checkpoint() is called before each prefill_chunk_size (4,096-token) chunk. On the first chunk (begin = 0), self.prompt_cache.tokens is empty, so zero checkpoints exist below 4,096 tokens. Any two requests whose shared prefix (system prompt + tool schemas + repo context) is < 4,096 tokens—or where a < 4,096-token prompt changes even its final user question—previously got checkpoint = None in select_prefix() and fell back to reused = 0 (cold prefill from scratch).
   • FixedSlidingKVCache now retains a bounded prompt history (_history_k / _history_v, capped at MAX_REWIND_TOKENS = 4096) and exposes can_rewind_to(target_pos) / rewind_to(target_pos). When prefix.checkpoints is empty (sub-4K prompts), select_prefix() rewinds all global and sliding layers to the exact common token offset (target = min(common, len(tokens) - 1)), lifting sub-4K prefix cache hit rates from 0% to 100% of the shared prefix with bit-exact KV restoration.

2. Multi-Branch LRU PrefixTreeCache & Batched Async KV Copies (rec_mlx/prefix.py, rec_mlx/cache.py, rec_mlx/quantized.py, rec_mlx/engine.py):

   • Adds PrefixTreeCache(max_branches=4) with compact per-layer global KV snapshots ([2, 4, L, 128], ~28 KB/token across all 14 global layers instead of copying the full 1.78 GiB KVStorage bank) plus a pinned base_branch for the shared system/tool preamble. Switching across files or sessions (Branch A -> Branch B -> Branch A) restores the full branch in < 10 ms (reused = len(branch_a)).
   • Adds copy_array_async and KVStorage.write_async so PromptCache.fork(), PromptCache.checkpoint(), and QATSession.prefill's 35 sliding-window KV writes dispatch in a single batched mx.eval instead of 35–70 synchronous CPU→GPU barriers.

3. Wired Metal Working Set, Bounded Cache Eviction & Idle Page-Residency Keep-Alive (rec_mlx/engine.py, rec_mlx/scheduler.py):

   • Previously, rec_mlx left MLX's wired limit at 0 (mx.set_wired_limit was never called) and called mx.clear_cache() unconditionally after every request. On 48 GB Macs, pausing for 1–3 minutes between interactive agent turns allowed macOS to compress/page out the unwired 22.9 GiB weight banks (RSS -> ~117 MB), causing a ~2.6s–3.5s page-in stall on the next prefill even when cached_prompt_tokens > 0.
   • Engine.__init__ now pins the Metal working set via mx.set_wired_limit(min(max_recommended_working_set_size, memory - 8 GiB)), retains the 2 GiB allocator cache across requests unless mx.get_cache_memory() > 2 GiB, and runs a lightweight touch_resident_pages() kernel (1 word per 16 KB page across weight and KV banks in ~3.5 ms) from Generator._next_request every 3.0s while idle.

4. Memoized ConstraintCompiler.tools Compilation (rec_mlx/grammar.py):

   • Caches compiled Constraint objects by canonical JSON tool-schema key (maxsize = 64), eliminating ~200–230 ms of redundant XGrammar compile_structural_tag CPU overhead on repeated tool-calling requests.

---

Benchmarks (RecursiveSI/flint-int4, Apple M5 Pro, 48 GB Unified Memory)

Interactive Agent & Multi-File Switching Latency (~1,800-token prompts, 917-token shared system + tools prefix)

Scenario	Prompt / Cached / New Tokens	main Prefill	This PR Prefill	Speedup
Step 2: Switch to File B (A -> B)	1,779 (cached=917, new=862)	~1,500–2,500 ms (reused=0)	508.9 ms (reused=917)	3.0–4.9× faster
Step 3: Switch back to File A (B -> A)	1,839 (cached=1,839, new=0)	~1,500–2,500 ms (reused=0)	7.0–11.7 ms (reused=1,839)	>120× faster
Step 4: File A + New User Question	1,840 (cached=1,826, new=14)	~1,500–2,500 ms (reused=0)	68.4 ms (reused=1,826)	>22× faster
Warm Prefill After 3 Min Idle	1,592 (cached=890, new=702)	3,131 ms (OS page-out)	~380–510 ms (wired)	~6.5× faster
Repeated Tool Schema Compilation	6 tool schemas (ConstraintCompiler.tools)	~230 ms / req	0.0 ms / req (memoized)	Eliminated

---

Testing

• cd engine && .venv/bin/pytest -v -rs: 71 passed in 15.26s (all 67 existing tests + 4 new unit tests):
  • PrefixRollbackTests.test_subchunk_exact_rewind_without_checkpoints: verifies 1-token-exact sub-chunk rewind (common = 9 -> reused = 9) and bit-exact sliding-window KV parity (mx.array_equal) against cold prefill.
  • PrefixRollbackTests.test_multi_branch_prefix_tree_cache_switch_and_restore: verifies Branch A -> Branch B -> Branch A switching reuses the shared prefix on A -> B and restores 100% of Branch A on B -> A with bit-exact global and sliding KV parity.
  • GrammarTests.test_tools_compilation_is_memoized: verifies ConstraintCompiler.tools memoizes compiled structural tags across equivalent tool schemas.
  • SharedKVTests.test_write_async_batches_writes_and_touch_resident_pages: verifies KVStorage.write_async + Engine.touch_resident_pages.
• uvx ruff check . & uvx ruff format --check .: All checks passed.