SPB Git

spb/wp33_uqo

Public MIT
5.8 KB · 69 lines latex
1\section{Introduction}2\label{sec:intro}34A sparse mixture-of-experts (MoE) transformer replaces the dense feed-forward block of each layer by a bank of $E$ expert5networks and a router that sends every token to $K \ll E$ of them \citep{rajbhandari2022deepspeedmoe}. The consequence that6matters for inference is arithmetic: a token activates $K/E$ of the expert parameters, so the \emph{working set} of a decode7step is small even when the \emph{model} is enormous. OLMoE-1B-7B activates 1.3 of its 6.9 billion parameters; in its84-bit MLX conversion a decode step touches 128 experts of 3.4 MB each, 432 MiB out of 3.6 GB of expert weights9(Section~\ref{sec:setup}). Engines such as colibri \citep{colibri2026} exploit exactly this to run models of hundreds of billions of10parameters from an NVMe disk on a laptop; the price is that the disk, not the GPU, sets the speed.1112Apple Silicon is an unusual target for this problem. Its unified memory removes the host-to-device copy that dominates13GPU offloading systems, its NVMe storage is fast (we measure 8.5~GB/s for random 3.4~MB reads on an M5~Max), and small14clusters of Macs are common in laboratories that cannot or do not want to buy datacenter accelerators. The MLX framework15\citep{hannun2023mlx} provides a quantized grouped matrix-multiplication kernel (\gatherqmm{}) that executes a batch of16expert products in one dispatch, and, since macOS~26.2, RDMA over Thunderbolt~5 between Mac Studios17\citep{apple2026mlxdistributed}. What does not exist is a runtime that decides, per expert and per token, \emph{where} an18expert should reside (local memory, a peer's memory, or the local disk), \emph{when} to move it (on demand or ahead of the19router), and \emph{which node} should execute it, using measured properties of the machines and links at hand.2021This working paper describes \titan{} (Topology-aware Intelligent exTert orchestrAtion for distributed MoE iNference), a22runtime built to answer one question stated as five testable hypotheses (Section~\ref{sec:hypotheses}): can topology-aware23placement, causal prefetching, hierarchical caching and specialised Metal kernels improve decode latency, prefill throughput,24energy per token and feasible model capacity over strong existing baselines on Apple Silicon? We insist on three25methodological commitments that shape the rest of the paper.2627\paragraph{Placement decides speed, never semantics.} The router's expert identities, gating weights and aggregation order28are those of the reference implementation (mlx-lm). A resident, a prefetched and a remotely executed expert must produce the29same bytes. Section~\ref{sec:results-correctness} shows that \titan{} with bounded residency reproduces mlx-lm's logits30exactly ($\max|\Delta| = 0$ in float16) on the test prompts; cache-aware routing, which trades expert identity for hits31\citep{skliar2025cacheaware}, is therefore out of scope by construction.3233\paragraph{Measure before optimising.} Every design decision in Section~\ref{sec:architecture} is preceded by a microbenchmark34recorded in the repository, from the cost of writing one expert into a stacked buffer (0.3~ms) to the per-expert fetch35latency of each tier (0.21~ms from the page cache, 0.43~ms from a cold NVMe, about 28~ms over the gigabit LAN that actually36connects the cluster). The inventory of the cluster is read live and is reported as read, including the seven nodes that did37not answer and the absence of any Thunderbolt link between the Mac Studios.3839\paragraph{Report losses.} Our first online experiment with one-layer-ahead prefetching \emph{raised} the expert hit rate and40\emph{lowered} throughput (Section~\ref{sec:results-prefetch}); the paper keeps that result and uses it to derive the41condition under which prefetching can pay on this hardware.4243\subsection{Hypotheses}44\label{sec:hypotheses}45\begin{description}[style=nextline]46  \item[H1 (causal prefetch)] Predicting upcoming expert demand from information available at decision time and fetching47    asynchronously reduces expert-loading stalls net of predictor overhead.48  \item[H2 (topology-aware placement)] Placing experts on heterogeneous nodes according to measured compute, memory and49    link costs outperforms uniform or static sharding.50  \item[H3 (hierarchical caching)] Coordinating local memory, remote memory and local NVMe under a fixed memory budget improves51    throughput over any single second tier.52  \item[H4 (Metal kernels)] Selected fused Metal kernels improve end-to-end inference over stock MLX kernels for matching53    shapes and precision.54  \item[H5 (adaptive orchestration)] Adapting placement and prefetching online improves the Pareto frontier of latency,55    throughput and energy, not only isolated microbenchmarks.56\end{description}5758\subsection{Contributions and status of this version}59Version~0.1 of this working paper contributes (i) a verified inventory and network characterisation of a 21-node Apple Silicon60cluster and the audit of four reference code bases (colibri, MLX, mlx-lm, llama.cpp) with a prior-art matrix of 59 works;61(ii) the \titan{} single-node runtime: an expert registry over safetensors byte ranges, a \code{pread}-based NVMe tier,62per-layer slot tables with LRU/LFU eviction, an exact execution block over \gatherqmm{}, and an asynchronous prefetcher63with four causal predictors; (iii) the correctness evidence (bit-identical logits, token-identical greedy decoding);64(iv) a trace-driven simulator calibrated on measured constants and the first controlled single-node experiments on65OLMoE-1B-7B; (v) a formal statement of the placement--prefetch problem (Section~\ref{sec:formulation}) and the deterministic66cost-aware heuristic that the distributed gates will use. Distributed execution (Gates E--F of the protocol), Metal kernel work67(Gate~D) and the energy Pareto study are implemented or planned as described in Section~\ref{sec:limitations} and reported only68to the extent that raw logs exist.69
Files · f419985eab99e178b5c98c31ba0a179ada784682
Menu

Theme