PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yujie Zhang, Bin Gao, Tulika Mitra

arXiv:2608.21614v1cs.AIcs.DC

TL;DR

CoT decoding creates substantial latency and memory pressure, while MoE expert weights can exceed GPU capacity and trigger costly transfers; existing runtimes miss coherent activation patterns across reasoning stages. SAEM detects stage boundaries and coordinates stage-aware caching, token repacking, and CPU execution. It delivers 1.33× average throughput improvement over caching and offloading baselines, increasing to 1.54× with calibration-matched data.

  • Problem

    Existing MoE inference runtimes treat CoT tokens uniformly despite coherent, predictable expert activation across consecutive reasoning stages, causing inefficient caching and data movement under constrained GPU memory.

  • Method

    SAEM detects reasoning-stage transitions and coordinates stage-aware caching, expert-aligned token repacking, and in-situ CPU execution for expert management.

  • Results

    1.33× average throughput improvement was achieved over state-of-the-art caching and offloading baselines, rising to 1.54× when calibration data matched the workload.

  • Takeaways & Limitations

    Stage-level activation locality can guide efficient MoE expert placement and improve CoT inference under constrained GPU memory.

  • Takeaways & Limitations

    The paper discusses runtime overhead and remaining limitations but does not specify a concrete limitation in the supplied passage.

Abstract

from arXiv · show

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.

I. INTRODUCTION

CoT reasoning increases latency and memory pressure, while MoE inference under constrained GPU memory requires costly expert movement. SAEM exploits coherent reasoning-stage activations through stage-aware management and coordinated execution mechanisms.

  • Motivation: Long CoT traces amplify inference latency and memory demand, especially as task difficulty increases.Evaluations may permit generation budgets up to 32,768 tokens.
  • Motivation: MoE inference often exceeds GPU capacity because full expert weights must move dynamically between CPU and GPU.Sparse routing also leaves many GPU-resident expert resources underused.
  • Motivation: Token-level caching and prefetching overlook stable, predictable expert subsets across consecutive CoT reasoning stages.This limits their ability to exploit semantic coherence in generated reasoning traces.
  • SAEM: SAEM detects stage transitions and aggregates expert usage at stage granularity to guide cache management.Stage-aware updates occur when reasoning semantics change rather than for every token.
  • SAEM: SAEM combines expert-aligned token repacking with in-situ CPU execution to improve locality and avoid unnecessary transfers.Repacking creates contiguous batches, while CPU execution handles infrequently activated experts without cache pollution.
  • Evaluation: 1.33× average throughput improvement was achieved across two MoE models and three reasoning benchmarks, rising to 1.54× with calibration-matched conditions.The evaluation varied batch sizes and GPU cache budgets against state-of-the-art caching and offloading baselines.

II. BACKGROUND

MoE models scale capacity through sparse expert activation, while CoT traces organize reasoning into semantically coherent stages. Existing expert-management systems generally operate at finer granularity and miss these stage-level activation regularities.

  • MoE Inference: MoE routing activates only a small subset of experts per token, enabling parameter scaling without proportional computation growth.This sparse activation keeps per-token computation relatively stable.
  • Existing Management: Prior systems use GPU caching, CPU execution, and sequence-level prediction or prefetching to manage non-resident experts.These strategies target transfer overhead and memory utilization.
  • CoT Structure: Reasoning stages exhibit stable, predictable expert activation patterns that finer-grained management fails to exploit.This regularity distinguishes structured CoT traces from uniform token streams.
  • CoT Structure: CoT traces contain distinct reasoning stages signaled by discourse transitions such as “alternatively,” “on second thought,” and “therefore.”These cues mark shifts in reasoning strategy, verification, or conclusion.

III. MOTIVATION

SAEM is motivated by the observation that CoT stages activate distinct expert subsets with stable routing across adjacent stages, while conventional systems manage experts uniformly at token granularity. This regularity supports stage-aware scheduling and motivates token repacking to reduce fragmented execution overhead.

  • CoT stages activate distinct expert subsets, so token-uniform expert management overlooks semantic regularities in generated reasoning.
  • 89.30% average TCseq indicates stable expert activation across reasoning-stage transitions on Qwen3 and ERNIE-4.5 across three datasets.Higher TCseq values indicate more stable routing across stage transitions.
  • Qwen3 MATH-500 traces contain 1–49 stages, with stage lengths spanning 55–7129 tokens, yet adjacent-stage activation regularity remains useful.The reported median and average are 5 and 8 stages, and 271 and 484 tokens per stage, respectively.
  • Qwen3’s top-8 routing quadruples scattered token-access kernel invocations relative to typical top-2 routing, producing more than 128 kernels per MoE-layer forward pass.
  • 54.2% of Bookkeeping execution time is spent on kernel launches, routing metadata, and token access/layout transformation, versus 40.0% with TokenRepacking.TokenRepacking reduces token access/layout transformation kernels by 4.7% and kernel launch overhead by 6.7%.

