Astrolabe: Balancing Load in LLM Serving with Randomized Prediction-Guided Scheduling

This paper presents Astrolabe, a randomized prediction-guided scheduler for one-shot request dispatch in multi-instance large language model (LLM) serving. Astrolabe improves load balancing without relying on migration-based rebalancing, whose KV-cache transfers can introduce substantial overhead and network contention under high load. Astrolabe combines response-length estimation, per-instance simulation-based latency prediction, and a power-of-two choices dispatch policy to improve load balancing and avoid request herding. On the default Llama-2-7B/ShareGPT setup, Astrolabe matches the best load-aware baseline's SLO capacity (31.6 vs 31.5 QPS) while improving latency metrics: 8–36% lower mean time-to-first-token (TTFT), 16–77% lower P99 TTFT, up to 5.6% lower mean end-to-end (E2E) latency, and 6× fewer preemptions once capacity is reached. Under configuration shifts, Astrolabe also improves SLO capacity by up to +6% on Qwen2-7B and +7.1% under tight batching, with 6–9% lower mean E2E under bursty arrivals, while achieving a 2.8-fold reduction in per-predictor CPU relative to full fanout. With migration enabled on A100, Astrolabe outperforms Llumnix by up to 2.6× in throughput with orders-of-magnitude lower per-token latency at saturation.

The paper notes that "the rise of large language models like GPT [1], Llama [62], Gemini [61], Qwen [72], and DeepSeek [41] has revolutionized modern applications, such as chatbots [1], virtual assistants [31], code generation [34], and creative writing [32]." This places pressure on LLM inference serving systems to meet latency requirements and maximize throughput. Key techniques such as continuous batching [74], Paged Attention [39], chunked prefill [4], and FlashAttention [19] have been proposed.

However, "LLM inference is often characterized as unpredictable [60]. For example, LLMs generate tokens autoregressively based on preceding tokens until a stop signal is reached, leading to variable response lengths and numbers of decoding steps [57]. Paged Attention [39], which dynamically allocates memory resources and allows request preemption, further contributes to the dynamic nature of runtime memory consumption. Decode-step latency also exhibits high variance due to batch-size variation and interference between requests in the same batch [60]."

These uncertainties make load balancing across LLM serving instances challenging: "common runtime metrics (latency, throughput, memory) no longer represent each request's E2E load. Moreover, unexpectedly long responses further load already-overloaded hosts and block subsequent requests (head-of-line blocking [16, 51]). In practice, production multi-instance frameworks [6, 47] dispatch with simple heuristics like round-robin, offering no performance guarantee."

Alternative solutions like Llumnix [60] combine dispatching with dynamic rebalancing through live migration, but "migration requires transferring request KV caches [38] across instances, adding network, memory, and orchestration overhead. Under high load, these costs can accumulate and contend with normal serving traffic, potentially degrading tail latency and throughput."

The paper identifies an opportunity: if length and latency can be estimated at dispatch time, the scheduler can proactively route long requests to less-loaded instances before imbalance manifests. A basic approach of querying every instance's predictor in parallel (full fanout) has three costs: "First, each dispatch costs O(N) predictor queries for N instances, which scales poorly in the scheduler control plane. Second, simulation runs on CPU to predict the scheduling signal (e.g., latency) for each candidate instance, so aggregate cluster-CPU grows linearly with N per request. Third, deterministic fanout creates a thundering-herd issue: under short correlated bursts, nearby requests compute similar 'best' candidates and route to the same instance, overloading it transiently [8, 44]."

Astrolabe comprises four services: a query length tagger, a global scheduler, per-instance Predictor sidecars, and inference framework backends. The query length tagger is the entry point and predicts the anticipated response length for each request. Requests are then forwarded to the global scheduler, which selects a target model instance. For each incoming request, the global scheduler issues predict RPCs to a sampled subset of Predictor sidecars (power-of-k choices; k=2 by default, and k=N recovers full fanout). Each Predictor queries its colocated backend via a status API to obtain the instance's current runtime state (e.g., running/waiting requests, batching state, and memory availability), and then runs a lightweight simulation to estimate the request's target metric (e.g., E2E latency). The global scheduler picks the target instance from these predictions per the configured policy, and the backend executes the request and returns the response.

