arXiv:2608.21614cs.AIcs.DC2026-08中稿 · DAC 2026

通过识别推理阶段提升MoE模型效率,降低内存占用和延迟。

SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning

论文配图:SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning
图 1 · 摘自论文原文
  • 根据推理阶段边界动态管理专家权重,利用激活模式的规律性
  • 在受限显存下平均提速1.33倍,匹配数据时达1.54倍
  • 适合需要高效链式推理的数学与科学任务场景

链式思维(CoT)提示通过将复杂问题分解为中间步骤提升大模型推理能力,但其串行特性增加了解码延迟和内存消耗。混合专家(MoE)模型通过稀疏激活扩展容量,但全量专家权重常超出GPU内存,需频繁进行昂贵的GPU-CPU数据迁移。现有运行时对所有标记统一处理,忽视了CoT轨迹的一个关键结构特征:连续推理阶段具有连贯且可预测的专家激活模式。忽略这一阶段级规律会导致缓存效率低下和不必要的数据移动。我们提出SAEM,一种面向阶段感知的MoE推理运行时,能够检测推理阶段边界,并利用阶段级激活一致性指导专家放置。SAEM结合阶段感知缓存、专家对齐的令牌重排以及就地CPU执行,减少数据传输与内核碎片化。在数学与科学推理任务上,当显存受限时,相较于最强的现有缓存与卸载基线,SAEM平均实现1.33倍吞吐量提升;当校准数据匹配工作负载时,提升达到1.54倍,验证了阶段感知、局部性驱动的MoE推理在CoT推理中的有效性。

原文摘要 · Abstract (English)

Chain-of-thought (CoT) prompting improves LLM reasoning by decomposing complex problems into intermediate steps, but its sequential nature increases decoding latency and memory usage. Mixture-of-Experts (MoE) models scale capacity through sparse expert activation, yet their full expert weights often exceed GPU memory and require costly GPU-CPU transfers. Existing runtimes treat all tokens uniformly, overlooking a key structural property of CoT traces: consecutive reasoning stages exhibit coherent and predictable expert activation patterns. Ignoring this stage-level regularity leads to inefficient caching and unnecessary data movement. We propose SAEM, a stage-aware MoE inference runtime that detects reasoning stage boundaries and exploits stage-level activation coherence to guide expert placement. SAEM combines stage-aware caching, expert-aligned token repacking, and in-situ CPU execution to reduce data transfer and kernel fragmentation. On mathematical and scientific reasoning workloads, SAEM achieves an average 1.33x throughput improvement over the strongest state-of-the-art caching and offloading baselines under constrained GPU memory, rising to 1.54x when calibration data matches the workload, demonstrating the effectiveness of stage-aware, locality-driven MoE inference for CoT reasoning.

MoE推理链式思维内存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。