"""MLX backend: real model loading, prefill to KV cache, logits, greedy generation, memory/timing instrumentation. Everything measured here is produced by actual mlx_lm model execution on Apple Silicon. """ from __future__ import annotations import copy, json, time, os from dataclasses import dataclass, field, asdict from pathlib import Path from typing import Any, Optional import mlx.core as mx import mlx.nn as nn from mlx.utils import tree_map import psutil from mlx_lm import load as _mlx_load from mlx_lm.models import cache as mcache from mlx_lm.sample_utils import make_sampler from ..util import sha256_text, sha256_file, Timer def synchronize(objs=None): """Force evaluation of lazy MLX arrays so timings include real computation.""" if objs is None: mx.synchronize() return mx.eval(objs) mx.synchronize() def memory_snapshot() -> dict: proc = psutil.Process(os.getpid()) mi = proc.memory_info() return { "mlx_active_bytes": int(mx.get_active_memory()), "mlx_peak_bytes": int(mx.get_peak_memory()), "mlx_cache_bytes": int(mx.get_cache_memory()), "process_rss_bytes": int(mi.rss), } def reset_peak(): mx.reset_peak_memory() @dataclass class ModelInfo: repo_id: str snapshot_path: str revision: str model_type: str num_layers: int num_kv_heads: int num_attention_heads: int head_dim: int hidden_size: int vocab_size: int max_position_embeddings: Optional[int] rope_theta: Optional[float] rope_scaling: Any sliding_window: Any weight_dtype: str weight_quantization: Any tokenizer_hash: str chat_template_hash: str config_hash: str weight_files_sha256: dict mlx_version: str mlx_lm_version: str bos_token_id: Any eos_token_ids: list def to_dict(self): return asdict(self) class MLXModel: """A loaded model with tokenizer and provenance metadata.""" def __init__(self, repo_id: str, lazy: bool = False): import mlx_lm self.repo_id = repo_id t0 = time.perf_counter() self.model, self.tokenizer = _mlx_load(repo_id, lazy=lazy) self.load_ms = (time.perf_counter() - t0) * 1000.0 from huggingface_hub import snapshot_download path = Path(snapshot_download(repo_id)) self.path = path cfg = json.loads((path / "config.json").read_text()) self.config = cfg tok_files = sorted([p for p in path.iterdir() if p.name.startswith("tokenizer") or p.name in ("vocab.json", "merges.txt", "special_tokens_map.json")]) tok_hash = sha256_text("|".join(f"{p.name}:{sha256_file(p)}" for p in tok_files)) tmpl = getattr(self.tokenizer, "chat_template", None) or "" weight_hashes = {p.name: sha256_file(p) for p in sorted(path.glob("*.safetensors"))} hd = cfg.get("head_dim") or (cfg["hidden_size"] // cfg["num_attention_heads"]) eos = self.tokenizer.eos_token_ids if hasattr(self.tokenizer, "eos_token_ids") else [self.tokenizer.eos_token_id] self.info = ModelInfo( repo_id=repo_id, snapshot_path=str(path), revision=path.name, model_type=cfg.get("model_type", "?"), num_layers=cfg["num_hidden_layers"], num_kv_heads=cfg.get("num_key_value_heads", cfg["num_attention_heads"]), num_attention_heads=cfg["num_attention_heads"], head_dim=hd, hidden_size=cfg["hidden_size"], vocab_size=cfg["vocab_size"], max_position_embeddings=cfg.get("max_position_embeddings"), rope_theta=cfg.get("rope_theta"), rope_scaling=cfg.get("rope_scaling"), sliding_window=cfg.get("sliding_window"), weight_dtype=cfg.get("torch_dtype", "?"), weight_quantization=cfg.get("quantization"), tokenizer_hash=tok_hash, chat_template_hash=sha256_text(tmpl), config_hash=sha256_text(json.dumps(cfg, sort_keys=True)), weight_files_sha256=weight_hashes, mlx_version=mx.__version__, mlx_lm_version=mlx_lm.__version__, bos_token_id=getattr(self.tokenizer, "bos_token_id", None), eos_token_ids=list(eos) if eos is not None else [], ) # ---- tokenization ------------------------------------------------- def chat_prompt_tokens(self, messages: list[dict], add_generation_prompt: bool = True) -> list[int]: """Apply the model's own chat template and return token ids (no double BOS).""" ids = self.tokenizer.apply_chat_template(messages, add_generation_prompt=add_generation_prompt, tokenize=True) if hasattr(ids, "input_ids"): ids = ids["input_ids"] return list(ids) def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: return list(self.tokenizer.encode(text, add_special_tokens=add_special_tokens)) def decode(self, ids: list[int]) -> str: return self.tokenizer.decode(ids) # ---- cache ops ------------------------------------------------------ def new_cache(self) -> list: return mcache.make_prompt_cache(self.model) def prefill(self, token_ids: list[int], cache: Optional[list] = None, step: int = 2048) -> tuple[list, dict]: """Run the model over token_ids, filling `cache`. Returns (cache, timing dict). Fully synchronized.""" if cache is None: cache = self.new_cache() toks = mx.array(token_ids, dtype=mx.int32) t = Timer(); t.start("prefill") n = 0 while n < len(token_ids): chunk = toks[n:n + step] self.model(chunk[None], cache=cache) mx.eval([c.state for c in cache]) n += chunk.shape[0] mx.synchronize() t.stop("prefill") return cache, {"prefill_ms": t.spans["prefill"], "prefill_tokens": len(token_ids)} def next_logits(self, cache: list, token_ids: list[int]) -> mx.array: """Logits for the last position after consuming token_ids on top of cache (cache mutated).""" toks = mx.array(token_ids, dtype=mx.int32) logits = self.model(toks[None], cache=cache)[:, -1, :] mx.eval(logits) mx.synchronize() return logits[0].astype(mx.float32) def generate_from_cache(self, cache: list, query_ids: list[int], max_tokens: int = 64, temp: float = 0.0, stop_on_eos: bool = True) -> dict: """Greedy (temp=0) generation: prefill query_ids into cache, then decode. Measures TTFT precisely (time to the first *content* token) and total decode time. Cache is mutated in place.""" sampler = make_sampler(temp) t = Timer() out_ids: list[int] = [] t.start("query_prefill") toks = mx.array(query_ids, dtype=mx.int32) # Heterogeneous retained lengths across layers (e.g. PyramidKV budgets) cannot share one attention mask # (mlx_lm builds the mask from cache[0]); process the query tokens one at a time in that case (mask=None). kept = {getattr(c, "n_kept", None) for c in cache} heterogeneous = len(kept) > 1 if heterogeneous and len(query_ids) > 1: for i in range(len(query_ids) - 1): self.model(toks[i:i + 1][None], cache=cache) mx.eval([c.keys for c in cache]) toks = toks[-1:] logits = self.model(toks[None], cache=cache)[:, -1, :] logprobs = logits - mx.logsumexp(logits, keepdims=True) y = sampler(logprobs) mx.eval(y) mx.synchronize() t.stop("query_prefill") ttft_ms = t.spans["query_prefill"] t.start("decode") tok = int(y.item()) eos = set(self.info.eos_token_ids) n = 0 while True: if stop_on_eos and tok in eos: break out_ids.append(tok) n += 1 if n >= max_tokens: break logits = self.model(y[None], cache=cache)[:, -1, :] logprobs = logits - mx.logsumexp(logits, keepdims=True) y = sampler(logprobs) mx.eval(y) tok = int(y.item()) mx.synchronize() t.stop("decode") text = self.decode(out_ids) return {"text": text, "output_ids": out_ids, "output_tokens": len(out_ids), "ttft_ms": ttft_ms, "sequential_query_prefill": heterogeneous, "query_prefill_ms": t.spans["query_prefill"], "decode_ms": t.spans["decode"], "query_tokens": len(query_ids)} def generate_text(self, token_ids: list[int], max_tokens: int = 64, temp: float = 0.0) -> dict: """Full fresh prefill + generation from a token list (no reusable cache).""" cache = self.new_cache() # prefill all but the last token, then let generate_from_cache consume the last one -> consistent TTFT boundary if len(token_ids) > 1: _, pre = self.prefill(token_ids[:-1], cache) else: pre = {"prefill_ms": 0.0, "prefill_tokens": 0} g = self.generate_from_cache(cache, token_ids[-1:], max_tokens=max_tokens, temp=temp) g.update(pre) return g from ..cache.sparse import copy_cache as deep_copy_cache, cache_tensor_bytes # noqa: E402