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

AI summary

Overview

• Research area: Efficient large-model inference systems — specifically Mixture-of-Experts (MoE) serving on consumer hardware, combining SSD weight streaming, learned routing prediction, and 4-bit quantization with adapter-based recovery.
• Technical level: Advanced. The paper assumes familiarity with MoE routing (softmax-top-k, sigmoid-group), quantization (int4 affine group-64, GPTQ/AWQ-style tradeoffs), LoRA adapters, distillation objectives, and memory-hierarchy/serving systems work.
• Scope in one sentence: The paper describes Edge0, a streaming MoE inference engine that keeps expert weights on SSD, predicts the next layer's routing one token ahead so expert reads overlap compute, and recovers the resulting quality loss with a distilled, unmerged recovery LoRA — serving a 35B-class MoE at ≈20 tok/s inside 3 GiB of peak active memory on a single 24 GB machine.

What This Paper Is About

MoE inference on consumer hardware is bounded not by compute but by weight memory: a 35B-class model is 19.5 GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Simply offloading weights to SSD does not solve the problem on its own, because layer N+1's experts must be chosen before layer N's output exists, so the disk reads cannot start early enough to hide behind compute. The goal is to break that dependency with a trained predictor so that expert reads overlap the forward pass, and then to pay back the quality cost of both the routing approximation and 4-bit quantization during training rather than at inference time.

Key Contributions

1. SSD as the weight tier. An expert executor that streams MoE expert weights from disk on demand, with routing math pinned bit-identical to the vendored models and every executor path verified element-wise against the dequantized reference, plus a slot mechanism (incremental sticky-slot stacks) that removes the per-step stack-rebuild tax.
2. Prediction as routing. A cross-token prerouter whose prediction is the routing at decode, so the staged expert set equals the routed set and nothing is dropped. This is worth +80% to +84% decode on a machine where the checkpoint does not fit, and the paper supplies the training recipe that makes a language model work under replaced routing.
3. Unmerged recovery LoRA. Merging an adapter into an int4 base and re-quantizing destroys most of its effect; serving it as a parallel delta avoids this at negligible cost (42 MB of adapter weights, no measurable decode time).
4. Two open tiers, end to end. A 35B and an 8B model released as checkpoint plus adapters that load together, each within a few points of its fp16 base on average, with the framework open source under Apache-2.0.

Main Findings

• Decode speed and memory on the 35B tier: On a Mac mini M4 Pro with 24 GB, edge0-35b (K=4) decodes at 20.4 tok/s with a peak active memory of 2.9 GiB, against a 19.5 GB checkpoint on disk. A vanilla mlx-lm server holding all 19.5 GB of 4-bit weights resident decodes at 3.9 tok/s occupying 18.2 GiB.
• Decode speed and memory on the 8B tier: edge0-8b (K=8) decodes at 28.0 tok/s inside 1.5 GiB of peak active memory, from a 4.5 GB checkpoint. Enlarging its shared LRU from 64 to 1024 expert bundles costs ≈0.9 GiB and lifts decode to 31.8 tok/s.
• Prefill is the mirror image: cold/warm prefill is 113/140 tok/s for the 35B tier on a 3.1k-token prompt, and 500/1102 tok/s for the 8B tier. Decode, not prefill, is the constrained phase.
• The prerouter's measured advantage: On a 16 GB MacBook M2 serving an 18.4 GiB checkpoint that does not fit, staged prefetch beats pure on-demand streaming by +80% at K=2 (8.6 vs 4.8 tok/s), +82% at K=4 (6.4 vs 3.5), and +84% at K=8 (3.3 vs 1.8). Disabling the prerouter on the 35B tier measures 19.9 tok/s.
• The gain is moved latency, not fewer bytes: Both arms read within a few percent of the same bytes per step at K=2 (29.5 vs 30.4 MiB) and K=8 (125.1 vs 126.9 MiB); only at K=4 does the prerouter arm read more (58.9 vs 50.9 MiB, 16% more). Time blocked on expert loads falls from 154.9 to 46.5 ms per step at K=2, 244.0 to 101.9 ms at K=4, and 575.0 to 211.6 ms at K=8.
• Prefetching shifts the cost to page cache: The prerouter issues fewer, larger, colder loads — 1.40 MiB per load at 90% cold versus 0.32 MiB per load at 20% cold at K=8. Per-load cost follows 1.17 ms + 1.33 ms × (cold fraction), fitted on the K=8 pair and reproducing every measured per-load cost within 0.13 ms. At K=8 the prerouter arm holds 1.43 GiB more unreclaimable MLX memory while the page cache loses 1.15 GiB, four fifths of what the residency takes.
• Resource headroom remains: The disk read is at most 12% of the step even at its widest, the process uses about one core of eight, and the GPU runs at 35–41%.
• Quality is close to fp16 across five benchmarks: Mean per-benchmark gaps of 3.9 points (35B) and 2.8 points (8B). For edge0-35b vs Qwen3.6 fp16: AIME 2026 86.6 vs 92.7, HumanEval 90.9 vs 95.1, GPQA-Diamond 79.8 vs 81.8, MMLU-Pro 81.0 vs 84.6, IFBench 57.9 vs 61.7, average 79.2 vs 83.2. For edge0-8b vs Ling 3.0 tiny fp16: 63.3 vs 73.3, 91.5 vs 92.7, 70.7 vs 71.2, 70.1 vs 65.8, 53.9 vs 60.6, average 69.9 vs 72.7.
• Loss concentrates in long-chain reasoning: 6.1 points on AIME for the 35B tier and 10.0 for the 8B tier; every other 8B benchmark is within 6.7 points, and MMLU-Pro is 4.3 points in the student's favour.
• Merging the LoRA destroys its effect: LoRA deltas (RMS 10⁻³) sit below the 4-bit group step; on requantization 34% of the effect survives on an attention projection and 2% on a dense projection, and 18% survives at the logits level, versus serving unmerged.
• Training order mattered: Distillation-only checkpoints with student routing produce repetitive, collapsed text, while the same heads plus SFT produce coherent output at identical speed and memory. The order was heads first, SFT second, on-policy distillation last. Phase 3 converges on roughly 200k rows against the ≈2M of Phase 2.
• Assembly cost was separately measured: Replacing per-step stacks with persistent sticky-slot tensors contributes a further +34% decode in a separate A/B. In a second same-session A/B on the K=8 configuration (native routing against the complete shipped pipeline), decode rises from 6.8 to 12.5 tok/s. The authors note these two pairs come from an earlier campaign and a different configuration, so their rates are not comparable with the prerouter table.
• Measurement discipline: Single-shot benchmarks on this hardware carry ±40% run-to-run spread and up to 2.3× across sessions on the 16 GB machine, so all speed comparisons use same-session rotated A/B with page-cache warmup, matched cache budgets, and medians over repeated runs.

