1"""MLX backend: real model loading, prefill to KV cache, logits, greedy generation, memory/timing instrumentation.23Everything measured here is produced by actual mlx_lm model execution on Apple Silicon.4"""5from __future__ import annotations6import copy, json, time, os7from dataclasses import dataclass, field, asdict8from pathlib import Path9from typing import Any, Optional1011import mlx.core as mx12import mlx.nn as nn13from mlx.utils import tree_map14import psutil15from mlx_lm import load as _mlx_load16from mlx_lm.models import cache as mcache17from mlx_lm.sample_utils import make_sampler1819from ..util import sha256_text, sha256_file, Timer202122def synchronize(objs=None):23 """Force evaluation of lazy MLX arrays so timings include real computation."""24 if objs is None:25 mx.synchronize()26 return27 mx.eval(objs)28 mx.synchronize()293031def memory_snapshot() -> dict:32 proc = psutil.Process(os.getpid())33 mi = proc.memory_info()34 return {35 "mlx_active_bytes": int(mx.get_active_memory()),36 "mlx_peak_bytes": int(mx.get_peak_memory()),37 "mlx_cache_bytes": int(mx.get_cache_memory()),38 "process_rss_bytes": int(mi.rss),39 }404142def reset_peak():43 mx.reset_peak_memory()444546@dataclass47class ModelInfo:48 repo_id: str49 snapshot_path: str50 revision: str51 model_type: str52 num_layers: int53 num_kv_heads: int54 num_attention_heads: int55 head_dim: int56 hidden_size: int57 vocab_size: int58 max_position_embeddings: Optional[int]59 rope_theta: Optional[float]60 rope_scaling: Any61 sliding_window: Any62 weight_dtype: str63 weight_quantization: Any64 tokenizer_hash: str65 chat_template_hash: str66 config_hash: str67 weight_files_sha256: dict68 mlx_version: str69 mlx_lm_version: str70 bos_token_id: Any71 eos_token_ids: list7273 def to_dict(self):74 return asdict(self)757677class MLXModel:78 """A loaded model with tokenizer and provenance metadata."""7980 def __init__(self, repo_id: str, lazy: bool = False):81 import mlx_lm82 self.repo_id = repo_id83 t0 = time.perf_counter()84 self.model, self.tokenizer = _mlx_load(repo_id, lazy=lazy)85 self.load_ms = (time.perf_counter() - t0) * 1000.086 from huggingface_hub import snapshot_download87 path = Path(snapshot_download(repo_id))88 self.path = path89 cfg = json.loads((path / "config.json").read_text())90 self.config = cfg91 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")])92 tok_hash = sha256_text("|".join(f"{p.name}:{sha256_file(p)}" for p in tok_files))93 tmpl = getattr(self.tokenizer, "chat_template", None) or ""94 weight_hashes = {p.name: sha256_file(p) for p in sorted(path.glob("*.safetensors"))}95 hd = cfg.get("head_dim") or (cfg["hidden_size"] // cfg["num_attention_heads"])96 eos = self.tokenizer.eos_token_ids if hasattr(self.tokenizer, "eos_token_ids") else [self.tokenizer.eos_token_id]97 self.info = ModelInfo(98 repo_id=repo_id, snapshot_path=str(path), revision=path.name,99 model_type=cfg.get("model_type", "?"), num_layers=cfg["num_hidden_layers"],100 num_kv_heads=cfg.get("num_key_value_heads", cfg["num_attention_heads"]),101 num_attention_heads=cfg["num_attention_heads"], head_dim=hd,102 hidden_size=cfg["hidden_size"], vocab_size=cfg["vocab_size"],103 max_position_embeddings=cfg.get("max_position_embeddings"),104 rope_theta=cfg.get("rope_theta"), rope_scaling=cfg.get("rope_scaling"),105 sliding_window=cfg.get("sliding_window"),106 weight_dtype=cfg.get("torch_dtype", "?"), weight_quantization=cfg.get("quantization"),107 tokenizer_hash=tok_hash, chat_template_hash=sha256_text(tmpl), config_hash=sha256_text(json.dumps(cfg, sort_keys=True)),108 weight_files_sha256=weight_hashes, mlx_version=mx.__version__, mlx_lm_version=mlx_lm.__version__,109 bos_token_id=getattr(self.tokenizer, "bos_token_id", None), eos_token_ids=list(eos) if eos is not None else [],110 )111112 # ---- tokenization -------------------------------------------------113 def chat_prompt_tokens(self, messages: list[dict], add_generation_prompt: bool = True) -> list[int]:114 """Apply the model's own chat template and return token ids (no double BOS)."""115 ids = self.tokenizer.apply_chat_template(messages, add_generation_prompt=add_generation_prompt, tokenize=True)116 if hasattr(ids, "input_ids"):117 ids = ids["input_ids"]118 return list(ids)119120 def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:121 return list(self.tokenizer.encode(text, add_special_tokens=add_special_tokens))122123 def decode(self, ids: list[int]) -> str:124 return self.tokenizer.decode(ids)125126 # ---- cache ops ------------------------------------------------------127 def new_cache(self) -> list:128 return mcache.make_prompt_cache(self.model)129130 def prefill(self, token_ids: list[int], cache: Optional[list] = None, step: int = 2048) -> tuple[list, dict]:131 """Run the model over token_ids, filling `cache`. Returns (cache, timing dict). Fully synchronized."""132 if cache is None:133 cache = self.new_cache()134 toks = mx.array(token_ids, dtype=mx.int32)135 t = Timer(); t.start("prefill")136 n = 0137 while n < len(token_ids):138 chunk = toks[n:n + step]139 self.model(chunk[None], cache=cache)140 mx.eval([c.state for c in cache])141 n += chunk.shape[0]142 mx.synchronize()143 t.stop("prefill")144 return cache, {"prefill_ms": t.spans["prefill"], "prefill_tokens": len(token_ids)}145146 def next_logits(self, cache: list, token_ids: list[int]) -> mx.array:147 """Logits for the last position after consuming token_ids on top of cache (cache mutated)."""148 toks = mx.array(token_ids, dtype=mx.int32)149 logits = self.model(toks[None], cache=cache)[:, -1, :]150 mx.eval(logits)151 mx.synchronize()152 return logits[0].astype(mx.float32)153154 def generate_from_cache(self, cache: list, query_ids: list[int], max_tokens: int = 64, temp: float = 0.0,155 stop_on_eos: bool = True) -> dict:156 """Greedy (temp=0) generation: prefill query_ids into cache, then decode. Measures TTFT precisely157 (time to the first *content* token) and total decode time. Cache is mutated in place."""158 sampler = make_sampler(temp)159 t = Timer()160 out_ids: list[int] = []161 t.start("query_prefill")162 toks = mx.array(query_ids, dtype=mx.int32)163 # Heterogeneous retained lengths across layers (e.g. PyramidKV budgets) cannot share one attention mask164 # (mlx_lm builds the mask from cache[0]); process the query tokens one at a time in that case (mask=None).165 kept = {getattr(c, "n_kept", None) for c in cache}166 heterogeneous = len(kept) > 1167 if heterogeneous and len(query_ids) > 1:168 for i in range(len(query_ids) - 1):169 self.model(toks[i:i + 1][None], cache=cache)170 mx.eval([c.keys for c in cache])171 toks = toks[-1:]172 logits = self.model(toks[None], cache=cache)[:, -1, :]173 logprobs = logits - mx.logsumexp(logits, keepdims=True)174 y = sampler(logprobs)175 mx.eval(y)176 mx.synchronize()177 t.stop("query_prefill")178 ttft_ms = t.spans["query_prefill"]179 t.start("decode")180 tok = int(y.item())181 eos = set(self.info.eos_token_ids)182 n = 0183 while True:184 if stop_on_eos and tok in eos:185 break186 out_ids.append(tok)187 n += 1188 if n >= max_tokens:189 break190 logits = self.model(y[None], cache=cache)[:, -1, :]191 logprobs = logits - mx.logsumexp(logits, keepdims=True)192 y = sampler(logprobs)193 mx.eval(y)194 tok = int(y.item())195 mx.synchronize()196 t.stop("decode")197 text = self.decode(out_ids)198 return {"text": text, "output_ids": out_ids, "output_tokens": len(out_ids), "ttft_ms": ttft_ms, "sequential_query_prefill": heterogeneous,199 "query_prefill_ms": t.spans["query_prefill"], "decode_ms": t.spans["decode"], "query_tokens": len(query_ids)}200201 def generate_text(self, token_ids: list[int], max_tokens: int = 64, temp: float = 0.0) -> dict:202 """Full fresh prefill + generation from a token list (no reusable cache)."""203 cache = self.new_cache()204 # prefill all but the last token, then let generate_from_cache consume the last one -> consistent TTFT boundary205 if len(token_ids) > 1:206 _, pre = self.prefill(token_ids[:-1], cache)207 else:208 pre = {"prefill_ms": 0.0, "prefill_tokens": 0}209 g = self.generate_from_cache(cache, token_ids[-1:], max_tokens=max_tokens, temp=temp)210 g.update(pre)211 return g212213214from ..cache.sparse import copy_cache as deep_copy_cache, cache_tensor_bytes # noqa: E402215