IV. SAEM: SYSTEM DESIGN

SAEM is a data-aware runtime that detects CoT stage boundaries and uses stage-level activation regularity to coordinate expert caching, prefetching, token repacking, and execution. Lightweight, model-customized cue matching identifies transitions without heavyweight semantic parsing.

  • SAEM combines a stage-boundary detector, stage-aware expert cache manager, and latency-hiding prefetch scheduler with token repacking and in-situ expert execution.
  • SAEM focuses on alternative-path exploration and detects transitions using discourse cues such as “Alternatively,” “Instead,” and “Another way is.”Broader reasoning-strategy detection is left for future work.
  • Cue-based detection uses compact lexical patterns rather than heavyweight semantic parsing or neural boundary classifiers on the decoding critical path.The detector targets coarse transition points sufficient for expert-placement decisions, not perfect human-interpretable segmentation.
  • Model-specific transition sets are pruned offline using marker frequency, segmentation noise, and adjacent-stage coherence measured by TCseq.Calibration data and a model checkpoint are used to construct the customized set Tmodel.
  • A sliding-window finite-state matcher tracks cue fragments across subword-token boundaries and emits a stage-transition event when a complete cue is matched.

2) Activation Pattern Aggregation at Boundaries:

At detected stage boundaries, SAEM aggregates expert usage over the completed stage and uses the resulting pattern to guide the next stage’s placement within a fixed GPU cache budget.

  • At each boundary, SAEM computes per-layer, per-expert usage frequencies across batch queries for the completed stage.The aggregation summarizes stage-level computational demand instead of reacting to individual token fluctuations.
  • Frequently activated experts are prioritized for GPU residency in the upcoming stage, while rarely activated experts are deprioritized.The policy relies on activation coherence between adjacent stages and updates the cache only at detected boundaries.
  • SAEM places non-MoE layers on the GPU and reserves a fixed-size per-layer expert cache whose capacity is determined by the Expert Cache Ratio.Fixed allocation provides predictable memory use and avoids fragmentation from dynamic allocation.

2) Dynamic Expert Placement Using Historical Patterns:

SAEM dynamically places experts using stage-history rankings, asynchronous prefetching, and a latency comparison between GPU transfer and direct CPU execution. This hybrid placement avoids unnecessary transfers while using both processors.

  • Experts in each layer are ranked by activation frequency from the just-completed stage, and residency is updated within the allocated cache budget.
  • Asynchronous DMA transfers overlap CPU–GPU expert movement with kernel execution, while GPU events coordinate completion to avoid pipeline stalls.Prefetch requests are issued immediately after a stage boundary is detected.
  • For non-resident experts, SAEM selects GPU transfer or direct CPU execution according to lower latency, favoring CPU execution when few tokens are assigned.In that regime, PCIe transfer latency can dominate.
  • After top-k routing, tokens are partitioned into GPU-bound and CPU-bound groups, with GPU tensor-core kernels and optimized CPU GEMM backends processing them in parallel.
  • NUMA-aware execution co-locates CPU threads, expert weights, and activation buffers on one socket and its corresponding NUMA node.This design minimizes cross-socket memory traffic and access latency.
  • GPU and CPU outputs are merged before the next layer, avoiding transfers for infrequently activated experts and using CPU computation alongside the GPU.

1) Expert-Aligned Token Repacking:

Expert-aligned token repacking groups routed tokens contiguously by expert, replacing scattered accesses with batched execution. This reduces layout-transformation and kernel-launch overhead, especially with imbalanced routing.

  • 1) Expert-Aligned Token Repacking:: Token repacking groups tokens by expert assignment and places them contiguously in memory across MoE layers.The temporary workspace is allocated once and reused, while grouping requires lightweight reindexing.
  • 1) Expert-Aligned Token Repacking:: Each expert can execute in a single batched kernel launch instead of multiple fragmented invocations.This reduces scattered memory accesses, layout transformations, and kernel-launch overhead.
  • 1) Expert-Aligned Token Repacking:: Token repacking operates independently of CoT stage boundaries and improves throughput particularly at small batch sizes or under imbalanced routing.It converts sparse token-to-expert mappings into dense expert-aligned batches.
  • 1) Expert-Aligned Token Repacking:: The mechanism is illustrated for a single MoE layer decoding four tokens, with routing assumed but not shown.The figure focuses on how tokens are reorganized for expert execution.

3) Implementation:

