A Universal Load Balancing Principle and Its Application to Large Language Model Serving

TL;DR

This paper develops a universal load-balancing principle, BF-IO, for barrier-synchronized, stateful parallel systems such as data-parallel LLM decoding. BF-IO uses a short-horizon, step-wise integer optimization to allocate work with sticky assignments, minimizing predicted imbalance over a few steps rather than predicting full remaining workloads. The authors prove worst-case, scale-sensitive guarantees showing BF-IO reduces long-run imbalance by a factor that scales as $oldsymbol{ ilde{ heta}}(\, ext{sqrt}(B\log G))$ with batch size $B$ and number of workers $G$, and they show the same scaling holds for a broad class of non-decreasing drift workloads. Empirically, BF-IO substantially improves throughput and latency while reducing energy consumption in both open and proprietary LLM workloads, with additional gains from longer lookahead windows. The results offer a theoretically grounded framework for load balancing in synchronization-gated resource allocation and have immediate implications for sustainable LLM serving and other parallel computing domains.

Abstract

Load balancing-the allocation of work across parallel resources to reduce delay, energy and cost-is a pervasive challenge in science and engineering, from large-scale simulation and data processing to cloud and manufacturing operations. Motivated by the emerging bottleneck in large language model (LLM) serving, we study a particularly stringent regime of load balancing that arises in barrier-synchronized, stateful systems: work cannot be freely migrated and progress is gated by the slowest participant at each step, so heterogeneity and temporal drift in workloads create persistent stragglers and substantial idle time. LLM serving under data-parallel decoding provides a prominent modern instance: in production traces, barrier-induced idle can exceed 40% of compute time per decode step. Here we develop a universal load-balancing principle, which admits a step-wise finite-horizon integer-optimization formulation and yields worst-case guarantees: across LLM decode models and a broader class of non-decreasing workload drift processes, it reduces long-run imbalance by a factor that grows with batch size and system scale. Extensive experiments corroborate the theory, showing substantial improvements in throughput and latency together with reductions in energy consumption. These results provide a general, theoretically grounded framework for load balancing, with immediate implications for sustainable LLM serving and broad relevance to other synchronization-gated resource-allocation problems.

Figures (7)

• Figure 1: Workload imbalance and idle time in $436$ decoding steps over $3$ minutes from a real industrial trace. (a): Dynamic total workload across 32 GPUs. At each step, the completion time for all GPUs is dictated by the heaviest workload unit, demonstrating significant and persistent workload imbalance. (b): Per-step idle time resulting from this imbalance. The yellow bars represent the average percentage of time that GPUs spend waiting during each decode step. With a mean (and median) idle time of 40% (and 41%), more than two-fifths of the aggregate computational resources are wasted on average per step.
• Figure 2: Instantaneous GPU power consumption over time under the default policy and the approach introduced in this paper (BF-IO). Total energy, computed as the time integral of instantaneous power, is 396 kJ for the default and 383 kJ for BF-IO—a 3.4% reduction. BF-IO draws slightly higher instantaneous power (reflecting improved GPU utilization from reduced barrier idling), but completes the same workload in less time; the resulting shorter integration window yields lower total energy consumption.
• Figure 3: Prefill–decode serving and the decode-stage load-balancing barrier. Left (Decode Step t): Prefill workers build each request’s KV cache and hand the request to the scheduler, which assigns it to a decode worker; the assignment is sticky because migrating the KV cache is impractical. Each decode worker processes a batch of requests during the local compute stage (Multi-head Latent Attention). Gray squares indicate the resident KV cache per request. After local compute, all workers must begin the EP stage synchronously. The step time is therefore governed by the slowest local stage: here, Worker 2 must wait for Worker 1, exposing the load-balancing bottleneck. Right (Decode Step t+1): Every active request produces one token, appending to its KV cache (red square). Completed requests leave the batch (e.g., Request 4), and newly prefetched requests may be admitted by the scheduler to a worker (e.g., Request 7 with green square). Because request lengths differ and assignments are sticky, per-worker loads drift from step to step, creating persistent imbalance that limits throughput and energy efficiency. Legend: gray = existing KV cache; red = KV added by the current step; green = newly admitted request after prefill.
• Figure 4: Intuition for selecting the window length $H$, evaluated across values from $0$ to $500$. This analysis uses an implementation of the Algorithm BF-IO in the LLM serving context. Further implementation and experimental details are provided in Sections \ref{['sec:num']}.
• Figure 5: Empirical distributions of decode lengths from production LLM traces wang2023openchatwang2024burstgptzheng2023lmsyszhao2024wildchat. Decode lengths exhibit the geometric (discrete-exponential) pattern.

Theorems & Definitions (24)

• Definition 1: Overloaded arrival instance family $\mathcal{I}$
• Theorem 1: Warm-up: Homogeneous decode lengths
• Theorem 2: Inhomogeneous decode lengths
• Definition 2: Non-decreasing Workload Drift
• Theorem 3: General non-decreasing workload drift
• proof : Proof of Theorem \ref{['thm:homog-o']}:
• Lemma 1
• proof : Proof of Lemma \ref{['clm:gap-le-smax']}
• proof : Proof of Theorem \ref{['thm:inhomog-o']}:
• Lemma 2: Separation and gap reduction