Scheduling LLM Inference with Uncertainty-Aware Output Length Predictions
Haoyu Zheng ⋅ Yongqiang Zhang ⋅ Fangcheng Fu ⋅ Xiaokai Zhou ⋅ Hao Luo ⋅ Hongchao Zhu ⋅ Yuanyuan Zhu ⋅ Hao Wang ⋅ Xiao Yan ⋅ Jiawei Jiang
Abstract
To schedule LLM inference, the \textit{shortest job first} (SJF) principle is favorable by prioritizing requests with short output lengths to avoid head-of-line (HOL) blocking. Existing methods usually predict a single output length for each request to facilitate scheduling. We argue that such a \textit{point estimate} does not match the \textit{stochastic} decoding process of LLM inference, where output length is \textit{uncertain} by nature and determined by when the end-of-sequence (EOS) token is sampled. Hence, the output length of each request should be fitted with a distribution rather than a single value. With an in-depth analysis of empirical data and the stochastic decoding process, we observe that output length follows a heavy-tailed distribution and can be fitted with the log-t distribution. On this basis, we propose a simple metric called Tail Inflated Expectation (TIE) to replace the output length in SJF scheduling, which adjusts the expectation of a log-t distribution with its tail probabilities to account for the risk that a request generates long outputs. To evaluate our TIE scheduler, we compare it with three strong baselines, and the results show that TIE reduces the per-token latency by $2.31\times$ for online inference and improves throughput by $1.42\times$ for offline data generation.
Lay Summary
**(Problem)** Modern AI chatbots serve billions of requests a day. A server can only process so many at once, so the order matters: if a request that generates a very long answer is served first, all the shorter requests behind it must wait, increasing everyone's response time. Existing systems try to avoid this by guessing how long each reply will be and serving the short ones first. **(Solution)** We argue that these guesses are fundamentally limited because language models write text by drawing words randomly one at a time — the same question can produce very different replies on different runs. So instead of guessing a single number, we predict a range of likely lengths together with a risk score that flags requests likely to turn out unexpectedly long. The scheduler combines these two signals, and automatically becomes more cautious when the server is under heavy load. **(Impact)** The per-token latency drops by up to $2.31\times$ for online chatbots, and bulk data-generation jobs finish $1.42\times$ more work in the same time, with no change to the underlying AI model.
Successful Page Load