SP-MoE: Speculative Decoding and Prefetching for Accelerating MoE-based Model Inference
Liangkun Chen, Zijian Wen, Tian Wu, Xiaoxi Zhang, Chuan Wu
cs.DC
2025-10-12
SP-MoE feeds draft attention to the target gate and prefetches experts during drafting, cutting TPOT 1.07x–3.5x vs SD-enabled offload systems.
MoE cuts compute by routing each token through a few experts, but the full parameter set still has to live somewhere. Mixtral 8×7B needs about 87 GB at inference, more than a 24 GB RTX 4090. The usual fix is to park idle experts in CPU memory and pull them over PCIe when the router asks. Loading one Mixtral expert layer over PCIe 4.0 takes about 80 ms; computing it on a 4090 takes about 3 ms. The GPU waits on the bus.
Speculative decoding (SD) should help: a small draft model proposes several tokens, the large model verifies them in one parallel pass. Combined with MoE offloading it often hurts. Multi-token verification inflates the union of activated experts and contends for CPU–GPU bandwidth. Experts loaded for tokens the target later rejects waste that bandwidth. Mixtral-Offloading, MoE-Infinity, and AdapMoE never use the draft stage, so PCIe sits idle while the draft model runs.
Neighboring draft tokens often share experts, which makes prediction feasible. Prefetching too many future layers thrashes the GPU cache: eviction rises past three layers, and prefetch starts failing past five. The draft stage is about 16.2% of a decode step, and that window usually leaves PCIe unused.
SP-MoE feeds each draft layer's attention output into the target model's already-trained gate at the same layer, scores experts, and prefetches only the top-k "critical" ones. Mixtral uses k=1 because one expert is 336 MB and a miss is expensive; Phi-3.5-MoE uses k=2; Deepseek-Lite uses k=6. Draft/target pairs are Mistral 7B with Mixtral 8×7B, Phi-mini-MoE with Phi-3.5-MoE, and Deepseek-Lite-AWQ with Deepseek-Lite, with HumanEval acceptance of 97.01% to 98.15%. Cross-model attention cosine similarity reaches 94.59% on DeepSeek; top-1 expert prediction is about 88% across the three pairs.
A cutoff layer L bounds how far ahead drafting may prefetch, so peak non-expert memory plus prefetched experts stay inside the GPU, and I/O finishes before drafting ends. An async worker on a dedicated CUDA stream batches per-layer copies. The cache is LRU. Compute prefers experts already on GPU. All reported runs use batch size 1.
Hardware: RTX 3090 24 GB, RTX 4090 24 GB, A100 40 GB. Models: Mixtral 8×7B, Phi-3.5-MoE, Deepseek-Lite. Datasets: HumanEval, BigBench, WikiText-103, MMLU-Pro. Baselines are the same offload systems with SD bolted on.
| Baseline | TPOT |
| Mixtral-Offloading+SD | −34% on average, up to 1.75× on HumanEval + 3090 |
| MoE-Infinity+SD | −19% on average |
| AdapMoE+SD | −12% on average |
On Deepseek-Lite, HumanEval, A100, the gain is 3.5× versus Mixtral-Offloading and 1.07× versus AdapMoE, the two ends of the paper's 1.07×–3.5× range. Deepseek hit rate averages 40.06% for SP-MoE versus 21.85% for AdapMoE. On Mixtral, AdapMoE actually hits more often (42.14% vs 20.89%), but its prefetch sync tax still loses end-to-end. Expert loading is about 69.4% of one decode step.
This is the first offload stack built around SD's two stages. On a consumer GPU the draft window is free PCIe time. If you already run speculative decoding, the change is the prefetch scheduler, not the weights. It is a systems paper. Gains are largest when VRAM is tight and the bus is slow; the A100 leaves less headroom.
Every number is batch size 1. The authors say a large batch mixes uncorrelated expert sets across requests, which breaks reuse; they leave that to future work. Drafting is greedy and sequential, not tree-based, and they do not sweep sampling temperature. Cutoff L and k come from offline profiles and must be remeasured if PCIe generation or expert size changes. Baselines are the authors' ports of Mixtral-Offloading, AdapMoE, and MoE-Infinity. There is no vLLM or SGLang comparison.