TLDR: Justitia is a novel scheduler for Large Language Model (LLM) applications designed to achieve both high efficiency and strong fairness guarantees. It operates by serving applications in a ‘saturated manner’ based on their fair completion order. Key techniques include a memory-centric cost model that focuses on GPU KV cache, a lightweight MLP-based neural network for accurate application demand prediction, and a virtual-time based fair queuing algorithm to efficiently determine scheduling priorities and prevent starvation. Implemented on vLLM, Justitia significantly reduces average application completion times (e.g., 57.5% better than VTC) while ensuring that applications are not unfairly delayed.
In the rapidly evolving landscape of Large Language Models (LLMs), a new challenge has emerged: efficiently and fairly managing multiple LLM applications running on shared computing resources, particularly GPUs. These LLM applications, which involve a series of correlated LLM inferences to solve complex problems like summarizing large files or intricate mathematical tasks, demand both fast completion times and guaranteed performance without being unfairly delayed by other users.
Traditional LLM schedulers, such as vLLM, often employ a First-Come-First-Serve (FCFS) approach. While simple, this can lead to a problem known as head-of-line blocking, where a long-running task can hold up many shorter, waiting tasks. Other schedulers, like VTC, attempt to ensure fairness by allocating equal resources instantly, but this can sometimes slow down the overall completion of applications. The core dilemma is achieving a balance between overall system efficiency (short average completion times) and fairness (ensuring no application is unduly delayed).
Introducing Justitia: A New Approach to LLM Scheduling
To address these challenges, researchers have developed Justitia, a novel scheduler designed specifically for LLM applications. Justitia’s core insight is to serve LLM applications in a “saturated manner following the fair completion order.” This means instead of trying to share resources equally at every instant, Justitia prioritizes applications one by one based on when they would ideally finish under a perfectly fair system, allowing the prioritized application to use all available resources. This approach can significantly reduce average completion times while still ensuring fairness over the long term.
How Justitia Works: Three Key Techniques
Justitia incorporates three key techniques to achieve its goals:
1. Memory-Centric Cost Modeling: For LLM inference, GPU memory, particularly the KV cache (used to store intermediate feature states), is often the primary bottleneck. Unlike previous methods that focused on computation, Justitia models the service cost of an LLM application in a memory-centric way. It quantifies this cost as “KV token-time,” which considers both the space and duration of memory occupation. This model recognizes that the relationship between cost and output length is quadratic, not linear, providing a more accurate representation of resource consumption.
2. MLP-based Demand Prediction: To determine the fair completion order, Justitia needs to predict an application’s total service cost upon its arrival. It uses a lightweight and accurate Multi-Layer Perceptron (MLP) neural network model for this prediction. Each type of LLM application has its own dedicated MLP model, which is trained on historical data and uses TF-IDF (Term Frequency-Inverse Document Frequency) to vectorize the input prompt. This application-specific approach ensures higher accuracy and minimal runtime overhead compared to more complex, general-purpose prediction models.
3. Application-level Fair Queuing: Justitia borrows a concept from network scheduling: virtual-time based fair queuing. This algorithm efficiently calculates the expected completion order of LLM applications under an idealized fair scheduler, known as Generalized Processor Sharing (GPS). By using a “virtual finish time” that remains constant once calculated, Justitia can determine scheduling priorities in a single step, without needing continuous updates or knowledge of future arrivals. This method also comes with a theoretical guarantee that an application’s maximum delay compared to an ideal fair scheduler is bounded by a constant, effectively preventing starvation for larger applications.
Also Read:
- Managing Large Language Model Traffic: Introducing beLLMan for Efficient AI Serving
- TokenTiming: Accelerating LLM Inference with Universal Speculative Decoding
Performance and Fairness
Justitia has been implemented on top of vLLM, a popular LLM serving framework. Experimental results with a diverse set of LLM applications demonstrate its effectiveness. Justitia significantly reduces the average application completion time, showing a 57.5% improvement over state-of-the-art fair schedulers like VTC and 61.1% over Parrot. Crucially, it achieves this while preserving fairness; 92% of applications complete no later than they would under VTC, with a worst-case delay of only 26.0%.
Further studies confirmed the importance of Justitia’s memory-centric cost modeling and MLP-based demand prediction. The scheduler also maintains a negligible overhead, even under high arrival rates, making it suitable for real-time deployment.
By combining efficiency with strong fairness guarantees, Justitia offers a promising solution for managing the complex and resource-intensive workloads of modern LLM applications. For more in-depth technical details, you can refer to the research paper.
