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 mx13from mlx_lm.models.activations import swiglu1415from .slots import SlotTable161718class TitanSwitchGLU:19 """Callable with the SwitchGLU signature: (x[..., 1, D], indices[..., k]) -> y[..., k, D]."""2021 def __init__(self, table: SlotTable, group_size: int, bits: int, activation=None,22 on_route: Optional[Callable[[int, List[int]], None]] = None):23 self.table = table24 self.group_size = group_size25 self.bits = bits26 self.on_route = on_route # hook(layer, unique_experts) used by predictors/prefetchers (Gate C)27 self.last_indices: Optional[List[int]] = None2829 def _qmm(self, x, m, idx, sorted_indices):30 b = self.table.buf31 return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx,32 transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices)3334 def _run(self, x, flat, shape):35 uniq = list(dict.fromkeys(flat))36 if self.on_route is not None:37 self.on_route(self.table.layer, uniq)38 slot_of = dict(zip(uniq, self.table.ensure(uniq)))39 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)40 x = mx.expand_dims(x, (-2, -3))41 x_up = self._qmm(x, "up", slot_idx, False)42 x_gate = self._qmm(x, "gate", slot_idx, False)43 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU44 y = self._qmm(h, "down", slot_idx, False)45 return y.squeeze(-2)4647 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:48 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.49 flat = indices.reshape(-1).tolist()50 self.last_indices = flat51 k = indices.shape[-1]52 n_rows = len(flat) // k53 cap = self.table.capacity54 x2 = x.reshape(n_rows, x.shape[-1])55 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).56 chunks, cur, cur_set = [], [], set()57 for r in range(n_rows):58 es = set(flat[r * k:(r + 1) * k])59 if cur and len(cur_set | es) > cap:60 chunks.append(cur); cur, cur_set = [], set()61 cur.append(r); cur_set |= es62 if cur:63 chunks.append(cur)64 self.n_chunks = len(chunks)65 if len(chunks) == 1:66 return self._run(x, flat, indices.shape)67 outs = []68 for rows in chunks:69 f = [flat[r * k + j] for r in rows for j in range(k)]70 outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k)))71 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))72