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

Yujie Zhang; Bin Gao; Tulika Mitra

cs.AIAug 21, 2026

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

Authors: Yujie Zhang, Bin Gao, Tulika Mitra

Organizations: School of Computing, National University of Singapore, Singapore

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.

Figures & tables

Fig. 1: CoT reasoning employs transition cues to structure multi-step problem solving, facilitating self-correction and the exploration of alternative solution paths.

Fig. 2: Illustration of CoT reasoning in a MoE model, where each token activates its top-2 selected experts.

Fig. 3: Layer-wise expert activation pattern in Qwen3 on MATH‑500 under CoT reasoning, averaged over 100 samples.

	MATH-500	AIME 2024	GPQA-Diamond
Qwen3	89.78%	89.50%	92.01%
ERNIE-4.5	89.35%	89.87%	85.27%

TABLE I: Temporal coherence scores (

TC_{seq}

) of expert activation patterns across reasoning stages, segmented by transition cues, measured on Qwen3 and ERNIE-4.5.

Panel of Figure 5 — Fig. 4: Operation-level breakdown of Bookkeeping and Token Repacking strategies in one MoE layer during decoding, isolating execution overhead (excluding routing) on Qwen3.

Fig. 5: Design Overview of SAEM.

Fig. 6: Lightweight pattern matching for stage transitions.

Fig. 7: Illustration of expert-aligned token repacking for a single MoE layer during decoding with a batch size of four. For clarity, the routing decision is assumed but not shown.

Model	#Blocks	#Shared Experts	#Routed Experts	Top-k	Experts Params.	Total Params.
Qwen3-30B-A3B (Qwen3)	48	0	128	8	29.0B	30.5B
ERNIE-4.5-21B-A3B (ERNIE-4.5)	28	2	64	6	21.0B	21.8B

TABLE II: Architectural details of MoE models optimized for multi-step reasoning tasks.

Panel of Figure 10 — Fig. 8: Throughput comparison over batch sizes and ECRs for MATH-500 . Overlaid lines denote speedup relative to the strongest baseline.

Panel of Figure 11 — Fig. 9: Throughput comparison over batch sizes and ECRs for AIME 2024 . Overlaid lines denote speedup relative to the strongest baseline.

Panel of Figure 12 — Fig. 10: Throughput comparison over batch sizes and ECRs for GPQA-Diamond . Overlaid lines denote speedup relative to the strongest baseline.

Batch Size	Full GPU (tokens/s)	ECR	SAEM (tokens/s)	CHR Eff. (%)	TP Eff. (%)
1	9.04	50.0%	7.19	90.36	87.87
		25.0%	6.45	81.40	90.82
		12.5%	6.10	76.09	92.80
8	20.90	50.0%	18.55	97.80	95.69
		25.0%	16.35	96.43	98.46
		12.5%	13.82	96.75	100.78

TABLE III: Prediction-oracle analysis on Qwen3. Full GPU is the ECR = 100% reference, and SAEM reports the free-running throughput from Fig. 8 . CHR and TP Eff. are defined in Eq. ( 2 ).

Batch Size	Technique	Throughput (tokens/s)	Speedup
1	Best-performing baseline	3.45	–
	Cache update only	4.81	1.39 $\times$
	In-situ CPU execution only	2.73	0.79 $\times$
	Token repacker only	5.02	1.46 $\times$
	All	6.10	1.77 $\times$
8	Best-performing baseline	8.51	–

TABLE IV: Inference speedup breakdown of proposed techniques under 12.5% cache ratio on Qwen3.

LLM Inference Efficiency

Large language model inference efficiency reduces the time, memory, energy, and cost of generating useful outputs.

See all 352 papers

MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration
Apr 16, 2026Xinyu Liu, Xin Liu, Bo Jin +8Multi-Token Prediction KV Caching
School of Computer Science and Engineering, Northeastern University, China · Meituan Inc. · NiuTrans Research, Shenyang, China
ReMoE: Boosting Expert Reuse through Router Fine-Tuning in Memory-Constrained MoE LLM Inference
May 26, 2026Xiongwei Zhu, Xiaojian Liao, Tianyang Jiang +3LLM Inference Efficiency Expert Offloading
School of Computer Science and Engineering, Beihang University, Beijing 100191, China · Huawei Technologies Ltd.
Speculating Experts Accelerates Inference for Mixture-of-Experts
Mar 9, 2026Vivan Madan, Prajwal Singhania, Abhinav Bhatele +2Mixture-of-Experts Inference LLM Inference Acceleration
University of Maryland · TogetherAI

Apr 16, 2026J/K move · Enter open · S save

MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration

Xinyu Liu, Xin Liu, Bo Jin, Runsong Zhao+7

School of Computer Science and Engineering, Northeastern University, China · Meituan Inc. · NiuTrans Research, Shenyang, China

Multi-Token Prediction KV Caching CoT Reasoning Memory-Efficient Optimization Memory-Efficient Inference LLM Inference Acceleration+1

Open paper AskPDF arXiv

While chain-of-thought (CoT) reasoning enables LLMs to solve challenging reasoning tasks, the linear growth of the KV cache leads to substantial memory and inference overhead. Existing approaches such as context compression and multi-token prediction (MTP) improve efficiency from two complementary directions by compressing historical tokens and generating future tokens in parallel. However, effectively combining them remains challenging due to their different training paradigms and architectural assumptions. In this work, we propose MemoSight (Memory-Foresight-Based Reasoning), a unified framework that integrates context compression and MTP to improve inference efficiency while preserving CoT performance. MemoSight adopts a shared minimalist design based on special tokens and token-specific positional layouts for both compression and parallel prediction. Experiments on four reasoning benchmarks show that, compared to the vanilla SFT baseline, MemoSight reduces KV cache usage by up to 66% and improves inference speed by 56%, while incurring less than a 3% drop in average reasoning accuracy, yielding a better efficiency-accuracy trade-off than existing CoT compression methods.

Abstract

LLM Inference Efficiency

Explore similar work

MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration

ReMoE: Boosting Expert Reuse through Router Fine-Tuning in Memory-Constrained MoE LLM Inference

Speculating Experts Accelerates Inference for Mixture-of-Experts