SPB Git

spb/wp33_uqo

Public MIT
2.6 KB · 41 lines latex
1\section{Algorithms}2\label{sec:algorithms}34\begin{algorithm}[t]5\caption{Decode step of one MoE layer in \titan{} (single node)}6\label{alg:layer}7\begin{algorithmic}[1]8\Require hidden rows $X$, router indices $I$ ($R \times K$), slot table $\mathcal{S}_l$, prefetcher $\mathcal{P}$, depth $d$9\If{$\mathcal{S}_l$ holds every expert} \State \Return $\gatherqmm(X, \mathcal{S}_l, \mathrm{slotmap}[I])$ \Comment{no host sync}10\EndIf11\State $U \gets$ unique experts of $I$ (host) \Comment{the unavoidable synchronisation}12\State $\mathcal{P}.\textsc{drain}(l, U)$ \Comment{install finished prefetches; wait for in-flight experts in $U$ (late)}13\State $\hat S \gets \mathcal{P}.\textsc{predict}(l + d)$; $\mathcal{P}.\textsc{issue}(l + d, \hat S)$ \Comment{before this layer's compute}14\State $M \gets U \setminus \mathrm{resident}(\mathcal{S}_l)$; read $M$ from the store as a batch; install into slots (lazy)15\State $J \gets \mathrm{slot}(I)$16\State \Return $\gatherqmm(X, \mathcal{S}_l, J)$ \Comment{same kernel, same bytes as mlx-lm}17\end{algorithmic}18\end{algorithm}1920\begin{algorithm}[t]21\caption{Deterministic cost-aware placement (Gate F heuristic)}22\label{alg:placement}23\begin{algorithmic}[1]24\Require expert popularity $\pi_e$ from fit traces, node costs $(c_n, M_n, \rho_{0n})$, expert size $b$25\State sort experts by $\pi_e$ descending; $\mathrm{free}_n \gets M_n$26\For{each expert $e$ in that order}27  \State $n^\star \gets \arg\min_{n:\ \mathrm{free}_n \ge b} \big(c_n + \ind[n \ne 0]\,\rho_{0n}\big) \cdot \pi_e$ \Comment{expected per-step cost}28  \State place $e$ on $n^\star$; $\mathrm{free}_{n^\star} \gets \mathrm{free}_{n^\star} - b$29\EndFor30\State replicate the $\rho$ most popular experts on node $0$ if $\mathrm{free}_0$ allows (hot set)31\State experts placed nowhere fall back to node $0$'s slot table (SSD tier, demand loading)32\end{algorithmic}33\end{algorithm}3435Algorithm~\ref{alg:layer} is the per-layer path of Section~\ref{sec:architecture}; the prefetch for layer $l+d$ is issued36\emph{before} the current layer's kernels are dispatched so that the reader threads overlap the GPU work of the next $d$ layers.37Algorithm~\ref{alg:placement} is the first placement heuristic: greedy by popularity, cost-aware by node, with replication of a38hot set on the coordinator. It runs in $\mathcal{O}(LE \log LE + LE N)$ time, negligible against inference, and will be the39static baseline against which the adaptive scheduler (online re-estimation of $\pi$ and of the node costs) is compared.40Predictor fitting is $\mathcal{O}(T L K^2)$ for transition counts and one logistic regression per layer for the probe.41
Files · 2d987b30cb427aa933cb485dd908495f7ea82b47
Menu

Theme