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

TL;DR

PROBE is a real-time inference system for Mixture-of-Experts models that proactively predicts and prefetches expert activations, balancing computation and communication to reduce latency and improve throughput in dynamic workloads.

Contribution

PROBE introduces a novel real-time predictive prefetching system with a lookahead predictor, dynamic expert balancing, and co-scheduling to optimize MoE inference performance.

Findings

Reduces prefill latency by up to 1.32X

Improves decoding throughput by up to 1.26X

Effective under extreme workload volatility

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.