aiqbench analysis
SAEM optimizes Mixture-of-Experts inference for Chain-of-Thought reasoning by leveraging predictable expert activation patterns across different reasoning stages to reduce GPU memory pressure and data transfer costs. This approach yields significant throughput gains on mathematical and scientific workloads, offering a practical solution for deploying large-scale MoE models under constrained hardware resources.
SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning
arXiv:2608.21614v1 Announce Type: new 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. 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.
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. 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.
Extended version of the paper accepted at DAC 2026. Extends the conference version with evaluations on AIME 2024 and GPQA-Diamond and a prediction-oracle upper-bound analysis
Artificial Intelligence (cs.AI); Distributed, Parallel, and Cluster Computing (cs.DC)
arXiv:2608.21614 [cs.AI]
(or arXiv:2608.21614v1 [cs.AI] for this version)
https://doi.org/10.48550/arXiv.2608.21614 arXiv-issued DOI via DataCite (pending registration)
Full text under the publisher's open licence · 216 words
Related coverage
Building Blocks for Foundation Model Training and Inference on AWS
Hugging Face · May 11, 2026
The hard physics and complex economics of AI’s insatiable hunger for power
SiliconANGLE · Sep 11, 2026
Scaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughput
AWS · Sep 25, 2026
[AINews] not much happened today
Latent Space · Sep 10, 2026