Project-Chronos 0.1.1__py3-none-any.whl

chronos/io/storage.py

@@ -0,0 +1,237 @@
+"""
+chronos/io/storage.py
+
+Cluster-aware safetensors backend for Chronos expert weights.
+
+A .ctsr file is a plain safetensors container packing every expert in one
+Louvain co-occurrence cluster. Tensor keys follow the pattern
+``l{layer_id}_e{expert_id}__{param_name}``.
+
+A sidecar ``cluster_manifest.json`` maps (cluster_id) → (file, expert_ids)
+and (expert_id) → cluster_id so the runtime can:
+
+1. Look up which .ctsr file owns an expert in O(1).
+2. Read an entire cluster in a single mmap + zero-copy tensor view, turning
+   N random .pt pickle decodes into one sequential pread.
+
+The format intentionally stays minimal — no compression, no encryption, no
+tensor sharding — so that it maps cleanly onto future io_uring / DirectIO
+backends without schema churn.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+try:
+    from safetensors.torch import save_file as _st_save_file
+    from safetensors import safe_open as _st_safe_open
+    _HAS_SAFETENSORS = True
+except ImportError:  # pragma: no cover
+    _HAS_SAFETENSORS = False
+
+
+MANIFEST_FILENAME = "cluster_manifest.json"
+MANIFEST_VERSION = 1
+KEY_SEP = "__"  # layer/expert prefix vs param-name separator
+
+
+def _key(layer_id: int, expert_id: int, param_name: str) -> str:
+    return f"l{layer_id}_e{expert_id}{KEY_SEP}{param_name}"
+
+
+def _parse_key(key: str) -> Optional[Tuple[int, int, str]]:
+    """Inverse of _key. Returns (layer_id, expert_id, param_name) or None."""
+    if KEY_SEP not in key:
+        return None
+    prefix, param_name = key.split(KEY_SEP, 1)
+    if not (prefix.startswith("l") and "_e" in prefix):
+        return None
+    try:
+        layer_str, expert_str = prefix[1:].split("_e", 1)
+        return int(layer_str), int(expert_str), param_name
+    except ValueError:
+        return None
+
+
+@dataclass
+class ClusterManifest:
+    """Parsed cluster_manifest.json."""
+    version: int
+    num_experts: int
+    num_layers: int
+    storage_format: str
+    # cluster_id -> (filename, expert_ids)
+    clusters: Dict[int, Tuple[str, List[int]]]
+    # expert_id -> cluster_id
+    expert_to_cluster: Dict[int, int]
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "ClusterManifest":
+        clusters: Dict[int, Tuple[str, List[int]]] = {}
+        for entry in data["clusters"]:
+            clusters[int(entry["cluster_id"])] = (
+                entry["file"],
+                [int(e) for e in entry["experts"]],
+            )
+        e2c = {int(k): int(v) for k, v in data["expert_to_cluster"].items()}
+        return cls(
+            version=int(data.get("version", MANIFEST_VERSION)),
+            num_experts=int(data["num_experts"]),
+            num_layers=int(data["num_layers"]),
+            storage_format=data.get("storage_format", "safetensors"),
+            clusters=clusters,
+            expert_to_cluster=e2c,
+        )
+
+    def to_dict(self) -> Dict:
+        return {
+            "version": self.version,
+            "num_experts": self.num_experts,
+            "num_layers": self.num_layers,
+            "storage_format": self.storage_format,
+            "clusters": [
+                {"cluster_id": cid, "file": fname, "experts": list(eids)}
+                for cid, (fname, eids) in sorted(self.clusters.items())
+            ],
+            "expert_to_cluster": {str(k): v for k, v in self.expert_to_cluster.items()},
+        }
+
+    def cluster_members(self, expert_id: int) -> List[int]:
+        cid = self.expert_to_cluster[expert_id]
+        return list(self.clusters[cid][1])
+
+    def cluster_file(self, cluster_id: int) -> str:
+        return self.clusters[cluster_id][0]
+
+
+class ClusterStorage:
+    """
+    Reader/writer for cluster-packed expert weights.
+
+    Write path (one-shot, offline):
+        ClusterStorage.write_clusters(
+            moe_layers, clusters, output_dir, num_layers, num_experts
+        )
+
+    Read path (runtime, hot loop):
+        storage = ClusterStorage(output_dir)
+        per_layer = storage.load_cluster(cluster_id, dtype=torch.float16)
+        # per_layer: {expert_id: {layer_id: {param_name: Tensor}}}
+    """
+
+    def __init__(self, root_dir: str):
+        if not _HAS_SAFETENSORS:
+            raise ImportError(
+                "ClusterStorage requires the `safetensors` package. "
+                "Install with `pip install safetensors>=0.4.0`."
+            )
+        self.root_dir = root_dir
+        manifest_path = os.path.join(root_dir, MANIFEST_FILENAME)
+        if not os.path.exists(manifest_path):
+            raise FileNotFoundError(
+                f"cluster_manifest.json not found at {manifest_path}. "
+                "Run offload_all_to_ssd() or repack_expert_weights_safetensors() first."
+            )
+        with open(manifest_path) as f:
+            self.manifest = ClusterManifest.from_dict(json.load(f))
+
+    # ── read path ────────────────────────────────────────────────
+
+    def load_cluster(
+        self,
+        cluster_id: int,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Dict[int, Dict[int, Dict[str, torch.Tensor]]]:
+        """
+        mmap a whole cluster file and return nested dict:
+            {expert_id: {layer_id: {param_name: Tensor}}}
+
+        Tensors are materialized by safetensors (CPU, target dtype). We keep
+        the logical layout compatible with the existing RAM buffer structure
+        used by ExpertStore so the refactor is a drop-in replacement.
+        """
+        file_name = self.manifest.cluster_file(cluster_id)
+        path = os.path.join(self.root_dir, file_name)
+        nested: Dict[int, Dict[int, Dict[str, torch.Tensor]]] = {}
+        with _st_safe_open(path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                parsed = _parse_key(key)
+                if parsed is None:
+                    continue
+                layer_id, expert_id, param_name = parsed
+                t = f.get_tensor(key)
+                if dtype is not None and t.dtype != dtype:
+                    t = t.to(dtype)
+                nested.setdefault(expert_id, {}).setdefault(layer_id, {})[param_name] = t
+        return nested
+
+    def cluster_for_expert(self, expert_id: int) -> int:
+        return self.manifest.expert_to_cluster[expert_id]
+
+    def experts_in_cluster(self, cluster_id: int) -> List[int]:
+        return list(self.manifest.clusters[cluster_id][1])
+
+    # ── write path ───────────────────────────────────────────────
+
+    @staticmethod
+    def write_clusters(
+        moe_layers,
+        clusters: List[List[int]],
+        output_dir: str,
+        num_layers: int,
+        num_experts: int,
+        dtype: torch.dtype = torch.float16,
+    ) -> str:
+        """
+        Serialize `moe_layers` (list of ChronosMOEFeedForward) into one
+        safetensors file per cluster and write the manifest.
+
+        Returns the path to cluster_manifest.json.
+        """
+        if not _HAS_SAFETENSORS:
+            raise ImportError(
+                "write_clusters requires the `safetensors` package. "
+                "Install with `pip install safetensors>=0.4.0`."
+            )
+        os.makedirs(output_dir, exist_ok=True)
+
+        manifest_clusters: Dict[int, Tuple[str, List[int]]] = {}
+        expert_to_cluster: Dict[int, int] = {}
+
+        for cid, cluster in enumerate(clusters):
+            tensors: Dict[str, torch.Tensor] = {}
+            for eid in cluster:
+                expert_to_cluster[int(eid)] = cid
+                for li in range(num_layers):
+                    if li >= len(moe_layers):
+                        continue
+                    expert = moe_layers[li].experts[eid]
+                    for pname, pval in expert.state_dict().items():
+                        tensors[_key(li, int(eid), pname)] = pval.detach().to(dtype=dtype).cpu().contiguous()
+
+            file_name = f"cluster_{cid}.ctsr"
+            file_path = os.path.join(output_dir, file_name)
+            _st_save_file(tensors, file_path)
+            manifest_clusters[cid] = (file_name, [int(e) for e in cluster])
+
+        manifest = ClusterManifest(
+            version=MANIFEST_VERSION,
+            num_experts=num_experts,
+            num_layers=num_layers,
+            storage_format="safetensors",
+            clusters=manifest_clusters,
+            expert_to_cluster=expert_to_cluster,
+        )
+        manifest_path = os.path.join(output_dir, MANIFEST_FILENAME)
+        with open(manifest_path, "w") as f:
+            json.dump(manifest.to_dict(), f, indent=2)
+        return manifest_path
+
+    @staticmethod
+    def has_manifest(root_dir: str) -> bool:
+        return os.path.exists(os.path.join(root_dir, MANIFEST_FILENAME))

chronos/mlx/__init__.py

@@ -0,0 +1,32 @@
+"""
+chronos/mlx/__init__.py
+
+MLX-native backend for Project Chronos (Apple Silicon / Metal).
+
+MLX uses unified memory — CPU and GPU share the same physical DRAM.
+There is no H2D copy, no pinned memory, and no CUDA streams.
+The "SSD prefetch" concept maps differently:
+  - Tier 0 (hot):  mx.array on GPU (materialised via mx.eval())
+  - Tier 1 (warm): mx.array in unified memory (not yet eval'd)
+  - Tier 2 (cold): .safetensors / .npy files on SSD
+
+Key differences from the CUDA backend:
+  - No ExpertStore.promote_to_vram() — mx.eval() pins to Metal on first use
+  - No H2D stream — Metal handles async via its own command queue
+  - torch.compile unavailable; mx.compile() used instead
+  - All modules are mlx.nn.Module, not torch.nn.Module
+"""
+from chronos.mlx.model import ChronosMLXModel
+from chronos.mlx.moe import ChronosMLXMOE
+from chronos.mlx.attention import MLAAttentionMLX, SlidingWindowAttentionMLX
+from chronos.mlx.expert_store import MLXExpertStore
+from chronos.mlx.inference import ChronosMLXInferenceEngine
+
+__all__ = [
+    "ChronosMLXModel",
+    "ChronosMLXMOE",
+    "MLAAttentionMLX",
+    "SlidingWindowAttentionMLX",
+    "MLXExpertStore",
+    "ChronosMLXInferenceEngine",
+]

chronos/mlx/attention.py

@@ -0,0 +1,230 @@
+"""
+chronos/mlx/attention.py — MLX-native MLA + Sliding Window Attention.
+
+Design notes vs. the PyTorch version:
+- No past_key_values list; instead we carry (c_kv_cache, k_rope_cache) as
+  plain mx.array and concatenate explicitly each step.
+- RoPE: computed with mx operations; no separate apply_rotary_pos_emb import.
+- SlidingWindow: slice oldest tokens with array indexing, not assignment.
+- mx.compile()-friendly: no Python branches on runtime tensor shapes.
+"""
+import math
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def _rope_freqs(dim: int, max_len: int = 4096, base: float = 10000.0) -> mx.array:
+    """Precompute RoPE cos/sin tables [max_len, dim/2].
+
+    `dim` is rounded down to the nearest even number internally so the
+    cos/sin half-dimension is well-defined. Callers that pass an odd
+    `dim` get tables that match the even portion only.
+    """
+    even = dim - (dim % 2)
+    half = max(even // 2, 1)
+    theta = 1.0 / (base ** (mx.arange(0, half, dtype=mx.float32) / half))
+    pos = mx.arange(max_len, dtype=mx.float32)
+    freqs = mx.outer(pos, theta)               # [T, half]
+    return mx.cos(freqs), mx.sin(freqs)        # each [T, half]
+
+
+def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+    """
+    Apply RoPE to x [B, T, nH, d].
+    cos/sin: [T, d_eff/2] where d_eff = d rounded down to even.
+
+    When d is odd, the trailing element is passed through unchanged so the
+    output shape matches the input. (Some Chronos checkpoints were saved
+    with rope_dim = head_dim - 1, which is odd whenever head_dim is even
+    — head_dim = hidden_size / num_heads = 256/8 = 32 → rope_dim = 31.)
+    """
+    B, T, nH, d = x.shape
+    d_eff = d - (d % 2)
+    half = d_eff // 2
+    x1 = x[..., :half]                          # [B, T, nH, half]
+    x2 = x[..., half:d_eff]                     # [B, T, nH, half]
+    c = cos[:T, :half][:, None, :]              # [T, 1, half]
+    s = sin[:T, :half][:, None, :]
+    rotated = mx.concatenate([x1 * c - x2 * s,
+                              x1 * s + x2 * c], axis=-1)  # [B, T, nH, d_eff]
+    if d_eff < d:
+        # Append the unrotated tail so the output keeps the original last dim.
+        tail = x[..., d_eff:]
+        rotated = mx.concatenate([rotated, tail], axis=-1)
+    return rotated
+
+
+class MLAAttentionMLX(nn.Module):
+    """
+    Multi-head Latent Attention for MLX.
+
+    KV cache stores (c_kv [B, T, kv_latent_dim], k_rope [B, T, nH, rope_dim]).
+    Both are concatenated each decode step — no copy in/out of VRAM.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads or config.n_heads
+        self.head_dim = config.head_dim
+        self.kv_latent_dim = getattr(config, 'kv_latent_dim', 64)
+        # rope_dim is clamped to ≤ head_dim−1 to leave room for nope. The
+        # MLX _apply_rope below handles odd rope_dim safely (trailing element
+        # passes through unchanged), so there's no evenness requirement.
+        self.rope_dim = min(
+            max(1, getattr(config, 'rope_dim', min(64, self.head_dim))),
+            self.head_dim - 1,
+        )
+        self.nope_dim = self.head_dim - self.rope_dim
+        H = config.hidden_size
+
+        self.q_nope_proj = nn.Linear(H, self.n_heads * self.nope_dim, bias=False)
+        self.q_rope_proj = nn.Linear(H, self.n_heads * self.rope_dim, bias=False)
+        self.kv_down_proj = nn.Linear(H, self.kv_latent_dim, bias=False)
+        self.k_nope_proj  = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.nope_dim, bias=False)
+        self.k_rope_proj  = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.rope_dim, bias=False)
+        self.v_proj       = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj       = nn.Linear(self.n_heads * self.head_dim, H, bias=False)
+
+        cos, sin = _rope_freqs(self.rope_dim)
+        self.cos = cos   # [max_len, rope_dim/2]  — not a parameter
+        self.sin = sin
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: mx.array = None,
+        cache=None,       # (c_kv [B, T_past, kv_latent_dim], k_rope [B, T_past, nH, rope_dim])
+    ):
+        B, S, H = x.shape
+
+        # ── Queries ───────────────────────────────────────────────
+        q_nope = self.q_nope_proj(x).reshape(B, S, self.n_heads, self.nope_dim)
+        q_rope = self.q_rope_proj(x).reshape(B, S, self.n_heads, self.rope_dim)
+
+        # ── KV latent (what gets cached) ──────────────────────────
+        c_kv = self.kv_down_proj(x)             # [B, S, kv_latent_dim]
+
+        # Append to cache
+        if cache is not None:
+            c_kv_past, k_rope_past = cache
+            c_kv_full = mx.concatenate([c_kv_past, c_kv], axis=1)
+        else:
+            c_kv_full = c_kv
+            k_rope_past = None
+
+        T_full = c_kv_full.shape[1]
+
+        # ── Keys / Values (recomputed from cache) ─────────────────
+        k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim)
+        k_rope = self.k_rope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.rope_dim)
+        v      = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim)
+
+        # ── RoPE ──────────────────────────────────────────────────
+        past_len = T_full - S
+        q_rope = _apply_rope(q_rope,
+                              self.cos[past_len:past_len + S],
+                              self.sin[past_len:past_len + S])
+        k_rope = _apply_rope(k_rope, self.cos[:T_full], self.sin[:T_full])
+
+        # Concat nope + rope
+        q = mx.concatenate([q_nope, q_rope], axis=-1)  # [B, S, nH, head_dim]
+        k = mx.concatenate([k_nope, k_rope], axis=-1)  # [B, T, nKV, head_dim]
+
+        # GQA repeat
+        if self.n_heads != self.n_kv_heads:
+            rep = self.n_heads // self.n_kv_heads
+            k = mx.repeat(k, rep, axis=2)
+            v = mx.repeat(v, rep, axis=2)
+
+        # ── Attention ─────────────────────────────────────────────
+        scale = 1.0 / math.sqrt(self.head_dim)
+        q = q.transpose(0, 2, 1, 3)    # [B, nH, S, d]
+        k = k.transpose(0, 2, 1, 3)    # [B, nH, T, d]
+        v = v.transpose(0, 2, 1, 3)
+        attn = (q @ k.transpose(0, 1, 3, 2)) * scale   # [B, nH, S, T]
+        if mask is not None:
+            attn = attn + mask[None, None, :, :]
+        attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype)
+        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, S, -1)
+
+        new_cache = (c_kv_full, k_rope)
+        return self.o_proj(out), new_cache
+
+
+class SlidingWindowAttentionMLX(nn.Module):
+    """
+    Standard multi-head attention with a sliding KV window for MLX.
+    KV cache is capped at window_size tokens — O(1) memory per step.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads or config.n_heads
+        self.head_dim = config.head_dim
+        self.window_size = getattr(config, 'sliding_window_size', 2048)
+        H = config.hidden_size
+
+        self.q_proj = nn.Linear(H, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(H, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(H, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False)
+
+        rope_dim = self.head_dim
+        cos, sin = _rope_freqs(rope_dim)
+        self.cos = cos
+        self.sin = sin
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: mx.array = None,
+        cache=None,  # (k [B, T_past, nKV, d], v [B, T_past, nKV, d])
+    ):
+        B, S, H = x.shape
+
+        q = self.q_proj(x).reshape(B, S, self.n_heads,    self.head_dim)
+        k = self.k_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim)
+        v = self.v_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim)
+
+        if cache is not None:
+            k_past, v_past = cache
+            k = mx.concatenate([k_past, k], axis=1)
+            v = mx.concatenate([v_past, v], axis=1)
+
+        T_full = k.shape[1]
+        past_len = T_full - S
+
+        # Apply RoPE to new queries and all keys (before GQA repeat)
+        q = _apply_rope(q, self.cos[past_len:past_len + S], self.sin[past_len:past_len + S])
+        k = _apply_rope(k, self.cos[:T_full], self.sin[:T_full])
+
+        # Slide window — keep only last window_size tokens
+        if T_full > self.window_size:
+            k = k[:, -self.window_size:, :, :]
+            v = v[:, -self.window_size:, :, :]
+
+        # GQA repeat for attention (k/v expand to n_heads)
+        k_attn = k
+        v_attn = v
+        if self.n_heads != self.n_kv_heads:
+            rep = self.n_heads // self.n_kv_heads
+            k_attn = mx.repeat(k, rep, axis=2)
+            v_attn = mx.repeat(v, rep, axis=2)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        # Transpose to [B, nH, T, d] for batched matmul
+        qt = q.transpose(0, 2, 1, 3)
+        kt = k_attn.transpose(0, 2, 1, 3)
+        vt = v_attn.transpose(0, 2, 1, 3)
+        attn = (qt @ kt.transpose(0, 1, 3, 2)) * scale
+        if mask is not None:
+            # mask is [S, T_k] (2D causal); broadcast over [B, nH, S, T_k]
+            attn = attn + mask[None, None, :, :]
+        attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype)
+        out = (attn @ vt).transpose(0, 2, 1, 3).reshape(B, S, -1)
+
+        # Cache stores [B, T, nKV, d] (pre-repeat, pre-transpose)
+        new_cache = (k, v)
+        return self.o_proj(out), new_cache

chronos/mlx/expert_store.py

@@ -0,0 +1,155 @@
+"""
+chronos/mlx/expert_store.py — Expert weight management for MLX (unified memory).
+
+On Apple Silicon there is no VRAM / RAM split — all memory is unified.
+"Prefetch" means loading from SSD into Python-side mx.array objects and
+calling mx.eval() to materialise them in Metal shared memory.
+
+Three logical tiers:
+  Tier 0 (hot):  mx.array already eval'd (resident in Metal)
+  Tier 1 (warm): mx.array allocated but not yet eval'd
+  Tier 2 (cold): weights stored as .npy / .pt files on SSD
+"""
+import os
+import threading
+import collections
+from typing import Dict, List, Optional
+
+import mlx.core as mx
+import numpy as np
+
+
+class MLXExpertStore:
+    """
+    Expert weight store for MLX (unified memory — no H2D stream needed).
+
+    Public API is intentionally similar to ExpertStore (PyTorch) so
+    CacheManager can swap backends.
+    """
+
+    def __init__(self, model, config, ssd_dir: str = "./expert_cache_mlx"):
+        self.config = config
+        self.ssd_dir = ssd_dir
+        self.num_experts = config.num_experts
+        os.makedirs(ssd_dir, exist_ok=True)
+
+        from chronos.mlx.moe import ChronosMLXMOE
+        self.moe_layers: List[ChronosMLXMOE] = [
+            layer.mlp for layer in model.layers
+            if isinstance(layer.mlp, ChronosMLXMOE)
+        ]
+        self.num_layers = len(self.moe_layers)
+
+        # LRU tracking which experts are "hot" (eval'd in Metal)
+        capacity = max(1, int(config.vram_budget_gb * 1.5))  # rough heuristic
+        self._hot_lru: collections.OrderedDict = collections.OrderedDict()
+        self._capacity = capacity
+        self._lock = threading.Lock()
+
+        # Warm buffers: {expert_id: {layer_idx: {param_name: mx.array}}}
+        self._warm: Dict[int, Dict[int, Dict[str, mx.array]]] = {}
+
+    # ── SSD serialisation ─────────────────────────────────────────
+
+    def offload_all_to_ssd(self):
+        """Save all expert weights to SSD as .npy files."""
+        for li, moe in enumerate(self.moe_layers):
+            for ei, expert in enumerate(moe.experts):
+                for name, param in vars(expert).items():
+                    if isinstance(param, mx.array):
+                        path = self._ssd_path(ei, li, name)
+                        if not os.path.exists(path):
+                            np.save(path, np.array(param))
+
+    def _ssd_path(self, expert_id: int, layer_idx: int, param_name: str) -> str:
+        return os.path.join(
+            self.ssd_dir, f"expert_l{layer_idx}_e{expert_id}_{param_name}.npy"
+        )
+
+    def _load_from_ssd(self, expert_id: int) -> Dict[int, Dict[str, mx.array]]:
+        layer_states = {}
+        for li in range(self.num_layers):
+            params = {}
+            for name in ("gate_proj", "up_proj", "down_proj"):
+                for subname in ("weight",):
+                    fname = f"{name}.{subname}"
+                    path = self._ssd_path(expert_id, li, fname)
+                    if os.path.exists(path):
+                        params[fname] = mx.array(np.load(path))
+            if params:
+                layer_states[li] = params
+        return layer_states
+
+    # ── Prefetch (SSD → unified memory) ──────────────────────────
+
+    def prefetch_to_ram(self, expert_ids: List[int]):
+        """
+        Load expert weights from SSD into mx.array (not yet eval'd).
+        Call from a background thread for async prefetch.
+        """
+        for eid in expert_ids:
+            with self._lock:
+                if eid in self._warm:
+                    continue
+            layer_states = self._load_from_ssd(eid)
+            if not layer_states:
+                continue
+            with self._lock:
+                self._warm[eid] = layer_states
+
+    # ── Promote to "hot" (materialise in Metal) ───────────────────
+
+    def promote_to_vram(self, expert_id: int) -> bool:
+        """
+        Materialise expert weights in Metal shared memory via mx.eval().
+        On MLX there is no actual copy — just lazy evaluation.
+        """
+        with self._lock:
+            if expert_id not in self._warm:
+                return False
+            layer_states = self._warm[expert_id]
+
+        # mx.eval() forces computation and pins arrays in Metal memory
+        for li, params in layer_states.items():
+            mx.eval(*params.values())
+            # Write back into the live model
+            expert = self.moe_layers[li].experts[expert_id]
+            for fname, arr in params.items():
+                parts = fname.split(".")
+                obj = expert
+                for part in parts[:-1]:
+                    obj = getattr(obj, part)
+                setattr(obj, parts[-1], arr)
+
+        with self._lock:
+            if expert_id in self._hot_lru:
+                self._hot_lru.move_to_end(expert_id)
+            else:
+                if len(self._hot_lru) >= self._capacity:
+                    self._hot_lru.popitem(last=False)
+                self._hot_lru[expert_id] = True
+        return True
+
+    def sync_h2d(self):
+        """No-op on MLX — Metal handles synchronisation internally."""
+        pass
+
+    # ── Availability mask ─────────────────────────────────────────
+
+    def vram_availability_mask(self) -> mx.array:
+        """[num_experts] float32 mask — 1.0 if expert is hot, 0.0 otherwise."""
+        mask = np.zeros(self.num_experts, dtype=np.float32)
+        with self._lock:
+            for eid in self._hot_lru:
+                if eid < self.num_experts:
+                    mask[eid] = 1.0
+        return mx.array(mask)
+
+    def stats(self) -> dict:
+        return {
+            "backend": "mlx",
+            "hot_experts": len(self._hot_lru),
+            "hot_capacity": self._capacity,
+            "warm_experts": len(self._warm),
+            "h2d_stream": "metal_unified",
+        }