perf(rec_mlx): exact sub-4K prefix rewind, multi-branch LRU cache, and wired Metal residency

engine/src/rec_mlx/cache.py

@@ -15,14 +15,18 @@
 
 
 # MLX deepcopy can share mutable Metal inputs; copy into independent storage.
-def copy_array(value):
-    snapshot = _copy(
+def copy_array_async(value):
+    return _copy(
         inputs=[value, mx.array([value.size], mx.uint32)],
         grid=(value.size, 1, 1),
         threadgroup=(256, 1, 1),
         output_shapes=[value.shape],
         output_dtypes=[value.dtype],
     )[0]
+
+
+def copy_array(value):
+    snapshot = copy_array_async(value)
     mx.eval(snapshot)
     return snapshot
 
@@ -47,6 +51,7 @@ def copy_array(value):
 # Four append-only blocks fit alongside the fused decoder's weight/scratch buffers.
 BLOCK_TOKENS = 65536
 MAX_BLOCKS = 4
+MAX_REWIND_TOKENS = 4096
 
 
 class KVStorage:
@@ -94,11 +99,12 @@ def layer(self, index, count=None):
             parts.append(part.reshape(2, 4, cap, 128)[:, :, :size])
         return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=2)
 
-    def write(self, index, offset, keys, values):
+    def write_async(self, index, offset, keys, values):
         start, cap = self.rows[index][61:63]
         count = keys.shape[-2]
         assert 0 <= offset < offset + count <= (self.capacity if self.rows[index][1] else cap)
         begin = 0
+        dones = []
         while begin < count:
             block, position = divmod(offset + begin, cap)
             size = min(count - begin, cap - position)
@@ -114,8 +120,12 @@ def write(self, index, offset, keys, values):
                 output_shapes=[(1,)],
                 output_dtypes=[mx.uint32],
             )
-            mx.eval(result)
+            dones.extend(result)
             begin += size
+        return dones
+
+    def write(self, index, offset, keys, values):
+        mx.eval(self.write_async(index, offset, keys, values))
 
 
 class FixedKVCache:
