TIE: Optimization of LLM Inference Scheduling Based on Uncertainty-Aware Output Length Prediction

TIE is an open-source project from an ICML 2026 paper. Addressing the problem of GPU idling caused by output length variations in LLM batch inference, it proposes an uncertainty-aware output length prediction method to optimize inference scheduling, effectively reducing GPU idle waiting time and improving inference throughput. The project is implemented based on the vLLM framework, with the open-source address at https://github.com/Hyzheng-code/TIE.

In the batch processing scenario of LLM inference services, the output lengths of different requests vary significantly (from dozens to thousands of tokens). Traditional scheduling strategies assume similar output lengths or process requests sequentially, leading to short requests waiting for long ones to complete, GPU computing resources being idle, and user waiting times being prolonged.

The core of TIE is uncertainty-aware output length prediction: it not only predicts the expected value of output length but also its uncertainty distribution (log-normal distribution). The technical architecture is based on the vLLM framework and includes four main components:

1. TIE Predictor (ua_predictor.py): Uses a DeBERTa encoder to predict the log mean and standard deviation of output length;
2. Score Calculator (ua_score_calculator.py): Computes scheduling scores based on predicted length and uncertainty, downweighting requests with high uncertainty;
3. Request Queue (request_queue.py): Prioritizes scheduling requests with deterministic predictions and efficient GPU utilization;
4. Scheduler (scheduler.py): Extends the core scheduling logic of vLLM v1 and integrates the above components.

• Implementation: Written in Python, training code is in the train directory, and inference is integrated into the vllm/v1/core/sched directory; requires configuration of predictor checkpoints, pre-trained encoders, and the path to training data CSV (including three columns: prompt, logt_mu, logt_sigma).
• Deployment: Started via start-server.sh, specifying scheduling strategy, GPU devices, and model path; when using the ua strategy, reserve one GPU to run the predictor, and use the rest for tensor parallel inference to avoid prediction overhead affecting performance.

• Academic: TIE was accepted by ICML 2026, introducing a machine learning-driven prediction mechanism to promote the evolution of LLM inference systems toward intelligent scheduling.
• Practical: Output length prediction can be applied to scenarios such as dynamic batch size adjustment, KV cache pre-allocation, user waiting time estimation, and resource quota planning.

• Usage Recommendations: Suitable for deployment by teams with sufficient historical request data; need to evaluate the additional computational overhead of the predictor to ensure that benefits outweigh costs.
• Summary: TIE introduces uncertainty quantification into LLM inference scheduling to solve resource wastage caused by output length variations, and it is easy to integrate based on the vLLM framework; as LLM applications expand, research on optimizing such system bottlenecks will become more important.