A model instance is the collection of services on a GPU host driving one or more GPUs: the inference framework itself plus a sidecar service called the Predictor. The Predictor's main role is to predict key performance metrics, such as E2E latency or TTFT, for incoming requests. "The Predictor runs locally on each instance, aggregates runtime data from the status API, and serves a scalar metric prediction via the predict API to the global scheduler. In cases where a request's decoded length exceeds its predicted length at runtime, the simulator extends the current prediction by assuming the request will complete within the next 10 tokens by default."

The prototype integrates Vidur [2], redesigned for single-instance prediction and encapsulated behind the Predictor API. Simulation proceeds in two stages: a local-scheduler simulator models the backend's batching strategy, then a linear model predicts execution times for the generated batches.

The global scheduler service is stateless and lightweight without maintaining a global table of instance statuses. Upon each request arrival, the scheduler issues a predict RPC to sampled Predictor sidecars and then dispatches the request to the instance with the lowest predicted metric. E2E latency is the default target metric.

"This differs from centralized predictors that need a global per-instance status view (e.g., the estimation-driven global schedulers in TetriServe [33] and DynamoLLM [59]). In Astrolabe, prediction runs on each instance in parallel using its local state, and each Predictor returns only a scalar metric (e.g., predicted latency). Decision time is bounded by the slowest Predictor reply, while avoiding the cost of exporting heavy instance state (e.g., request lists) to a central simulator."

"To avoid the O(N) per-dispatch messaging cost of full fanout and the scheduler-herding issue, Astrolabe's default samples k=2 random instances (power-of-two choices) and queries only those Predictors. Setting k=N recovers full fanout."

"The query length tagger is an online service that estimates the response length of each request from its prompt and the target serving model. It runs in parallel with request handling and uses a lightweight proxy model to keep overhead low." The predicted length is used only for metric prediction and scheduling decisions and does not change the inference parameters and results.

The prototype integrates with vLLM 0.7.2 and comprises about 4,000 LoC. The Predictor internally invokes a backendspecific simulator based on Vidur's offline simulator codebase, substantially adapted for online, per-request use. The authors rewrote the primary simulation hot paths and added a per-Predictor cache that memoizes latency predictions per (batch size, token count), substantially reducing the per-request simulation cost.

For the query length tagger, the authors fine-tuned a lightweight RoBERTa-base [42] regression model (125M parameters). The length estimation model achieves MAE 78.8 and MAPE 24.4% on the 10k test split, which is comparable to the 7B prompt-based model reported in Sequence Scheduling [77].

The primary cluster consists of 12 d7525 nodes on CloudLab [22], each with two 16-core AMD 7302 CPUs, 128 GB ECC memory, one NVIDIA A30 GPU (24 GB), and 25 Gbps NICs. A second 2-node CloudLab d8545 cluster has 4× A100-40GB GPUs per node and 100 Gbps NICs. The main experiments use ShareGPT [71] with Llama-2-7B [62] in FP16 precision. The model weights occupy 12.5 GB of GPU memory, split into 1056 vLLM memory blocks for KV cache. The maximum request batch size is 48, and the Chunked-Prefill chunk size is 512.

The length prediction model reaches MAE 78.8 and MAPE 24.4% on the 10k test split. For simulation-based metric prediction, Enabling Chunked Prefill yields much lower error than vLLM's default Prioritized Prefill. The scheduler maintains a high probability (40 to 80%) of selecting the Rank-1 instance among candidates. For Llama-2-70B on the A100 cluster with TP=4, the simulator achieves a latency-prediction error of 16.6%.

Across the 17-point QPS sweep, "Astrolabe's peak reductions over the best load-aware baseline are 38% mean TTFT, 78% TTFT-P99, 7% mean E2E, and 12% E2E-P99, with up to 5% higher throughput; relative to the worst baseline, the gains reach 97% lower TTFT and 53% higher throughput." Astrolabe's per-request overhead stays under 3% of E2E within capacity.