@@ -172,6 +182,8 @@ def __init__(
         self._valid = 0
         self._slot = 0
         self._ring_index = 0
+        self._history_k = None
+        self._history_v = None
 
     def grow(self, capacity):
         capacity = min(self.window_size, capacity)
@@ -183,6 +195,8 @@ def grow(self, capacity):
 
     @property
     def state(self):
+        if self._history_k is not None:
+            return self.key_buffers, self.value_buffers, self._history_k, self._history_v
         return self.key_buffers, self.value_buffers
 
     @property
@@ -198,6 +212,56 @@ def size(self) -> int:
     def is_trimmable(self) -> bool:
         return False
 
+    def _record_history(self, keys: mx.array, values: mx.array, prev_offset: int):
+        if prev_offset >= MAX_REWIND_TOKENS:
+            return
+        keep = min(keys.shape[-2], MAX_REWIND_TOKENS - prev_offset)
+        if keep <= 0:
+            return
+        k_slice = keys[..., :keep, :]
+        v_slice = values[..., :keep, :]
+        if prev_offset == 0 or self._history_k is None:
+            self._history_k = k_slice
+            self._history_v = v_slice
+        else:
+            self._history_k = mx.concatenate(
+                (self._history_k[..., :prev_offset, :], k_slice), axis=-2
+            )
+            self._history_v = mx.concatenate(
+                (self._history_v[..., :prev_offset, :], v_slice), axis=-2
+            )
+
+    def can_rewind_to(self, target_pos: int) -> bool:
+        if target_pos <= 0 or target_pos > self.offset:
+            return False
+        if self.offset <= self.capacity and self._ring_index == 0:
+            return True
+        return self._history_k is not None and target_pos <= self._history_k.shape[-2]
+
+    def rewind_to(self, target_pos: int) -> bool:
+        if not self.can_rewind_to(target_pos):
+            return False
+        retained = min(self.capacity, target_pos)
+        next_slot = 1 - self._slot
+        if self.offset <= self.capacity and self._ring_index == 0:
+            src_k = self.key_buffers[self._slot][..., :retained, :]
+            src_v = self.value_buffers[self._slot][..., :retained, :]
+        else:
+            src_k = self._history_k[..., target_pos - retained : target_pos, :]
+            src_v = self._history_v[..., target_pos - retained : target_pos, :]
+        if self._history_k is not None and target_pos < self._history_k.shape[-2]:
+            self._history_k = self._history_k[..., :target_pos, :]
+            self._history_v = self._history_v[..., :target_pos, :]
+        self.key_buffers[next_slot][..., :retained, :] = src_k
+        self.value_buffers[next_slot][..., :retained, :] = src_v
+        self.key_buffers = tuple(mx.contiguous(value) for value in self.key_buffers)
+        self.value_buffers = tuple(mx.contiguous(value) for value in self.value_buffers)
+        self._slot = next_slot
+        self._valid = retained
+        self._ring_index = 0
+        self.offset = target_pos
+        return True
+
     def make_mask(self, n_tokens: int, window_size=None, return_array: bool = False):
         window = self.window_size if window_size is None else window_size
         if n_tokens == 1 and not return_array:
@@ -225,6 +289,7 @@ def _prior(self):
         return keys[..., : self._valid, :], values[..., : self._valid, :]
 
     def update_and_fetch(self, keys: mx.array, values: mx.array):
+        self._record_history(keys, values, self.offset)
         if keys.shape[-2] == 1:
             active_keys = self.key_buffers[self._slot]
             active_values = self.value_buffers[self._slot]

engine/src/rec_mlx/engine.py

@@ -75,6 +75,10 @@ def __init__(
         self.context_limit = context_limit
         self._prompt_cache = None
         self._prefix_preparation = None
+        from .prefix import PrefixTreeCache
+
+        self._branch_cache = PrefixTreeCache()
+        self._touch_kernel = None
         self.mx = mx
         self.tokenizer = load_tokenizer(directory)
         self._constraints = None
@@ -84,6 +88,12 @@ def __init__(
         if not self.eos:
             raise ValueError("Checkpoint has no EOS tokens")
         start = time.perf_counter()
+        wired_limit = min(
+            int(device.get("max_recommended_working_set_size", 0)),
+            max(0, memory - 8 * 1024**3),
+        )
+        if wired_limit > 0:
+            mx.set_wired_limit(wired_limit)
         mx.set_cache_limit(0)
         mx.clear_cache()
         mx.reset_peak_memory()
@@ -101,6 +111,7 @@ def __init__(
                 self.tokenizer.chat_template.encode()
             ).hexdigest(),
             "allocation_cache_limit_bytes": cache_limit,
+            "wired_limit_bytes": wired_limit,
             "load_seconds": time.perf_counter() - start,
             "load_peak_mlx_bytes": mx.get_peak_memory(),
             "process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
@@ -202,6 +213,49 @@ def prepare_prefix(self, *, should_stop=lambda: False):
             mx.synchronize()
         return {"prepared_tokens": prepared, "preparation_seconds": time.perf_counter() - start}
 
+    def touch_resident_pages(self):
+        mx = self.mx
+        if getattr(self, "_touch_kernel", None) is None:
+            self._touch_kernel = mx.fast.metal_kernel(
+                name="flint_touch_resident_pages",
+                input_names=["buf", "num_pages"],
+                output_names=["out"],
+                source="""
+                    uint tid = thread_position_in_grid.x;
+                    uint n = num_pages[0];
+                    uint acc = 0u;
+                    for (uint p = tid; p < n; p += 1024u) {
+                        acc ^= ((const device uint*)buf)[p * 4096u];
+                    }
+                    if (tid == 0u && acc == 0xdeadbeefu) {
+                        out[0] = acc;
+                    } else if (tid == 0u) {
+                        out[0] = 1u;
+                    }
+                """,
+            )
+        arrays = list(getattr(self.runtime, "banks", ()))
+        if self._prompt_cache is not None and hasattr(self._prompt_cache.storage, "banks"):
+            arrays.extend(self._prompt_cache.storage.banks)
+        if not arrays:
+            return
+        with mx.stream(mx.new_stream(mx.gpu)):
+            outs = []
+            for arr in arrays:
+                if arr is not None and arr.nbytes >= 65536:
+                    pages = max(1, arr.nbytes // 16384)
+                    outs.append(
+                        self._touch_kernel(
+                            inputs=[arr, mx.array([pages], dtype=mx.uint32)],
+                            grid=(min(1024, pages), 1, 1),
+                            threadgroup=(min(256, pages), 1, 1),
+                            output_shapes=[(1,)],
+                            output_dtypes=[mx.uint32],
+                        )[0]
+                    )
+            if outs:
+                mx.eval(outs)
+
     def _schedule_prefix(self, ids, generated):
         from .prefix import PrefixPreparation
 
@@ -211,6 +265,7 @@ def _schedule_prefix(self, ids, generated):
     def _take_prefix(self, ids):
         prefix, self._prompt_cache = self._prompt_cache, None
         pending, self._prefix_preparation = self._prefix_preparation, None
+        branch_cache = getattr(self, "_branch_cache", None)
         if pending is not None and pending.cache is not None:
             prepared = pending.cache
             if ids[: len(prepared.tokens)] != prepared.tokens:
@@ -219,6 +274,8 @@ def _take_prefix(self, ids):
                 prepared, _ = select_prefix(prepared, ids)
             if prepared is not None and len(prepared.tokens) > len(pending.origin.tokens):
                 prefix = prepared
+        if branch_cache is not None:
+            prefix = branch_cache.select_best(ids, prefix)
         return prefix
 
     def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_logprobs=False):
@@ -261,6 +318,9 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l
                 raise
             del prefix, preparation
             self._prompt_cache = session.prompt_cache
+            branch_cache = getattr(self, "_branch_cache", None)
+            if branch_cache is not None:
+                branch_cache.save(self._prompt_cache)
             text_stream = StreamingText(self.tokenizer)
             if sampling.presence_penalty:
                 session._state["presence_mask"] = mx.zeros(logits.size, mx.float32)
@@ -394,4 +454,5 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l
                 self._schedule_prefix(ids, generated)
             mx.synchronize()
             del session
-            mx.clear_cache()
+            if mx.get_cache_memory() > 2 * 1024**3:
+                mx.clear_cache()

engine/src/rec_mlx/grammar.py

@@ -1,3 +1,5 @@
+import json
+import threading
 import time
 from dataclasses import dataclass
 
@@ -29,17 +31,34 @@ def __init__(self, tokenizer, *, enforce_tool_schemas=False):
             tokenizer, vocab_size=VOCAB, stop_token_ids=[tokenizer.eos_token_id]
         )
         self.compiler = xgr.GrammarCompiler(info, cache_limit_bytes=128 * 1024**2)
+        self._tools_cache = {}
+        self._tools_lock = threading.Lock()
+        self._stops = tuple(
+            self.tokenizer.convert_tokens_to_ids(x)
+            for x in ("<|end_of_text|>", "<|user|>", "<|observation|>")
+        )
 
     def tools(self, tools=None):
+        try:
+            key = (
+                None if tools is None else json.dumps(tools, sort_keys=True, separators=(",", ":"))
+            )
+        except TypeError:
+            key = repr(tools)
+        with self._tools_lock:
+            cached = self._tools_cache.get(key)
+            if cached is not None:
+                return cached
         tag = build_flint_glm_grammar(
             self.tokenizer, tools, enforce_tool_schemas=self.enforce_tool_schemas
         )
         compiled = self.compiler.compile_structural_tag(tag)
-        stops = tuple(
-            self.tokenizer.convert_tokens_to_ids(x)
-            for x in ("<|end_of_text|>", "<|user|>", "<|observation|>")
-        )
-        return Constraint(compiled, stops)
+        constraint = Constraint(compiled, self._stops)
+        with self._tools_lock:
+            if len(self._tools_cache) >= 64:
+                self._tools_cache.pop(next(iter(self._tools_cache)))
+            self._tools_cache[key] = constraint
+        return constraint
 
 
 class ConstraintSession: