PROBE: Co-Balancing Computation and Communication in MoE Inference via Real-Time Predictive Prefetching

TL;DR

Mixture-of-Experts (MoE) inference faces latency bottlenecks from spatial skew and temporal hotspot shifts under expert parallelism. PROBE introduces Continuous Lookahead Pipelining with a Gate-Initialized Lookahead Predictor, Hardware-Aware Balance Planning, and Phase-Locked Co-Scheduling to proactively predict, plan, and prefetch for upcoming layers while hiding overheads behind computation. The system achieves up to 1.32x prefill latency reduction and up to 1.26x decoding throughput gains, and demonstrates robustness to volatile workloads through high predictor fidelity (Top-K accuracy around 90%+) and effective latency breakdown that shows near-zero overhead on the critical path. This work enables more scalable and responsive trillion-parameter MoE deployments in latency-sensitive serving scenarios by shifting balancing from reactive adjustments to predictive, real-time execution.

Abstract

Mixture-of-Experts models have become a dominant architecture for scaling Large Language Models by activating only a sparse subset of experts per token. However, latency-critical MoE inference faces a fundamental tension: while expert parallelism improves memory efficiency, it also amplifies execution stragglers. In real-world serving, continuous batching and diverse concurrent requests induce rapid semantic shifts, causing expert hotspots to migrate abruptly across GPUs and triggering the 'double penalty' of coupled computational skew and network congestion. We propose PROBE, an inference system that co-balances computation and communication in real time. PROBE introduces Continuous Lookahead Pipelining, which proactively predicts, plans, and prefetches for upcoming layers while keeping all control overheads off the critical path. PROBE consists of: (1) a Gate-Initialized Lookahead Predictor that distills the target router to forecast next-layer expert activation with high fidelity; (2) a Hardware-Aware Balance Planning solver that jointly optimizes dynamic expert replication and token assignment under strict hiding-window constraints; and (3) a Phase-Locked Co-Scheduling policy that uses split-phase transmission to hide bandwidth-intensive expert transfers behind computation without contending with All-to-All collectives. Experiments show that PROBE reduces prefill latency by up to 1.32X and improves decoding throughput by up to 1.26X over state-of-the-art baselines, especially under extreme workload volatility.

Figures (10)

• Figure 1: Spatial and Temporal Challenges in MoE Inference. Under Expert Parallelism, efficiency is constrained by: (1) Spatial imbalance, where skewed token-to-expert routing creates computational stragglers and communication bottlenecks; and (2) Temporal volatility, where expert hotspots shift rapidly over time under continuous batching. This motivates a system that can handle both instantaneous skew and continuous distribution shifts.
• Figure 2: Expert activation patterns across prefill and decoding. Measurements use $ep=8$ with a standard sharded expert placement policy. Subfigures (a) and (b) show the concentrated, bursty skew during prefill ($\approx$32K tokens). Subfigures (c) and (d) show rapid load shifts during decoding ($\approx$8K tokens), where expert popularity changes with semantic transitions. Comparing GPT-OSS-120B (Top-4 out of 128 experts) and Qwen3-235B (Top-8 out of 128 experts) illustrates that model sparsity patterns further modulate imbalance severity.
• Figure 3: MoE compute latency. Profiling via SGLang on GPT-OSS-120B (128 experts, Top-4). We compare EP (Max/Avg/Min) with DP and EP + 4 extra experts. DP is bottlenecked by fragmentation (low arithmetic intensity and padding), while modest EP redundancy mitigates stragglers with minimal memory overhead.
• Figure 5: Skew hurts All-to-All efficiency. Benchmarked on 8$\times$H800, GPT-OSS-120B with DeepEP zhao2025deepep. Top: effective All-to-All Dispatch bandwidth. Bottom: max per-rank traffic volume. Compared to a manually balanced top-$K$ baseline, real workloads create receiver hotspots and reduce effective bandwidth; Combine phase shows similar behavior.
• Figure 6: Overview of PROBE. The system implements a dual-track execution that overlaps control-plane operations with the main stream on complementary resources: predictor/planner run during network-bound All-to-All, while P2P prefetching is overlapped with compute-bound GEMM and attention to hidden management overhead.