PaperlayerView exact revision on arXiv

Source-linked AI summary

Cache-Aware Joint Router Adaptation for Memory-Efficient MoE Inference

Zhenhe Wu, Yaping Jin, Qinghua Xing, Hang Zhou, Wei He, Xianjie Wu, Xianfu Cheng, Jian Yang, Hanting Chen

arXiv:2609.04895v1cs.CL

TL;DR

MoE expert sets can exceed GPU memory, causing repeated decode-time transfers. The paper introduces cache-aware post-training with native Top-K routing and auxiliary Temporal and Spatio Routers. Across two backbones and three benchmarks, the methods improve cache coverage and reduce traffic, with joint training stronger than auxiliary-only adaptation, while decode-stage metrics do not establish end-to-end speedup.

  • Problem

    Growing MoE expert sets can exceed GPU memory, forcing repeated expert-weight transfers during decoding.

  • Method

    The framework jointly post-trains the MoE backbone and lightweight cache routers while preserving native inference-time Top-K selection; Temporal retains experts after access and Spatio-Temporal adds causal pre-access refinement.

  • Results

    Across two backbones and three benchmarks, the configurations improve cache coverage and reduce decode-stage traffic, with joint post-training substantially stronger than auxiliary-only training.

  • Takeaways & Limitations

    Temporal Router provides an update-only low-traffic operating point, while Spatio-Temporal Router trades bounded proactive traffic for improved pre-access coverage.

  • Takeaways & Limitations

    The study evaluates algorithmic decode-stage traffic rather than end-to-end serving latency or throughput, which also depend on scheduling, overlap, bandwidth, and batching.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) models activate only a small subset of experts per token, but the full expert set often exceeds GPU memory, causing repeated weight transfers during decoding. We formulate expert-cache management as a model-side algorithmic problem and propose a cache-aware post-training framework that jointly adapts the MoE backbone and lightweight auxiliary cache routers while preserving the native Top-K expert-selection rule at inference. Its update-only mode, Temporal Router, predicts same-layer reuse and retains experts for future tokens without proactive loading. The full Spatio-Temporal Router adds a Spatio Router that uses the causal predecessor's hidden state to refine the temporal cache before target-layer access. We evaluate both modes on Qwen3 and GPT-OSS across GSM8K, MATH, and CommonsenseQA. Temporal Router consistently improves cache hit rate and reduces expert-weight traffic over matched LM-only baselines. On Qwen3, Spatio-Temporal Router achieves the best load-adjusted efficiency across three tasks, improving adjusted hit rate by 1.15--18.03 points and reducing traffic by 4.6--53.3% relative to the strongest evaluated prefetching baseline; results on GPT-OSS are competitive but task-dependent. An auxiliary-only ablation preserves baseline accuracy but yields modest cache gains, whereas joint post-training produces larger improvements. Sensitivity analyses show that cache capacity controls transfer demand, while the refinement budget governs the trade-off between pre-access coverage and proactive traffic.

1 INTRODUCTION

MoE inference reduces per-token computation through sparse expert activation, but expert sets can exceed GPU memory and force repeated decode-time transfers. The paper addresses this gap with cache-aware post-training that jointly adapts native routing and auxiliary cache routers while preserving inference-time Top-K selection.

  • MoE increases model capacity without proportional per-token computation by routing each token to a small expert subset.
  • GPU-memory limits can still force repeated expert-weight transfers during decoding as expert sets grow.
  • Heuristic cache policies and runtime prefetchers do not directly learn which resident experts will remain useful at later decoding steps.
  • The framework jointly adapts native routing behavior and auxiliary cache routers under a cache-aware objective, while leaving scheduling and hardware realization outside the algorithm.
  • Temporal Router performs update-only cache retention, while Spatio-Temporal Router adds bounded pre-access refinement using causal information.
  • Experiments compare replacement and prefetching baselines across two sparse MoE backbones and three reasoning benchmarks using adaptation, capacity, budget, and traffic-aware analyses.

2 PROBLEM FORMULATION

The formulation models each MoE layer as a bounded GPU expert cache whose resident set must support the native routed experts. Cache management can retain experts after access or refine residency before access using only causally available information.

  • At each decoding step and MoE layer, the native router selects an active expert set from the layer’s N routed experts.
  • The cache stores a bounded resident expert set, with capacity B and the assumption B ≥ K for the K activated experts.
  • Any selected expert absent from the cache must be demand-loaded before its computation can proceed.
  • Post-access retention chooses experts for later decoding but cannot prefetch for the access that has already occurred.
  • Temporal Router uses only post-access retention, whereas Spatio-Temporal Router performs pre-access refinement before applying the same temporal update.
  • Pre-access refinement uses causally available information to increase residency of experts likely to be routed at the target layer.

3 METHOD

