Justitia: Fair and Efficient Scheduling of Task-parallel LLM Agents with Selective Pampering

Conference: NeurIPS 2026
arXiv: 2510.17015
Area: LLM Safety
Keywords: LLM Agent Scheduling, Task Parallelism, Memory Bottleneck Modeling, Selective Pampering, Virtual Time Fair Queuing

TL;DR

To resolve the fundamental conflict between scheduling efficiency and fairness for task-parallel LLM agents on shared GPU clusters, Justitia introduces "selective pampering" based on virtual-time fair completion orders, combined with memory-centric quadratic KV token-time cost modeling and lightweight agent-specific MLP predictors, reducing average job completion time by 57.5% while mathematically guaranteeing a constant worst-case delay bound.

Background & Motivation

With the rapid deployment of large language models for complex real-world workflows, task-parallel LLM agents—which coordinate multiple concurrent inference requests to accomplish a composite objective—have become an essential workload paradigm. Prominent instances include MapReduce Summarization for large-scale documents, Document Merging and scoring, Fact Verification executing parallel claim check inferences, and Self-Consistency generating multi-branch reasoning paths for majority voting. In shared multi-tenant GPU clusters, a task-parallel agent only delivers its final value once all constituent inference tasks finish; hence, minimizing end-to-end Job Completion Time (JCT) is the foremost optimization goal, while robust multi-tenant fairness must also be preserved to prevent starvation and cross-tenant interference.

However, existing serving infrastructures fail to reconcile efficiency and fairness when handling task-parallel agents. Standard frameworks such as vLLM and SGLang rely on First-Come-First-Serve (FCFS) request scheduling, which triggers severe head-of-line blocking when long agent tasks arrive early. Conversely, efficiency-focused heuristic schedulers like Shortest-Job-First (SJF) or Multi-Level Feedback Queues (MLFQ) prioritize short tasks, accelerating overall throughput but inevitably starving large "elephant" agents. Meanwhile, tenant-fair schedulers such as VTC and FairServe enforce instantaneous fair sharing by strictly constraining each agent's concurrent footprint to an equal resource slice; this artificial throttling directly elongates the critical path of multi-task agents and degrades user-perceived JCT.

A critical insight underlying this tension is that users primarily demand long-term finish-time fairness (i.e., deterministic worst-case completion guarantees relative to a dedicated cluster slice) rather than instantaneous resource parity at every millisecond. When multiple agents compete, serving them sequentially in their idealized fair completion order with saturated resources allows prioritized agents to complete much faster and yield GPU memory earlier—dramatically reducing mean JCT without delaying any single agent beyond its fair-sharing deadline. Core idea: Propose selective pampering, which schedules task-parallel LLM agents in a saturated, non-preemptive manner prioritized by their Generalized Processor Sharing (GPS) virtual finish times, underpinned by memory-centric quadratic KV token-time cost modeling and lightweight agent-specific MLP demand predictors to achieve near-optimal efficiency with constant worst-case delay bounds.

Method

Overall Architecture

Justitia is architected as an end-to-end scheduling pipeline atop the vLLM serving framework, tailored to multi-task parallel agent workloads. When an agent comprising multiple concurrent inference requests arrives, Justitia first performs TF-IDF vectorization on the prompt and forwards it to an agent-specific 4-layer MLP to accurately predict its cumulative memory service cost. Next, the scheduler updates the global virtual time \(V(t)\) and derives a static virtual finish timestamp \(\bar{f}_j\) representing the agent's completion time under an idealized fair scheduler. All inference tasks belonging to the agent inherit this timestamp as their queue priority. Finally, the non-preemptive execution engine dispatches inference tasks to the vLLM instance in strict timestamp order, granting prioritized agents saturated access to KV cache blocks without suffering expensive cache swapping.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["User Submits Task-parallel Agent<br/>MRS / DM / FV / SC Workloads"] --> B["Agent-Specific MLP Demand Predictor<br/>TF-IDF Vectorization + 4-Layer MLP"]
    B --> C["Memory-Centric Cost Modeling<br/>Cumulative KV Token-Time pd + d^2/2"]
    C --> D["Virtual-Time Fair Queuing<br/>Compute Virtual Finish Time f_j = V(a_j) + C_j"]
    D --> E["Selective Pampering Dispatcher<br/>Saturated Non-preemptive Execution"]
    E --> F["Underlying vLLM Engine Execution<br/>Saturated Running Queue / Swapped Backup"]
    F --> G["All Sub-tasks Complete & Emit Final Output"]

