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_device(self, x, slot_idx):35 x = mx.expand_dims(x, (-2, -3))36 x_up = self._qmm(x, "up", slot_idx, False)37 x_gate = self._qmm(x, "gate", slot_idx, False)38 h = swiglu(x_gate, x_up)39 y = self._qmm(h, "down", slot_idx, False)40 return y.squeeze(-2)4142 def _run(self, x, flat, shape):43 uniq = list(dict.fromkeys(flat))44 if self.on_route is not None:45 self.on_route(self.table.layer, uniq)46 slot_of = dict(zip(uniq, self.table.ensure(uniq)))47 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)48 x = mx.expand_dims(x, (-2, -3))49 x_up = self._qmm(x, "up", slot_idx, False)50 x_gate = self._qmm(x, "gate", slot_idx, False)51 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU52 y = self._qmm(h, "down", slot_idx, False)53 return y.squeeze(-2)5455 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:56 t = self.table57 if t.complete():58 # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.59 self.n_chunks = 160 self.last_indices = None61 t.hits_device += indices.size62 return self._run_device(x, t.slot_map[indices])63 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.64 flat = indices.reshape(-1).tolist()65 self.last_indices = flat66 k = indices.shape[-1]67 n_rows = len(flat) // k68 cap = self.table.capacity69 x2 = x.reshape(n_rows, x.shape[-1])70 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).71 chunks, cur, cur_set = [], [], set()72 for r in range(n_rows):73 es = set(flat[r * k:(r + 1) * k])74 if cur and len(cur_set | es) > cap:75 chunks.append(cur); cur, cur_set = [], set()76 cur.append(r); cur_set |= es77 if cur:78 chunks.append(cur)79 self.n_chunks = len(chunks)80 if len(chunks) == 1:81 return self._run(x, flat, indices.shape)82 outs = []83 for rows in chunks:84 f = [flat[r * k + j] for r in rows for j in range(k)]85 outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k)))86 return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))87