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

Abstract Large language model (LLM) inference workloads are bursty, have output lengths that are unknown at arrival, expose two resource-distinct phases (prefill and decode), and serve requests carrying heterogeneous service-level objectives (SLOs) on time-to-first-token (TTFT) and time-per-output-token (TPOT). Prevailing schedulers rely on static heuristics that cannot adapt to this uncertainty, while online reinforcement learning (RL) is unsafe and sample-inefficient in production. We present URSA, an uncertainty-aware, risk-sensitive scheduler trained entirely offline , without interacting with a live serving system, on scheduling transitions generated by replaying historical request traces under existing heuristic schedulers. URSA (i) predicts a calibrated distribution over each request’s output length and derived key–value (KV) cache footprint rather than a point estimate; (ii) encodes the joint state of the waiting queue, heterogeneous GPU instances, and KV occupancy as a request–instance heterogeneous graph processed by a graph neural network; (iii) learns a dispatch policy with conservative, distributional offline RL whose reward jointly optimizes SLO attainment, goodput, and KV-overflow risk; and (iv) reserves memory at a class-dependent high quantile of the predicted footprint to bound out-of-memory events. Across the BurstGPT, Azure LLM inference, and Alibaba GPU-cluster traces, URSA raises overall SLO attainment from $$71.3\%$$ (FCFS) and $$84.1\%$$ (learning-to-rank) to $$90.4\%$$ , cuts P99 TTFT by up to $$3.3\times $$ , and reduces recompute/eviction events below $$1\%$$ , approaching an oracle that knows true lengths. We analyze limitations, robustness to prediction error, and transfer to other scheduling domains.