SageSched: Efficient LLM Scheduling Confronting Demand Uncertainty and Hybridity

SageSched is an efficient LLM scheduler that addresses demand uncertainty and hybrid resource requirements by integrating lightweight output-length prediction with a comprehensive cost model and an uncertainty-aware policy, achieving over 28.7% efficiency improvement in diverse testbed experiments.

The Gist

Imagine you run a very busy, high-end bakery called "The LLM Bakery." Customers (users) walk in and ask for custom cakes (LLM responses). Some orders are for a tiny cupcake, while others are for a massive, 10-tier wedding cake.

The problem is, you don't know how big the cake will be until the baker finishes baking it. You also don't know if the baker will need more oven space (compute) or more counter space (memory) for a specific order.

This is exactly the challenge the paper SageSched solves for Large Language Models (LLMs). Here is the story of how they fixed the chaos, explained simply.

The Problem: The "Guessing Game" Bakery

In the past, bakeries (LLM servers) handled orders in two bad ways:

1. The "First-Come, First-Served" Line: If a customer asks for a 10-tier cake, they get in line first. While they wait, 50 people asking for cupcakes are stuck behind them. Everyone waits forever. This is called Head-of-Line Blocking.
2. The "Shortest Job First" Guess: Some smart bakeries tried to guess the cake size. They'd say, "That looks like a small order, let's do it first!" But their guesses were often wrong because they relied on heavy, complicated training to guess the size. Plus, they only looked at the oven (compute) and forgot about the counter space (memory). If a "small" cake needed a huge counter, it would still clog up the kitchen.

The result? The bakery was slow, customers were angry, and the bakers were confused.

The Solution: Enter "SageSched"

The authors built a new manager for the bakery called SageSched. It uses three clever tricks to run the kitchen perfectly.

1. The "Memory Match" Trick (Predicting the Future)

Instead of trying to magically predict the future with a complex crystal ball (a heavy AI model), SageSched looks at the history.

• The Analogy: If a customer walks in and says, "I want a chocolate cake with sprinkles," the manager doesn't guess. They look at their logbook and say, "Ah, last Tuesday, three people asked for exactly that. They all got 500 grams of cake. Let's assume this one will be similar."
• Why it's better: It doesn't need to be retrained every time. It just matches the "flavor" (prompt) of the current request to past requests. It predicts a range of possibilities (e.g., "It's likely between 400 and 600 grams") rather than just one guess.

2. The "Full Kitchen" Cost (Computing the Real Price)

Old managers only counted how long the cake would take to bake (compute time). SageSched knows that baking isn't just about time; it's about space.

• The Analogy: Imagine two orders:
  • Order A: A tiny cake that takes 1 minute but requires a giant, custom mold that takes up the whole counter.
  • Order B: A huge cake that takes 10 minutes but fits on a small tray.
  • If the kitchen is running out of counter space, Order A is actually the "expensive" one because it blocks everyone else, even though it's fast.
• The Fix: SageSched calculates the "True Cost" by weighing both the baking time and the counter space needed. It knows when the oven is the bottleneck and when the counter is the bottleneck.

3. The "Fair Dice" Scheduler (Handling Uncertainty)

This is the magic sauce. Since the manager only has a probability (a range) of how big the cake will be, they can't just pick the "shortest" one. They need a strategy that handles the risk.

• The Analogy: Imagine you have two customers.
  • Customer X: You are 90% sure they want a tiny cupcake, but there's a 10% chance they want a giant cake.
  • Customer Y: You are 50% sure they want a medium cake, and 50% sure they want a huge one.
  • Who do you serve first?
• The Fix: SageSched uses a mathematical tool called the Gittins Index. Think of this as a "Fairness Dice." It doesn't just look at the average size; it calculates the best possible outcome if you serve them now. It prioritizes the order that is most likely to finish quickly and free up resources for everyone else. It constantly updates this "dice roll" as the cake is being baked, ensuring the line keeps moving efficiently.

The Result

When the researchers tested this new manager in a real kitchen (using powerful GPUs):

• Customers waited 28.7% less time for their final cake (Time-to-Last-Token).
• The system handled heavy traffic without crashing.
• It worked great even when the "guesses" weren't perfect, because it was designed to handle uncertainty.

In a Nutshell

SageSched is like a super-smart restaurant manager who:

1. Remembers what similar customers ordered in the past to guess the size.
2. Counts both the cooking time and the table space needed.
3. Uses math to decide who to serve next, ensuring that even if the future is uncertain, the whole restaurant runs as fast as possible.

It turns a chaotic, guessing game into a smooth, efficient operation.

Technical Summary

Here is a detailed technical summary of the paper "SageSched: Efficient LLM Scheduling Confronting Demand Uncertainty and Hybridity".

1. Problem Statement

The paper addresses the inefficiency of existing Large Language Model (LLM) schedulers when handling two unique characteristics of LLM inference workloads: Demand Uncertainty and Demand Hybridity.

