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

• URL: http://arxiv.org/abs/2608.21614v1
• Date: Fri, 21 Aug 2026 20:26:32 GMT
• Title: SAEM: Stage-Aware Expert Management for Memory-Efficient MoE Inference in Chain-of-Thought Reasoning
• Abstract summary: Chain-of-thought (CoT) prompting improves LLM reasoning by decomposing complex problems into intermediate steps.<n>MoE models scale capacity through sparse expert activation, yet their full expert weights often exceed GPU memory.<n>We propose SAEM, a stage-aware MoE inference runtime that detects reasoning stage boundaries and exploits stage-level activation coherence to guide expert placement.
• Score: 26.672933821270885
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• LLM Inference in a Flash! [66.43896490006965]
  Large Language Models (LLMs) have shown impressive capabilities across a range of natural language processing tasks.<n>We present an end-to-end integer-only quantization approach to eliminate expensive floating-point computations.
  arXiv  Detail & Related papers  (2026-09-14T18:04:30Z)
• FLARE: Diffusion for Hybrid Language Model [72.60770374799634]
  FLARE is a systematic conversion framework for hybrid-attention large language models.<n>It combines a token-equal AR-and-diffusion objective, hardware-aware kernels, and unified inference, enabling one checkpoint to support both AR-style verified decoding and diffusion-style parallel denoising.<n>Our results suggest practical dLLMs are limited not only by decoding algorithms, but also by transfer data quality and the training inefficiency of current block-diffusion objectives.
  arXiv  Detail & Related papers  (2026-06-01T06:58:15Z)
• CoX-MoE: Coalesced Expert Execution for High-Throughput MoE Inference with AMX-Enabled CPU-GPU Co-Execution [6.04230666444406]
  CoX-MoE is an Advanced Matrix Extensions (AMX)-enabled CPU- GPU collaborative system.<n>It comprehensively optimize MoE inference by combining coalesced expert execution with strategic workload orchestration for higher throughput.<n>Compared to state-of-the-art frameworks, CoX-MoE delivers significant gains, achieving up to 7.1x and 2.4x higher throughput.
  arXiv  Detail & Related papers  (2026-05-18T05:54:30Z)
• DexWorldModel: Causal Latent World Modeling towards Automated Learning of Embodied Tasks [54.32016216994156]
  We introduce the Causal Latent World Model (CLWM), which employs DINOv3 features as generative targets to disentangle interaction semantics from visual noise.<n>To overcome memory scaling, CLWM features a Dual-State Test-Time Training (TTT) Memory that guarantees a strict $mathcalO(1)$ footprint for long-horizon tasks.<n>We present EmbodiChain, an online framework that establishes the Efficiency Law by injecting an infinite flow of physics-grounded trajectories during training.
  arXiv  Detail & Related papers  (2026-04-13T03:19:36Z)
• Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts [68.79341332280062]
  Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time.<n>We introduce OOMB, a highly memory-efficient training system that directly confronts this barrier.<n>Our approach employs a chunk-recurrent training framework with on-the-fly activation recomputation, which maintains a constant activation memory footprint.
  arXiv  Detail & Related papers  (2026-02-02T13:52:40Z)
• MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs [9.086910335841772]
  "Memory wall" bottleneck is significantly amplified in modern large-scale Mixture-of-Experts (MoE) architectures.<n>We present MoEBlaze, a memory-efficient MoE training framework that addresses these issues through a co-designed system approach.<n>We demonstrate that MoEBlaze can achieve over 4x speedups and over 50% memory savings compared to existing MoE frameworks.
  arXiv  Detail & Related papers  (2026-01-08T08:38:23Z)
• Training-free Context-adaptive Attention for Efficient Long Context Modeling [57.703159205740185]
  Training-free Context-adaptive Attention (TCA-Attention) is a training-free sparse attention mechanism that selectively attends to only the informative tokens for efficient long-context inference.<n>TCA-Attention achieves a 2.8$times$ speedup and reduces KV cache by 61% at 128K context length while maintaining performance comparable to full attention.
  arXiv  Detail & Related papers  (2025-12-10T01:54:57Z)
• Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference [16.71963410333802]
  Large language models (LLMs) have demonstrated remarkable success across diverse artificial intelligence tasks.<n>MoC substantially reduces activation memory during pre-training.<n>MoC delivers significant memory savings and throughput gains while maintaining competitive model performance.
  arXiv  Detail & Related papers  (2025-11-12T13:30:57Z)
• Sparse-dLLM: Accelerating Diffusion LLMs with Dynamic Cache Eviction [72.27673320976933]
  Diffusion Large Language Models (dLLMs) enable breakthroughs in reasoning and parallel decoding.<n>Current caching techniques accelerate decoding by storing full-layer states, yet impose substantial memory usage.<n>We propose Sparse-dLLM, the first training-free framework integrating dynamic cache eviction with sparse attention.
  arXiv  Detail & Related papers  (2025-08-04T16:14:03Z)
• FlashForge: Ultra-Efficient Prefix-Aware Attention for LLM Decoding [44.47821531299985]
  Prefix-sharing among multiple prompts presents opportunities to combine the operations of the shared prefix.<n>Decoding is a memory-intensive process requiring heavy memory access on the key-value ( KV) cache of the prefixes.<n>We propose a dedicated attention kernel to combine the memory access of shared KV cache in the decode stage, namely FlashForge.
  arXiv  Detail & Related papers  (2025-05-23T10:03:28Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.