SAEM is evaluated against established MoE caching and offloading baselines across batch sizes, cache ratios, and reasoning benchmarks. It outperforms the strongest baseline across nearly every tested configuration, with larger gains on calibrated workloads.

  • 3) Implementation:: SAEM outperforms the strongest baseline in nearly every configuration across MATH-500, AIME 2024, and GPQA-Diamond.Average speedups are reported separately for Qwen3 and ERNIE-4.5 across the three benchmarks.
  • 3) Implementation:: Average speedups are 1.60× and 1.47× on MATH-500, 1.20× and 1.21× on AIME 2024, and 1.14× and 1.34× on GPQA-Diamond for Qwen3 and ERNIE-4.5, respectively.The comparisons span varying batch sizes and Expert Cache Ratios.
  • 3) Implementation:: MATH-500 produces the largest margin because it calibrates both dominant-expert statistics and the pruned transition set Tmodel.AIME 2024 and GPQA-Diamond are held out, so their results measure calibration mismatch.
  • 3) Implementation:: Without target-domain calibration, SAEM retains a 1.14–1.34× advantage across reasoning domains.The paper contrasts this transfer with the additional headroom available when calibration data matches the workload.

1) Single-batch Regime:

In single-batch execution, SAEM can outperform token-level migration despite a lower cache hit rate by recovering efficiency through token repacking. Its oracle analysis separates prediction quality from cache-capacity effects.

  • 1) Single-batch Regime:: 1.62× throughput improvement over Mixtral-Offloading is achieved despite SAEM’s lower cache hit rate under 50% ECR on Qwen3.The cache hit rates are 90.08% for SAEM and 94.53% for Mixtral-Offloading.
  • 1) Single-batch Regime:: At batch size 8 and 3.125% ECR, SAEM achieves a 2.10× throughput gain over Mixtral-Offloading.Its cache hit ratio is higher by 187.95% at that setting.
  • 1) Single-batch Regime:: SAEM increases cache hit ratio relative to Mixtral-Offloading by 53.06% at 12.5% ECR and 187.95% at 3.125% ECR for Qwen3 with batch size 8.The gains occur under multi-batch execution and constrained cache capacity.
  • 1) Single-batch Regime:: The oracle replaces the previous stage’s activation profile with perfect knowledge of the upcoming stage while preserving cache capacity and update events.Both policies use identical boundary events and cache-enforcement mechanisms.
  • 1) Single-batch Regime:: At batch size 1, cache-hit-ratio efficiency falls from 90.36% to 76.09% as ECR decreases, while throughput efficiency rises from 87.87% to 92.80%.The oracle throughput improvement remains limited to 1.08–1.14×.
  • 1) Single-batch Regime:: At batch size 8, cache-hit-ratio and throughput efficiencies remain above 96% and 95%, respectively, indicating near-oracle operation.The larger SAEM-to-full-GPU gap is attributed primarily to limited GPU cache capacity rather than prediction inaccuracy.

D. Ablation Study

The ablation shows that stage-aware caching and token repacking address complementary bottlenecks, while in-situ CPU execution is most effective when combined with those mechanisms. Avoiding on-demand transfers is essential under constrained GPU memory.

  • D. Ablation Study: Stage-aware cache updates and expert-aligned token repacking each accelerate inference independently by reducing cache churn or improving memory locality and kernel efficiency.The two mechanisms address complementary bottlenecks in expert placement and token execution.
  • D. Ablation Study: In-situ CPU execution alone reaches 0.79× at batch size 1 and matches the baseline at batch size 8.Without stage-aware placement, the GPU cache retains a suboptimal expert set and sends many tokens to the slower CPU path.
  • D. Ablation Study: Combining all three components yields the highest throughput, with speedups of 1.77× and 1.62× at the evaluated settings.The results support complementarity among stage-aware caching, token repacking, and in-situ CPU execution.
  • D. Ablation Study: Disabling CPU-side expert execution reduces throughput from 6.10 to 3.22 tokens/s at batch size 1 and from 13.82 to 9.00 tokens/s at batch size 8.Non-resident experts are instead transferred to the GPU on demand through a reserved temporary slot.

VI. DISCUSSION & LIMITATIONS

SAEM improves MoE inference efficiency by exploiting reasoning-stage locality, while incurring bounded runtime overhead and remaining sensitive to the presence of explicit transition cues. Future robustness could come from expert-activation signals that detect implicit stage changes.

  • Discussion: SAEM improves MoE inference throughput by exploiting reasoning-stage locality for expert management.Its coordinated caching, scheduling, and token repacking reduce memory pressure and accelerate decoding for CoT reasoning.
  • Practical Overhead: Transition matching, token repacking, and stage-level placement add bounded overhead that remains lightweight relative to MoE computation and data movement.The design amortizes placement updates over reasoning stages and uses short-window matching with preallocated workspaces.
  • Stage Boundary Detection: Explicit linguistic transition cues make boundary detection effective when reasoning traces contain clear structural markers.The current detector relies on cues such as “Alternatively,” and “Instead,”.
  • Stage Boundary Detection: When generated traces lack explicit cues, boundary detection may become less precise, reducing the effectiveness of stage-level cache updates.
  • Future Directions: Future work could supplement lexical cues with entropy shifts in expert activation distributions to detect implicit reasoning-stage transitions.Such statistical signals may improve generality across models, reasoning styles, and workloads.
Loading 2608.21614v1…