PaperlayerView exact revision on arXiv

Source-linked AI summary

S2-MoE: Enabling Efficient Self-Speculative Decoding for Mixture-of-Experts on Edge Devices

Haochen Huang, Shengxuan Qiu, Meng Li

arXiv:2608.15018v1cs.AI

TL;DR

Edge-device LLM inference is constrained by memory bandwidth and limited MoE expert reuse. S2-MoE addresses this with routing-aware speculation, reuse-aware gating, and context alignment, achieving 1.3×–5.3× speedup on Jetson Orin and 1.2×–2.9× on RTX 4090 over autoregressive decoding.

  • Problem

    Edge LLM inference is constrained by memory bandwidth, while diverse MoE expert activations limit parameter reuse during verification.

  • Method

    S2-MoE combines routing-aware adaptive speculative expansion, reuse-aware expert gating, and context-aligned self-speculative decoding for MoE edge inference.

  • Results

    1.3×–5.3× speedup on Jetson Orin and 1.2×–2.9× on RTX 4090 over autoregressive decoding, with consistent improvements across models, tasks, and memory constraints.

  • Takeaways & Limitations

    S2-MoE improves end-to-end MoE inference efficiency in memory-constrained edge environments while maintaining comparable accuracy.

  • Takeaways & Limitations

    PPL/KL evaluation is omitted for GPT-OSS because standard next-token likelihood is unreliable outside its Harmony chat format and post-training objectives.

Abstract

from arXiv · show

Deploying large language models (LLMs) for inference on edge devices is challenging due to severe memory and bandwidth constraints. While speculative decoding and Mixture-of-Experts (MoE) have been proposed to improve inference efficiency, naively combining them often incurs excessive verification overhead and poor expert reuse, limiting their effectiveness in memory-bound edge settings. In this work, we propose S2-MoE, an efficient self-speculative decoding framework for MoE inference on edge devices. S2-MoE reduces redundant verification through routing-aware adaptive speculative expansion, improves verification efficiency with reuse-aware expert gating, and aligns draft and target execution via shared context. Implemented in llama.cpp, S2-MoE achieves up to 5.3x speedup (about 2.0x on average) over standard autoregressive de?coding across diverse MoE models and datasets on edge devices.Code is available at https://github.com/angerybob/S2-MoE.

1 Introduction

Edge LLM inference is constrained primarily by memory bandwidth, while speculative decoding and MoE improve output density and parameter density respectively but combine poorly because MoE reduces expert reuse during verification. S2-MoE addresses these challenges with routing-aware expansion and reuse-aware gating, achieving substantial speedups over standard autoregressive decoding.

  • Motivation: Edge LLM inference faces severe resource constraints, with memory bandwidth becoming a key bottleneck under small-batch or single-batch execution.Limited parameter reuse increases memory-bandwidth pressure on edge devices.
  • Efficiency factors: Throughput depends on output density—the tokens generated per iteration—and parameter density—the model parameters accessed per iteration.Bandwidth limits on resource-constrained devices constrain this tradeoff.
  • Existing acceleration: Speculative decoding increases output density, whereas MoE reduces parameter density by activating only a subset of model parameters per iteration.These mechanisms are complementary in principle, but their combination is difficult in memory-bound edge settings.
  • Challenge: Naïve speculative decoding on MoE models incurs high verification cost because speculative tokens activate diverse experts, reducing parameter reuse and causing redundant verification.Self-speculative decoding improves draft-target alignment, but lightweight drafts can still suffer limited acceptance, rejected-draft overhead, and exacerbated expert-diversity costs.
  • S2-MoE: S2-MoE introduces routing-aware adaptive speculative expansion and reuse-aware expert gating to reduce redundant verification and improve expert reuse for edge MoE inference.The framework is designed to align draft and target behavior while avoiding low-quality candidates with high verification cost.
  • Results: 1.3× to 5.3× speedup on Jetson Orin and 1.2× to 2.9× on RTX 4090 are achieved over standard autoregressive decoding.The implementation is released at https://github.com/angerybob/S2-MoE.

2 Background