At QPS=32, Llumnix- accumulates 1.8k preemptions vs Astrolabe's 300 (6× fewer). The serving capacity is 31.6 QPS for Astrolabe vs 31.5 for Llumnix-.

On the A100 cluster with full migration enabled, Llumnix collapses at QPS=20, where mean request latency jumps from 6.3 s to 106 s and mean token latency from 21 ms to 1.15 s. At QPS=36, "Astrolabe without CP delivers 2.57× the throughput (18.7k vs 7.3k tok/s), 36× lower mean request latency (8.2 s vs 300 s), and 123× lower mean token latency (26 ms vs 3.2 s) than Llumnix. Enabling CP widens these to 2.58×/40×/135×."

Capacity is effectively tied from k=2 upward: Po2-est reaches 31.6 QPS vs Po4-est 31.9 (the maximum), Po8-est 31.7, and Fanout-est 31.7, all within a ±0.3 QPS band at SLO = 10s. The paper notes that k=4 gives the best empirical tradeoff: it has the highest capacity and throughput, while retaining the second-lowest E2E latency and scheduling overhead.

Oracle adds 0.8 QPS for Po2 (31.6 → 32.4) and 0.9 QPS for Fanout (31.7 → 32.6) at SLO = 10s, a 3% capacity gain.

Under high burstiness (α=0.25), Astrolabe's mean E2E latency is about 9% lower than Llumnix-; under near-regular arrivals (α=2.0), the gap narrows but Astrolabe retains a small edge. Moving from α=2.0 to α=0.25 raises mean E2E by 13% for Astrolabe vs 22% for Llumnix-.

The largest mean degradation is +6.3% at (100%, 25%), while the extreme (100%, 100%) cell is +5.1%; most cells remain within the ±4% single-deploy noise floor. The paper notes that P99 degradation remains below Llumnix- in every cell (maximum +3.7%).

At QPS=20 mean per-predictor CPU is 20.3% for Po2 and 56.3% for Fanout, a 2.8-fold reduction. Per-predictor RSS is Po2 1225 MB, Fanout 1773 MB at QPS=20.

Across configuration variations: with batch size reduced 48→24, Astrolabe reaches 28.5 QPS vs Llumnix- 26.6 QPS (+7.1%); with Qwen2-7B, Astrolabe reaches 73.9 QPS vs Llumnix- 69.7 QPS (+6.0%, +12.6% with oracle lengths); with BurstGPT [67], Astrolabe-oracle reaches 64.5 QPS vs Llumnix- 61.6 QPS (+4.7%).

"Astrolabe combines response-length estimation, simulation-based latency prediction, and randomized dispatch for multi-instance LLM serving. Across our experiments, prediction-guided one-shot dispatch matches or outperforms migration-based rebalancing while avoiding KV-cache transfer overhead. The paper leaves larger-scale multi-GPU evaluation, richer failure handling, and length-model retraining under workload drift as future work."

Based on the paper, here are the specific improvements I can implement in an AI system:

• Add a lightweight RoBERTa-base regression model (125M parameters) that predicts the expected response length from the input prompt before generation begins

• Use this prediction to inform scheduling decisions, not to alter generation behavior

• Achieve MAE of 78.8 tokens and MAPE of 24.4% on ShareGPT-style prompts

• The model runs in parallel with request handling, adding only 4.8ms per request

