diff --git a/engine/src/rec_mlx/cache.py b/engine/src/rec_mlx/cache.py index b9a9a5b..66fd3dd 100644 --- a/engine/src/rec_mlx/cache.py +++ b/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] diff --git a/engine/src/rec_mlx/engine.py b/engine/src/rec_mlx/engine.py index 6ba06f1..9e14bf9 100644 --- a/engine/src/rec_mlx/engine.py +++ b/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() diff --git a/engine/src/rec_mlx/grammar.py b/engine/src/rec_mlx/grammar.py index 1a49410..1bc0f3b 100644 --- a/engine/src/rec_mlx/grammar.py +++ b/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: diff --git a/engine/src/rec_mlx/prefix.py b/engine/src/rec_mlx/prefix.py index 4e292ed..d92ea62 100644 --- a/engine/src/rec_mlx/prefix.py +++ b/engine/src/rec_mlx/prefix.py @@ -4,7 +4,7 @@ import mlx.core as mx -from .cache import FixedSlidingKVCache, copy_array +from .cache import MAX_REWIND_TOKENS, FixedSlidingKVCache, copy_array_async class PrefillCancelled(Exception): @@ -18,6 +18,184 @@ class Checkpoint: logits: mx.array +@dataclass +class BranchSnapshot: + tokens: list[int] + global_kv: dict[int, mx.array] + sliding: list[tuple] + logits: mx.array | None + + +class PrefixTreeCache: + """Multi-branch LRU prefix cache for fast switching across files and sessions.""" + + def __init__(self, max_branches: int = 4): + self.max_branches = max_branches + self.branches: list[BranchSnapshot] = [] + self.base_branch: BranchSnapshot | None = None + + def clear(self): + self.branches.clear() + self.base_branch = None + + @staticmethod + def _lcp(a: list[int], b: list[int]) -> int: + common = 0 + for x, y in zip(a, b): + if x != y: + break + common += 1 + return common + + def _snapshot(self, cache: "PromptCache", target_len: int) -> BranchSnapshot | None: + storage = cache.storage + if not (hasattr(storage, "layer") and hasattr(storage, "write_async")): + return None + global_kv = {} + sliding = [] + pending = [] + for index, layer in enumerate(cache.layers): + if isinstance(layer, FixedSlidingKVCache): + if target_len == len(cache.tokens): + k_buf = copy_array_async(layer.key_buffers[layer._slot]) + v_buf = copy_array_async(layer.value_buffers[layer._slot]) + valid = layer._valid + ring_idx = layer._ring_index + hist_k = ( + copy_array_async(layer._history_k[..., :target_len, :]) + if layer._history_k is not None + else None + ) + hist_v = ( + copy_array_async(layer._history_v[..., :target_len, :]) + if layer._history_v is not None + else None + ) + else: + if not layer.can_rewind_to(target_len): + return None + retained = min(layer.capacity, target_len) + if layer.offset <= layer.capacity and layer._ring_index == 0: + src_k = layer.key_buffers[layer._slot][..., :retained, :] + src_v = layer.value_buffers[layer._slot][..., :retained, :] + else: + src_k = layer._history_k[..., target_len - retained : target_len, :] + src_v = layer._history_v[..., target_len - retained : target_len, :] + k_buf = mx.zeros_like(layer.key_buffers[0]) + v_buf = mx.zeros_like(layer.value_buffers[0]) + k_buf[..., :retained, :] = src_k + v_buf[..., :retained, :] = src_v + k_buf = mx.contiguous(k_buf) + v_buf = mx.contiguous(v_buf) + valid = retained + ring_idx = 0 + hist_k = ( + copy_array_async(layer._history_k[..., :target_len, :]) + if layer._history_k is not None + else None + ) + hist_v = ( + copy_array_async(layer._history_v[..., :target_len, :]) + if layer._history_v is not None + else None + ) + pending.extend([k_buf, v_buf]) + if hist_k is not None: + pending.extend([hist_k, hist_v]) + sliding.append((k_buf, v_buf, valid, ring_idx, hist_k, hist_v)) + else: + kv = copy_array_async(storage.layer(index, target_len)) + global_kv[index] = kv + pending.append(kv) + logits = cache.logits if target_len == len(cache.tokens) else None + if logits is not None: + pending.append(logits) + if pending: + mx.eval(pending) + return BranchSnapshot( + tokens=cache.tokens[:target_len], + global_kv=global_kv, + sliding=sliding, + logits=logits, + ) + + def save(self, cache: "PromptCache"): + if cache is None or not cache.tokens: + return + n_tok = len(cache.tokens) + if n_tok > MAX_REWIND_TOKENS: + return + for i, branch in enumerate(self.branches): + if branch.tokens == cache.tokens: + branch.logits = cache.logits + self.branches.append(self.branches.pop(i)) + return + snap = self._snapshot(cache, n_tok) + if snap is None: + return + self.branches.append(snap) + if len(self.branches) > self.max_branches: + self.branches.pop(0) + if self.base_branch is None and n_tok >= 256: + base_len = min(n_tok, 1024) + self.base_branch = snap if base_len == n_tok else self._snapshot(cache, base_len) + + def select_best(self, tokens: list[int], active_cache: "PromptCache | None"): + if active_cache is None: + self.clear() + return None + active_common = self._lcp(active_cache.tokens, tokens) + best_snap = None + best_common = active_common + candidates = list(reversed(self.branches)) + if self.base_branch is not None: + candidates.append(self.base_branch) + for snap in candidates: + common = self._lcp(snap.tokens, tokens) + if common > best_common: + best_common = common + best_snap = snap + if best_snap is None: + return active_cache + shared_with_active = self._lcp(active_cache.tokens, best_snap.tokens) + target_len = len(best_snap.tokens) + storage = active_cache.storage + dones = [] + if shared_with_active < target_len: + for index, kv in best_snap.global_kv.items(): + k_slice = kv[0:1, :, shared_with_active:target_len, :] + v_slice = kv[1:2, :, shared_with_active:target_len, :] + dones.extend(storage.write_async(index, shared_with_active, k_slice, v_slice)) + sliding_iter = iter(best_snap.sliding) + eval_tensors = list(dones) + for layer in active_cache.layers: + layer.offset = target_len + if isinstance(layer, FixedSlidingKVCache): + k_buf, v_buf, valid, ring_idx, hist_k, hist_v = next(sliding_iter) + k_copy = copy_array_async(k_buf) + v_copy = copy_array_async(v_buf) + layer.capacity = k_copy.shape[-2] + layer.key_buffers = (k_copy, mx.zeros_like(k_copy)) + layer.value_buffers = (v_copy, mx.zeros_like(v_copy)) + layer._slot = 0 + layer._valid = valid + layer._ring_index = ring_idx + layer._history_k = copy_array_async(hist_k) if hist_k is not None else None + layer._history_v = copy_array_async(hist_v) if hist_v is not None else None + eval_tensors.extend([layer.key_buffers[0], layer.value_buffers[0]]) + if layer._history_k is not None: + eval_tensors.extend([layer._history_k, layer._history_v]) + if eval_tensors: + mx.eval(eval_tensors) + active_cache.tokens = best_snap.tokens[:] + active_cache.logits = best_snap.logits + active_cache.checkpoints = [] + if best_snap in self.branches: + self.branches.remove(best_snap) + self.branches.append(best_snap) + return active_cache + + @dataclass class PromptCache: tokens: list[int] @@ -28,13 +206,21 @@ class PromptCache: def fork(self): layers = [copy.copy(layer) for layer in self.layers] + pending = [] for layer in layers: if isinstance(layer, FixedSlidingKVCache): - keys = copy_array(layer.key_buffers[layer._slot]) - values = copy_array(layer.value_buffers[layer._slot]) + keys = copy_array_async(layer.key_buffers[layer._slot]) + values = copy_array_async(layer.value_buffers[layer._slot]) layer.key_buffers = (keys, mx.zeros_like(keys)) layer.value_buffers = (values, mx.zeros_like(values)) layer._slot = 0 + pending.extend([keys, values, layer.key_buffers[1], layer.value_buffers[1]]) + if layer._history_k is not None: + layer._history_k = copy_array_async(layer._history_k) + layer._history_v = copy_array_async(layer._history_v) + pending.extend([layer._history_k, layer._history_v]) + if pending: + mx.eval(pending) return PromptCache(self.tokens[:], layers, self.logits, self.storage, self.checkpoints[:]) def advance(self, tokens, logits): @@ -53,8 +239,8 @@ def checkpoint(self): if isinstance(layer, FixedSlidingKVCache): sliding.append( ( - copy.deepcopy(layer.key_buffers[layer._slot]), - copy.deepcopy(layer.value_buffers[layer._slot]), + copy_array_async(layer.key_buffers[layer._slot]), + copy_array_async(layer.value_buffers[layer._slot]), layer._valid, layer._ring_index, ) @@ -90,9 +276,7 @@ def select_prefix(prefix, tokens): checkpoint = next( (item for item in reversed(prefix.checkpoints) if item.position <= common), None ) - if checkpoint is None: - prefix = None - else: + if checkpoint is not None: sliding = iter(checkpoint.sliding) for layer in prefix.layers: layer.offset = checkpoint.position @@ -102,11 +286,39 @@ def select_prefix(prefix, tokens): layer.key_buffers = (keys, mx.zeros_like(keys)) layer.value_buffers = (values, mx.zeros_like(values)) layer._slot = 0 + if ( + layer._history_k is not None + and checkpoint.position < layer._history_k.shape[-2] + ): + layer._history_k = layer._history_k[..., : checkpoint.position, :] + layer._history_v = layer._history_v[..., : checkpoint.position, :] prefix.tokens = prefix.tokens[: checkpoint.position] prefix.logits = checkpoint.logits prefix.checkpoints = [ item for item in prefix.checkpoints if item.position < checkpoint.position ] + elif not prefix.checkpoints: + target = min(common, len(tokens) - 1) + sliding_layers = [ + layer for layer in prefix.layers if isinstance(layer, FixedSlidingKVCache) + ] + if ( + sliding_layers + and 0 < target <= MAX_REWIND_TOKENS + and all(layer.can_rewind_to(target) for layer in sliding_layers) + ): + for layer in prefix.layers: + if isinstance(layer, FixedSlidingKVCache): + layer.rewind_to(target) + else: + layer.offset = target + mx.eval([layer.state for layer in sliding_layers]) + prefix.tokens = prefix.tokens[:target] + prefix.logits = None + else: + prefix = None + else: + prefix = None cached = 0 if prefix is None else len(prefix.tokens) logging.info( "Prefill: prompt=%d common_prefix=%d reused=%d new=%d", diff --git a/engine/src/rec_mlx/quantized.py b/engine/src/rec_mlx/quantized.py index a204421..85391b5 100644 --- a/engine/src/rec_mlx/quantized.py +++ b/engine/src/rec_mlx/quantized.py @@ -138,10 +138,17 @@ def prefill( raise PrefillCancelled() # Global attention already uses decode storage. Only the bounded sliding # rings need copying; retain their native prefix for subsequent requests. + write_async = getattr(state["kv"], "write_async", None) + dones = [] for index, entry in enumerate(cache): if not state["kv"].rows[index][1]: keys, values = entry._prior() - state["kv"].write(index, 0, keys, values) + if write_async is not None: + dones.extend(write_async(index, 0, keys, values)) + else: + state["kv"].write(index, 0, keys, values) + if dones: + mx.eval(dones) state["capacity"] = state["kv"].capacity state["position"] = state["initial"] = len(tokens) self._state = state diff --git a/engine/src/rec_mlx/scheduler.py b/engine/src/rec_mlx/scheduler.py index 17e5a94..b6889db 100644 --- a/engine/src/rec_mlx/scheduler.py +++ b/engine/src/rec_mlx/scheduler.py @@ -1,6 +1,7 @@ import json import sys import threading +import time from dataclasses import asdict from queue import Empty, Queue from types import SimpleNamespace @@ -58,8 +59,20 @@ def __init__(self, engine, metrics_path, defaults): self.busy = threading.Event() self._request_lock = threading.Lock() self.active_context = None + self._last_touch = time.monotonic() super().__init__(ModelProvider(engine, self.tokenizer, cli_args), prompt_cache=None) + def _next_request(self, timeout=None): + request = super()._next_request(timeout=timeout) + if request is None and not self.stopping.is_set(): + touch = getattr(self.engine, "touch_resident_pages", None) + if touch is not None: + now = time.monotonic() + if now - self._last_touch >= 3.0: + self._last_touch = now + touch() + return request + def cancel(self): with self._request_lock: self.stopping.set() @@ -133,6 +146,7 @@ def _serve_single(self, queued_request): should_stop=lambda: self.stopping.is_set() or not self.requests.empty() ) finally: + self._last_touch = time.monotonic() self.busy.clear() def _start_request(self, request, args): diff --git a/engine/tests/test_grammar.py b/engine/tests/test_grammar.py index 2795fc3..f8223fe 100644 --- a/engine/tests/test_grammar.py +++ b/engine/tests/test_grammar.py @@ -271,6 +271,18 @@ def rejected(token, active): native.close(identifier) self.check_fixture(self.ebnf('root ::= "ok"\n'), "ok") + def test_tools_compilation_is_memoized(self): + tools = self.shell_tools(True) + first = self.compiler.tools(tools) + with unittest.mock.patch.object( + self.compiler.compiler, + "compile_structural_tag", + wraps=self.compiler.compiler.compile_structural_tag, + ) as compile_tag: + second = self.compiler.tools(deepcopy(tools)) + compile_tag.assert_not_called() + self.assertIs(first, second) + if __name__ == "__main__": unittest.main() diff --git a/engine/tests/test_prefix_rollback.py b/engine/tests/test_prefix_rollback.py index ee6e8d5..9fe7fc5 100644 --- a/engine/tests/test_prefix_rollback.py +++ b/engine/tests/test_prefix_rollback.py @@ -4,9 +4,9 @@ import mlx.core as mx -from rec_mlx.cache import FixedSlidingKVCache +from rec_mlx.cache import FixedKVCache, FixedSlidingKVCache, KVStorage from rec_mlx.engine import Engine, Sampling -from rec_mlx.prefix import PrefillCancelled, PromptCache, select_prefix +from rec_mlx.prefix import PrefillCancelled, PrefixTreeCache, PromptCache, select_prefix from rec_mlx.quantized import QATSession @@ -124,6 +124,70 @@ def test_generation_selects_prefix_once_before_cancelled_prefill(self): ["INFO:root:Prefill: prompt=11 common_prefix=9 reused=8 new=3"], ) + def test_subchunk_exact_rewind_without_checkpoints(self): + prefix = self.prefix() + self.append(prefix, list(range(12))) + self.assertEqual(prefix.checkpoints, []) + diverged = list(range(9)) + [100, 101] + restored, common = select_prefix(prefix, diverged) + self.assertIs(restored, prefix) + self.assertEqual(common, 9) + self.assertEqual(len(restored.tokens), 9) + self.assertEqual(restored.layers[0].offset, 9) + self.append(restored, [100, 101]) + fresh = self.prefix() + self.append(fresh, diverged) + for actual, wanted in zip( + restored.layers[1]._prior(), fresh.layers[1]._prior(), strict=True + ): + self.assertTrue(mx.array_equal(actual, wanted).item()) + + @unittest.skipUnless(mx.metal.is_available(), "requires Apple Metal") + def test_multi_branch_prefix_tree_cache_switch_and_restore(self): + metadata = [[0] * 64 for _ in range(50)] + metadata[0][1] = 1 + storage = KVStorage(metadata, 32) + global_layer = FixedKVCache(storage, 0) + sliding_layer = FixedSlidingKVCache( + max_sequence_tokens=16, num_kv_heads=4, head_dim=128, window_size=5 + ) + cache = PromptCache([], [global_layer, sliding_layer], None, storage) + tree = PrefixTreeCache(max_branches=4) + + def run_prompt(tokens): + nonlocal cache + selected = tree.select_best(tokens, cache) + restored, _ = select_prefix(selected, tokens) + cached = 0 if restored is None else len(restored.tokens) + suffix = tokens[cached:] + if suffix: + k = mx.broadcast_to( + mx.array(suffix, dtype=mx.bfloat16).reshape(1, 1, len(suffix), 1), + (1, 4, len(suffix), 128), + ) + k = mx.contiguous(k) + v = mx.contiguous(-k) + cache.layers[0].update_and_fetch(k, v) + cache.layers[1].update_and_fetch(k, v) + cache.tokens.extend(suffix) + cache.logits = mx.array([suffix[-1]], dtype=mx.bfloat16) + mx.eval(cache.logits, [layer.state for layer in cache.layers]) + tree.save(cache) + return cached + + branch_a = list(range(1, 11)) + branch_b = list(range(1, 7)) + [201, 202, 203] + self.assertEqual(run_prompt(branch_a), 0) + expected_global_a = mx.array(storage.layer(0, len(branch_a))) + expected_sliding_a = [mx.array(x) for x in cache.layers[1]._prior()] + mx.eval(expected_global_a, expected_sliding_a) + + self.assertEqual(run_prompt(branch_b), 6) + self.assertEqual(run_prompt(branch_a), len(branch_a)) + self.assertTrue(mx.array_equal(storage.layer(0, len(branch_a)), expected_global_a).item()) + for actual, wanted in zip(cache.layers[1]._prior(), expected_sliding_a, strict=True): + self.assertTrue(mx.array_equal(actual, wanted).item()) + if __name__ == "__main__": unittest.main() diff --git a/engine/tests/test_shared_kv.py b/engine/tests/test_shared_kv.py index bf3eb42..2a49c2d 100644 --- a/engine/tests/test_shared_kv.py +++ b/engine/tests/test_shared_kv.py @@ -78,6 +78,25 @@ def test_packet_growth_and_limit(self): with self.assertRaisesRegex(RuntimeError, "invalidated"): FullModel.ensure_capacity(engine, state, 1) + @unittest.skipUnless(mx.metal.is_available(), "requires Apple Metal") + def test_write_async_batches_writes_and_touch_resident_pages(self): + from rec_mlx.engine import Engine + + storage = self.storage() + ring = mx.arange(4 * 513 * 128).reshape(1, 4, 513, 128).astype(mx.bfloat16) + dones = storage.write_async(1, 0, ring, -ring) + mx.eval(dones) + self.assertTrue(mx.array_equal(storage.layer(1)[0:1], ring).item()) + self.assertTrue(mx.array_equal(storage.layer(1)[1:2], -ring).item()) + + engine = Engine.__new__(Engine) + engine.mx = mx + engine._touch_kernel = None + engine.runtime = SimpleNamespace(banks=[storage.banks[0]]) + engine._prompt_cache = SimpleNamespace(storage=storage) + engine.touch_resident_pages() + self.assertIsNotNone(engine._touch_kernel) + if __name__ == "__main__": unittest.main()