Key Designs

1. Memory-Centric Cost Modeling: Capturing Quadratic KV Cache Spatiotemporal Footprints Prior LLM fair schedulers (e.g., VTC) gauge serving cost primarily through compute tokens (e.g., \(p + 2d\)). However, in production serving systems, GPU memory (specifically KV cache capacity) is the dominant throughput bottleneck: new sequences can be scheduled as long as KV space permits, and memory exhaustion immediately forces sequence preemption. Crucially, as auto-regressive decoding proceeds, sequence length expands dynamically over time. For a sub-task with prefill length \(p\) and decode length \(d\), its cumulative memory occupation (KV token-time) over all iterations is defined as: $\(c = \sum_{i=1}^{d} (p + i) = pd + \frac{d^2}{2}\)$ where the quadratic term \(d^2/2\) accurately captures the compounding inflation of decode length across both spatial memory allocation and temporal execution duration. The aggregate service cost \(C_j\) of agent \(j\) is the sum of KV token-times over all its constituent inference tasks, \(\sum c\). This formulations prevents underestimating memory-intensive long decoding sequences.

2. Lightweight Agent-Specific MLP Demand Prediction: Millisecond Inference with High Fidelity Selective pampering requires knowing the aggregate cost \(C_j\) at agent arrival. Existing prediction techniques (such as \(S^3\)) fine-tune a 66M-parameter Distillbert language model, incurring an unacceptable inference overhead of ~55 ms and failing to generalize across heterogeneous agent categories. Justitia observes that within a specific agent class (e.g., query generation in Fact Verification, chunk summarization in MapReduce), token length distributions exhibit strong statistical stability. Consequently, Justitia trains an independent, compact 4-layer Multi-Layer Perceptron (MLP) for each agent type. Runtime prompts are rapidly converted via TF-IDF into numerical frequency vectors, with the first layer scaled proportionally to average input size. This approach cuts prediction latency down to 2.16 ms and achieves convergence within 1 minute using only 100 profiling samples per agent.

3. Virtual-Time Fair Queuing and Constant Delay Bounds: Starvation-Free Saturated Scheduling To determine fair completion orders without simulating dynamic runtime arrivals or continuously refreshing residual allocations, Justitia draws inspiration from network packet scheduling under Generalized Processor Sharing (GPS). The system tracks a real-time virtual time \(V(t)\), advancing at a rate inversely proportional to the number of active agents \(N_t\) and proportional to total KV cache capacity \(M\): $\(V(0) = 0, \quad \frac{\mathrm{d}}{\mathrm{d}t} V(t) = \frac{M}{N_t}\)$ When agent \(j\) arrives at physical time \(a_j\) with predicted cost \(C_j\), its virtual finish time \(\bar{f}_j\) in GPS is calculated once and for all: $\(\bar{f}_j = V(a_j) + C_j\)$ Because all active tenants share resources at identical marginal rates, arrivals of subsequent agents do not alter the relative order of existing \(\bar{f}_j\) timestamps. Justitia assigns \(\bar{f}_j\) as the priority for all inference requests of agent \(j\), allowing top-priority agents to monopolize available KV blocks until completion. Crucially, the authors prove (Theorem B.1) that the actual completion time \(f_j\) under Justitia relative to the idealized GPS completion time \(\bar{f}_j\) is strictly bounded by a constant: $\(f_j - \bar{f}_j \le 2 c_{\max} + \frac{C_{\max}}{M}\)$ where \(c_{\max}\) is the maximum cost of any single inference and \(C_{\max}\) is the maximum total cost of any agent. This mathematical guarantee prevents short tasks from starving large ones, eliminating the catastrophic tail latencies typical of heuristic SJF schemes.

A Worked Example

Consider a GPU server with total KV cache capacity \(M = 100\). At time \(t=0\), the system is idle when an elephant agent \(A\) arrives (with 10 sub-inferences, total cost \(C_A = 200\), and max sub-task cost \(c_{\max} = 20\)). With \(N_0 = 1\), the virtual time rate is \(\mathrm{d}V/\mathrm{d}t = 100/1 = 100\). The virtual finish timestamp is \(\bar{f}_A = V(0) + 200 = 200\). The scheduler allocates all 100 KV blocks to saturate the first batch of agent \(A\)'s sub-tasks.

