1"""TITAN expert execution block, drop-in replacement for mlx-lm's `SwitchGLU`.23The router (gate linear, softmax, top-k, gating weights, weighted sum) stays in the host model's4code so that expert identity and aggregation are exactly those of the reference implementation;5this module only decides *where* the selected experts' weights come from and executes them with6`mx.gather_qmm` on the layer's slot table.7"""8from __future__ import annotations910from typing import Callable, List, Optional1112import time13import mlx.core as mx14import numpy as np15from mlx_lm.models.activations import swiglu1617from .slots import SlotTable181920class TitanSwitchGLU:21 """Callable with the SwitchGLU signature: (x[..., 1, D], indices[..., k]) -> y[..., k, D]."""2223 def __init__(self, table: SlotTable, group_size: int, bits: int, activation=None,24 on_route: Optional[Callable[[int, List[int]], None]] = None, prefetcher=None):25 self.table = table26 self.prefetcher = prefetcher27 self.current_probs = None28 self.remote = None # RemoteExecutor (Gate E); when set, experts placed on workers are executed remotely29 self.t_sync = 0.0; self.t_prefetch = 0.0; self.t_ensure = 0.0; self.t_build = 0.0 # per-layer host-time attribution (seconds)30 self.group_size = group_size31 self.bits = bits32 self.on_route = on_route # hook(layer, unique_experts) used by predictors/prefetchers (Gate C)33 self.last_indices: Optional[List[int]] = None3435 def _qmm(self, x, m, idx, sorted_indices):36 b = self.table.buf37 return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx,38 transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices)3940 def _run_device(self, x, slot_idx):41 x = mx.expand_dims(x, (-2, -3))42 x_up = self._qmm(x, "up", slot_idx, False)43 x_gate = self._qmm(x, "gate", slot_idx, False)44 h = swiglu(x_gate, x_up)45 y = self._qmm(h, "down", slot_idx, False)46 return y.squeeze(-2)4748 def _run(self, x, flat, shape, probs_np=None):49 uniq = list(dict.fromkeys(flat))50 if self.on_route is not None:51 self.on_route(self.table.layer, uniq)52 pf = self.prefetcher53 if pf is not None:54 _t = time.perf_counter()55 pf.before_layer(self.table.layer, uniq) # install completed prefetches (wait for late ones we need)56 pf.after_layer(self.table.layer, uniq, probs_np) # issue prefetch for the next layer BEFORE this layer's compute57 self.t_prefetch += time.perf_counter() - _t58 _t = time.perf_counter()59 slot_of = dict(zip(uniq, self.table.ensure(uniq)))60 self.t_ensure += time.perf_counter() - _t61 _t = time.perf_counter()62 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)63 x = mx.expand_dims(x, (-2, -3))64 x_up = self._qmm(x, "up", slot_idx, False)65 x_gate = self._qmm(x, "gate", slot_idx, False)66 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU67 y = self._qmm(h, "down", slot_idx, False)68 self.t_build += time.perf_counter() - _t69 return y.squeeze(-2)7071 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:72 t = self.table73 if t.complete():74 # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.75 self.n_chunks = 176 self.last_indices = None77 t.hits_device += indices.size78 return self._run_device(x, t.slot_map[indices])79 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.80 _t0 = time.perf_counter()81 probs_np = None82 if self.prefetcher is not None and self.current_probs is not None:83 mx.eval(indices, self.current_probs)84 probs_np = np.array(self.current_probs[-1]) # last row: decode token (prefill rows are not used for prediction)85 flat = indices.reshape(-1).tolist()86 self.t_sync += time.perf_counter() - _t087 self.last_indices = flat88 k = indices.shape[-1]89 n_rows = len(flat) // k90 cap = self.table.capacity91 x2 = x.reshape(n_rows, x.shape[-1])92 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).93 chunks, cur, cur_set = [], [], set()94 for r in range(n_rows):95 es = set(flat[r * k:(r + 1) * k])96 if cur and len(cur_set | es) > cap:97 chunks.append(cur); cur, cur_set = [], set()98 cur.append(r); cur_set |= es99 if cur:100 chunks.append(cur)101 self.n_chunks = len(chunks)102 if self.remote is not None:103 outs = []104 for ci, rows in enumerate(chunks):105 f = [flat[r * k + j] for r in rows for j in range(k)]106 outs.append(self.remote.run_layer(self, x2[rows[0]:rows[-1] + 1], f, (len(rows), k), probs_np if ci == len(chunks) - 1 else None))107 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))108 if len(chunks) == 1:109 return self._run(x, flat, indices.shape, probs_np)110 outs = []111 for ci, rows in enumerate(chunks):112 f = [flat[r * k + j] for r in rows for j in range(k)]113 outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k), probs_np if ci == len(chunks) - 1 else None))114 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))115116117class TitanMoEBlock:118 """Drop-in replacement for the host model's sparse MoE block (OLMoE / Qwen3-MoE routing: softmax over all experts,119 top-k by argpartition, gating weights = softmax probabilities, renormalised iff norm_topk_prob). The router ops are120 exactly mlx-lm's (verified bit-identical in tests/test_numerical.py); the only additions are the prefetcher hooks."""121122 def __init__(self, host_mlp, block: "TitanSwitchGLU", model_type: str, prefetcher=None):123 self.gate = host_mlp.gate124 self.top_k = host_mlp.top_k125 self.norm_topk_prob = getattr(host_mlp, "norm_topk_prob", False)126 self.switch_mlp = block127 self.model_type = model_type128 self.prefetcher = prefetcher129 self.layer = block.table.layer130131 def __call__(self, x: mx.array) -> mx.array:132 B, L, D = x.shape133 x_flat = x.reshape(-1, D)134 k = self.top_k135 if self.model_type == "olmoe":136 router_logits = self.gate(x_flat)137 rw = mx.softmax(router_logits, axis=1, precise=True)138 indices = mx.stop_gradient(mx.argpartition(-rw, kth=k - 1, axis=-1)[..., :k])139 scores = mx.take_along_axis(rw, indices, axis=-1)140 else: # qwen3_moe141 gates = self.gate(x_flat)142 rw = mx.softmax(gates, axis=-1, precise=True)143 indices = mx.argpartition(rw, kth=-k, axis=-1)[..., -k:]144 scores = mx.take_along_axis(rw, indices, axis=-1)145 if self.norm_topk_prob:146 scores = scores / scores.sum(axis=-1, keepdims=True)147 self.switch_mlp.current_probs = rw if self.prefetcher is not None else None148 y = self.switch_mlp(x_flat, indices)149 y = (y * scores[..., None]).sum(axis=-2)150 return y.reshape(B, L, D)151