MoE reduces inference cost through sparse expert activation, making it attractive for edge deployment, while speculative decoding accelerates generation by proposing and parallelly verifying future tokens. However, applying speculative decoding to MoE models is difficult because suitable external drafts are challenging to construct and existing approaches may overlook verification inefficiencies.

  • Mixture-of-Experts: MoE models sparsely activate a small subset of experts per token, reducing effective computation and parameter access for memory-constrained edge inference.This sparse execution supports large-capacity models under lower average memory and compute cost.
  • Speculative Decoding: Speculative decoding accelerates autoregressive generation by proposing future tokens with a lightweight mechanism and verifying them in parallel with the target model.The approach is particularly beneficial in memory-bound settings.
  • Challenges for MoE: External-draft speculative decoding is difficult for MoE models because a draft must be lightweight and well aligned despite already sparse target activation.MoE sparsity leaves limited room to construct a substantially smaller, well-aligned draft model.
  • Self-Speculative Decoding: Self-speculative decoding derives a lightweight draft from the target model, avoiding an external draft while maintaining strong behavioral alignment.Existing approaches simplify the target through quantization, expert sparsity, or layer sparsity.
  • Prior MoE Studies: Prior MoE studies use speculative decoding to predict future expert activations for prefetching and offloading but primarily address caching and scheduling rather than intrinsic verification inefficiencies.These studies focus on system-level mechanisms without resolving verification costs under sparse activation.

3 Motivation

Section 3 identifies three obstacles to efficient self-speculative MoE decoding: redundant verification, limited expert reuse, and draft divergence. It motivates routing-aware expansion, reuse-aware gating, and shared context as corresponding remedies.

  • Challenge 1 and Opportunity 1: Confidence-only pruning is weak because retained-token confidence concentrates near 1.0 while verification costs vary with expert activation patterns.This motivates expansion that balances acceptance potential against expected verification overhead.
  • Challenge 1 and Opportunity 1: Draft-time expert routing predicts target routing accurately in many layers, providing a proxy for verification cost during adaptive speculative expansion.The proposed motivation is to jointly consider token-level acceptance potential and routing-derived verification overhead.
  • Challenge 2 and Opportunity 2: MoE models activate different experts for neighboring tokens, limiting parameter reuse during parallel verification.Reuse is measured by the unique experts activated within an interval, normalized by autoregressive expert activations over that interval.
  • Challenge 2 and Opportunity 2: Experts with similar low gating scores can be functionally redundant, enabling tokens to select already activated experts with slightly lower scores.This trades a small loss in per-token gating score for improved expert reuse.
  • Challenge 3 and Opportunity 3: Independent draft contexts accumulate prediction errors, causing progressive divergence and lower acceptance rates in long-form generation.Sharing draft and target context confines errors to the current speculative step instead of propagating across iterations.

4 S2-MoE Design

S2-MoE combines routing-aware adaptive speculative expansion, reuse-aware expert gating, and context-aligned self-speculative decoding to improve MoE inference efficiency on edge devices. These components jointly target redundant verification, excessive expert activation and memory traffic, and draft–target context misalignment.

  • Overall design: S2-MoE combines routing-aware adaptive speculative expansion, reuse-aware expert gating, and context-aligned self-speculative decoding.The three components are described in Sections 4.1–4.3.
  • Routing-aware adaptive speculative expansion: Routing-aware expansion ranks draft candidates by expected acceptance benefit against verification cost, accounting for the experts activated by each candidate.The expected benefit is B_i = p_i · T_AR, while verification cost includes draft expansion and newly introduced expert activation costs.
  • Routing-aware adaptive speculative expansion: Expansion greedily admits candidates with U_i ≥ 1, unlike confidence-only rules that overlook marginal expert activation cost.The criterion continues until no remaining candidate satisfies it.
  • Reuse-aware expert gating: Reuse-aware expert gating aggregates confidence-weighted expert importance across speculative tokens and softly biases a bounded set of preferred experts.The cap preserves the original top-k routing structure while promoting cross-token expert reuse.
  • Context-aligned self-speculative decoding: Context-aligned self-speculative decoding addresses draft–target fidelity loss by sharing a unified KV cache.The method targets errors caused by independently maintained decoding states and accumulating context misalignment.

