1"""Build a TITAN model from an MLX checkpoint: the host architecture (attention, norms, router) is2mlx-lm's own module tree; every layer's `switch_mlp` is replaced by a TitanSwitchGLU over a slot table.3Expert weights are never materialised from the checkpoint by mlx-lm: they are read on demand by the4ExpertStore. `capacity` = number of resident experts per layer (num_experts = everything resident).5"""6from __future__ import annotations78import os9from typing import Dict, List, Optional1011import mlx.core as mx12import mlx.nn as nn13from mlx_lm.utils import load_model, load_tokenizer1415from .moe import TitanSwitchGLU16from .registry import ModelSpec, load_spec17from .slots import EvictionPolicy, LFUPolicy, SlotTable18from .store import ExpertStore1920POLICIES = {"lru": EvictionPolicy, "lfu": LFUPolicy}212223class TitanModel:24 def __init__(self, model_dir: str, capacity: Optional[int] = None, policy: str = "lru", nocache: bool = False,25 on_route=None):26 self.model_dir = os.path.realpath(model_dir)27 self.spec: ModelSpec = load_spec(self.model_dir)28 self.store = ExpertStore(self.spec, nocache=nocache)29 self.capacity = capacity or self.spec.num_experts30 self.policy_name = policy31 # Load the host model lazily so the stacked expert tensors are never evaluated.32 from pathlib import Path33 self.model, cfg = load_model(Path(self.model_dir), lazy=True)34 self.tokenizer = load_tokenizer(Path(self.model_dir))35 q = cfg.get("quantization") or cfg.get("quantization_config")36 self.tables: List[SlotTable] = []37 self.blocks: List[TitanSwitchGLU] = []38 for l, layer in enumerate(self.model.model.layers):39 table = SlotTable(l, self.capacity, self.store, POLICIES[policy]())40 block = TitanSwitchGLU(table, q["group_size"], q["bits"], on_route=on_route)41 # drop the lazy stacked weights and splice in the TITAN block42 mlp = layer.mlp43 del mlp.switch_mlp # removes the lazy stacked SwitchGLU (never evaluated)44 mlp.switch_mlp = block # plain attribute: not an nn.Module, not in parameters()45 self.tables.append(table)46 self.blocks.append(block)47 mx.eval(self.model.parameters()) # dense weights only4849 # convenience -----------------------------------------------------------------50 @property51 def layers(self):52 return self.model.layers5354 def __call__(self, inputs, cache=None):55 return self.model(inputs, cache=cache)5657 def warm_all(self):58 """Load every expert (only meaningful when capacity == num_experts)."""59 for t in self.tables:60 t.ensure(list(range(self.spec.num_experts)))61 t.reset_stats()62 self.store.reset_stats()6364 def cache_stats(self) -> Dict:65 per = [t.stats() for t in self.tables]66 hits = sum(p["hits"] + p["hits_device"] for p in per)67 misses = sum(p["misses"] for p in per)68 return {"layers": per, "hits": hits, "misses": misses,69 "hit_rate": hits / (hits + misses) if hits + misses else None,70 "stall_time_s": sum(p["stall_time_s"] for p in per), "store": self.store.stats()}7172 def reset_stats(self):73 for t in self.tables:74 t.reset_stats()75 self.store.reset_stats()76