The method combines native MoE routing with learned cache-priority routers and a shared bounded-cache update. Temporal Router retains experts after access, while Spatio-Temporal Router adds causal pre-access refinement; joint post-training supports both deployment modes.

  • 3.1 UNIFIED CACHE-ROUTING FRAMEWORK: The unified framework carries a temporal cache across decoding positions and optionally refines it before target access using a causal predecessor’s Spatio Router output.
  • 3.1 UNIFIED CACHE-ROUTING FRAMEWORK: Auxiliary routers assign cache priorities, while the native router still selects the executed Top-K expert set.
  • 3.2 TEMPORAL ROUTER: POST-ACCESS RETENTION: The shared update inserts demand-loaded selected experts and evicts the lowest-priority resident outside the current selected set when capacity is full.
  • 3.1 UNIFIED CACHE-ROUTING FRAMEWORK: The framework uses dense supervision from full target distributions, while hard Top-B membership is enforced by the cache algorithms.
  • 3.2 TEMPORAL ROUTER: POST-ACCESS RETENTION: Temporal Router predicts next-step same-layer expert demand and updates retention after execution without proactive transfers.
  • 3.3 SPATIO-TEMPORAL ROUTER: Spatio-Temporal Router predicts demand from a causal predecessor, refines the cache before access, bounds candidates with Top-R, and then applies temporal retention.
  • 3.3 SPATIO-TEMPORAL ROUTER: The full mode counts each accepted pre-access insertion as proactive loading, while remaining misses are demand-loaded before MoE execution.
  • 3.4 TRAINING AND DEPLOYMENT MODES: Joint training updates backbone and auxiliary-router parameters, while auxiliary-only training freezes the LM-only backbone and trains only the new routers.

4 EXPERIMENTS

Experiments evaluate cache-aware routing on two sparse MoE backbones and three reasoning benchmarks using matched baselines, ablations, and cache-efficiency metrics. Temporal Router improves update-only retention, while Spatio-Temporal Router trades bounded proactive traffic for stronger pre-access coverage and load-adjusted efficiency.

  • Experimental setup: Qwen3 and GPT-OSS are evaluated on GSM8K, MATH, and CommonsenseQA against replacement and prefetching baselines.The study also examines adaptation scope, router composition, cache capacity, refinement budget, and cache-loss weight.
  • Metrics: Hard hit rate measures resident routed experts at access time, while adjusted hit rate penalizes proactive traffic and Load/token counts expert-weight transfers.Update-only methods have no proactive traffic, so adjusted hit rate is omitted for them; neither metric models overlap or wall-clock latency.
  • Update-only retention: Temporal Router leads update-only methods across all six backbone–dataset settings, with Qwen3 gains of 10.46, 10.70, and 33.34 points over the strongest classical hit-rate baselines.On Qwen3, Load falls from 1353 to 974, 1294 to 906, and 1512 to 304 MB/token on GSM8K, MATH, and CommonsenseQA, respectively.
  • Prefetch-enabled refinement: On Qwen3, Spatio-Temporal Router achieves the best adjusted hit rate and lowest Load across all three tasks, improving adjusted hit rate over ProMoE by 1.15–18.03 points.It reduces traffic by 4.6–53.3% relative to ProMoE; GPT-OSS results are task-dependent.
  • Operating and parameter trade-offs: Temporal Router has lower total Load in all six settings, whereas Spatio-Temporal Router spends bounded proactive traffic for higher pre-access hard hit rate and earlier expert availability.The full method adds 25.2M parameters on Qwen3 and 4.4M on GPT-OSS, compared with ProMoE’s 96.0M and 48.0M.
  • Ablations: Joint post-training is substantially stronger than auxiliary-only training, which preserves baseline accuracy but produces modest cache gains.The ablation compares frozen-backbone auxiliary-only training with jointly adapted routing and cache routers.
  • Metric caveat: Perfect raw hit rate is insufficient evidence of efficient prefetching because proactive-load-aware metrics also account for transferred traffic and task accuracy.Temporally Extended MoE reaches 100% hard hit rate while adjusted hit rate, Load, and accuracy can be less favorable.

5 ANALYSIS AND DISCUSSION

