SPB Git

spb/wp33_uqo

Public MIT
2.3 KB · 30 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{Reproducibility.} Every benchmark writes JSON lines with the command, git commit, host, chip, memory, macOS,25library versions, seed and timestamps; raw logs are append-only under \code{benchmarks/raw/} and tables and figures are26generated from them by scripts. Model revisions are pinned by commit hash in \code{benchmarks/configs/models.yaml};27weights are never stored in the repository. A node is prepared with \code{scripts/cluster\_setup.sh} (rsync of the code,28pinned virtual environment, import check), and the sanitized inventory \code{artifacts/environment.json} is generated from29raw probes that are kept out of version control.30
Files · 27ccb9ab24f62ffe6b74e804d21dba925de70970
Menu

Theme