The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Reading guide

Problem and motivation

Mixture-of-Experts (MoE) models scale total parameters while keeping active parameters low per token (e.g., activating ~3B parameters per token in a 35B model). However, inference on consumer hardware remains constrained by static weight memory: a 35B MoE requires roughly 19.5 GB even when quantized to int4. On a 24 GB machine, this completely crowds out the operating system and KV cache, leading to severe memory paging and degraded decode throughput (e.g., 3.9 tok/s for a resident model on an M4 Pro Mac mini).

Offloading expert weights to SSD reduces memory footprint, but naive on-demand loading introduces a hard serialization bottleneck: expert selection for layer $N+1$ depends on the output of layer $N$. The SSD reads cannot start early enough to hide behind compute, stalling execution on disk latency once per layer per token.

Key ideas

• SSD as the Weight Tier: MoE expert weights stay on SSD in int4 format and are memory-mapped (mmap) into a compact streaming pool. Peak memory is determined by the active working set rather than total model parameter size.
• Prediction as Routing: Instead of waiting for layer $N$ to complete, a per-layer prerouter predicts layer $N+1$'s routing a full token in advance. Crucially, the model treats the prediction as the routing itself at decode time. This eliminates dropped tokens or fallback loads by construction because the staged set identically equals the routed set.
• Unmerged Recovery LoRA: 4-bit quantization and approximate routing incur quality loss. Edge0 trains a low-rank adapter on the student routing path via distillation. Because merging the LoRA into the int4 base and re-quantizing erases most of the delta (only 18% logits-level retention survives), Edge0 serves the adapter as a parallel unmerged delta at negligible overhead.

Method

Prerouter Head Architecture

Each prerouter head owned by layer $N$ runs at token $t$ on the post-attention layer norm output and predicts routing for layer $N+1$ at token $t+1$. The input features concatenate the hidden state with two top-$k$ one-hot routing vectors (current and previous token):

$$\text{dim}=d_{\text{model}}+2E$$

For the 35B tier, this is $2048+2\times 256=2560$. The head computes:

$$\text{head}(x)={\mathrm{fc}}_2(\mathrm{erf}\text{-}\mathrm{gelu}({\mathrm{fc}}_1(x)))+\ell (x)$$

where $\ell (x)$ is a linear residual path warm-started from the next layer's router weights.

Unmerged Adapter Execution

The served MoE layer preserves the quantized base bytes and runs the parallel LoRA delta:

$$y=W_{int4}(x)+\frac{\alpha }{r}BAx$$

Three-Phase Training

1. Head Distillation: Prerouter heads are trained to imitate the teacher's next-layer true router decisions while keeping the rest of the network frozen.
2. Student-Path SFT: Recovery LoRA adapters (attached to attention, linear-attention, and shared-expert projections, but not streaming experts) are trained using teacher-generated text while forcing routing through the prerouter.
3. On-Policy Distillation: The student generates text under prerouted execution, the fp16 base teacher scores the generations, and the LoRA parameters are optimized via reverse Kullback–Leibler divergence with a tail-aware top-$k$ probability objective.

Results

• Throughput & Memory (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Achieves 20.4 tok/s decode using only 2.9 GiB peak active memory from a 19.5 GB int4 checkpoint (compared to 3.9 tok/s and 18.2 GiB when fully resident).
  • edge0-8b ($K=8$): Achieves 28.0 tok/s decode inside 1.5 GiB peak memory.
• Prefetching Speedup (MacBook M2, 16 GB):
  • Compared to purely reactive on-demand streaming, the 1-token prerouter prefetch increases decode speed by +80% to +84% across routing widths ($K=2$: 4.8 to 8.6 tok/s; $K=4$: 3.5 to 6.4 tok/s; $K=8$: 1.8 to 3.3 tok/s).
  • Main-thread blocked load time falls dramatically (e.g., from 244.0 ms to 101.9 ms per step at $K=4$).
• Model Quality (OpenCompass):
  • Across five benchmarks (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), edge0-35b scores an average of 79.2 vs. 83.2 for the fp16 Qwen3.6 base (3.9-point gap).
  • edge0-8b achieves 69.9 vs. 72.7 for the Ling 3.0 tiny fp16 base (2.8-point gap).

Limitations and open questions

• Serving Concurrency: Edge0 operates FIFO-serialized for single requests; batching changes the expert working set and caching dynamics in ways not addressed by the single-stream profiles.
• CPU Graph-Building Overhead: Decode speed is bounded by CPU-side framework overhead (~44 ms/step in MLX graph building across 40 layers), which SSD optimizations cannot alleviate.
• Quality Drop on Reasoning: Performance degradation is concentrated on complex reasoning benchmarks (e.g., a 6.1-point drop on AIME for 35B and 10.0 points for 8B).
• Hardware Dependency: The prerouter's relative throughput gains depend on disk latency and shrink when weight caches are warm or when running on ultra-fast storage tiers.