SPB Git

spb/wp33_uqo

Public MIT
5.5 KB · 73 lines latex
1\section{Architecture}2\label{sec:architecture}34Figure~\ref{fig:arch} shows the components of \titan{}. The host model (embeddings, attention, norms, routers, KV cache) is5mlx-lm's own module tree, loaded lazily so that the stacked expert tensors are never materialised; each layer's expert block6is replaced by a \titan{} block that owns a \emph{slot table}. Everything else is orchestration around that table.78\begin{figure}[t]9\centering10\begin{tikzpicture}[node distance=4mm and 6mm, every node/.style={font=\small},11  box/.style={draw, rounded corners=2pt, minimum height=7mm, minimum width=24mm, align=center, fill=gray!8},12  tier/.style={draw, dashed, rounded corners=2pt, minimum height=7mm, minimum width=24mm, align=center},13  arr/.style={-{Latex[length=2mm]}, thick}]14  \node[box] (reg) {Expert registry\\\scriptsize safetensors byte ranges};15  \node[box, right=of reg] (store) {Expert store\\\scriptsize \code{pread}, F\_NOCACHE};16  \node[box, right=of store] (slots) {Slot tables ($L$)\\\scriptsize capacity $C$, LRU/LFU};17  \node[box, right=of slots] (exec) {Execution block\\\scriptsize \gatherqmm{} on slots};18  \node[box, below=of store] (pred) {Predictors\\\scriptsize prev-rank, transition$_d$, probe};19  \node[box, right=of pred] (pf) {Prefetcher\\\scriptsize reader pool, drain/install};20  \node[box, right=of pf] (host) {Host model (mlx-lm)\\\scriptsize attention, router, KV};21  \node[tier, below=of pf] (net) {Transport (Gate E)\\\scriptsize TCP, authenticated};22  \node[tier, right=of net] (sched) {Scheduler (Gate F)\\\scriptsize cost model, placement};23  \draw[arr] (reg) -- (store);24  \draw[arr] (store) -- (slots);25  \draw[arr] (slots) -- (exec);26  \draw[arr] (host) -- node[above,font=\scriptsize]{routes} (exec);27  \draw[arr] (exec) -- node[left,font=\scriptsize]{routes, probs} (pf);28  \draw[arr] (pf) -- (pred);29  \draw[arr] (pf) -- node[left,font=\scriptsize]{install} (slots);30  \draw[arr] (pred) -- node[above,font=\scriptsize]{$\hat S$} (pf);31  \draw[arr, dashed] (net) -- (slots);32  \draw[arr, dashed] (sched) -- (net);33\end{tikzpicture}34\caption{Components of \titan{} on one node. Solid boxes are implemented and tested in version 0.1; dashed boxes are the35distributed components whose status is given in Section~\ref{sec:limitations}.}36\label{fig:arch}37\end{figure}3839\paragraph{Expert registry.} The registry parses \code{config.json} and the safetensors headers of a checkpoint directory40and records, for every (layer, expert) and every matrix (gate, up, down) and component (packed weight, scales, biases), the41absolute byte offset and length in the file, together with the routing semantics (softmax-then-top-$k$, renormalisation flag,42shared experts) and the quantization (bits, group size). It accepts only verified formats (OLMoE and Qwen3-MoE in this43version) and refuses the others rather than guessing.4445\paragraph{Expert store (SSD tier).} An expert is read with one \code{pread} per tensor at the registry's offsets, wrapped46into MLX arrays without further copies on the Python side; bytes, read count and read time are accounted. The file descriptor47can be opened with \code{F\_NOCACHE} so that reads bypass the unified buffer cache, which is how the cold-NVMe tier is48measured (after \code{purge}); without it the same code measures the page-cache tier.4950\paragraph{Slot tables.} Each layer owns $C$ slots in stacked MLX buffers of the exact shapes mlx-lm uses for its stacked51experts, so that the same \gatherqmm{} kernel can run on them. A dictionary maps resident expert ids to slots; eviction is52LRU by default (LFU available). Installing an expert is a slice assignment into the buffers, which MLX performs in place.53Misses are read as a batch per layer and installed lazily so that the layer's kernel consumes the new buffers in the same54graph: there is no per-expert synchronisation. The table records hits, misses, evictions, loads and stall time.5556\paragraph{Execution block.} Given the router's indices, the block materialises them on the host (the one synchronisation a57demand-loading design cannot avoid), ensures residency, remaps expert ids to slot ids and runs \gatherqmm{} for the up, gate58and down projections with mlx-lm's compiled SwiGLU. When a layer's table holds every expert (the all-resident control) a59device-side slot map is used and no synchronisation occurs, which recovers mlx-lm's throughput. Prefill rows whose union of60experts exceeds $C$ are chunked greedily by rows, which is exact because rows are independent.6162\paragraph{Predictors and prefetcher.} A predictor proposes up to $m$ experts for a layer from information available at the63decision point: the previous token's routing and router probabilities (token horizon), or the routing of layer $l-d$ at the64current token (layer horizon, depth $d$). The prefetcher submits reads to a small thread pool, installs completed fetches when65the target layer is reached, waits for fetches that are needed but still in flight (counted as \emph{late}, with the residual66wait charged as stall), and never changes which experts run.6768\paragraph{Hierarchical memory and distributed execution.} The remote tier does not move weights during decode (Section69\ref{sec:prefetch-condition}): a remote node that holds an expert executes it on the coordinator's activations and returns the70$d$-dimensional result, exactly as colibri's cluster mode and as expert parallelism do. The transport (Gate E) is a persistent71authenticated TCP connection per worker with a per-layer batched request, timeouts and local fallback; placement (Gate F)72follows the cost model of Section~\ref{sec:formulation} with measured $c_n$, $M_n$, $\rho_{0n}$ and $\lambda$ per node.73
Files · 2d987b30cb427aa933cb485dd908495f7ea82b47
Menu

Theme