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
中山大学把投机解码的草稿阶段改成专家预取窗口,用草稿attention喂给目标门控,相对AdapMoE等offload系统TPOT最高快3.5倍。
MoE把大模型拆成一堆专家,单token只走少数FFN,算力下来了,显存却涨上去。Mixtral 8×7B推理大约要87GB,一张24GB的RTX 4090塞不下。常见做法是把暂时不用的专家放到CPU,用时再经PCIe搬回GPU。PCIe 4.0搬一层Mixtral专家大约80ms,4090上算一层只要3ms,时间几乎全耗在搬数据上。
投机解码(speculative decoding, SD)本来能加速生成:小草稿模型先写若干token,大模型一次并行验证。但SD碰上MoE offload会更糟。多token一起验证时激活的专家集合膨胀,CPU-GPU带宽争用更狠;最后被拒绝的草稿token,对应的专家加载还白占带宽。Mixtral-Offloading、MoE-Infinity、AdapMoE这类offload系统都不吃SD的两阶段结构,草稿那一段时间PCIe常常空转。
相邻草稿token经常激活同一批专家,这给预测留了空间。预取层数太深会把GPU缓存打爆:超过三层驱逐率明显上升,超过五层预取开始失败。草稿阶段大约占一次decode的16.2%,这段PCIe往往闲着。
SP-MoE把草稿模型每一层的attention输出,喂给目标模型同一层已经训好的门控网络,给专家打分,只预取top-k「关键专家」。Mixtral取k=1,因为单个专家336MB,错预取代价太大;Phi-3.5-MoE取k=2,Deepseek-Lite取k=6。草稿-目标对分别是Mistral 7B配Mixtral 8×7B、Phi-mini-MoE配Phi-3.5-MoE、Deepseek-Lite-AWQ配Deepseek-Lite,HumanEval上接受率97.01%到98.15%。跨模型attention余弦相似度在DeepSeek上最高94.59%,三对模型的top-1专家预测准确率大约88%。
为防缓存抖动,只预取前L层,L由分析模型给出:预取专家加上非专家峰值显存不能超GPU,且预取必须在草稿结束前完成。运行时用独立CUDA stream和worker线程做异步预取,同层专家批量I/O,缓存用LRU。计算时优先跑已经在GPU上的专家。评测固定batch size=1。
三套硬件(3090 24GB、4090 24GB、A100 40GB),三个模型,HumanEval、BigBench、WikiText-103、MMLU-Pro四个数据集。对照是把SD接到现有offload系统上。
| 对照 | TPOT |
| Mixtral-Offloading+SD | 平均降34%,HumanEval+3090最高1.75× |
| MoE-Infinity+SD | 平均降19% |
| AdapMoE+SD | 平均降12% |
Deepseek-Lite在A100加HumanEval上,相对Mixtral-Offloading达到3.5×,相对AdapMoE只有1.07×,这是全文区间的两端。Deepseek上SP-MoE平均命中率40.06%,AdapMoE是21.85%。Mixtral上AdapMoE命中率反而更高,42.14%对20.89%,但AdapMoE的预取同步把端到端优势吃掉了。专家加载占单次decode延迟约69.4%。
这是第一套按SD两阶段来设计的MoE offload。消费级单卡跑Mixtral、Phi-MoE这类模型,草稿阶段那段空闲PCIe是现成的,不用加卡。已经在用投机解码的人,改动主要在预取调度,不改模型权重。它是系统活,不是新训练算法。收益在显存紧、PCIe慢的卡上最大,A100上空间更小。
评测锁死batch=1。论文自己写,大batch会让不同请求的专家集合几乎不相关,缓存和预取都会失效,这块留给后续。草稿只走贪心顺序解码,没有树状草稿,也没有采样温度扫描。cutoff和k靠离线profile,换PCIe代际或专家尺寸要重测。基线是作者把Mixtral-Offloading、AdapMoE、MoE-Infinity扩到三个模型后的复现,跟vLLM、SGLang没有直接对比。