• Integrate a simulator (based on Vidur's approach) that predicts end-to-end latency for a request on each candidate instance

• The simulator accounts for:

• Current batch composition and size

• Pending queue depth

• Available KV cache memory blocks

• The backend's batching strategy (e.g., chunked prefill vs. prioritized prefill)

• Predicted response length

• Maintain a cache that memoizes latency predictions per (batch size, token count) to reduce per-request simulation cost

• Achieve 10-15% absolute prediction error, with 40-80% probability of selecting the true Rank-1 instance

• For each incoming request, sample exactly 2 random instances (not all instances)

• Query both instances' predictors in parallel

• Route the request to the instance with the lower predicted latency

• This reduces per-request predictor fanout from O(N) to O(2), cutting per-predictor CPU by 2.8x compared to full fanout

• Randomization breaks herding: simultaneous requests are unlikely to sample the same pair, preventing transient overload

• Implement the auto-extension rule: if a request's actual length exceeds prediction, extend the simulation assuming completion within the next 10 tokens

• Use relative ranking (not absolute latency values) for instance selection, so common-mode errors cancel out

• The system tolerates up to ±100% independent Gaussian noise on latency predictions with at most +6.3% mean E2E degradation and +3.7% P99 E2E degradation

• Default to E2E latency as the scheduling objective, but keep the interface configurable

• Support alternative objectives (e.g., TTFT, throughput) via a pluggable interface

• Allow tuning the number of sampled instances (k) from 2 to N, where k=N recovers full fanout

For a multi-instance LLM serving cluster (e.g., 12 GPUs):

1. Achieve 31.6 QPS SLO capacity (P99 TTFT ≤ 10s) on Llama-2-7B/ShareGPT, matching the best load-aware baseline (31.5 QPS) and full-fanout prediction (31.7 QPS) while using 6x fewer predictor queries

2. Reduce latency metrics:

• 8-36% lower mean TTFT vs. best baseline

• 16-77% lower P99 TTFT

• Up to 5.6% lower mean E2E latency

• Up to 97% lower mean TTFT vs. round-robin

3. Reduce preemptions by 6x at overload (300 vs. 1,800 preemptions at QPS=32), avoiding KV-cache recomputation stalls

4. Handle configuration shifts automatically:

• +7.1% capacity gain with tight batching (batch size 48→24)

• +6.0% capacity gain on Qwen2-7B

• +4.7% capacity gain on BurstGPT dataset (oracle lengths)

5. Maintain performance under bursty arrivals: 6-9% lower mean E2E latency vs. Llumnix- across gamma-distributed arrival patterns (shape parameter 0.25 to 2.0)

6. Outperform migration-based rebalancing (Llumnix) on A100 clusters:

• Up to 2.6x higher throughput at saturation

• Orders-of-magnitude lower per-token latency (26ms vs. 3.2s at QPS=36)

• Avoids KV-cache transfer overhead and network contention

7. Keep control-plane overhead low: per-request scheduling overhead stays under 3% of E2E latency within capacity, with per-predictor CPU at 20% (vs. 56% for full fanout)

8. Scale to larger models: validated on Llama-2-70B with tensor parallelism (TP=4), achieving 16.6% latency-prediction error

In summary: The improved AI system can serve LLM requests across a cluster with near-optimal load balancing, avoiding the overhead of live migration while matching or exceeding its performance, and remaining robust to prediction errors and workload shifts.

This paper presents Astrolabe, a randomized prediction-guided scheduler for one-shot request dispatch in multi-instance large language model (LLM) serving. Astrolabe improves load balancing without relying on migration-based rebalancing, whose KV-cache transfers can introduce substantial overhead and network contention under high load. It combines response-length estimation, per-instance simulation-based latency prediction, and a power-of-two-choices dispatch policy to balance load while avoiding request herding. On the default Llama-2-7B/ShareGPT setup, Astrolabe matches the SLO capacity of the best load-aware baseline (31.6 versus 31.5 QPS), while reducing mean time-to-first-token (TTFT) by 8 to 36 percent, P99 TTFT by 16 to 77 percent, and mean end-to-end (E2E) latency by up to 5.6 percent, with approximately six times fewer preemptions once capacity is reached. Under configuration shifts, Astrolabe improves SLO capacity by up to 6 percent on Qwen2-7B and 7.1 percent under tight batching, reduces mean E2E latency by 6 to 9 percent under bursty arrivals, and achieves an approximately 2.8-fold reduction in per-predictor CPU usage relative to full fanout. With migration enabled on A100 GPUs, Astrolabe outperforms Llumnix by up to 2.6 times in throughput while achieving orders-of-magnitude lower per-token latency at saturation.