Methodology in Plain English

The authors start from the observation that the 35B model's expert weights — 453 MB per layer at int4, or 18 GB (16.9 GiB) of routed experts inside the 19.5 GB checkpoint — cannot live in 24 GB of RAM alongside an operating system and a KV cache. They keep those experts on SSD as int4 per-layer stacked safetensors, memory-map them, and read byte ranges on demand, so peak memory tracks the active set rather than the parameter count.

On-demand loading alone stalls, because a layer's experts must be chosen before the previous layer has finished producing the hidden state the router would use. Their fix is to train a small per-layer head, the prerouter, that predicts layer N+1's routing at token t+1 from layer N's post-attention norm output at token t. The prediction is then used as the routing, so the staged expert set and the routed set are identical by construction and no expert is dropped. The head is a small MLP (fc1 → erf-gelu → fc2) with a linear residual path warm-started from the next layer's router weight; its input concatenates the hidden state with two top-k one-hots (the experts this layer actually routed to at this token and at the previous token).

Because a single flush per step can predict every staged layer at once, the head evaluation moves off the per-layer critical path — the paper reports that same-token, per-layer variants cost 30–100 ms per step in their engine and every one they measured fell below a plain LRU baseline. The engine has four executor paths (exact, staged, hot, whole-layer) that all compute the same math and are verified element-wise against a dequantized reference. Quantization and routing replacement cost quality, so a recovery LoRA is distilled from the fp16 teacher on the student path (routing by the prerouter) in three phases, and shipped unmerged as a parallel delta.

Why This Matters

• Research impact: The paper reframes the "memory wall" for inference as two problems — dynamic KV state and static weights — and argues that the field has attacked only the dynamic half. It shows that the weight half is a storage placement decision, and that a learned, cross-token routing predictor is a principled way to make that placement work. The transfer of the quality cost from inference-time fallback to one-time training is the distinctive move.
• Real-world applications:
  • Running a 35B-class assistant locally on a 24 GB consumer desktop, where the alternative is a datacenter or a smaller model.
  • Privacy-sensitive or offline deployments where prompts and outputs cannot leave the machine.
  • Laptop-class inference (16 GB) of checkpoints that exceed RAM, where the prerouter's +80% to +84% decode gain is measured.
  • Tiered product lines, since retraining a tier for a different routing width is a file swap rather than a training campaign.
• Industry relevance: The released artifacts (framework, checkpoints, adapters, Apache-2.0) lower the barrier to reproducing and extending the approach. If routing prediction becomes a standard component, the economics of local serving change: memory becomes the scarce resource that is bounded, and disk becomes the tier that is tolerated, which reshapes how models are packaged and how serving stacks are designed.

Future Directions

• Close the CPU-side assembly gap. The authors state the streaming engine's decode is CPU-side, not storage-side: 44 ms per step of graph building in the 40-layer forward is the floor, and no storage-side optimization moves it. Kernel-level graph amortization, or a smaller served model, is the suggested route.
• Pull the two identified levers. Dropping the duplicate bundle/stacked-tensor copy in the stager would return about 0.45 GiB at K=8, and moving the incremental-stack fill off the main thread would remove a 62.5 ms per-step cost at K=8. Neither change is in the measured configuration.
• Add concurrency and batching. Edge0 currently serves one request at a time, FIFO-serialized; the authors note that batching changes the expert working set in ways their per-request profiles do not model.
• Address the reasoning gap and reuse limits. Quality loss concentrates in long-chain reasoning (6.1 points on AIME for 35B, 10.0 for 8B), and the prerouter's ceiling is bounded by reuse: adjacent tokens agree on only about a quarter of a layer's expert set, so a prefetched expert is often read once and paid for twice. Better reuse modeling or a better head is the open question.
• Broaden backends. The MLX backend is the only implementation; the paper describes the CUDA slot as architecture, not code.

Target Audience

Systems and inference engineers building local or edge serving stacks; MoE and quantization researchers interested in routing-level approximations and adapter-serving arithmetic; practitioners who need to run 30B-class models on consumer hardware with limited VRAM or unified memory; and anyone studying memory-hierarchy tradeoffs in LLM deployment. Readers should already be comfortable with MoE routing variants, post-training quantization, and LoRA/distillation training, since the paper's central claims interleave all three.

Authors’ abstract

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.