• Demand Uncertainty: Unlike traditional workloads (e.g., OS processes or big data) with stable resource demands, LLM inference is auto-regressive. The output token length is non-deterministic and cannot be known until generation is complete. Existing schedulers often rely on single-point predictions (e.g., expected mean length) or simple heuristics (First-Come-First-Served), failing to capture the probabilistic nature of the workload.
• Demand Hybridity: LLM inference is both compute-intensive (matrix operations) and memory-intensive (heavy reliance on KVCache). Existing schedulers often treat service cost purely as computation time (output length), ignoring memory constraints. In memory-bound scenarios, prioritizing short outputs can be suboptimal if they consume excessive memory, causing contention.
• Current Limitations:
  • FCFS (e.g., vLLM, SGLang): Suffers from head-of-line blocking.
  • Prediction-based (e.g., SSJF, TRAIL): Rely on heavy-weight fine-tuned models that are expensive to train and often inaccurate. They typically predict a single value rather than a distribution, losing critical uncertainty information. They also fail to model memory costs accurately.

2. Methodology: SageSched

SageSched is a novel scheduler designed to handle uncertainty and hybridity through three core techniques:

A. Semantic-Aware History-Based Predictor

Instead of training heavy models to emulate generation, SageSched leverages the correlation between prompt similarity and output-length similarity.

• Mechanism: It maintains a history of recently served requests (prompt content + actual output length). For a new incoming request, it searches for historical requests with high semantic similarity (using cosine similarity of prompt embeddings).
• Output: Rather than predicting a single token count, it aggregates the output lengths of similar historical requests to generate a probability distribution of the output length.
• Advantages: This approach is training-free (lightweight), generalizable across different LLMs, and captures the inherent uncertainty of the workload.

B. Resource-Bound-Based Cost Modeling

SageSched moves beyond simple output-length metrics to model the true service cost considering both compute and memory.

• Analysis: The system analyzes whether the backend is compute-bound or memory-bound.
  • Memory-bound: Cost is driven by cumulative KVCache consumption.
  • Compute-bound: Cost is driven by accumulated computation time.
• Unified Model: The authors derive a unified cost formula that applies to both scenarios:
  $$C = \frac{O^2}{2} + IO$$
  Where $I$ is input length and $O$ is output length. This quadratic relationship reflects the fact that both memory usage and compute time grow non-linearly with sequence length during the auto-regressive decoding process. This model accurately captures the "true" cost regardless of the bottleneck type.

C. Uncertainty-Aware Scheduling Policy (Gittins Index)

To optimize the Time-to-Last-Token (TTLT) given cost distributions, SageSched employs the Gittins Index, a policy known to be optimal for multi-armed bandit problems with unknown durations but known distributions.

• Mechanism: Instead of sorting by the mean cost, the scheduler calculates the Gittins index for each request based on its cost distribution. Requests with lower indices (indicating a higher probability of early completion or lower amortized cost) are prioritized.
• Dynamic Refreshing: To handle runtime changes, the Gittins index is periodically refreshed at specific "bucket boundaries" (e.g., every 200 tokens) rather than at every step. This balances scheduling timeliness with system stability, avoiding thrashing while adapting to remaining service costs.

3. Key Contributions

1. Problem Identification: Empirically demonstrated that existing schedulers fail to handle the dual challenges of uncertainty (variable output length) and hybridity (compute/memory interplay) in LLM workloads.
2. Novel Architecture: Designed SageSched, integrating:
   • A lightweight, distribution-based predictor using semantic history.
   • A unified cost model ($O^2/2 + IO$) covering both compute and memory bounds.
   • A Gittins-index-based scheduling policy that optimizes for distribution-aware queuing.
3. Comprehensive Evaluation: Validated the system through testbed experiments on real hardware (A40, H800 GPUs) and large-scale simulations (up to 64 nodes).

4. Experimental Results

The authors evaluated SageSched against state-of-the-art baselines (FCFS, FastServe, SSJF, LTR, TRAIL) using diverse datasets (SharedGPT, Alpaca, Document-Write) and models (Llama3.1-8B, Qwen3-32B).

• Performance Gain: SageSched achieved an average TTLT improvement of over 28.7% compared to the best existing scheduler (TRAIL) under high-load conditions (8 requests/sec).
• Robustness: The system maintained superior performance across different datasets, particularly excelling on the Alpaca dataset where input lengths were long and existing length-based models failed.
• Component Analysis:
  • The Semantic-Aware History Predictor outperformed LLM-based predictors in both accuracy and latency (prediction time < 0.5ms vs. ~3.6ms).
  • The Resource-Bound Cost Model significantly outperformed simple output-length or weighted-sum models.
  • The Gittins Policy with dynamic refreshing proved superior to Mean-value sorting, even when prediction noise was introduced.
• Scalability: In simulations with up to 64 GPU nodes, the scheduling overhead remained negligible (linear growth, ~100ms per request at scale), confirming its viability for large clusters.

5. Significance

This work represents a significant step forward in LLM infrastructure optimization. By shifting from deterministic, single-value heuristics to probabilistic, distribution-aware scheduling, SageSched effectively mitigates the "unknown" nature of LLM generation. It provides a practical framework for balancing compute and memory constraints, which is critical as LLMs become more complex and memory-bound. The adoption of the Gittins index offers a theoretically grounded approach to minimizing latency in stochastic environments, setting a new standard for efficient LLM serving.