5 System Implementation

S2-MoE is integrated into llama.cpp for edge inference and implements expert-level offloading tailored to sparse MoE activation. The evaluation uses established speculative-decoding baselines and sparsity-based draft construction across representative MoE families.

  • Framework: S2-MoE fully integrates its proposed mechanisms into llama.cpp’s existing decoding pipeline.llama.cpp is used as the implementation framework for edge and resource-constrained deployment.
  • Expert-level offloading: On Jetson Orin, parameters exceeding device memory are offloaded to SSDs, whereas RTX 4090 transfers cold experts from CPU memory on demand.Jetson Orin uses unified CPU/GPU DRAM, while RTX 4090 has separate CPU and GPU memories.
  • Expert-level offloading: Expert-level offloading dynamically fetches individual expert parameters while keeping non-expert and shared parameters resident in device memory.This avoids the unnecessary data movement caused by offloading entire layers, since only a small subset of experts activates per token.
  • Baseline models and implementations: S2-MoE is evaluated against state-of-the-art speculative-decoding baselines across GPT-OSS, Qwen3, OLMoE, and DeepSeek model families.Public EAGLE-3 checkpoints are used for GPT-OSS, Qwen3, and OLMoE, while DeepSeek decoding is implemented using the open-source SpecForge framework built on sglang.
  • Draft configuration: The draft is constructed using expert sparsity because quantization can add substantial draft overhead and aggressive layer sparsity can degrade draft fidelity.When memory permits, quantization is additionally applied to the sparse draft to reduce computation and memory footprint.

6 Experiments · 6.1 Experimental Setup

The experiments evaluate S2-MoE across four MoE language models, diverse generation tasks, and memory-constrained edge and GPU platforms. Comparisons include autoregressive, self-speculative, and MoE-aware speculative decoding baselines.

  • 6.1 Experimental Setup: Four MoE LLMs span lightweight edge-friendly to large-scale systems: DeepSeek-V2-Lite-Chat, OLMoE-1B-7B, Qwen3-30B-A3B, and GPT-OSS-120B.The models cover diverse parameter sizes, expert configurations, and routing behaviors.
  • 6.1 Experimental Setup: Seven generation-task categories—MT, RG, SU, TR, QA, MA, and HE—are evaluated with identical speculative decoding hyperparameters across tasks.The tasks include conversation, retrieval-augmented generation, summarization, translation, question answering, mathematical reasoning, and code generation.
  • 6.1 Experimental Setup: Two platforms are tested: NVIDIA Jetson Orin edge devices with 16 GB, 32 GB, and 64 GB constraints, and an NVIDIA RTX 4090 under constrained GPU memory.Experts exceeding device capacity are offloaded to SSD on Jetson Orin, while cold experts are offloaded to CPU memory on RTX 4090.
  • 6.1 Experimental Setup: The baseline suite includes standard autoregressive decoding (AR), expert-sparsity self-speculative decoding (ES), and layer-sparsity self-speculative decoding (LS).ES activates fewer experts per layer, whereas LS skips transformer layers selected via Bayesian optimization.
  • 6.1 Experimental Setup: LS is evaluated only on D because of its consistently poor efficiency in experiments.The supplied passage truncates the remainder of the dataset name.
  • 6.1 Experimental Setup: Cascade provides an MoE-aware speculative decoding baseline that selectively enables speculation to balance draft benefit and verification overhead.For a stronger comparison, Cascade is instantiated with EAGLE3 as its underlying speculative decoding engine.

6.2 Main Results

