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 _upgate(self, x, slot_idx, sorted_indices=False):41 """Fused up+gate: one gather_qmm over the [C, 2H, D] buffer, split into (up, gate)."""42 y = self._qmm(x, "upgate", slot_idx, sorted_indices)43 H = y.shape[-1] // 244 return y[..., :H], y[..., H:]4546 def _run_device(self, x, slot_idx):47 x = mx.expand_dims(x, (-2, -3))48 if self.table.fused:49 x_up, x_gate = self._upgate(x, slot_idx)50 else:51 x_up = self._qmm(x, "up", slot_idx, False)52 x_gate = self._qmm(x, "gate", slot_idx, False)53 h = swiglu(x_gate, x_up)54 y = self._qmm(h, "down", slot_idx, False)55 return y.squeeze(-2)5657 def _run(self, x, flat, shape, probs_np=None):58 uniq = list(dict.fromkeys(flat))59 if self.on_route is not None:60 self.on_route(self.table.layer, uniq)61 pf = self.prefetcher62 if pf is not None:63 _t = time.perf_counter()64 pf.before_layer(self.table.layer, uniq) # install completed prefetches (wait for late ones we need)65 pf.after_layer(self.table.layer, uniq, probs_np) # issue prefetch for the next layer BEFORE this layer's compute66 self.t_prefetch += time.perf_counter() - _t67 _t = time.perf_counter()68 slot_of = dict(zip(uniq, self.table.ensure(uniq)))69 self.t_ensure += time.perf_counter() - _t70 _t = time.perf_counter()71 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)72 x = mx.expand_dims(x, (-2, -3))73 if self.table.fused:74 x_up, x_gate = self._upgate(x, slot_idx)75 else:76 x_up = self._qmm(x, "up", slot_idx, False)77 x_gate = self._qmm(x, "gate", slot_idx, False)78 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU79 y = self._qmm(h, "down", slot_idx, False)80 self.t_build += time.perf_counter() - _t81 return y.squeeze(-2)8283 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:84 t = self.table85 if t.complete():86 # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.87 self.n_chunks = 188 self.last_indices = None89 t.hits_device += indices.size90 return self._run_device(x, t.slot_map[indices])91 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.92 _t0 = time.perf_counter()93 probs_np = None94 if self.prefetcher is not None and self.current_probs is not None:95 mx.eval(indices, self.current_probs)96 probs_np = np.array(self.current_probs[-1]) # last row: decode token (prefill rows are not used for prediction)97 flat = indices.reshape(-1).tolist()98 self.t_sync += time.perf_counter() - _t099 self.last_indices = flat100 k = indices.shape[-1]101 n_rows = len(flat) // k102 cap = self.table.capacity103 x2 = x.reshape(n_rows, x.shape[-1])104 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).105 chunks, cur, cur_set = [], [], set()106 for r in range(n_rows):107 es = set(flat[r * k:(r + 1) * k])108 if cur and len(cur_set | es) > cap:109 chunks.append(cur); cur, cur_set = [], set()110 cur.append(r); cur_set |= es111 if cur:112 chunks.append(cur)113 self.n_chunks = len(chunks)114 if self.remote is not None:115 outs = []116 for ci, rows in enumerate(chunks):117 f = [flat[r * k + j] for r in rows for j in range(k)]118 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))119 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))120 if len(chunks) == 1:121 return self._run(x, flat, indices.shape, probs_np)122 outs = []123 for ci, rows in enumerate(chunks):124 f = [flat[r * k + j] for r in rows for j in range(k)]125 outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k), probs_np if ci == len(chunks) - 1 else None))126 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))127128129class TitanMoEBlock:130 """Drop-in replacement for the host model's sparse MoE block (OLMoE / Qwen3-MoE routing: softmax over all experts,131 top-k by argpartition, gating weights = softmax probabilities, renormalised iff norm_topk_prob). The router ops are132 exactly mlx-lm's (verified bit-identical in tests/test_numerical.py); the only additions are the prefetcher hooks."""133134 def __init__(self, host_mlp, block: "TitanSwitchGLU", model_type: str, prefetcher=None):135 self.gate = host_mlp.gate136 self.top_k = host_mlp.top_k137 self.norm_topk_prob = getattr(host_mlp, "norm_topk_prob", False)138 self.switch_mlp = block139 self.model_type = model_type140 self.prefetcher = prefetcher141 self.layer = block.table.layer142143 def __call__(self, x: mx.array) -> mx.array:144 B, L, D = x.shape145 x_flat = x.reshape(-1, D)146 k = self.top_k147 if self.model_type == "olmoe":148 router_logits = self.gate(x_flat)149 rw = mx.softmax(router_logits, axis=1, precise=True)150 indices = mx.stop_gradient(mx.argpartition(-rw, kth=k - 1, axis=-1)[..., :k])151 scores = mx.take_along_axis(rw, indices, axis=-1)152 else: # qwen3_moe153 gates = self.gate(x_flat)154 rw = mx.softmax(gates, axis=-1, precise=True)155 indices = mx.argpartition(rw, kth=-k, axis=-1)[..., -k:]156 scores = mx.take_along_axis(rw, indices, axis=-1)157 if self.norm_topk_prob:158 scores = scores / scores.sum(axis=-1, keepdims=True)159 self.switch_mlp.current_probs = rw if self.prefetcher is not None else None160 y = self.switch_mlp(x_flat, indices)161 y = (y * scores[..., None]).sum(axis=-2)162 return y.reshape(B, L, D)163