"""TITAN expert execution block, drop-in replacement for mlx-lm's `SwitchGLU`. The router (gate linear, softmax, top-k, gating weights, weighted sum) stays in the host model's code so that expert identity and aggregation are exactly those of the reference implementation; this module only decides *where* the selected experts' weights come from and executes them with `mx.gather_qmm` on the layer's slot table. """ from __future__ import annotations from typing import Callable, List, Optional import mlx.core as mx import numpy as np from mlx_lm.models.activations import swiglu from .slots import SlotTable class TitanSwitchGLU: """Callable with the SwitchGLU signature: (x[..., 1, D], indices[..., k]) -> y[..., k, D].""" def __init__(self, table: SlotTable, group_size: int, bits: int, activation=None, on_route: Optional[Callable[[int, List[int]], None]] = None, prefetcher=None): self.table = table self.prefetcher = prefetcher self.current_probs = None self.remote = None # RemoteExecutor (Gate E); when set, experts placed on workers are executed remotely self.group_size = group_size self.bits = bits self.on_route = on_route # hook(layer, unique_experts) used by predictors/prefetchers (Gate C) self.last_indices: Optional[List[int]] = None def _qmm(self, x, m, idx, sorted_indices): b = self.table.buf return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx, transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices) def _run_device(self, x, slot_idx): x = mx.expand_dims(x, (-2, -3)) x_up = self._qmm(x, "up", slot_idx, False) x_gate = self._qmm(x, "gate", slot_idx, False) h = swiglu(x_gate, x_up) y = self._qmm(h, "down", slot_idx, False) return y.squeeze(-2) def _run(self, x, flat, shape, probs_np=None): uniq = list(dict.fromkeys(flat)) if self.on_route is not None: self.on_route(self.table.layer, uniq) pf = self.prefetcher if pf is not None: pf.before_layer(self.table.layer, uniq) # install completed prefetches (wait for late ones we need) pf.after_layer(self.table.layer, uniq, probs_np) # issue prefetch for the next layer BEFORE this layer's compute slot_of = dict(zip(uniq, self.table.ensure(uniq))) slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape) x = mx.expand_dims(x, (-2, -3)) x_up = self._qmm(x, "up", slot_idx, False) x_gate = self._qmm(x, "gate", slot_idx, False) h = swiglu(x_gate, x_up) # identical compiled activation to mlx-lm SwitchGLU y = self._qmm(h, "down", slot_idx, False) return y.squeeze(-2) def __call__(self, x: mx.array, indices: mx.array) -> mx.array: t = self.table if t.complete(): # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync. self.n_chunks = 1 self.last_indices = None t.hits_device += indices.size return self._run_device(x, t.slot_map[indices]) # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design. probs_np = None if self.prefetcher is not None and self.current_probs is not None: mx.eval(indices, self.current_probs) probs_np = np.array(self.current_probs[-1]) # last row: decode token (prefill rows are not used for prediction) flat = indices.reshape(-1).tolist() self.last_indices = flat k = indices.shape[-1] n_rows = len(flat) // k cap = self.table.capacity x2 = x.reshape(n_rows, x.shape[-1]) # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent). chunks, cur, cur_set = [], [], set() for r in range(n_rows): es = set(flat[r * k:(r + 1) * k]) if cur and len(cur_set | es) > cap: chunks.append(cur); cur, cur_set = [], set() cur.append(r); cur_set |= es if cur: chunks.append(cur) self.n_chunks = len(chunks) if self.remote is not None: outs = [] for ci, rows in enumerate(chunks): f = [flat[r * k + j] for r in rows for j in range(k)] 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)) return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],)) if len(chunks) == 1: return self._run(x, flat, indices.shape, probs_np) outs = [] for ci, rows in enumerate(chunks): f = [flat[r * k + j] for r in rows for j in range(k)] outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k), probs_np if ci == len(chunks) - 1 else None)) return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],)) class TitanMoEBlock: """Drop-in replacement for the host model's sparse MoE block (OLMoE / Qwen3-MoE routing: softmax over all experts, top-k by argpartition, gating weights = softmax probabilities, renormalised iff norm_topk_prob). The router ops are exactly mlx-lm's (verified bit-identical in tests/test_numerical.py); the only additions are the prefetcher hooks.""" def __init__(self, host_mlp, block: "TitanSwitchGLU", model_type: str, prefetcher=None): self.gate = host_mlp.gate self.top_k = host_mlp.top_k self.norm_topk_prob = getattr(host_mlp, "norm_topk_prob", False) self.switch_mlp = block self.model_type = model_type self.prefetcher = prefetcher self.layer = block.table.layer def __call__(self, x: mx.array) -> mx.array: B, L, D = x.shape x_flat = x.reshape(-1, D) k = self.top_k if self.model_type == "olmoe": router_logits = self.gate(x_flat) rw = mx.softmax(router_logits, axis=1, precise=True) indices = mx.stop_gradient(mx.argpartition(-rw, kth=k - 1, axis=-1)[..., :k]) scores = mx.take_along_axis(rw, indices, axis=-1) else: # qwen3_moe gates = self.gate(x_flat) rw = mx.softmax(gates, axis=-1, precise=True) indices = mx.argpartition(rw, kth=-k, axis=-1)[..., -k:] scores = mx.take_along_axis(rw, indices, axis=-1) if self.norm_topk_prob: scores = scores / scores.sum(axis=-1, keepdims=True) self.switch_mlp.current_probs = rw if self.prefetcher is not None else None y = self.switch_mlp(x_flat, indices) y = (y * scores[..., None]).sum(axis=-2) return y.reshape(B, L, D)