At time \(t=1\), agent \(A\) has processed 100 units of memory-time when two mice agents, \(B\) and \(C\), arrive with costs \(C_B = 30\) and \(C_C = 40\). At \(t=1\), the system virtual time has progressed to \(V(1) = 100\). The virtual finish timestamps for the newcomers are \(\bar{f}_B = 100 + 30 = 130\) and \(\bar{f}_C = 100 + 40 = 140\). Because \(\bar{f}_B < \bar{f}_C < \bar{f}_A\) (\(130 < 140 < 200\)), once agent \(A\)'s running sub-tasks complete non-preemptively, the engine prioritizes agents \(B\) and \(C\) rather than continuing with \(A\). Agents \(B\) and \(C\) complete in rapid succession under saturated resources, after which full capacity returns to agent \(A\). Both \(B\) and \(C\) experience near-instantaneous completion, while agent \(A\) is guaranteed to finish within its theoretical constant bound \(2 c_{\max} + C_{\max}/M\).

Loss & Training

The 4-layer MLP predictor for each agent category is trained offline using Mean Squared Error (MSE) with \(L_2\) weight regularization: $\(\mathcal{L}(\theta) = \frac{1}{K} \sum_{k=1}^{K} \left( \hat{C}_k(\theta) - C_k \right)^2 + \gamma \|\theta\|_2^2\)$ where \(K=100\) profiling runs are collected per agent type, \(\hat{C}_k\) is the predicted cost, \(C_k\) is the profiled ground truth computed at temperature 0 to ensure determinism, and \(\gamma\) is the regularization coefficient. Training converges in approximately 1 minute on a single NVIDIA V100 GPU.

Key Experimental Results

Main Results

Justitia was evaluated on an empirical benchmark of 300 diverse LLM agents drawn from 9 representative workflows (MapReduce Summarization, Document Merging, Code Checking, Knowledge QA Verification, Equation Verification, Fact Verification, ALFWorld, etc., sampled at 72% small, 26% medium, and 2% large). Hardware platforms spanned LLaMA-7B on an A100-40GB GPU, LLaMA-13B across 4× V100-16GB GPUs via tensor parallelism, and Qwen2.5-32B on an H800-80GB GPU. Workload arrival traces mirrored Alibaba Mooncake production cluster logs across 1×, 2×, and 3× submission densities.

The table below summarizes the quantitative comparisons on the LLaMA-7B testbed across mainstream scheduling paradigms:

Scheduler	Granularity & Strategy	Mean JCT (s)	Improvement vs. VTC (%)	Delay at 92% Agents	Starvation Vulnerability	Scheduling Overhead
vLLM (FCFS)	Request-level FCFS	388.2	-	Severe delay (>80%)	None	Minimal (<1 ms)
vLLM-SJF	Request-level Distillbert SJF	215.4	39.5%	Moderate delay (>50%)	Present	High (~55 ms)
Parrot	Agent-level FCFS	412.5	-6.7%	Extreme HoL blocking	None	Minimal (<1 ms)
VTC	Agent-level instantaneous fair share	386.4	Baseline (0.0%)	1.0× (reference)	None	Low (~5 ms)
SRJF	Agent-level Shortest Remaining Job First	148.6	61.5%	High tail latency	Severe (infinite starvation)	Low (~3 ms)
Justitia (Ours)	Agent-level virtual-time selective pampering	164.2	57.5%	Only 8% delayed (max 26.0%)	None (mathematically bounded)	Minimal (<10 ms)

Ablation Study

To isolate the benefits of memory-centric cost modeling and lightweight MLP prediction, the authors conducted two targeted ablations: 1. Predictor Ablation: Comparing the agent-specific MLP against \(S^3\)'s Distillbert predictor on a V100 GPU under 2× workload density. 2. Cost Model Ablation: Replacing the quadratic KV token-time model with VTC's compute-token formula \(p + 2d\) (denoted Justitia/C).

Metric / Variant	Average Relative Error (%)	Inference Overhead (ms)	Average JCT (s)	Offline Training Time	Performance Degradation
Distillbert (\(S^3\) Baseline)	452.0%	55.70 ms	366.7 s	~2 hours	Baseline comparison
Justitia Dedicated MLP	53.0%	2.16 ms	151.1 s	~1 minute	Superior speed & accuracy
Justitia Full (KV Token-Time)	-	-	164.2 s	-	Optimal performance
Justitia/C (Compute-Token Cost)	-	-	233.7 s	-	Degrades by 42.3%

