跳转至

文章背景与核心概要

本文介绍了一种专为混合专家（MoE）大语言模型在执行思维链（CoT）推理时设计的先进推理运行时——SAEM（阶段感知专家管理，Stage-Aware Expert Management）。虽然思维链提示词通过分解复杂问题增强了推理能力，但其顺序特性显著增加了 GPU 内存需求和解码延迟。由于 MoE 模型拥有完整的专家权重且通常超出 GPU 显存限制，传统系统往往面临高昂的 GPU-CPU 数据传输和低效的缓存问题。

SAEM 通过利用 CoT 轨迹的一个关键结构特征解决了这一痛点：连续的推理阶段表现出连贯且可预测的专家激活模式。通过检测阶段边界并利用这种局部性，SAEM 实现了阶段感知缓存、专家对齐的Token重新打包以及原位 CPU 执行。该方法大幅减少了数据移动和内核碎片化，在受限的 GPU 内存下，相比最先进的基线取得了 1.33倍至1.54倍的吞吐量提升。

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

Summary

Summary

SAEM (Stage-Aware Expert Management) is an advanced inference runtime designed for Mixture-of-Experts (MoE) Large Language Models performing Chain-of-Thought (CoT) reasoning. While CoT prompting enhances reasoning by breaking down complex problems, its sequential nature drastically increases GPU memory demands and decoding latency. Because MoE models feature full expert weights that often exceed GPU limits, traditional systems experience expensive GPU-CPU data transfers and inefficient caching.

SAEM (Stage-Aware Expert Management) is an advanced inference runtime designed for Mixture-of-Experts (MoE) Large Language Models performing Chain-of-Thought (CoT) reasoning. While CoT prompting enhances reasoning by breaking down complex problems, its sequential nature drastically increases GPU memory demands and decoding latency. Because MoE models feature full expert weights that often exceed GPU limits, traditional systems experience expensive GPU-CPU data transfers and inefficient caching.

SAEM solves this by capitalizing on a key structural trait of CoT traces: consecutive reasoning stages exhibit coherent and predictable expert activation patterns. By detecting stage boundaries and leveraging this locality, SAEM implements stage-aware caching, expert-aligned token repacking, and in-situ CPU execution. This approach drastically minimizes data movement and kernel fragmentation, achieving a 1.33x to 1.54x throughput improvement over state-of-the-art baselines under constrained GPU memory.

SAEM solves this by capitalizing on a key structural trait of CoT traces: consecutive reasoning stages exhibit coherent and predictable expert activation patterns. By detecting stage boundaries and leveraging this locality, SAEM implements stage-aware caching, expert-aligned token repacking, and in-situ CPU execution. This approach drastically minimizes data movement and kernel fragmentation, achieving a 1.33x to 1.54x throughput improvement over state-of-the-art baselines under constrained GPU memory.

Paper Metadata

Paper Metadata

Attribute Details
arXiv ID arXiv:2608.21614 [cs.AI]
Title SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning
Authors Yujie Zhang, Bin Gao, Tulika Mitra
Submitted August 21, 2026
Primary Subject Artificial Intelligence (cs.AI)
Secondary Subjects Distributed, Parallel, and Cluster Computing (cs.DC)
Status Extended version of a paper accepted at DAC 2026 (includes evaluations on AIME 2024 and GPQA-Diamond, plus a prediction-oracle upper-bound analysis).
Attribute Details
arXiv ID arXiv:2608.21614 [cs.AI]
Title SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning
Authors Yujie Zhang, Bin Gao, Tulika Mitra
Submitted August 21, 2026
Primary Subject Artificial Intelligence (cs.AI)
Secondary Subjects Distributed, Parallel, and Cluster Computing (cs.DC)
Status Extended version of a paper accepted at DAC 2026 (includes evaluations on AIME 2024 and GPQA-Diamond, plus a prediction-oracle upper-bound analysis).

Abstract

Abstract

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.

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.

现有运行时将所有 Token 同等对待，忽略了 CoT 轨迹的一个关键结构属性：连续的推理阶段表现出连贯且可预测的专家激活模式。忽视这种阶段级别的规律性会导致缓存低效和不必要的数据移动。

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.

我们提出了 SAEM，这是一个阶段感知的 MoE 推理运行时，它能够检测推理阶段边界并利用阶段级别的激活连贯性来指导专家放置。SAEM 结合了： * 阶段感知缓存 * 专家对齐的 Token 重新打包 * 原位 CPU 执行

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 * In-situ CPU execution

这些机制协同工作，减少了数据传输和内核碎片化。在数学和科学推理工作负载下，在受限的 GPU 内存下，SAEM 相比最强大的最先进缓存和卸载基线取得了平均 1.33倍的吞吐量提升；当校准数据与工作负载匹配时，提升幅度可达 1.54倍。这证明了面向 CoT 推理的、基于局部性的阶段感知 MoE 推理具有极高的有效性。

These mechanisms work together 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. This demonstrates the high effectiveness of stage-aware, locality-driven MoE inference for CoT reasoning.

Full-Text & Access Links

Full-Text & Access Links