Source-linked AI summary
The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction
Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng
TL;DR
Consumer hardware cannot conveniently hold 35B MoE weights, while naive SSD offloading cannot hide expert-load latency because routing depends on prior-layer outputs. Edge0 uses trained one-token-ahead routing as the actual routing decision and an unmerged recovery LoRA; it serves a 35B MoE at 20tok/s within 3GiB on a 24GB machine, within a few points of its fp16 teacher across five benchmarks.
Problem
35B MoE inference on consumer hardware is constrained by 19.5GB of 4-bit weight memory, while sparsity reduces computation rather than stored bytes and naive SSD offloading cannot begin next-layer reads early enough.
Method
Edge0 predicts each next layer’s routing one token ahead, uses that prediction as the routing itself, and trains an unmerged recovery LoRA on the resulting student path.
Results
20tok/s at 3GiB of peak active memory lets Edge0 serve a 35B MoE on a 24GB machine within a few points of its fp16 teacher across five public benchmarks.
Takeaways & Limitations
The framework supports both 35B and 8B tiers on consumer hardware, with framework, checkpoints, and adapters released open source.
Takeaways & Limitations
Quality loss remains concentrated in long-chain reasoning, reaching 6.1 points on AIME for the 35B tier and 10.0 points for the 8B tier.
Abstract
from arXiv · showhide
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.
1 Introduction
Edge0 targets the weight-memory bottleneck that makes 35B MoEs impractical on consumer hardware. It combines SSD-resident expert streaming, trained prediction-as-routing, and an unmerged recovery LoRA to preserve quality while enabling two open model tiers.
- 35B MoE sparsity reduces per-token computation but leaves 19.5 GB of 4-bit weights to fit in memory, exhausting a 24 GB desktop alongside the OS and other workloads.
- Edge0 streams expert weights from SSD so peak memory follows the active set, then uses a per-layer prerouter to predict the next layer’s routing one token ahead.The prediction overlaps SSD reads with computation and is consumed as the routing itself, making staged and routed expert sets identical.
- The recovery LoRA is trained on the student routing path with a frozen int4 base, recovering quality lost to quantization and routing replacement without merging into the base.
- The framework releases both 35B and 8B tiers as checkpoints plus adapters, with each averaging within a few points of its fp16 base.
2 Background and Related Work
The background frames Edge0 as addressing the static-weight half of the memory problem rather than dynamic state. Existing approaches compress or relocate memory, but do not make disk reads aware of the next routed experts.
- MoE routing: MoE sparsity selects only K of E experts per layer, while Edge0 supports both softmax-top-k and sigmoid-group routing using the models’ original routing rules.
- Related routing: Pre-gated MoE chooses the next block’s experts within the same token, but streaming requires Edge0’s full-token lead to begin reads early enough.
- KV-side compression: KV-side methods compress dynamic state and are orthogonal to Edge0, whose 8B MLA-plus-MoE tier leaves more memory for active experts.
- Quantization: Quantization makes 4-bit the practical working point, so Edge0 accepts int4 loss and uses distillation-based adapters to recover it at serving time.
- Offloading: Prior offloading systems move expert footprints across CPU, GPU, or disk, but disk accesses are not informed by which experts the next step will need.
3 System Design
Edge0 is a backend-independent streaming framework that mmaps quantized experts, overlaps SSD reads with decode using a one-token prerouter, and serves an unmerged LoRA alongside the frozen int4 model. Its executor supports multiple correctness-preserving paths and outperforms a fully resident baseline when the checkpoint does not fit.
- Framework: Edge0 describes models with one MoESpec and routes all core logic through a backend facade, allowing one generic streaming layer to serve multiple models.
- SSD streaming: Quantized expert safetensors are mmapped and read by byte range, while the page cache carries hot data and active memory tracks the staged set rather than total parameters.
- Prerouter: A one-token prerouter predicts layer N+1 routing from layer N state, allowing SSD reads to overlap the current forward pass.
- Executor paths: The executor provides exact, staged, hot, and whole-layer paths, with staged decode using fixed-slot double buffering and whole-layer loading serving prefill.
- Correctness: All paths share the same quantized gather computation and are checked element-wise against a dequantized reference, enabling path switching without changing outputs.
- Performance rationale: 20.4 tok/s at 2.9 GiB versus 3.9 tok/s at 18.2 GiB for a fully resident baseline on the same 24 GB machine.The resident path cannot fit comfortably with the KV cache and operating system, while Edge0 pays streaming and assembly costs instead.
- Prediction-as-routing: At decode, prerouter predictions replace the original router, so staged and routed expert sets match exactly and no experts or tokens are dropped.
- Recovery LoRA: The LoRA is trained on the deployed student path and served as a parallel delta beside the int4 base, while merging and requantization erase most of its effect.
4 Training
Training first teaches prerouter heads to predict routing, then trains the recovery LoRA under deployed student routing and finally adapts it through on-policy distillation. The released routing width can be changed by swapping adapter files and retraining the tier.
- Training setup: The recipe trains on the dequantized reconstruction of the 4-bit deployment checkpoint with the base frozen throughout.
- Phase 1: Phase 1 distills prerouter heads against the next layer’s true router, teaching routing prediction rather than text modeling.
- Phase 2: Phase 2 performs SFT on roughly two million teacher-generated rows with prerouting active, and is necessary for coherent output after distillation-only student routing collapses.
- Training order: The training order is heads first, SFT second, and on-policy distillation last because otherwise SFT gradients overwhelm the small head gradients.
- Phase 3: Phase 3 trains on the student’s own generations using teacher-scored top-k tokens and reverse KL, rather than requiring coverage of the teacher’s full support.
- Routing width: Retraining for a different routing width only swaps adapter files; narrowing from K=8 to K=4 nearly doubles decode from 3.3 to 6.4 tok/s while retaining the released tier’s quality.
5 Evaluation
On machines where MoE weights do not fit, Edge0 overlaps SSD reads with decoding through prerouter staging, improving throughput while keeping active memory small. The evaluation also measures quality, memory trade-offs, storage costs, and phase-specific performance across released tiers.
- 5 Evaluation: Mean per-benchmark quality gaps versus fp16 bases are 3.9 points for 35B and 2.8 points for 8B after int4 quantization, routing replacement, and recovery LoRA.The quality comparison uses OpenCompass under identical settings.
- 5.3 The Advantage of the Prerouter: Blocked expert-load time falls from 154.9 to 46.5 ms at K=2, 244.0 to 101.9 ms at K=4, and 575.0 to 211.6 ms at K=8.The predictor moves serialized load latency off the critical path while reads overlap the forward pass.
- 5.3 The Advantage of the Prerouter: +80%, +82%, and +84% decode throughput at K=2, 4, and 8 comes from prefetching the routed expert set one token ahead.Throughput rises from 4.8, 3.5, and 1.8 tok/s on demand to 8.6, 6.4, and 3.3 tok/s with staging.
- 5.4 Memory: What the Machine Actually Pays: The prerouter reads nearly the same bytes at K=2 and K=8 but 16% more at K=4, in fewer, larger, colder loads that increase unreclaimable residency.At K=8, active memory rises by 1.43 GiB while page cache falls by 1.15 GiB; per-load size increases from 0.32 MiB to 1.40 MiB.
- 5.4 Memory: What the Machine Actually Pays: 20.4 tok/s at 2.9 GiB peak active memory lets the 35B tier outperform a fully resident server at 3.9 tok/s and 18.2 GiB.Edge0 uses one-sixth to one-seventh of the resident server’s memory and decodes five times faster.
- 5.4 Memory: What the Machine Actually Pays: Decode is the constrained phase: cold/warm prefill reaches 113/140 tok/s for 35B and 500/1102 tok/s for 8B, while the mechanisms target decode.The experiments use a 16 GB MacBook M2 with an 18.4 GiB checkpoint that does not fit, and a 24 GB Mac mini for tier profiles.
6 Limitations
Edge0’s gains are bounded by reuse, serving scope, and CPU-side assembly costs, while quality losses remain concentrated in long-chain reasoning.
- Serving scope: Edge0 serves one request at a time, FIFO-serialized, and does not model how batching changes the expert working set.Concurrency is delegated to a serving layer rather than handled by the engine.
- Compute boundary: 44 ms per step of CPU-side graph building is a floor that storage-side optimization cannot remove.Closing this residual gap requires kernel-level graph amortization or a smaller model.
- Reuse boundary: Adjacent tokens agree on only about a quarter of a layer’s expert set, so many prefetched experts are read once without reuse.Prerouter gains therefore depend on the storage tier and reuse, while unreclaimable residency displaces page-cache capacity.
- Quality boundary: 6.1 points on AIME for 35B and 10.0 for 8B remain the main quality losses, while every other 8B benchmark is within 6.7 points and MMLU-Pro favors the student by 4.3 points.The recovery LoRA recovers most of the pipeline’s loss elsewhere, but int4 quantization and routing replacement remain visible on reasoning tasks.
- Implementation scope: The MLX backend is the only implementation, while the backend facade defines the intended abstraction for additional backends.The CUDA slot is architecture rather than implemented code.
7 Conclusion
Edge0 places expert weights on SSD, predicts routing one token ahead, and uses an unmerged LoRA to recover quality under 4-bit quantization. It serves both model tiers on consumer hardware with low active memory and open-source artifacts.
- Conclusion: Experts live on SSD, a trained prerouter predicts routing one token ahead, and prediction is consumed as routing so staged reads overlap compute without dropping experts.A distilled, unmerged LoRA addresses quality loss from int4 quantization and routing replacement.
- Conclusion: 20 tok/s at 3 GiB on a 24 GB consumer desktop supports the 35B-class MoE, while the 8B hybrid reaches 28 tok/s at 1.5 GiB; both remain within a few points of fp16 teachers.The framework, checkpoints, and adapters are open.
- Conclusion: Framework, checkpoints, and adapters are released as open artifacts.The listed framework and model repositories are distributed under Apache-2.0.
A Routing Math
Edge0 preserves the model’s routing mathematics across resident, streaming, and prerouter paths by implementing the relevant selection functions verbatim and testing bit-identical parity.
- Parity: The routing functions are extracted verbatim from vendored model implementations and shared by resident, streaming, and prerouter paths.Parity tests pin their behavior bit-identically.
- Softmax-top-k: Softmax-top-k computes exact softmax expert logits, selects the top K experts, and renormalizes their weights.This routing function is one of the model implementations reused across all execution paths.
- Sigmoid-group: Sigmoid-group ranks groups by the sum of their top two sigmoid scores, keeps the best G groups, masks the rest, and selects top K experts within survivors.Selection uses biased scores, while expert weights use raw sigmoid values that are renormalized and scaled.
B The Assembly Tax and Incremental Stacks
Per-step tensor assembly remains a distinct bottleneck after prefetching, and persistent sticky-slot tensors reduce that assembly tax.
- Incremental stacks: +34% decode comes from replacing per-step stack rebuilding with persistent sticky-slot tensors updated in place.The A/B isolates assembly overhead rather than prefetching gains.
- Assembly tax: 6.8 to 12.5 tok/s is the reported same-session increase for native routing versus the complete shipped pipeline in a K=8 configuration.This pair comes from an earlier campaign and is not comparable with Table 3’s rates.