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 mlx.core as mx13import numpy as np14from mlx_lm.models.activations import swiglu1516from .slots import SlotTable171819class TitanSwitchGLU:20 """Callable with the SwitchGLU signature: (x[..., 1, D], indices[..., k]) -> y[..., k, D]."""2122 def __init__(self, table: SlotTable, group_size: int, bits: int, activation=None,23 on_route: Optional[Callable[[int, List[int]], None]] = None, prefetcher=None):24 self.table = table25 self.prefetcher = prefetcher26 self.current_probs = None27 self.remote = None # RemoteExecutor (Gate E); when set, experts placed on workers are executed remotely28 self.group_size = group_size29 self.bits = bits30 self.on_route = on_route # hook(layer, unique_experts) used by predictors/prefetchers (Gate C)31 self.last_indices: Optional[List[int]] = None3233 def _qmm(self, x, m, idx, sorted_indices):34 b = self.table.buf35 return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx,36 transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices)3738 def _run_device(self, x, slot_idx):39 x = mx.expand_dims(x, (-2, -3))40 x_up = self._qmm(x, "up", slot_idx, False)41 x_gate = self._qmm(x, "gate", slot_idx, False)42 h = swiglu(x_gate, x_up)43 y = self._qmm(h, "down", slot_idx, False)44 return y.squeeze(-2)4546 def _run(self, x, flat, shape, probs_np=None):47 uniq = list(dict.fromkeys(flat))48 if self.on_route is not None:49 self.on_route(self.table.layer, uniq)50 pf = self.prefetcher51 if pf is not None:52 pf.before_layer(self.table.layer, uniq) # install completed prefetches (wait for late ones we need)53 pf.after_layer(self.table.layer, uniq, probs_np) # issue prefetch for the next layer BEFORE this layer's compute54 slot_of = dict(zip(uniq, self.table.ensure(uniq)))55 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)56 x = mx.expand_dims(x, (-2, -3))57 x_up = self._qmm(x, "up", slot_idx, False)58 x_gate = self._qmm(x, "gate", slot_idx, False)59 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU60 y = self._qmm(h, "down", slot_idx, False)61 return y.squeeze(-2)6263 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:64 t = self.table65 if t.complete():66 # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.67 self.n_chunks = 168 self.last_indices = None69 t.hits_device += indices.size70 return self._run_device(x, t.slot_map[indices])71 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.72 probs_np = None73 if self.prefetcher is not None and self.current_probs is not None:74 mx.eval(indices, self.current_probs)75 probs_np = np.array(self.current_probs[-1]) # last row: decode token (prefill rows are not used for prediction)76 flat = indices.reshape(-1).tolist()77 self.last_indices = flat78 k = indices.shape[-1]79 n_rows = len(flat) // k80 cap = self.table.capacity81 x2 = x.reshape(n_rows, x.shape[-1])82 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).83 chunks, cur, cur_set = [], [], set()84 for r in range(n_rows):85 es = set(flat[r * k:(r + 1) * k])86 if cur and len(cur_set | es) > cap:87 chunks.append(cur); cur, cur_set = [], set()88 cur.append(r); cur_set |= es89 if cur:90 chunks.append(cur)91 self.n_chunks = len(chunks)92 if self.remote is not None:93 outs = []94 for ci, rows in enumerate(chunks):95 f = [flat[r * k + j] for r in rows for j in range(k)]96 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))97 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))98 if len(chunks) == 1:99 return self._run(x, flat, indices.shape, probs_np)100 outs = []101 for ci, rows in enumerate(chunks):102 f = [flat[r * k + j] for r in rows for j in range(k)]103 outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k), probs_np if ci == len(chunks) - 1 else None))104 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))105106107class TitanMoEBlock:108 """Drop-in replacement for the host model's sparse MoE block (OLMoE / Qwen3-MoE routing: softmax over all experts,109 top-k by argpartition, gating weights = softmax probabilities, renormalised iff norm_topk_prob). The router ops are110 exactly mlx-lm's (verified bit-identical in tests/test_numerical.py); the only additions are the prefetcher hooks."""111112 def __init__(self, host_mlp, block: "TitanSwitchGLU", model_type: str, prefetcher=None):113 self.gate = host_mlp.gate114 self.top_k = host_mlp.top_k115 self.norm_topk_prob = getattr(host_mlp, "norm_topk_prob", False)116 self.switch_mlp = block117 self.model_type = model_type118 self.prefetcher = prefetcher119 self.layer = block.table.layer120121 def __call__(self, x: mx.array) -> mx.array:122 B, L, D = x.shape123 x_flat = x.reshape(-1, D)124 k = self.top_k125 if self.model_type == "olmoe":126 router_logits = self.gate(x_flat)127 rw = mx.softmax(router_logits, axis=1, precise=True)128 indices = mx.stop_gradient(mx.argpartition(-rw, kth=k - 1, axis=-1)[..., :k])129 scores = mx.take_along_axis(rw, indices, axis=-1)130 else: # qwen3_moe131 gates = self.gate(x_flat)132 rw = mx.softmax(gates, axis=-1, precise=True)133 indices = mx.argpartition(rw, kth=-k, axis=-1)[..., -k:]134 scores = mx.take_along_axis(rw, indices, axis=-1)135 if self.norm_topk_prob:136 scores = scores / scores.sum(axis=-1, keepdims=True)137 self.switch_mlp.current_probs = rw if self.prefetcher is not None else None138 y = self.switch_mlp(x_flat, indices)139 y = (y * scores[..., None]).sum(axis=-2)140 return y.reshape(B, L, D)141