Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

XShare: Collaborative in-Batch Expert Sharing for Faster MoE Inference

Daniil Vankov, Nikita Ivkin, Kyle Ulrich, Xiang Song, Ashish Khetan, George Karypis

TL;DR

XShare addresses the batch-induced explosion of activated experts in production MoE inference by formulating batch-aware expert selection as a modular optimization problem and solving it with efficient greedy methods that require no retraining. The approach aggregates gating mass per layer, enforces a global budget, and refines routing to maintain per-token accuracy, with warm-up and refinement steps; it extends to speculative decoding via a hierarchical proxy and to multi-GPU deployments with GPU-aware load balancing. Empirical results show up to 30% reduction in activated experts, up to 3x reduction in peak GPU load in expert-parallel deployments, and up to 14% end-to-end throughput gains, while keeping accuracy within roughly 1% across diverse benchmarks and batch patterns. The methods are designed to integrate with existing serving frameworks and demonstrate robust performance on heterogeneous batches, providing practical improvements for production MoE serving without retraining.

Abstract

Mixture-of-Experts (MoE) architectures are increasingly used to efficiently scale large language models. However, in production inference, request batching and speculative decoding significantly amplify expert activation, eroding these efficiency benefits. We address this issue by modeling batch-aware expert selection as a modular optimization problem and designing efficient greedy algorithms for different deployment settings. The proposed method, namely XShare, requires no retraining and dynamically adapts to each batch by maximizing the total gating score of selected experts. It reduces expert activation by up to 30% under standard batching, cuts peak GPU load by up to 3x in expert-parallel deployments, and achieves up to 14% throughput gains in speculative decoding via hierarchical, correlation-aware expert selection even if requests in a batch drawn from heterogeneous datasets.

XShare: Collaborative in-Batch Expert Sharing for Faster MoE Inference

TL;DR

XShare addresses the batch-induced explosion of activated experts in production MoE inference by formulating batch-aware expert selection as a modular optimization problem and solving it with efficient greedy methods that require no retraining. The approach aggregates gating mass per layer, enforces a global budget, and refines routing to maintain per-token accuracy, with warm-up and refinement steps; it extends to speculative decoding via a hierarchical proxy and to multi-GPU deployments with GPU-aware load balancing. Empirical results show up to 30% reduction in activated experts, up to 3x reduction in peak GPU load in expert-parallel deployments, and up to 14% end-to-end throughput gains, while keeping accuracy within roughly 1% across diverse benchmarks and batch patterns. The methods are designed to integrate with existing serving frameworks and demonstrate robust performance on heterogeneous batches, providing practical improvements for production MoE serving without retraining.

Abstract

Mixture-of-Experts (MoE) architectures are increasingly used to efficiently scale large language models. However, in production inference, request batching and speculative decoding significantly amplify expert activation, eroding these efficiency benefits. We address this issue by modeling batch-aware expert selection as a modular optimization problem and designing efficient greedy algorithms for different deployment settings. The proposed method, namely XShare, requires no retraining and dynamically adapts to each batch by maximizing the total gating score of selected experts. It reduces expert activation by up to 30% under standard batching, cuts peak GPU load by up to 3x in expert-parallel deployments, and achieves up to 14% throughput gains in speculative decoding via hierarchical, correlation-aware expert selection even if requests in a batch drawn from heterogeneous datasets.
Paper Structure (24 sections, 2 theorems, 10 equations, 8 figures, 4 tables, 6 algorithms)

This paper contains 24 sections, 2 theorems, 10 equations, 8 figures, 4 tables, 6 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. MoE Layer Definition
  4. Gating and Routing Mechanism
  5. Inference Setting
  6. Problem Formulation
  7. Batch-Level Expert Activation
  8. Optimization Objective
  9. Modularity and Optimal Greedy Solution
  10. Practical Algorithm
  11. Speculative Decoding Aware Pruning
  12. Hierarchical Proxy for Speculative Decoding
  13. Experts Selection for Multi-GPU Inference
  14. GPU-Aware Proxy Function
  15. Experiments
  16. ...and 9 more sections

Key Result

Proposition 3.2

For each layer $l$, the function $f_l(S_l; G^{(l)})$ is modular. That is, for any subset of experts $S \subseteq E$ and expert $e \in E \setminus S$: independent of $S$.

Figures (8)

  • Figure 1: Average number of activated experts
  • Figure 2: Batch utility expert pruning
  • Figure 3: Overlap between the top-$k$ ($k=5,10,15,30$) experts ($|E_i\cap E_j|$) for (1) two speculative tokens; (2) two tokens from the same dataset; and (3) two tokens from different datasets.
  • Figure 4: Trade-off between OTPS improvement and accuracy degradation across three datatsets for GPT-OSS 120B (BS=16, speculations off). Each datapoint is accompanied with the pair of parameters defining the algorithm settings (budget $m_l$, warm-up $k_0$).
  • Figure 5: Trade-off between OTPS improvement and accuracy degradation across three datatsets for GPT-OSS 120B (BS=4, speculation length 3). Each datapoint has a set of parameters defining the algorithm settings (warm-up $k_0$, batch budget $m_l$, per-request budget $m_r$).
  • ...and 3 more figures

Theorems & Definitions (2)

  • Proposition 3.2: Modularity of $f_l$
  • Corollary 3.3