通过联合优化路由器与缓存,提升MoE模型推理时的内存效率。
Cache-Aware Joint Router Adaptation for Memory-Efficient MoE Inference

- 设计缓存感知的后训练框架,联合调整主干模型与轻量缓存路由器。
- 在Qwen3上使缓存命中率提升1.15至18.03点,专家权重传输减少4.6%至53.3%。
- 适合需要高效推理的MoE大模型部署场景,尤其关注内存瓶颈的系统
Mixture-of-Experts(MoE)模型每令牌仅激活少数专家,但完整专家集常超出GPU内存,导致解码时重复传输权重。本文将专家缓存管理建模为模型侧算法问题,提出一种缓存感知的后训练框架,联合适配MoE主干与轻量辅助缓存路由器,同时保留推理时的原生Top-K选择规则。其仅更新模式Temporal Router预测同层复用,保留专家以供后续令牌使用,无需主动加载。全量Spatio-Temporal Router增加空间路由器,利用因果前驱隐藏状态在目标层访问前优化时间缓存。在Qwen3与GPT-OSS上评估GSM8K、MATH和CommonsenseQA三个任务。Temporal Router持续提升缓存命中率并减少专家权重流量。在Qwen3上,Spatio-Temporal Router在三项任务中均实现最优负载调整效率,命中率提升1.15–18.03点,流量减少4.6%–53.3%(相较最强预取基线);在GPT-OSS上表现竞争,但任务依赖性强。仅辅助模块消融实验保持基线准确率但缓存增益有限,而联合后训练带来更大改进。敏感性分析表明,缓存容量控制传输需求,精炼预算决定预访问覆盖与主动流量间的权衡。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) models activate only a small subset of experts per token, but the full expert set often exceeds GPU memory, causing repeated weight transfers during decoding. We formulate expert-cache management as a model-side algorithmic problem and propose a cache-aware post-training framework that jointly adapts the MoE backbone and lightweight auxiliary cache routers while preserving the native Top-K expert-selection rule at inference. Its update-only mode, Temporal Router, predicts same-layer reuse and retains experts for future tokens without proactive loading. The full Spatio-Temporal Router adds a Spatio Router that uses the causal predecessor's hidden state to refine the temporal cache before target-layer access. We evaluate both modes on Qwen3 and GPT-OSS across GSM8K, MATH, and CommonsenseQA. Temporal Router consistently improves cache hit rate and reduces expert-weight traffic over matched LM-only baselines. On Qwen3, Spatio-Temporal Router achieves the best load-adjusted efficiency across three tasks, improving adjusted hit rate by 1.15--18.03 points and reducing traffic by 4.6--53.3% relative to the strongest evaluated prefetching baseline; results on GPT-OSS are competitive but task-dependent. An auxiliary-only ablation preserves baseline accuracy but yields modest cache gains, whereas joint post-training produces larger improvements. Sensitivity analyses show that cache capacity controls transfer demand, while the refinement budget governs the trade-off between pre-access coverage and proactive traffic.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。