Key Findings

• Memory-Centric Modeling is Indispensable: Reverting from quadratic KV token-time to linear compute tokens caused average JCT to deteriorate by 42.3%, demonstrating that memory allocation over time dominates actual GPU throughput.
• Selective Pampering Preserves Finish-Time Fairness: Under 3× arrival density, 92% of agents finished earlier than or equal to their completion time under VTC. For the 8% of delayed agents, their average delay was below 10%, with a worst-case delay capped at 26.0%.
• Robust Resistance to Starvation: In a stress test where short mice agents arrived continuously at 1 req/s alongside a large MapReduce Summarization elephant agent, SRJF stalled the elephant indefinitely. In contrast, Justitia's elephant completion time remained strictly bounded by the theoretical constant delay ceiling.
• Resilience Against Prediction Errors: When artificial prediction noise up to factor \(\lambda = 3\) (random scaling in \([1/3, 3]\)) was injected, Justitia's average JCT degraded by only 9.5%, underscoring exceptional error tolerance.

Highlights & Insights

• Conceptual Shift to Finish-Time Fairness: Shifting from rigid millisecond-level fair sharing to finish-time fairness underpins the insight that sequential saturated execution accelerates all competing tenants without penalizing worst-case delay.
• Quadratic KV Token-Time Formalization: The \(pd + d^2/2\) formulation accurately captures the escalating memory pressure of auto-regressive decoding, providing a principled metric for LLM system schedulers.
• Pragmatic Lightweight Prediction: By leveraging the behavioral consistency of structured business agents, microsecond TF-IDF vectorization paired with a shallow MLP sidesteps the heavy runtime overhead of language model predictors.

Limitations & Future Work

• Prompt Size Assumptions: The mathematical delay bound assumes individual prompt sizes are modest relative to total KV capacity \(M\). Under extreme long-context workloads (e.g., 128k+ tokens), memory fragmentation may challenge this assumption.
• Cold-Start for Unseen Agents: The system assumes agent workloads belong to registered categories with pre-profiled MLPs; completely novel, unstructured agent workflows may require online adaptation.
• Prefix Cache Agnosticism: Multi-task agents often share system instructions or context prefixes. Extending the virtual-time metric to account for prompt-cache hits across sub-tasks represents a promising avenue for further optimization.

Related Work & Insights

• vs. VTC (OSDI 2024): VTC enforces instantaneous token fair sharing, leaving multi-task agents partially throttled; Justitia reinterprets fairness via virtual finish times, allowing saturated execution that improves mean JCT by 57.5%.
• vs. Parrot (OSDI 2024) / Ayo (ASPLOS 2025): Parrot and Ayo optimize agent graphs but rely on FCFS scheduling at runtime, suffering head-of-line blocking; Justitia achieves dynamic, starvation-free fair queuing.
• vs. FastServe / S³: FastServe incurs preemptive swapping overheads, while \(S^3\) requires costly LLM-based sequence length predictors; Justitia achieves sub-3 ms prediction overhead with non-preemptive saturated scheduling.

Rating

• Novelty: ⭐⭐⭐⭐⭐ [Bridges network GPS fair queuing and selective pampering to solve the agent-level scheduling bottleneck with sound theoretical proofs]
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ [Evaluated on 300 diverse agents, 3 model backends, real Mooncake production traces, comprehensive ablations, and microbenchmarks]
• Writing Quality: ⭐⭐⭐⭐⭐ [Clear motivation, rigorous mathematical derivations, and clean visual mapping between architecture and mechanisms]
• Value: ⭐⭐⭐⭐⭐ [Provides actionable, highly impactful scheduling guidance for production multi-agent serving clusters]

Related Papers

• [NeurIPS 2026] InTAct: Interval-based Task Activation Consolidation for Continual Learning
• [ICLR 2026] DPQuant: Efficient and Private Model Training via Dynamic Quantization Scheduling
• [NeurIPS 2026] Evaluating and Understanding Scheming Propensity in LLM Agents
• [NeurIPS 2026] ASPI: Seeking Ambiguity Clarification Amplifies Prompt Injection Vulnerability in LLM Agents
• [NeurIPS 2026] POLAR-Bench: A Diagnostic Benchmark for Privacy-Utility Trade-offs in LLM Agents