SPB Git

spb/wp33_uqo

Public MIT
2.9 KB · 36 lines latex
1\section{Implementation}2\label{sec:implementation}34\titan{} is about 1\,100 lines of Python over MLX 0.32.4 and mlx-lm 0.32.2 (Python 3.12), organised as5\code{titan/\{registry, store, slots, moe, prefetch, predictors, runtime\}.py}, with benchmark drivers in6\code{benchmarks/scripts/} and tests in \code{tests/}. Three implementation findings determined the design.78\paragraph{In-place slot updates.} Writing the nine tensors of one OLMoE expert (3.4~MB) into stacked buffers of 64 slots9(216~MB per matrix) costs 0.305~ms, which shows that MLX reuses the buffer rather than copying it; a per-expert10\code{quantized\_matmul} loop costs 0.43~ms per layer for one token against 0.26~ms for \gatherqmm{}, and restacking eight11experts per call 0.50~ms. The slot table with \gatherqmm{} is therefore both the fastest and the most faithful option12(\code{benchmarks/raw/gateB/design/}).1314\paragraph{No per-miss synchronisation.} The first version evaluated the table after every miss; removing that15synchronisation and batching the reads of a layer's misses raised decode throughput at 16 slots from 28 to 51~\tps{}16(general prompts) and from 41 to 73~\tps{} (code), with the stall halved17(\code{benchmarks/raw/gateC/prefetch\_online\_laptop/}). The per-miss cost on the page-cache tier is now 0.21~ms: 0.13~ms18of \code{pread}, 0.06~ms of array creation and the remainder in the lazy install.1920\paragraph{Activation parity.} An unfused float16 $\mathrm{up}\cdot\sigma(\mathrm{gate})\cdot\mathrm{gate}$ gave21$\max|\Delta\mathrm{logit}| = 0.0195$ against mlx-lm; using mlx-lm's compiled SwiGLU gives $0.0$. Numerical parity in float16 is22a property of the exact kernel sequence, not only of the weights, and the test suite pins it.2324\paragraph{Threads and dtypes.} Two defects were found by the exactness tests rather than by inspection. MLX arrays created in reader25threads cannot be evaluated in another thread (streams are thread-local), so reader threads now return NumPy views and the consuming26thread creates the MLX arrays. And the first transport cast activations to float16 on the wire, which is exact for OLMoE (float16) but27not for Qwen3 (bfloat16): the first two-node sweep on Qwen3 showed $\max|\Delta\mathrm{logit}| = 0.65$ and was discarded; activations and28results now travel in the model's dtype and the Qwen3 two-node path is bit-identical again.2930\paragraph{Reproducibility.} Every benchmark writes JSON lines with the command, git commit, host, chip, memory, macOS,31library versions, seed and timestamps; raw logs are append-only under \code{benchmarks/raw/} and tables and figures are32generated from them by scripts. Model revisions are pinned by commit hash in \code{benchmarks/configs/models.yaml};33weights are never stored in the repository. A node is prepared with \code{scripts/cluster\_setup.sh} (rsync of the code,34pinned virtual environment, import check), and the sanitized inventory \code{artifacts/environment.json} is generated from35raw probes that are kept out of version control.36
Files · 2d987b30cb427aa933cb485dd908495f7ea82b47
Menu

Theme