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.group_size = group_size28 self.bits = bits29 self.on_route = on_route # hook(layer, unique_experts) used by predictors/prefetchers (Gate C)30 self.last_indices: Optional[List[int]] = None3132 def _qmm(self, x, m, idx, sorted_indices):33 b = self.table.buf34 return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx,35 transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices)3637 def _run_device(self, x, slot_idx):38 x = mx.expand_dims(x, (-2, -3))39 x_up = self._qmm(x, "up", slot_idx, False)40 x_gate = self._qmm(x, "gate", slot_idx, False)41 h = swiglu(x_gate, x_up)42 y = self._qmm(h, "down", slot_idx, False)43 return y.squeeze(-2)4445 def _run(self, x, flat, shape, probs_np=None):46 uniq = list(dict.fromkeys(flat))47 if self.on_route is not None:48 self.on_route(self.table.layer, uniq)49 pf = self.prefetcher50 if pf is not None:51 pf.before_layer(self.table.layer, uniq) # install completed prefetches (wait for late ones we need)52 pf.after_layer(self.table.layer, uniq, probs_np) # issue prefetch for the next layer BEFORE this layer's compute53 slot_of = dict(zip(uniq, self.table.ensure(uniq)))54 slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)55 x = mx.expand_dims(x, (-2, -3))56 x_up = self._qmm(x, "up", slot_idx, False)57 x_gate = self._qmm(x, "gate", slot_idx, False)58 h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU59 y = self._qmm(h, "down", slot_idx, False)60 return y.squeeze(-2)6162 def __call__(self, x: mx.array, indices: mx.array) -> mx.array:63 t = self.table64 if t.complete():65 # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.66 self.n_chunks = 167 self.last_indices = None68 t.hits_device += indices.size69 return self._run_device(x, t.slot_map[indices])70 # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.71 probs_np = None72 if self.prefetcher is not None and self.current_probs is not None:73 mx.eval(indices, self.current_probs)74 probs_np = np.array(self.current_probs[-1]) # last row: decode token (prefill rows are not used for prediction)75 flat = indices.reshape(-1).tolist()76 self.last_indices = flat77 k = indices.shape[-1]78 n_rows = len(flat) // k79 cap = self.table.capacity80 x2 = x.reshape(n_rows, x.shape[-1])81 # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).82 chunks, cur, cur_set = [], [], set()83 for r in range(n_rows):84 es = set(flat[r * k:(r + 1) * k])85 if cur and len(cur_set | es) > cap:86 chunks.append(cur); cur, cur_set = [], set()87 cur.append(r); cur_set |= es88 if cur:89 chunks.append(cur)90 self.n_chunks = len(chunks)91 if len(chunks) == 1:92 return self._run(x, flat, indices.shape, probs_np)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._run(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],))9899100class TitanMoEBlock:101 """Drop-in replacement for the host model's sparse MoE block (OLMoE / Qwen3-MoE routing: softmax over all experts,102 top-k by argpartition, gating weights = softmax probabilities, renormalised iff norm_topk_prob). The router ops are103 exactly mlx-lm's (verified bit-identical in tests/test_numerical.py); the only additions are the prefetcher hooks."""104105 def __init__(self, host_mlp, block: "TitanSwitchGLU", model_type: str, prefetcher=None):106 self.gate = host_mlp.gate107 self.top_k = host_mlp.top_k108 self.norm_topk_prob = getattr(host_mlp, "norm_topk_prob", False)109 self.switch_mlp = block110 self.model_type = model_type111 self.prefetcher = prefetcher112 self.layer = block.table.layer113114 def __call__(self, x: mx.array) -> mx.array:115 B, L, D = x.shape116 x_flat = x.reshape(-1, D)117 k = self.top_k118 if self.model_type == "olmoe":119 router_logits = self.gate(x_flat)120 rw = mx.softmax(router_logits, axis=1, precise=True)121 indices = mx.stop_gradient(mx.argpartition(-rw, kth=k - 1, axis=-1)[..., :k])122 scores = mx.take_along_axis(rw, indices, axis=-1)123 else: # qwen3_moe124 gates = self.gate(x_flat)125 rw = mx.softmax(gates, axis=-1, precise=True)126 indices = mx.argpartition(rw, kth=-k, axis=-1)[..., -k:]127 scores = mx.take_along_axis(rw, indices, axis=-1)128 if self.norm_topk_prob:129 scores = scores / scores.sum(axis=-1, keepdims=True)130 self.switch_mlp.current_probs = rw if self.prefetcher is not None else None131 y = self.switch_mlp(x_flat, indices)132 y = (y * scores[..., None]).sum(axis=-2)133 return y.reshape(B, L, D)134