Scheduling LLM Inference with Uncertainty-Aware Output Length Predictions

TL;DR

This paper introduces a distribution-based approach to predict output lengths in LLM inference scheduling, improving latency and throughput by accounting for uncertainty in output length predictions.

Contribution

It proposes a novel distribution fitting method and a metric called TIE for better scheduling of LLM inference based on output length uncertainty.

Findings

TIE reduces per-token latency by 2.31 times in online inference.

TIE improves throughput by 1.42 times in offline data generation.

Output length follows a heavy-tailed distribution, well modeled by log-t distribution.

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.