The analysis shows that joint cache-aware adaptation improves cache efficiency more effectively than auxiliary-only training, while cache capacity, refinement budget, and cache-loss weight govern traffic–quality trade-offs. Higher supervision increases expert concentration and locality, but excessive concentration can reduce task accuracy.

  • Adaptation scope: Joint post-training substantially outperforms auxiliary-only training on cache metrics, while auxiliary-only training preserves baseline accuracy.With sw = 0.1, joint Temporal Router reaches 73.13% hard hit and 973.86 MB/token; the full mode reaches 90.62% hard hit and 69.03% adjusted hit, while auxiliary-only gains are modest.
  • Router composition: Spatio-only achieves higher raw hit than the full mode but lower adjusted hit and worse traffic efficiency because replacements occur during prefetching.Spatio-only reaches 93.04% versus 90.62% hard hit, but adjusted hit is 66.94% versus 69.03% and Load is 1665.51 versus 1474.00 MB/token.
  • Cache capacity and refinement budget: Increasing cache capacity raises hit rates and reduces Load for both router modes.For Temporal Router, increasing B from 12 to 30 raises hard hit from 62.91% to 80.00% and reduces Load from 1344 to 725 MB/token; at R = 8, the full mode improves hard and adjusted hit while reducing Load from 1435 to 761 MB/token.
  • Cache capacity and refinement budget: Increasing refinement budget improves raw hit but can reduce adjusted hit and increase Load through over-prefetching.For B = 20, moving from R = 8 to 20 raises hard hit from 81.84% to 93.49%, while adjusted hit falls from 74.51% to 57.93% and Load rises from 1015 to 2461 MB/token.
  • Cache-loss weight and expert concentration: Higher cache-loss weight increases expert-use concentration and cache efficiency, but excessive supervision correlates with lower task accuracy.For Temporal Router, sw from 0.1 to 2.0 raises maximum expert frequency from 6.82% to 9.40% and top-8 mass from 36.76% to 48.56%, while entropy falls from 3.96 to 3.51.
  • Cache-loss weight and expert concentration: For Spatio-Temporal Router, stronger supervision raises hard hit from 90.62% to 97.41% and adjusted hit from 69.03% to 79.67%, while moderate weights offer a better quality–traffic balance than large weights.At sw = 0.3 and 2.0, Load falls relative to MoE/LRU, but accuracy drops by 1.97–2.04 and 8.94 points, respectively.

6 RELATED WORK

Prior work addresses expert serving through memory management, replacement, prediction, and hardware support. The paper distinguishes its approach by learning cache priorities during post-training while preserving native Top-K routing and supporting update-only retention or bounded pre-access refinement.

  • Efficient LLM and MoE serving: Large-model and MoE serving systems reduce memory pressure through tensor placement, KV-cache management, batching, and parallel execution.Examples include FlexGen, vLLM/PagedAttention, DeepSpeed-MoE, and Tutel.
  • Expert caching and prefetching: Classical LRU, LFU, and LRFU policies are training-free but cannot learn task- or model-specific expert reuse.MoE-specific methods use richer activation, stale-access, or prompt-level signals for caching and prefetching.
  • Spatio-temporal expert prefetching: Spatio-temporal prefetching methods emphasize inference-time prediction or system support, whereas this method learns in-model cache priorities through post-training.The proposed framework preserves native Top-K selection and supports both update-only retention and bounded pre-access refinement.

7 CONCLUSION

The framework manages expert caches during memory-constrained MoE inference while preserving native Top-K routing. Across two backbones and three benchmarks, joint post-training improves cache coverage and reduces decode-stage traffic, with Spatio-Temporal refinement offering the strongest evaluated configuration.

  • 7 CONCLUSION: Joint post-training improves cache coverage and reduces decode-stage traffic across two backbones and three benchmarks.The framework adds at most 0.083% inference-time parameters.
  • 7 CONCLUSION: Temporal Router retains experts after access, while Spatio-Temporal Router adds bounded pre-access refinement; both preserve the native inference-time Top-K rule.
  • 7 CONCLUSION: Auxiliary-only training yields modest gains, whereas joint post-training is substantially stronger.
  • 7 CONCLUSION: Spatio-only refinement raises raw hit but incurs more proactive traffic without temporal updates.
  • 7 CONCLUSION: Cache size, refinement budget, and cache-loss weight control memory, pre-access coverage, traffic, and quality.

LIMITATIONS

The evaluation is limited to algorithmic cache management and decode-stage traffic rather than end-to-end serving. Its empirical scope and adaptation requirements also constrain how broadly the results can be interpreted.

  • LIMITATIONS: The study evaluates algorithmic cache management rather than an end-to-end serving stack.Realized latency and throughput also depend on scheduling, overlap, bandwidth, and batching.
  • LIMITATIONS: Because prefill traffic is excluded, the metrics characterize decode-stage rather than request-level efficiency and do not imply end-to-end speedup.
  • LIMITATIONS: Main experiments update the full backbone, requiring full-model post-training despite small inference-time model-size overhead.The auxiliary-only ablation does not isolate native-router adaptation from other backbone changes.
  • LIMITATIONS: The study focuses on academic reasoning benchmarks and pure MoE models, leaving broader workloads, memory budgets, pretraining settings, and hybrid architectures for future work.
  • LIMITATIONS: Baseline comparisons may change under other memory hierarchies or runtime policies because the baselines use different task definitions and system assumptions.
Loading 2609.04895v1…