S2-MoE delivers consistent end-to-end speedups over autoregressive decoding and other speculative-decoding baselines across MoE models, tasks, and memory budgets. Its routing- and reuse-aware design is especially beneficial under tight memory constraints, while measured output quality remains comparable to the original model.

  • Effectiveness across settings: 1.3–5.3× speedup on Jetson Orin and 1.2–2.9× on RTX 4090 over standard autoregressive decoding.Speedup Ratio is the primary metric for memory-constrained edge deployment.
  • Effectiveness across settings: Naïve self-speculative decoding often fails to beat autoregressive decoding because low expert reuse and redundant verification offset draft-target alignment.S2-MoE addresses these bottlenecks by reducing verification cost and improving parameter reuse.
  • Baseline comparisons: S2-MoE achieves more stable, consistently higher speedups than EAGLE-3 across diverse modern MoE models and tasks without training.EAGLE-3’s effectiveness varies with draft prediction quality and training characteristics.
  • Baseline comparisons: Routing-aware, token-level speculative control avoids Cascade’s extra online test-and-set overhead and yields more consistent speedup gains.S2-MoE estimates verification utility from routing-aware costs at the current decoding step rather than historical information.
  • Memory and quality impact: Tighter memory budgets amplify S2-MoE’s advantages because more experts are offloaded and verification becomes more sensitive to redundant activation and parameter reuse.Reuse-aware expert gating introduces a bounded routing bias, while measured PPL ratios remain 1.012–1.013 and Top-1 agreement exceeds 89%.
  • Memory and quality impact: GPT-OSS omits PPL/KL evaluation because standard next-token likelihood is unreliable for Harmony-trained models optimized for chat-format and CoT/RL objectives.The passage identifies anomalously high raw-corpus PPL and poorly calibrated token log-probabilities as known evaluation artifacts.

6.3 Ablation Study

The ablation study finds that context alignment, adaptive expansion, and reuse-aware gating each improve S2-MoE and provide complementary gains. Their benefits reflect higher draft fidelity, less redundant verification, greater expert reuse, and sensitive hyperparameter trade-offs.

  • Component contributions: Each S2-MoE component consistently improves performance across Qwen3 and DeepSeek, with their combination producing the highest overall speedup.The results characterize context alignment, adaptive expansion, and reuse-aware gating as complementary rather than redundant.
  • Higher Draft Fidelity: Context alignment consistently improves acceptance rates across DeepSeek and OLMoE, increasing acceptance by up to 79% on reasoning-heavy GSM8K.The method shares an identical KV cache between draft and target execution.
  • Less Redundant Verification: Utility-guided adaptive expansion matches or exceeds fixed-width acceptance length while approaching aggressive confidence pruning in verification efficiency.It achieves a strictly better acceptance–verification trade-off than confidence-based pruning by jointly considering acceptance likelihood and verification cost.
  • Higher Expert Reuse: Reuse-aware expert gating consistently reduces the reuse ratio across datasets, indicating substantially higher expert reuse during DeepSeek verification.The gating aligns expert selection across speculative tokens and improves parameter locality.
  • Hyperparameter sensitivity: Draft expert top-k trades off acceptance length against draft overhead, while the reuse-aware gating cap must avoid being too small or excessively large.Increasing top-k improves draft quality but can reduce overall speedup; cap extremes weaken expert-reuse improvement and the reuse signal.

7 Conclusion

S2-MoE is an efficient self-speculative decoding framework for MoE models in memory-constrained edge environments, using context alignment, utility-guided speculation, and reuse-aware gating. It achieves substantial speedups over autoregressive decoding while maintaining comparable accuracy and outperforming state-of-the-art speculative baselines.

  • 7 Conclusion: S2-MoE alleviates memory-bandwidth bottlenecks through context alignment, utility-guided speculation, and reuse-aware gating.The framework targets memory-constrained edge environments for Mixture-of-Experts models.
  • 7 Conclusion: S2-MoE consistently outperforms state-of-the-art speculative baselines while maintaining comparable accuracy.The conclusion reports this comparison alongside the device-specific speedups.

A End-to-End Raw Measurements

This appendix presents the raw measurements underlying Fig. 7, covering throughput, acceptance length, and speedup across four hardware configurations. Tables 5–8 organize results for Jetson Orin NX, Jetson AGX Orin, and RTX 4090 devices.

  • Measurement scope: The appendix reports raw throughput (Tok/s), acceptance length (Acc.), and speedup measurements underlying Fig. 7.These measurements are organized in Tables 5–8.
  • Hardware configurations: The measurements cover NVIDIA Jetson Orin NX 16 GB, Jetson AGX Orin 32 GB and 64 GB, and RTX 4090.Tables 5–8 correspond to these four hardware configurations.
Loading 2608.15018v1…