DraftExpert: Expansion-Aware Self-Speculative Decoding for End-Device MoE Inference

2026-07-27

AI Summary

This work addresses the high inference latency of on-device Mixture-of-Experts (MoE) large language models in speculative decoding, which stems from the trade-off between expert loading overhead and draft acceptance rate. The authors propose an expansion-aware speculative decoding framework that enhances efficiency through lightweight resident draft experts. Innovatively, they integrate residual, logit/token, and routing consistency signals into a self-distillation training scheme and design a shared + Top-1 + draft expert architecture. Coupled with confidence-driven expansion truncation and a target-expert prefetching mechanism, the approach significantly accelerates inference while preserving verification accuracy. Experiments on DeepSeek-V2-Lite and Moonlight-16B-A3B demonstrate an average 1.45× throughput improvement, with draft acceptance rates of 84–87% and prefetch hit rates of 86–88%.

Abstract

Large Mixture-of-Experts (MoE) language models are attractive for end-device deployment because only a small subset of experts is active per token, but their routed expert weights often exceed accelerator memory. We target latency-critical single-user settings where routed experts are staged on demand from CPU memory to a GPU or from Flash to a mobile NPU. In this setting, self-speculative decoding faces a new bottleneck: increasing the draft expert set improves accuracy but triggers extra expert loading, while cheap small-footprint drafts have low acceptance; moreover, verifying a multi-token block activates the union of target experts and is no longer close to one target step. We propose DraftExpert, an expansion-aware self-speculative decoding framework for expert-offloaded MoE inference. DraftExpert trains one lightweight accelerator-resident draft expert per layer by self-distilling residual, logit/token, and router-agreement signals from the frozen target MoE. At inference time, it uses a fixed-footprint shared+top-1+draft-expert drafter together with confidence--expansion truncation and target-expert prefetching, while final tokens are still exactly verified by the target model. On DeepSeek-V2-Lite and Moonlight-16B-A3B across CPU-GPU and Flash-NPU offload, DraftExpert improves decode throughput by 1.45x on average, raises draft acceptance to 84~87%, and achieves 86~88% prefetch hit rates.

Problem

• Mixture-of-Experts
• self-speculative decoding
• expert offloading
• end-device inference
• memory bottleneck

Innovation

• Mixture-of-Experts
• self-speculative decoding
• expert offloading
• draft expert
• prefetching