1"""Model and expert registry.23Parses a Hugging Face / MLX MoE checkpoint directory (config.json + safetensors) into an explicit4description of layers, experts, routing semantics and the *byte ranges* of every expert tensor,5so that the expert store can read a single expert from disk with `pread` without touching the rest6of the file. Only formats that we have verified are accepted (see `SUPPORTED`).7"""8from __future__ import annotations910import json11import os12import struct13from dataclasses import dataclass, field14from typing import Dict, List, Tuple1516DTYPE_ITEMSIZE = {"F32": 4, "F16": 2, "BF16": 2, "U32": 4, "I32": 4, "U8": 1, "I8": 1, "U16": 2}1718# model_type -> (expert tensor key template, matrices, routing description)19SUPPORTED = {20 "olmoe": {21 "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",22 "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",23 "matrices": ("gate", "up", "down"),24 "router_key": "model.layers.{l}.mlp.gate.{s}",25 "routing": "softmax_then_topk", # softmax over all experts, then top-k, weights = softmax probs (renormalised iff norm_topk_prob)26 "shared_experts": 0,27 },28 "qwen3_moe": {29 "expert_key": "model.layers.{l}.mlp.experts.{e}.{m}_proj.{s}",30 "stacked_key": "model.layers.{l}.mlp.switch_mlp.{m}_proj.{s}",31 "matrices": ("gate", "up", "down"),32 "router_key": "model.layers.{l}.mlp.gate.{s}",33 "routing": "softmax_then_topk",34 "shared_experts": 0,35 },36}373839@dataclass(frozen=True)40class TensorLoc:41 key: str42 dtype: str43 shape: Tuple[int, ...]44 file: str45 offset: int # absolute byte offset in file46 nbytes: int474849@dataclass50class ExpertSpec:51 layer: int52 expert: int53 tensors: Dict[str, TensorLoc] # e.g. "gate.weight" -> loc5455 @property56 def nbytes(self) -> int:57 return sum(t.nbytes for t in self.tensors.values())585960@dataclass61class ModelSpec:62 path: str63 model_type: str64 num_layers: int65 num_experts: int66 top_k: int67 norm_topk_prob: bool68 hidden_size: int69 intermediate_size: int70 quantization: Dict[str, int]71 routing: str72 shared_experts: int73 experts: Dict[Tuple[int, int], ExpertSpec] = field(default_factory=dict)74 dense: Dict[str, TensorLoc] = field(default_factory=dict)75 layout: str = "per_expert" # or "stacked"7677 @property78 def bytes_per_expert(self) -> int:79 return next(iter(self.experts.values())).nbytes8081 @property82 def expert_bytes_total(self) -> int:83 return sum(e.nbytes for e in self.experts.values())8485 @property86 def dense_bytes_total(self) -> int:87 return sum(t.nbytes for t in self.dense.values())8889 def summary(self) -> dict:90 return {91 "path": self.path,92 "model_type": self.model_type,93 "num_layers": self.num_layers,94 "num_experts": self.num_experts,95 "top_k": self.top_k,96 "norm_topk_prob": self.norm_topk_prob,97 "hidden_size": self.hidden_size,98 "intermediate_size": self.intermediate_size,99 "quantization": self.quantization,100 "routing": self.routing,101 "shared_experts": self.shared_experts,102 "layout": self.layout,103 "bytes_per_expert": self.bytes_per_expert,104 "expert_bytes_total": self.expert_bytes_total,105 "dense_bytes_total": self.dense_bytes_total,106 "experts_per_token": self.num_layers * self.top_k,107 "bytes_per_token_if_uncached": self.num_layers * self.top_k * self.bytes_per_expert,108 }109110111def read_safetensors_header(path: str) -> Dict[str, TensorLoc]:112 """Parse the safetensors header and return absolute byte locations for every tensor."""113 locs = {}114 with open(path, "rb") as f:115 (n,) = struct.unpack("<Q", f.read(8))116 header = json.loads(f.read(n))117 base = 8 + n118 for key, meta in header.items():119 if key == "__metadata__":120 continue121 a, b = meta["data_offsets"]122 locs[key] = TensorLoc(key, meta["dtype"], tuple(meta["shape"]), path, base + a, b - a)123 return locs124125126def titan_layout_dir(model_dir: str):127 import hashlib128 d = os.path.expanduser(f"~/.cache/titan/{hashlib.sha1(os.path.realpath(model_dir).encode()).hexdigest()[:12]}")129 return d if os.path.exists(os.path.join(d, "experts.titan.json")) else None130131132def apply_titan_layout(spec: "ModelSpec", layout_dir: str) -> "ModelSpec":133 """Point every expert's tensors at the contiguous blob produced by scripts/convert_layout.py."""134 idx = json.load(open(os.path.join(layout_dir, "experts.titan.json")))135 blob = os.path.join(layout_dir, "experts.titan")136 for key, ent in idx["experts"].items():137 l, e = map(int, key.split(","))138 es = spec.experts[(l, e)]139 es.tensors = {name: TensorLoc(f"titan:{l},{e}:{name}", p["dtype"], tuple(p["shape"]), blob, ent["offset"] + p["off"], p["n"])140 for name, p in ent["tensors"].items()}141 spec.layout = spec.layout + "+titan_contiguous"142 return spec143144145def load_spec(model_dir: str, use_titan_layout: bool = True) -> ModelSpec:146 model_dir = os.path.realpath(model_dir)147 with open(os.path.join(model_dir, "config.json")) as f:148 cfg = json.load(f)149 mt = cfg["model_type"]150 if mt not in SUPPORTED:151 raise ValueError(f"model_type {mt!r} is not in the verified set {sorted(SUPPORTED)}")152 fmt = SUPPORTED[mt]153 locs: Dict[str, TensorLoc] = {}154 for fn in sorted(os.listdir(model_dir)):155 if fn.endswith(".safetensors"):156 locs.update(read_safetensors_header(os.path.join(model_dir, fn)))157 q = cfg.get("quantization") or cfg.get("quantization_config") or {}158 spec = ModelSpec(159 path=model_dir, model_type=mt, num_layers=cfg["num_hidden_layers"], num_experts=cfg["num_experts"],160 top_k=cfg["num_experts_per_tok"], norm_topk_prob=bool(cfg.get("norm_topk_prob", False)),161 hidden_size=cfg["hidden_size"], intermediate_size=cfg["intermediate_size"],162 quantization={"group_size": q.get("group_size"), "bits": q.get("bits")}, routing=fmt["routing"],163 shared_experts=fmt["shared_experts"],164 )165 per_expert_keys = set()166 for l in range(spec.num_layers):167 for e in range(spec.num_experts):168 tensors = {}169 for m in fmt["matrices"]:170 for s in ("weight", "scales", "biases"):171 k = fmt["expert_key"].format(l=l, e=e, m=m, s=s)172 if k in locs:173 tensors[f"{m}.{s}"] = locs[k]174 per_expert_keys.add(k)175 if not tensors:176 # stacked layout: switch_mlp.<m>_proj.<s> has shape [E, ...]; expert e is the contiguous slice e of the first axis177 for m in fmt["matrices"]:178 for s in ("weight", "scales", "biases"):179 k = fmt["stacked_key"].format(l=l, m=m, s=s)180 if k in locs:181 loc = locs[k]182 if loc.shape[0] != spec.num_experts:183 raise ValueError(f"{k}: first axis {loc.shape[0]} != num_experts {spec.num_experts}")184 per = loc.nbytes // spec.num_experts185 tensors[f"{m}.{s}"] = TensorLoc(f"{k}[{e}]", loc.dtype, tuple(loc.shape[1:]), loc.file, loc.offset + e * per, per)186 per_expert_keys.add(k)187 spec.layout = "stacked"188 if not tensors:189 raise ValueError(f"no expert tensors found for layer {l} expert {e} (neither per-expert nor stacked keys)")190 spec.experts[(l, e)] = ExpertSpec(l, e, tensors)191 spec.dense = {k: v for k, v in locs.items() if k not in per_expert_keys}192 if use_titan_layout:193 d = titan_layout_dir(model_dir)194 if d:195 apply_titan_layout(spec, d)196 return spec197198199if __name__ == "__main__":200 import sys201 s = load_spec(sys.argv[1])202 print(json.dumps(s.summary(), indent=2))203