AI Scheduler Learns Offline to Meet Conflicting LLM Latency Goals

Large language model inference is a scheduling nightmare. Requests arrive in bursts, their output lengths are unknown until generation finishes, and each one carries its own service-level objectives for two different latency measures: time to first token and time per output token. A new study published in Scientific Reports introduces URSA, an uncertainty-aware scheduler trained entirely offline that raises overall SLO attainment to 90.4 percent across three real-world request traces, up from 71.3 percent under first-come-first-served and 84.1 percent under a learning-to-rank baseline.

The core problem with prevailing schedulers, the authors argue, is that they rely on static heuristics that cannot adapt to uncertainty. Online reinforcement learning, which could in principle learn adaptive policies, is unsafe and sample-inefficient in production, since experimenting on a live serving system risks violating latency guarantees for real users. URSA sidesteps this by training only on scheduling transitions generated by replaying historical request traces under existing heuristic schedulers, never touching a live system.

URSA's first innovation is probabilistic rather than point prediction. Instead of guessing a single output length for each incoming request, it predicts a calibrated distribution over possible lengths and the derived key-value cache footprint that the request will occupy on GPU memory. This distribution becomes the raw material for risk-sensitive decisions downstream.

The second component is a graph neural network that encodes the joint state of the system: the waiting queue, heterogeneous GPU instances, and current KV cache occupancy are represented as a request-instance heterogeneous graph. This lets the policy reason jointly about which request should go where, rather than treating queues and machines in isolation. Training uses conservative, distributional offline reinforcement learning.

The reward jointly optimizes SLO attainment, goodput, and the risk of KV cache overflow, discouraging actions that look good on average but occasionally cause out-of-memory events. At dispatch time, URSA reserves memory at a class-dependent high quantile of the predicted footprint, bounding how often memory runs out. The evaluation draws on three production-scale datasets: the BurstGPT trace, the Azure LLM inference trace, and an Alibaba GPU-cluster trace.

Beyond the headline SLO gains, URSA cuts the 99th-percentile time to first token by up to 3.3 times and pushes recompute and eviction events below 1 percent, approaching the performance of an oracle scheduler that knows the true output lengths in advance. The comparison with the oracle is telling. The authors also report analyses of robustness to prediction error and transfer to other scheduling domains, acknowledging that the approach depends on the quality of the predicted length distributions.

For operators running LLM services, the practical appeal is that the scheduler learns from logs of past traffic without any live experimentation, making deployment far less risky than online learning.

Subject of Research: Uncertainty-aware offline reinforcement learning for scheduling LLM inference under heterogeneous service-level objectives

Article Title: Uncertainty-aware offline reinforcement learning for LLM inference serving with heterogeneous SLOs

Article References: Feng, E., Wang, G., Li, P., Wang, Y., & Huang, S. (2026). Uncertainty-aware offline reinforcement learning for LLM inference serving with heterogeneous SLOs. Scientific Reports. https://doi.org/10.1038/s41598-026-76614-6

Image Credits: AI Generated

DOI: 10.1038/s41598-026-76614-6

Keywords: LLM inference, reinforcement learning, offline RL, scheduling, SLO attainment, TTFT, TPOT, KV cache, graph neural network, uncertainty quantification, GPU clusters, goodput