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 TitanMoEBlock, TitanSwitchGLU16from .prefetch import Prefetcher17from . import predictors as _pred18from .registry import ModelSpec, load_spec19from .slots import EvictionPolicy, LFUPolicy, SlotTable20from .store import ExpertStore2122POLICIES = {"lru": EvictionPolicy, "lfu": LFUPolicy}232425class TitanModel:26 def __init__(self, model_dir: str, capacity: Optional[int] = None, policy: str = "lru", nocache: bool = False,27 on_route=None, predictor: str = "none", m: int = 8, prefetch_workers: int = 2, traces_dir: Optional[str] = None,28 fused_upgate: bool = False):29 self.model_dir = os.path.realpath(model_dir)30 self.spec: ModelSpec = load_spec(self.model_dir)31 self.store = ExpertStore(self.spec, nocache=nocache)32 self.capacity = capacity or self.spec.num_experts33 self.policy_name = policy34 # Load the host model lazily so the stacked expert tensors are never evaluated.35 from pathlib import Path36 self.model, cfg = load_model(Path(self.model_dir), lazy=True)37 self.tokenizer = load_tokenizer(Path(self.model_dir))38 q = cfg.get("quantization") or cfg.get("quantization_config")39 self.tables: List[SlotTable] = []40 self.blocks: List[TitanSwitchGLU] = []41 for l, layer in enumerate(self.model.model.layers):42 table = SlotTable(l, self.capacity, self.store, POLICIES[policy](), fused_upgate=fused_upgate)43 block = TitanSwitchGLU(table, q["group_size"], q["bits"], on_route=on_route)44 # drop the lazy stacked weights and splice in the TITAN block45 mlp = layer.mlp46 del mlp.switch_mlp # removes the lazy stacked SwitchGLU (never evaluated)47 mlp.switch_mlp = block # plain attribute: not an nn.Module, not in parameters()48 self.tables.append(table)49 self.blocks.append(block)50 mx.eval(self.model.parameters()) # dense weights only51 # optional predictive prefetcher: replaces each layer's MoE block with the hooked exact-routing block52 self.prefetcher = None53 self.predictor_name = predictor54 if predictor != "none":55 pred = _pred.make(predictor, self.spec.num_layers, self.spec.num_experts)56 if traces_dir is not None and hasattr(pred, "fit"):57 pred.fit(_load_fit_traces(traces_dir))58 self.prefetcher = Prefetcher(self.tables, self.store, pred, m=m, workers=prefetch_workers)59 for l, layer in enumerate(self.model.model.layers):60 self.blocks[l].prefetcher = self.prefetcher61 layer.mlp = TitanMoEBlock(layer.mlp, self.blocks[l], self.spec.model_type, self.prefetcher)6263 # convenience -----------------------------------------------------------------64 @property65 def layers(self):66 return self.model.layers6768 def __call__(self, inputs, cache=None):69 return self.model(inputs, cache=cache)7071 def attach_workers(self, workers: Dict[str, "WorkerClient"], placement: Dict, fallback_local: bool = True):72 """Enable distributed execution (Gate E). placement: {(layer, expert): worker_name or None}."""73 from .distributed import RemoteExecutor74 self.remote = RemoteExecutor(workers, placement, fallback_local)75 for b in self.blocks:76 b.remote = self.remote77 # a layer with every expert resident locally would take the sync-free fast path and bypass the remote split:78 # the capacity is bounded by construction in distributed runs (capacity < num_experts), so no change is needed here.79 return self.remote8081 def warm_all(self):82 """Load every expert (only meaningful when capacity == num_experts)."""83 for t in self.tables:84 t.ensure(list(range(self.spec.num_experts)))85 t.reset_stats()86 self.store.reset_stats()8788 def cache_stats(self) -> Dict:89 per = [t.stats() for t in self.tables]90 if self.prefetcher is not None:91 per_pf = self.prefetcher.stats()92 else:93 per_pf = None94 hits = sum(p["hits"] + p["hits_device"] for p in per)95 misses = sum(p["misses"] for p in per)96 return {"layers": per, "hits": hits, "misses": misses,97 "hit_rate": hits / (hits + misses) if hits + misses else None,98 "stall_time_s": sum(p["stall_time_s"] for p in per), "store": self.store.stats(), "prefetch": per_pf,99 "host_time_s": {"sync": sum(b.t_sync for b in self.blocks), "prefetch": sum(b.t_prefetch for b in self.blocks),100 "ensure": sum(b.t_ensure for b in self.blocks), "build": sum(b.t_build for b in self.blocks)}}101102 def reset_stats(self):103 for t in self.tables:104 t.reset_stats()105 for b in self.blocks:106 b.t_sync = b.t_prefetch = b.t_ensure = b.t_build = 0.0107 self.store.reset_stats()108 if self.prefetcher is not None:109 self.prefetcher.reset_stats()110111112def _load_fit_traces(traces_dir: str):113 import json, numpy as np114 idx = json.load(open(os.path.join(traces_dir, "index.json")))115 out = []116 for p in idx["prompts"]:117 if p["split"] != "fit":118 continue119 z = np.load(os.path.join(traces_dir, p["name"] + ".npz"))120 out.append((z["topk"].astype(np.int64), z["probs"]))121 return out122