DON’T STOP ME NOW: EMBEDDING BASED SCHEDULING FOR LLMS
Authors: Rana Shahout
eran malach
Chunwei Liu
Weifan Jiang
Minlan Yu
Michael Mitzenmacher
Venue: ICLR 2025 Poster (2025)
Keywords: LLM serving
scheduling
algorithms with predictions
Primary area: optimization
TL;DR: TRAIL improves LLM inference efficiency by using LLM embeddings for accurate output length predictions and a memory-aware SRPT scheduling policy with limited preemptions, reducing response time while managing GPU memory overhead.

Abstract: Efficient scheduling is crucial for interactive Large Language Model (LLM) applications, where low request completion time directly impacts user engagement. Size-based scheduling algorithms like Shortest Remaining Process Time (SRPT) aim to reduce average request completion time by leveraging known or estimated request sizes and allowing preemption by incoming jobs with shorter service times. However, two main challenges arise when applying size-based scheduling to LLM systems. First, accurately predicting output lengths from prompts is challenging and often resource-intensive, making it impractical for many systems. As a result, the state-of-the-art LLM systems default to first-come, first-served scheduling, which can lead to head-of-line blocking and reduced system efficiency. Second, preemption introduces extra memory overhead to LLM systems as they must maintain intermediate states for unfinished (preempted) requests.
In this paper, we propose TRAIL, a method to obtain output predictions from the target LLM itself. After generating each output token, we recycle the embedding of its internal structure as input for a lightweight classifier that predicts the remaining length for each running request. Using these predictions, we propose a prediction-based SRPT variant with limited preemption designed to account for memory overhead in LLM systems. This variant allows preemption early in request execution when memory consumption is low but restricts preemption as requests approach completion to optimize resource utilization. On the theoretical side, we derive a closed-form formula for this SRPT variant in an M/G/1 queue model, which demonstrates its potential value. In our system, we implement this preemption policy alongside our embedding-based prediction method. Our refined predictions from layer embeddings achieve 2.66x lower mean absolute error compared to BERT predictions from sequence prompts. TRAIL achieves 1.66x to 2.01x lower mean latency on the Alpaca dataset and 1.76x to 24.07x lower mean time to the first token compared to the state-of-the-art serving system.

Published as a conference paper at ICLR 2025
DON'T STOP ME Now: EMBEDDING BASED SCHEDULING FOR LLMS
Rana Shahout Harvard University
Eran Malach
Chunwei Liu MIT
Weifan Jiang
Harvard University
Harvard University
Minlan Yu Harvard University
Michael Mitzenmacher Harvard University
ABSTRACT
Efficient scheduling is crucial for interactive Large Language Model (LLM) ap- plications, where low request completion time directly impacts user engagement. Size-based scheduling algorithms like Shortest Remaining Process Time (SRPT) aim to reduce average request completion time by leveraging known or estimated request sizes and allowing preemption by incoming requests with shorter service times. However, two main challenges arise when applying size-based schedul- ing LLM systems. First, accurately predicting output lengths from prompts is challenging and often resource-intensive, making it impractical for many systems. As result, the state-of-the-art LLM systems default to first-come, first-served scheduling, which can lead to head-of-line blocking and reduced system effi- ciency. Second, preemption introduces extra memory overhead LLM systems as they must maintain intermediate states for unfinished (preempted) requests. In this paper, we propose TRAIL, novel approach that improves the response time in LLM inference through two key contributions: (1) obtain output predic- tions from the target LLM itself. After generating each output token, we recycle the embedding of its internal structure as input for a lightweight classifier that pre- dicts the remaining length for each rumning request. (2) prediction based SRPT variant with limited preemption designed to account for memory overhead in L.LM systems. This variant allows preemption early in request execution when memory consumption is low but restricts preemption as requests approach completion to optimize resource utilization. On the theoretical side, derive closed-form formula for this SRPT variant in an M/G/1 queue model, which demonstrates its potential value. In our system, we implement this preemption policy alongside our embedding based prediction method. Our refined predictions from layer em- beddings achieve 2.66x lower mean absolute error compared to BERT predictions from sequence prompts. TR ATL. achieves 1.66x to 2.01x lower mean latency the Alpaca dataset and 1.76x to 24.07x lower mean time to the first token compared to the state-of-the--e serving system.
INTRODUCTION
Recent advances in large language models (LL.Ms) have sparked new generation of interactive AI applications. OpenAl's ChatGPT (OpenAl 2022) exemplifies this trend, facilitating conversational interaction across diverse tasks. The interactive nature of these applications requires low request completion time to ensure user engagement. Users expect near-instant responses, making efficient inference serving critical for LLM-based interactive AI applications.
LLM inference requests exhibit a distinct autoregressive pattern comprising multiple iterations in which each output token is appended the input to generate the next token. Orca (Yu et al.) 2022) and vLLM (Kwon et al, 2023) are state-of-the-an solutions for LLM inference systems. These systems ieeattion-level scheduling, which allows requests be added or completed requests

---

Published as a conference paper at ICLR 2025
to be removed at the end of each iteration, providing greater flexibility in managing the processing batch than request-level scheduling. However, these works still process requests using a first-come- first-served (FCFS) policy, leading to head-of-line blocking (Kaffes et al., 2019) under high load. This problem is particularly severe for LLM inference, where many short requests may wait for other long requests to finish.
Size-based scheduling aims to reduce request completion time by leveraging known or estimated request sizes. For LLM requests, the execution time depends on both input and output lengths, with the latter being unknown and potentially variable. Preemption allows for dynamic scheduling: when one scheduled request finishes generating an output token, we can decide whether continue this request or preempt it with another (newly arrived) request. The Shortest Remaining Process Time (SRPT) policy is an example of preemption scheduling. However, preemption here introduces ex- tra memory overhead to maintain an intermediate state for started but unfinished requests. L.MMs maintain a key-value (KV) cache for each Transformer layer to store the intermediate state. In pre- emption scheduling policies, the cache must keep the intermediate state for all started but unfinished requests in the waiting queue, which may cause memory overflow, given the large size of LLMs and the limited memory capacity of GPUs. Non-preemptive policies, e.g. FCFS, do not face this issue, as there are no unfinished requests in the waiting queue.)
Recent works propose methods to predict request sizes. 53 (Jin et al. 2023) fine-tunes BERT model (Sanh, 2019) to predict output sequence lengths from input prompts. While this model is lightweight in resources, its accuracy diminishes for requests with varying execution times. Zheng et al. (Zheng et al., 2024) employ a separate, lighter LLM to predict output lengths before schedul- ing, achieving higher precision than the BERT model. However, this approach consumes more resources, raising questions about its efficiency and cost-effectivenes in practical applications. An- other method proposed by the same work Perception in Advance, where models are asked to predict the length of the responses they are about to generate. However, this approach is limited by potential dependencies between the prediction and the generated output, which may affect the quality of the output.
In this paper, we introduce TRAIL (Token Response and Embedding Adaptive Inference Layer), a novel approach that improves the response time in LLM inference through two key contributions: (I) iteration-level predictions for request lengths with low overhead and high accuracy to enable size-based scheduling, and (2) preemption scheduling at token granularity that accounts for memory overhead. For predictions, our approach takes advantage of the autoregressive nature of LLM output generation to predict the output size and facilitate the prediction of request length the granularity of individual tokens. We begin with an initial prediction based on the prompt and iteratively refine these predictions as the sequence progresses by innovatively using LLM layer embeddings. Specifically, we "recycle" the embedding from intermediate Transformer layers of the LLM, feeding them into a lightweight linear classifier to predict the remaining sequence length. This method, which involves analyzing the outputs of the internal model layer by using them as input for a separate machine leamingg predictor, is commonly referred to as probing (Belinkov) [2022] Hewitt & Liang/ 2019 Hewitt& & Manning, 2019). This approach combines the advantages of LLM-based prediction with computational efficiency, eliminating the need for an additional L.M dedicated solely to length predictions. To address the memory constraints for preemption scheduling, TRAII. limits the number of times each request can be preempted and introduces SRPT with limited preemptions. Early in the request's 's execution, preemption is allowed since its KV cache memory usage is small, making preemption less costly in terms of memory. However, as the request progresses and its memory consumption in the KV cache grows, preemption becomes more expensive. To manage this, turn the preemption off in the later stages of the request to avoid heavy memory overhead.
Our contributions in this paper are (1) an output length prediction framework using recycled LLM embeddings, with dynamic refinement at each iteration. In our experiments, refined predictions from layer embeddings achieve 2.66x lower mean absolute error compared BERT predictions from sequence prompts (2) a proposed variant of Shortest Remaining Processing Time using these predictions with limited preemptions (3) derivation of closed-form formula for this SRPT vari- ant in an M/G/1 queue model. We also show via experiments that our prediction and scheduling approaches result in mitigating head-of-line blocking. When TRAIL is integrated with a ofate-of- the-art LLM system, it has I. 1,66x to 2.01x lower mean latency and 1.76x to 24.07x lower mean time to first token compared to the state-of-the-an serving system tested with the Alpaca dataset. Our
2

---

Published as a conference paper at ICLR 2025
Sthnckling
LLV
ta HillOioii haggl 2 1902
Lageen
Figure 1: TRAIL architecture. The system (1) initially orders requests using BERT model, (2) schedules requests using modified SPRPT with limited preemption, and (3) refines predictions during token generation using embeddings from the LLM's internal layers. At every iteration, steps 2 and 3 are repeated (represented as red dashed lines), which allows preemption at iteration-level granularity and refined predictions. We focus on identifying the LLM layer that best predicts output length rather than using multi-layer embeddings (i = 11).
TRAIL approach has the potential to significantly enhance LLM inference efficiency and reliability, paving the way for more adaptive L.M serving systems.
BACKGROUND AND MOTIVATION
Transformer-Based Generative Models and Key-Value Cache. At each step, Transformer model predicts the next most likely token based on the sequence of tokens generated so far. A generative model of length n must perform n iterations, with each token passing through multiple transformer layers composed of self-attention and feed-forward networks.
In typical iteration the i-th step, the model computes over all previously generated tokens (to, f1........ f4-1) via self-attention. This can be expressed as:
2 K Vdn
hour softmax
V
Where q is the query vector representing the hidden state of the current token t,. and K, V Rixd, are matrices representing the keys and values derived from all previous tokens.
To enhance efficiency, LLMs cache the key and value matrices (KV cache) throughout the sequence generation process, eliminating the need to recompute them at every iteration. This caching mech- anism significantly reduces computation time but requires substantial memory proportional to number of layers and hidden dimensions. As more tokens generated, the cache stores informa- tion from all previously generated tokens, and its memory demand grows linearly with the length of the sequence. This makes the memory consumption of the KV cache considerable, especially for long sequences. For example, in the GPT-3 175B model, a single request with a sequence length of 512 tokens requires least 2.3 GB of memory to store the key-value pairs. The limited capacity of GPU memory restricts the size of the KV cache and poses a challenge to efficiently implementing preemptive scheduling policies.
Iteration-Level Scheduling. Iteration-lever scheeuuuuingn differs from the conventional request-level scheduling approach. In request-level scheduling, the system processes a batch of requests to com- pletion, forcing requests that finish earlier to wait until the entire batch completes, while new re- quests must remain in a queue until the next batch begins. contrast, iteration-level scheduling (Yu) et al., 2022) offers flexibility by processing only single iteration (i.e., generating one output token) for each request in the batch. After each iteration, the scheduler is called, and the completed requests exit the batch while newly arrived requests can be added, allowing dynamic batch adjustments. However, the batch size remains constrained by the GPU memory capacity.
3 METHOD
TRAIL leverages the autoregressive nature of LLM inference and iteration-level scheduling, em- ploying primary strategies to reduce L.LM inference response time: iteration-level prediction
3

---

Published as a conference paper at ICLR 2025
Prediction w/ per-layer embeddings
150
MAE
50
0
0
10
15
20
25
30
Layer index
Figure 2: MAE for length prediction using embeddings vs. layer (1,000 prompts).
and preemption that accounts for memory overhead. TRAIL uses low-overhead, high-accurage pre- dictions by utilizing the intermediate layer embeddings of the LLM itself to predict output length while implementing loken-granulat limited preemption scheduling account for the memory over- head associated with preemption. Figure iillustrates the TRAIL architecture. Users submit requests to the request pool. TRAIL initially orders requests using BERT model based input prompts (Step 1). This method relies exclusively on the prompt using BERT suggested in (Jin et al. 2023). Given the predictions of the output lengths, the scheduler implements a variant of the Shortest Pre- dicted Remaining Process Time (SPRPT) with limited preemption (Step 2, Section 3.3). Although SPRPT traditionally designed for single-server, sequential scheduling, we adapt handle multi- ple concurrent requests while still prioritizing those with the shortest predicted remaining time. The number of requests that can be scheduled simultaneously limited by the available GPU memory. As the LLM generates each output token, we refine predictions using embeddings from the L.M''s intermediate layers. linear classifier processes these embeddings and informs the scheduler, which adjusts based on updated predictions (Section 3.1
3.1 REFINED OUTPUT LENGTH PREDICTION
In this section, we describe our method for estimating the prompt length based on intermediate layer embeddings in the LLM.
Problem Definition. Denote by u ( ". E Rd the output of some intermediate layer for the N output tokens generated by the model (where df is the hidden dimension of the Transformer). Additionally, we denote by (0) the embedding of the input (during the prefilling phase), which is generated by averaging the embeddings of all the input tokens at the given layer. Let Bi,... Bg C be a choice of bins for the length prediction, i.c. B, {b,, b,. bs+1 1} for some choice of bin boundaries bj b2
b We train linear classifier that, given some embedding (*), outputs vector (1) E [0, 1] where (j) indicates the probability that the number of remaining tokens (i.e., N - () is in B,.
Our approach focuses identifying which LLM layer provides the most accurate predictions for the output length. One potential extension would be to select multiple layers and estimate the prediction using either weighted average of their outputs training linear classifier based on embeddings from several layers. We leave this multi-layer approach for future work. To do so, we have to profilc numerous parameters across all LLM layers for each request.
Using LLama3-8b-ing 8b-instruct as our model and the Alpaca dataset TTori et al. 2023), our profil- ing process begins by extracting embedded elements from each layer during the prefilling and de- coding phases. We profile embeddings across all 32 layers for each token, using 1,000 prompts from the dataset, and retain the embedding tensors along with the remaining token counts as train- ing data. The shape of the embedding tensor after the prefilling phase and before decoding is [1, input token size, 4096), while the embedding tensor for each decoding iteration is [1, 1, 4096).
Predictor architecture. To develop our predictor, we train MLP using the embeddings from each layer and evaluate prediction accuracy. We use a neural network model consisting of fully connected network with two linear layers. The first layer maps the input embedding to 512- dimensional space, followed by a ReLU activation. The second layer outputs predictions for the number of output tokens by classifying the predicted length into one of k 10 equal-width bins, rep-
4

---

Published as a conference paper at ICLR 2025
resenting output lengths between 0 and 512 tokens. The i-th bin b; covers the range
512s
512(++1)
10
10
For the final bin, 610. it includes the upper boundary and covers the range [460.8, 512].
The model is trained over 30 epochs with a batch size of 32, using the AdamW optimizer to control for overfitting We employ cosine annealing schedule to reduce the learning rate from 0.01 gradually. The loss function is CrossEntropy" Here, we present one example a predictor; our approach can, therefore, used with any suitable and efficient learning scheme that yields a predictor. For the first token prediction, we compute the average of all token embeddings the prompt, and this averaged embedding serves as input to the predictor. Preliminary evaluations across all 32 layers (Figure 2 indicate that layers 10- 15 provide the most accurate length predictions.
appropriate for our multi-class classification task.
Focused profiling. Based on this result, we concentrate profiling on these layers, gathering over 7 million training pairs. Subsequently, we train the predictor, employing smoothing techniques to improve the accuracy of our length predictions further.
Smoothing. We observe that while the predicted probability vector in each iteration is reasonable, there can be large variance between iterations. We, therefore, employ Bayesian inference (Särkkä & Svensson 2023) to update the probability estimate as new predictions become available, maintain- ing more accurate estimate of the probability distribution over time. In the context of Bayesian inference, the estimate from previous iterations provides a starting point for the current prediction. After observing the current prediction, the initial estimate updated to reflect the new information. This update is done by adjusting the previous estimate based on how well it aligns with the current prediction and then normalizing it. The updated estimate from this iteration will then be used as the starting point for the next one.
At each iteration 4, the prior is updated because the predicted length may shift between bins over time. Specifically, let T E [0, 1| represent the transition matrix, where T; denotes the probabil- ity that a value in bin B3 at iteration t moves to bin B, at iteration t. In our case, transitions can only occur between neighboring bins, specifically from Bi+1 to B,. We assume that lengths within each bin are uniformly distributed. Therefore, the diagonal entry Te reflects the probability that the value remains in the same bin, which is bin size while the off- off-diagonal entry T2,+11 represents the probability of moving from Bi+1 to B,, which is bin T is calculated once from bin sizes, with its structure detailed in Appendix A The estimated probability g(t) at iteration t is computed as follows:
1. Initialize Sprior (0))
2. At each iteration t. update the prior: q( T gft-o1)
( (+)
3. Update the current probability: (i) Σ, Sprar (f)p()()))
The predicted length at iteration & is L, -,,=,,0 namely my 1 128(21++1 This iterative process refines the prior distribution over time. Figure Bumneesshww shows the mean absolute error for predictions from different model layers.
E0 (1) m,, where my is the average length in bin Bi,
3.2 COMPUTING OUTPUT LENGTH PREDICTION PER ITERATION
There are two approaches to computing predictions based on the embedding extracted from interme- diate layer(s) (in our case, layer 11) of the LLM. In the first approach, we can compute the prediction directly on the GPU using the embedding available at the selected layer. While this method avoids data transfer, introduces slowdown the decoding phase, as both the prediction computation and the model's forward pass share the same GPU. In the second approach, we can compute the prediction on the CPU. For this, we extract the embedding from layer II and transfer it to the CPU. The prediction is then computed in parallel while the LLM processes the remaining layers (in our case, layers 12-32) on the GPU.
We note that whether implement the length prediction the GPU the CPU, the overhead in terms of computational cost is minimal. Specifically, the size the MLP that we train for the length prediction has around 2.1 million parameters, while the overall Llama model we use has 8 billion
5

---

Published as a conference paper at ICLR 2025
Device
Batch
Mean (118
Std (118)
Mean Absolute Error (MAE)
CPU
0.155
0.0043
8.
CPU
512
9.43
3,75
CPU
1024
6.19
1,46
CPU
2048
5.94
1.09
Input
CUDA
0.087
0.00017 7
60-
Refined
CUDA
512
0.615
0.093
BERT
50
CUDA
1024
0.497
0.078
CUDA
2048
0.429
0.084
40
LLM Forward Pass
30
Mcan
101.10 ms
L11
L15
Layer
Table 1: Mean and standard deviation of inference time per sample (TPS) for CPU and CUDA with different hatch sizes, in microseconds (JzS). Addition- ally, the LLM forward pass latency per token for is provided.
Figure 3: Mean Absolute Error for the pre- dicted length, comparing BERT input em- bedding (dashed red), average token embed- ding without refinement (blue), and with re- finement (orange), for different layers.
parameters. Since the computation (number of FLOPs) per token is roughly scaled with the number of parameters, the overhead the length prediction is around 0.03%. This overhead becomes even more negligible for larger model sizes or when using smaller MLP for length prediction. Table shows the mean inference time for length prediction of the two approaches.
3.3 SCHEDULING POLICY
background, traditional scheduling typically assumes that request completion times are either completely unknown, in which case FCFS is natural strategy, or completely known in advance, allowing strategies such as shortest jot first (SJF) to minimize average wait time. many practical systems, exact request completion times unknown. Several studies have explored queueing systems with predicted, rather than exact, service times, generally to minimize the average time request spends in the system. Building previous works, (Mitzenmacheri 2019 Shahout & Mitzenmacher] 2024) analyzed the shortest predicted job first (SPJF) and the shortest predicted remaining processing time (SPRPT) policies in the context of single server where only one request executes at a time. SPRPT tracks the predicted remaining processing time for each request, allowing preemption if new request arrives with a shorter predicted remaining service time.
Such previous work, while motivating using predictions, does not take into account the challenge in LLM systems that preemptions introduce memory overhead, particularly in managing the key value (KV) cache (as explained in Section While preemptions can reduce response time, they increase memory consumption, which may ultimately degrade performance due memory constraints. This is because, in LLM systems, when memory is full, we either discard requests from memory and recompute them once memory available, or we swap KV cache entries from the GPU to CPU. Both approaches impact response time, as discarding requires recomputation, and swapping interrupts the forward-pass of the model, causing delays for the entire running batch. Intuitively, should limit preemptions avoid exceeding memory constraints; in particular, recognize that preemptions made earlier in request's 's execution consume fewer resources, while preemptions closer to completion should be avoided.
In TRAIL, given the initial prediction r for a request (which we treat as a number corresponding to the middle of its predicted bin), only allow preemption for the first re iterations for fixed constant c. This restricts preemption as request nears completion and takes substantial memory.
analyze corresponding theoretical model for single-server SPRPT with lim- ited preemptions, which we believe is interesting in its own right. While this model does not capture the complexity of LLM systems (as it has no notion of memory, and single-server), we think also provides insight into the potential of limited preemption. Using standard queueing theory as-
Interestingly,
'The terms joh and request used interchangeably.

---

Published as a conference paper at ICLR 2025
sumptions (M/G/1 queuing system with Poisson arrivals of rate À < 1 and independent service and prediction times), we can derive a closed- expression for the mean response time (Lemmal I.
Specifically, the processing times for each arriving request independent and drawn based on the cumulative distribution F' ((x)) with an associated density function f(x). Predictions, independent over requests, follow the density function g(x;, r), where x is the actual size and is the predicted size. Thus, s,= g(T, r)dr ((x). A request is described by triple (z, T, a), where is the actual size. r is the predicted size, and is request age (time spent serving the request). We set threshold do = c. r, with as a tunable parameter. The system allows preemption when ag and disables when a 2 dg. balancing early resource use with timely completion. When c = 1, the system becomes the same as SPRPT.
The proof of Lemma ] below is presented in Appendix C Simulations showing that SPRPT with limited preemption improves memory usage compared traditional SPRPT are presented in Ap- pendix D Additionally, in Appendix we provide another theoretical model in the same framework for the setting of refined predictions, that may also be of independent interest.
Lemma 1. For SPRPT with limited preemption, where at age do the requests become non- preemptable, the expected mean response time for request of true size x and predicted size r
B[T(2r))] = (Syno x7 g( (101,y)d + Searaa Sei=t-, g(x,,t) (Ør - (t - r))² drydt)
2(1 -²))
n0
+
()r da + (x ao).
where dl. X 21 g(11, 9) y) dy.
4 EVALUATION
Setup and implementation. Our implementation is based on the open-source vLLM system (v0.5.0), with chunked prefill enabled in both our scheduler and the baseline scheduling methods. We set the out-of-memont mode to discard jobs and recompute them once memory becomes avail- able. The evaluation is conducted on a server with a single NVIDIA A100 80GB GPU and 64 AMD EPYC 7313 16-Core Processor cores, with 503 GiB of memory and running CUDA version 12.3. We use -aimna3----intrucc as a serving model with single-GPU. The workload generated us- ing the Alpaca dataset (Taori et al.) 2023), derived from open-source conversational exchanges and originally used to fine-tune the Alpaca model. We sample 10k unique prompts from the dataset for model serving, distinct from those used to train the length predictor.
Benchmark and Metrics. The benchmark we use simulates a chatbot serving in a client-server setup, where the server hosts LLM model, and the client sends requests specified request rate. The server runs the vLLM OpenAl API, while the client sends prompts from the Alpaca dataset to mimic real-world scenario. The benchmark measures the mean and median latency and Time To First Token (TTFT), providing insights into the LLM'' responsiveness. under varying request loads. In addition regular load, evaluate performance under burst conditions.
Baselines. We compare our approach with four baselines: (1) vanilla vLLM (Kwon et al., 2023) using First Come First Served; (2) vLLM with Shortest Job First based on BERT predictions and chunked prefill; (3) TRAIL with refined embedding predictions; and (4) TRAIL with BERT predic- tions.
4.1 PREDICTIONS ACCURACY
Figure 3 compares the mean absolute error (MAE) of the BERT predictions from the prompts, our linear classifier predictions from the embedding layers, and the refined predictions after using Bayesian inference the different layers choices. We further evaluated the accuracy of this refined prediction from layer 11 against the ground truth and compared it to BERT predictions. A heatmap was used to compare the ground truth remaining length with the predicted remaining length across multiple sequences. Both the x-axis (ground truth) and y-axis (predictions) were divided into ten bins, representing the remaining lengths as the sequence progresses: Where the i-th bin by covers
7

---

Published as conference paper at ICLR 2025
512: 512(++1)
Each request is counted multiple times the heatmap, once for cach prediction made during its execution. Each cell in the heatmap contains the logarithm of the num- ber of occurrences for a given pair of ground truth and predicted length at different stages. For BERT, where only prediction is made, subtract each generated token to create a comparable heatmap. As shown in Figure 4 the refined predictions from layer-11 embedding are more accurate than the BERT predictions, as they exhibit higher val- ues along the diagonal and lower values off the diagonal (indicating greater prediction accuracy). Meanwhile, the BERT predictions show lower values along the diagonal and higher values off the diagonal, reflecting larger differences between predicted and actual lengths.
the range 10
10
from the initial predicted length after
1 90
8
6
6
9 w
Predcteed
Prodicted
to
21
by
An bc
be
Li
0: 25 Di W. As bo Di De Groundkuuth flemiiiing Length
(a) embedding- predictions
(h) BERT predictions
Figure 4: Log-scaled comparison of ground-truth vs predicted lengths bins. The i-th bin bj covers
the range
5124
512(++1)
10
4.2 LLM SERVING
We begin by comparing the mean latency and TTFT of TRAIL across different values of e (c 0.5, 0.8. 1) while setting the request rate 14. (c mimics SRPT without limiting preemption.)
1, TRAIL exhibits higher mean latency and TTFT, while = 0.8 results in the lowest values, closely followed by 0.5. We also evaluated the system with 0.2, which produced higher latency and TTFT compared to the larger values (though not shown here focus more closely on this range of values). These results highlight that while preemption in scheduling benefits LLM systems, it should be limited due to the additional memory overhead it introduces, which can affect serving efficiency. Based on this, we set the default e value for TRAIL to 0.8 in the remainder of the evaluation. The optimal c, however, may vary depending system memory and workload, as it affects the trade-off between preemption benefits and memory overhead.
Figure 5 shows that when
14
200
Mean Latency
10
150
100
50-
C-0.5
0.8.8
c=1
C=0.5
0.8
(a) Mean Latency
Mean TTFT (ms)
Figure 5: Comparison of mean latency and TTFT across different values of (c 0.5, 0.8, 1) at
request rate of 14.
Figure 6 presents the mean and median latency, along with TTFT, function of the request rate for four LLM serving systems: (1) vanilla vLLM (Kwon et al. 2023), which uses First Come First Served (FCFS) the scheduling policy (VLLM-FCFS); (2) vLLM with Shortest Job First (SJF) scheduling policy for newly scheduled sequences based on BERT predictions from sequence prompts (yLLM-SJJ BERT), which also prioritizes new sequences and implements chunked prefill;
8

---

Published as a conference paper at ICLR 2025
(3) TRAIL with 0.8 using refined predictions from embeddings with chunked prefill enabled; (4) TRAIL with c = 0.8 using BERT predictions (TRAIL-BERT); and (5)TRAIL+ oracle baseline where requests scheduled based their exact length. The comparison between TRAIL-BERT and TRAII. evaluates the effectiveness of embedding- based predictions, while the comparison be- Tween vLL.M (with both scheduling policies) and TRAIL. (with both prediction methods) evaluates the impact of limited preemption in LLM systems. The figure shows that SJF based on BERT predic- tions provides minimal benefits over VLLM-FCFS, as vLLM implementation prioritizes incoming requests over existing running requests. Additionally, both TRAIL and TRAIL-BERT, which im- plement limited preemption, achieve lower mean and median latency (Figures [lb and lower mean and median TTFT (Figures 6b 6d) compared to vLLM, with TRAIL using refined embedding predictions showing the lowest latency and TTFT due to more accurate predictions (Figures
Mean Latency
Mean TTET (5)
252
1.00
1:0
100
50-
52
60
60
40
Request Rate (qp/s)
Request Rate (ap/s)
(a) Mean Latency
(b) Mean TTFT
vLLM-FCFS
Trail (c=0.8)
VLLM-SJF_BER BBRTTT
Trail (c=0.8)
Trail-BERT
400
Median Latency 191
350
3:0
250
200
Median
350
100
50-
0-
40
Request Rate (ap/s)
Request Rate (apis)
(c) Median Latency
(d) Median TTFT
Figure 6: Mean and median latency, along with TTFT, as function of request rate for four LLM systems: (I) vLLM-FCFS, vLLM using FCFS; (2) vLLM-SJF BEETTT,FBBERT,,, vLLM using SJF based on
BERT; (3) TRAIL with e 0.8 and refined embedding predictions; and (4) TRAIL-BERT with = 0.8 using BERT predictions.
Figure shows the mean and median latency and TTFT for vLLM, vLLM-SJF BERT, and TRAIL with two different e values (c 0.8 and 1) under a burst scenario. In this scenario, all requests arrive within a very short time interval at the beginning of the experiment, simulating sudden spike demand. The results show that TRAIL. continues to offer benefits with lower latency and TTFT, as implementation ranks all requests (running and waiting) based length predictions and prioritizes them accordingly, while vLLM prioritizes new requests over existing running ones. However, since new requests arrive during processing, preemption has advantage, leading to similar performance between TRAIL with c = 0.8 and c
RELATED WORKS
Recent efforts to optimize the serving of LLMs have been explored. ORCA (Yu et al. 2022) in- troduced token-level scheduling, assuming that batches are processed at the token level rather than the sequence level. vLLM (Kwon et al. 2023) applies PagedAttention to reduce memory overhead. However, these works rely on first-come, first-served scheduling, which could face head-of-ling oi-l blocking.

---

Published as a conference paper at ICLR 2025
Mean TIF Ime
Saat S
400
Noen Labeg
3000
300
300
300
200
200
200
100
100
Sp
Reil e
Teai crl
Trai col
Hit
Rea calc
(a) Mean Lateney (8)
(b) Mcan TTFT (ms)
(c) Median Latency (s)
(d) Median TTFT (ms)
Figure 7: Mean, Median of Latency and TTFT when we have burst of requests
In traditional scheduling, methods like shortest job first (SJF) and shortest remaining processing time (SRPT) aim to minimize response times, while multi-level feedback queues (MLFQ) adjust
job priorities dynamically without prior knowledge of job sizes. FastServe (Wu et al. 2023) uses MLFQ to avoid head-of-line blocking. However, this leads to frequent preemptions, which increase the cost of managing KV cache memory and offloading to the CPU. Several works have proposed prediction-based scheduling to address these challenges, Zhen et al. (Zheng et al.] (2024) enhanced LLM inference by predicting response lengths with an additional LLM model and scheduling them according to length predictions. While this approach optimizes inference, it introduces prediction overhead. Other works predict output lengths using machine learning models such as DistilBERT and OPT, improving resource allocation and reducing memory issues during inference. The works in (Tin et al. 2023; Stojkovic et al. 2024; Cheng et al. 2024) approach length prediction as a clas- sification problem, while other works (Qiu et al. 2024b adopt regression- based techniques.
approach similar to TRAIL. Instead of predicting the absolute output size of a request, LTR ranks requests based their output size, allowing the system to prioritize those with fewer remaining tokens. Although ranking is a simpler task than absolute size prediction, it requires training ranking model in offline phase. A limitation of LTR is that it ignores the size of the prompt when ranking, considering only the output size.
LTR (Fu et 2024) employs a Learning-to-KE
6 LIMITATIONS AND CONCLUSTON
We have presented TRAIL, novel approach that significantly improves response time in LLM inference. Our method is built on two key contributions. First, we obtained low-overhead, high- accuracy predictions from the target LLM itself. After generating each output token, we recycle the embedding of its internal structure as input for a lightweight classifier. This classifier predicts the remaining length for each running request. We demonstrated the high accuracy of this approach and its low overhead compared to existing prediction methods that rely solely on sequence prompts. Specifically, our refined predictions from layer embeddings achieve 2.66x lower mean absolute error compared to BERT predictions from sequence prompts. Second, utilizing these predictions for output lengths, we proposed a Shortest Remaining Processing Time (SRPT) variant with limited preemption. This variant is specifically designed to account for memory overhead in LLM systems. We derived a closed-form formula for this SRPT variant in an M/G/1 queue model and simulated it to demonstrate its potential. By integrating TRAIL with vLLM, we have demonstrated improvements in overall latency and fime-t-ffist-tokenn metrics using real-world datasets. Our experimental results show that TRAIL achieves 1.66x 2.0lx lower mean latency and 1.76x to 24.07x lower mean time to first token compared to vLLM tested using the Alpaca dataset.
Limitations. One key challenge is to select the most informative layer embedding to use as input for the linear classifier. Due to time and resource constraints, we focused only on the Alpaca dataset, as profiling embeddings across all layers is computationally demanding. There are several promis- ing directions for future research. We plan to explore the relationship between layer selection and prediction accuracy across multiple datasets. Another direction leveraging multiple-layer embed- dings through weighted averaging to enhance prediction accuracy. Additionally, experimenting with logarithmic bin sizes for the linear classifier could offer further benefits. A potential optimization is to compute embedding predictions at specific intervals rather than every iteration, reducing com- putational overhead. Lastly, to address data drifts and maintain performance over time, we aim to explore the dynamic retraining of the linear classifier.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We thank Sham Kakade for providing access to the Kempner cluster for evaluation. We also thank Yonatan Belinkov for the helpful discussions. Rana Shahout was supported in part by the Schmidt Futures Initiative and Zuckerman Institute. Michael Mitzenmacher was supported in part by NSF grants CCF-2101144 140 and DMS RMS-2023528. Rana Shahout, Minlan Yu, and Michael Mitzenmacher are partially supported by NSF CNS NeTS 2107078. This work was supported in part by ACE, one of the seven centers in JUMP 2.0, Semiconductor Research Corporation (SRC) program sponsored by DARPA. We are grateful for support from the DARPA ASKEM project (Award HR001 12220042) and the ARPA-H Biomedical Data Fabric project.

---

Published as a conference paper at ICLR 2025
REFERENCES
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207-219, 2022.
Ke Cheng, Wen Hu, Zhi Wang, Peng Du, Jianguo Li, and Sheng Zhang. Enabling efficient batch serving for Imaas via generation length prediction. arXiv preprint arXiv:2406.04780 2024.
Yichao Fu, Siqi Zhu, Runlong Su, Aurick Qiao, lon Stoica, and Hao Zhang. Efficient Ilm scheduling by learning rank. arXiv preprint arXiv:2408. 2024.
John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. arXiv preprint arXiv: /90900339, 2019.
John Hewitt and Christopher Manning. A structural probe for finding syntax in word representa- tions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4129-4138, 2019.
Yunho Jin, Chun-Feng Wu, David Brooks, and Yeon Wei. s3: Increasing gpu utilization during generative inference for higher throughput. Advances in Neural Information Processing Systems, 36:18015-18027, 2023.
Kostis Kaffes, Timothy Chong, Jack Tigar Humphries, Adam Belay, David Mazières, and Christos Kozyrakis. Shinjuku: Preemptive scheduling for {psecond-scale) tail latency. In 16th USENIX Symposium on Networked Systems Design and Implementation (NSDI 19), pp. 345-360, 2019.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 61 611-626, 2023.
Michael Mitzenmacher. Scheduling with predictions and the price of misprediction. arXiv preprint arXiv:190..00732, 2019.
OpenAl. Introducing
https:
2022.
Haoran Qiu, Weichao Mao, Archit Patke, Shengkun Cui, Saurabh Jha, Chen Wang. Hubertus Franke, Zbigniew Kalbarczyk, Tamer Başar, and Ravishankar K Iyer. Power-aware deep learning model serving with {μ-Serve). In 2024 USENIX Annual Technical Conference (USENIX ATC 24), pp. 75-93, 2024a.
Haoran Qiu, Weichao Mao, Archit Patke, Shengkun Cui, Saurabh Jha, Chen Wang, Hubertus Franke, Zbigniew T Kalbarczyk, Tamer Başar, and Ravishankar K Iyer. Efficient interactive 11m serving with proxy model-based sequence length prediction. arXiv preprint arXiv:2404.00 085509, 2024b.
Sanh. Distilbert, distilled version of bert: Smaller, faster, cheaper and lighter. arXiv preprint 1000 20119.
Simo Särkkä and Lennart Svensson. Bayesian filtering and smoothing, volume 17. Cambridge university press, 2023.
Ziv Scully and Mor Harchol-Balter. Soap bubbles: Robust scheduling under adversarial noise. In 2018 56th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 144-154, 2018.
Rana Shahout and Michael Mitzenmacher. Skippredict: When to invest in predictions for schedul- ing. arXiv preprint arXiv:2402.0356/ 2024.
Jovan Stojkovic, Chaojie Zhang, fñigo Goiri, Josep Torrellas, and Esha Choukse. Dynamollm: Designing I1m inference clusters for performance and energy efficiency. arXiv preprint 00741,, 2024.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An iasruction-following llama model. https: //ts thub. www.otttttt..ooooooooooooooooo lab/ sabbssaaffor_all alpacf 2023.
12

---

Published as a conference paper at ICLR 2025
Bingyang Wu, Yinmin Zhong, Zili Zhang, Gang Huang, Xuanzhe Liu, and Xin Jin. Fast distributed inference serving for large language models. arXiv preprint arXiv:2305.05920 2023.
Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A distributed serving system for (Transformer-Based) generative models. In 16th USENIX Sympo- sium on Operating Systems Design and Implementation (OSDI 22). pp. 521-538, 2022.
Zangwei Zheng, Xiaozhe Ren, Fuzhao Xue, Yang Luo, Xin Jiang, and Yang You. Response length perception and sequence scheduling: An llm-empowered Ilm inference pipeline. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Bayesian Inference Transition matrix
14
B SOAP Analysis
14
SPRPT with Limited Preemption
15
D Simulation of SPRPT with Limited Preemption
E Refined Prediction Framework
18
E.1 SPRPT with refined predictions
19
F Dataset Analysis
21
G Multi-GPU Evaluation
21
Ministral--b-nnsiie Predictions
22
A BAYESIAN INFERENCE TRANSITION MATRIX
The transition matrix captures transitions between neighboring bins based on the assumption of uniform distribution within each bin. The diagonal and sub- sub-diagonal entries reflect the probabilities of remaining in the same bin or moving to the adjacent lower bin, respectively. Below is the explicit form of the transition matrix for general number of bins.
0
bis size
Jun size
T
bun size
bus size
size
B SOAP ANALYSIS
We utilize the SOAP framework (Scully & Hlrchhol-Blted 2018) to derive precise expressions for the mean response time. SOAP also provides the Laplace-Stielipee transform of the response time distribution, but we focus on mean response time for ease of comparison throughout the paper.
The SOAP framework is designed analyze scheduling policies for M/G/1 queues that can be characterized by rank functions. Recent work by Scully and Harchol-Balter (Scully & Harchol- Balter 2018) has categorized a broad range of scheduling policies as belonging the SOAP class. These policies determine job scheduling based on ranks, always prioritizing the job with the lowest rank. (If multiple jobs share the same rank, First Come First Served used as the tiebreaker.) The rank function itself is key component, assigned to each job based on certain static properties, typically referred to as the job's type or descriptor. For instance, the descriptor could indicate the job's class in multi-class models, as well as static characteristics like its service time (job size). The rank may also depend on the job's age, or the time it has already been served.
SOAP policies assume that job's rank is influenced only by its inherent characteristics and its age, aligning well with the model and scheduling algorithm used in TRAIL. For more detailed information, the reader is referred to (Scully & Harchol- 2018).
14

---

Published as a conference paper at ICLR 2025
In SOAP analysis, the tagged-job technique is employed. Specifically, track tagged job J with size zJ and descriptor d., and use a, to represent the time J has already been served. The mean response time for J is calculated as the sum of its waiting time (the period from its arrival to the start of service) and its residence time (the duration from the start of service to its completion). To calculate the waiting time, SOAP evaluates the delays caused by both old jobs that arrived before J and new jobs that arrive after J. A key concept in this analysis is the worst future rank of a job, as job ranks may change over time, The worst future rank for a job with descriptor d, and age a, is denoted by rankay (a.)) When a. = O, the rank function is represented as Tworst = ra (0).
wowsst
In the SOAP framework, the waiting time is equivalent the transformed busy period in an M/G/1 queue, where the arrival rate is X and the job size is described by X"ee [rankay (a)]. Where X (rankraa)) is random variable representing the time a newly arrived job will be served until it either completes or its rank exceeds rankay (a). The initial workload during this period corresponds the delays caused by old jobs. To account for the delays from old jobs, SOAP transforms the system by categorizing jobs according their rank. Old jobs that exceed the rank threshold Tworst are classified as discarded and are not included in the transformed system. Those with ranks or below Tworst, referred to original jobs, treated arriving rate À with specific size distribution [T Where [Toorst] is random variable representing the ser- vice time of an original job with respect rank Tworst- Recycled jobs, which were once above the threshold but have now fallen below, are treated as server vacations of length X [[worse]. Where x [[awrrst] a random variable representing the service time recycled job for the i-th time with respect rank Tworst for i 1. In TRAII., jobs are recycled only once, so we only consider X [Tworst).
SOAP shows that, due to Poisson arrivals time averages, the stationary delay caused by old jobs follows the same distribution as the queuing time in the transformed M/G/I/FCFS system, characterized by sparse server vacations where original jobs arrive rate X and follow the size distribution [ (rooret]
Theorem (Theorem 5.5 of (Scully & Harchol-Balter 2018 Under any SOAP policy, the mean
response time of a job with descriptor d and size is:
A E. 0 B worrst 121
([T d)] 2(1 X old (woret)) (1 ALE X new (worst))
/
XE new frank (a)l] da.
C SPRPT WITH LIMITED PREEMPTION
We begin by describing the rank function for job with descriptor (x, r, a). where is the actual size. r is the predicted size, and a is the job's 's age. The threshold where we stop preemption from occurring set as ag c- The rank function is:
rank(I, fr-a if < 00
-00 otherwise
As this rank function is monotonic, A job's worst future rank is its initial prediction:
d. (a) r if < ag
rank worst
(1)
-00 otherwise
When = 0, the rank function is denoted by Tworst rank (0)
Lemma 1. For SPRPT with limited preemption, where at age ag the requests become non- preemptable, the expected mean response time for request of true size and predicted size
15

---

Published as a conference paper at ICLR 2025
is
(Jyzo -2 o9 g( /dd + Startan fo 9(11,) (ør - (t r))2 dr
[T)
2(1 -,))
00
P(r-a)+ da + (x an).
where Py = X ر Sn 2I g(31, y) ddd.
Proof. To analyze SPRPT with limited preemption using SOAP, based on the worst future rank (equation we calculate Xuoe (rank (a)), [ and X ["worst] for tagged job with descriptor az)..
C.0.1 XNE (rank (a)] COMPUTATION:
Suppose that new job K of predicted size т arrives when J has age If dg and K has predicted job size less than J's predicted remaining process time (r - a). K will always outrank J. Thus
X 2x [r (r- al 2kl(Tk <r- a)-a) ao)
E new
a]] =
HK g(0K dy
=x=0
C.0.2 X Tworss] COMPUTATION:
Let be an old job in the system. Whether job I is an original or recycled job depends on its predicted size relative to J's predicted size. Ifrr then I is original until its completion because its rank never exceeds r.
Xold 0,21
[r] = 1r s))
E E [r]] = I p
Dn yxzd.y
y= =1=0
/ at g)zr,
B[ [r)=]
C.0.3 XOLD I [worst] COMPUTATION:
ry then starts discarded but becomes recycled when r7 This means age
a = Fi and served till completion only a < ao, which will be xy aj = 21 - (TI r):
Thus, we have
xold 1, [r] = I1 - (rI - r).
For 2 2,
oold [r] 0.
100
B [
g0r, (x1 - (rr - r)) dxydri-
16

---

Published as a conference paper at ICLR 2025
Applying Theorem 5.5 of SOAP (Scully & Harchol-Bas Balter 2018) yields that the response time of jobs with descriptor (r) and size x is as follows. Let
0 = f 2r g(zr,
y= 0
Then
Uysoler 27 g( ddr,yy + Searta 1 9 (ær - (t r)) dx dt
r))
2(1 M,)
00
da + (x an).
P(r-a)
Let fp(y) = - y)d.. Then the mean response time for job with size æ, and the mean response time all jobs are given by
g(,
f y-
E (T))
ET
g(x, ET(( dydd.
==0 y=0
D SIMULATION OF SPRPT WITH LIMITED PREEMPTION
In simulation, memory usage is modeled proportional to the age of each job: the amount of service it has received so far. This reflects the idea that jobs consume more memory the longer they remain in the system. We measure the response time to evaluate scheduling efficiency along with memory consumption.
We consider single-queue setting with Poisson arrivals, where the job service time follows
ponential distribution with mean (f(x) --)) Two prediction models are used: 1) Exponential predictions (Mitzenmacher) 2019), where the prediction for job with service time æ is also expo- nentially distributed with mean x, given by g 3) = =z-v 2) A "perfect predictor", where the prediction is always accurate, given by g(..
Figure shows the mean response time (blue line) and the peak memory usage (green line) under various arrival rates and with different values. The key takeaway is that limiting preemption; reducing how often jobs are interrupted and rescheduled, can lead to better memory utilization while maintaining reasonable response time.
17

---

Published as a conference paper at ICLR 2025
10-
23
0.04
Response Time
Peak Memory
Response Time
Memory
0.02
5
18
-..2
3
-0.04
2
15
0.00
0.70
6.75
0.00
3.00
0.55
0.00
0.70
0.50
0.55
0.50
0.15
Aetioa Kate
Ariia Rade
(a) arrival rate when
0
(h) arrival when c 0.3
2.3
17.5
10
15.0
Response Time
Peak Memory
Response Time
2.2-
12.5
Peak Memory
2.1
10.0
20
7.5
2.0
5.0
16
0.0
6.0
02
3.1
0.6
6A
1.
C
Ar Rotz
(c) vs. arrival rate when 0.5
(d) vs. e when 0.7
4.24
2.9
20
Response Time
2.8-
20
Response Time
4.1
Peak Memory
Peak Memory
2.7
15
4.0
15
2.6
10
10
3.8
3.7
3.6
0b
04
0.6
6.0
6.
3.6
64
(e) vs. c when X = 0.8
(I) vs. c when À = 0.9
Figure 8: Comparing memory usage and response time across different arrival rates and values of c.
The parameter c, used to manage the trade-offs memory utilization, is influenced by the KV cache's available memory, which depends on the model size, batch sizes, and incoming request sizes. In scenarios with a mix of long and short requests, an appropriately chosen e value can prevent memory monopolization by long requests, thereby improving throughput. The value of can also be dynamically adjusted in real time.
E REFINED PREDICTION FRAMEWORK
We consider M/G/1 queueing systems with arrival rate X. The processing times for each arriving job are independent and drawn based on the cumulative distribution F(x), with an associated density function (xr)
18

---

Published as a conference paper at ICLR 2025
Given job of size X z, and sequence of refined predictions Yo, Yi...., Y,, where Y, represents the prediction after processing units of 2, we model the refined predictions as follows:
The probability of the refined prediction yi step i, given the actual size z and the entire history of previous predictions, is denoted as:
P(Y: ys X x -i- Y1- y1- Yo = y0)
Under the Markovian assumption, the refined prediction in step i depends only on the actual size and the previous prediction 1/-11
P(Y; yı X = Y. 91-1, Y 2:-2. Yo =y0) P(Y; yi
X Y. y/-1)
The joint probability of the sequence of predictions Yo, YI, Yu can be modeled as:
P(Y = yo. Y1 = 91 Y = Yo X = x) = P(Y = yo X x) II P(Y, - 34 = Y, 91-1)
The density function g(z, y) corresponding to the list of predictions y 1o ,............................... Mr I.. where each ys is obtained after processing one unit of job x, is obtained by summing over all possible refined predictions:
g(z,y)=
P(Yo Vo. Y1 y1..... Yo-1 Mx-1 X = x)
Ve -1=0
E. SPRPT WITH REFINED PREDICTIONS
The relevant attributes are size and the refined prediction list r. We can model the system using descriptor (size, list of refined predictions, age).
Let J/ be job with a descriptor (z, T, Q), where x is J's size, r is a list of refined predictions for each generated after processing one unit the job, and a is J's age. Thus, SPRPT with refined predictions has rank function rank(1;, r,a) r(a] for job of size x and predicted size r r[a] at time a. We denote the maximum rank function as Tman-
I's worst future rank is:
rankaa
wursi
sup (rlb) - 6) = max (r[b) - b) <b<
a<ber
When = 0, the rank function is denoted by Tworst rank
Tmazz
To analyze SPRPT with refined predictions using SOAP, the calculation involves solving multidi- mensional summation. While obtaining a closed-form expression for such sums may be challeng- ing, in this section, we propose a method to compute the relevant components (X (Xuoww [rank (xoo"rankk) worst (a))),
X [worss]] and ( [Tworst]) for the SOAP analysis without requiring a closed-form solution.
Let be a tagged job with descriptor (I,, a). based its worst future rank.
X prankk d,z
(a) computation: Suppose that a new job K of predicted size vx arrives when J has age a. K has a lower rank than the worst possible rank of job J is given by:
X [ morss (a))
[rank
Σ
1 [ < max (r(b) - b)) -1 Vj < ar, a<b<z
max (rlb) -b)
a<b<z
=0
where 1(.) is the indicator function, which equals I if the condition inside is true and otherwise.
19

---

Published as a conference paper at ICLR 2025
00
ØK (Tk(ak] ak
E new
rank d,z
(a)
max (r(b]
acbez
Nak
fj < or, rzli] -j max (r(M) - b)
a<b<æ
g()x,
where the density function g( y) is defined as:
g(rk,))=
Σ
Σ
P(Y = you Y1 = -1. y1...., Yux-1 = 92x-1 X = IK)
yo = in
Meg-1=t
X [worss] computation: Let I be an old job in the system. Whether job is an original depends on its predicted size relative to J's predicted size. I original till it is exceeds I'ma=- terminate the summing once you encounter age where rr rr(ar] 01 > Tmax:
O.ay [[moz]
(r/[ar)) as <maw)) 1 1(j 01 r1li) - Tmaz) 61=0
Σ
X [woree] computation: If > then starts discarded but becomes recycled when rr[a] [a]-c -a Tworst: This means at age ry [ar]]] Tworst and served till it will be again about r Tworsts
Let i be an interval index. The i-interval for job I is defined as the interval of ages during which the rank of job is less than Tmax- Specifically:
bo
This indicates that the first interval starts at age O.
co[r] inf ry rmaz
Here, [r] is the smallest age greater than 0 where the rank function ry (a) becomes greater than Tmaz, marking the end of the first interval where the rank is less than rmax.
For subsequent intervals > 1:
[r] = inf G-1lr] (a( Tmaz)
This is the starting point of the i-th interval, which begins the rank ry(a) drops below Tmaz again after the previous interval (b.-1[r),c,-1111
cr(r] = inf {a > by[rl (2(a)) - > Tmar]
Here, ci(r] is the endpoint of the i-th interval, where the rank rr(a) - becomes greater than Tmoz again.
For 2 0, the i-interval work is a random variable, written X.LD [r]. representing the sum of the tokens where job I has rank less than Tmax within its i-interval.
0
if X by[r] ifb, < Xa cg(r] of < Xa
OLD [r]
Xa 6,(r]
cs|ri [rl
If [:(r]] [r(r] no, define X [r] 0.
20

---

Published as a conference paper at ICLR 2025
F DATASET ANALYSIS
The first figure illustrates the cumulative distribution function (CDF) of the Alpaca (Taori et al. 2023) dataset's output lengths, showcasing the variation output sizes.
1.00
0.75
C
0.50
0.25
0.00
100
200
300
400
500
Output length
Figure Cumulative Distribution Function (CDF) of output lengths in the dataset.
The second figure presents a histogram of the dataset binned into uniform-sized bins described in Section 3.1 highlighting the number of prompts within each bin.
Num. prompts
1000
500
Bin (uniform size)
Figure 10: Histogram of the dataset binned into sefform-sized bins.
G MULTI-GPU EVALUATION
Testbed. For testing multi-GPU settings, we used machine with dual AMD EPYC 7313 CPUs (16 cores per CPU, totaling 64 threads), 503 of RAM, and two NVIDIA A100 GPUs with GB memory each connected via NVLink.
For the first experiment, we used the same serving model (L.Lama3-8b-ing 8b-itsrrutnn distributed across two GPUs using tensor parallelism. As shown in Figure II TRAIL continues to perform effectively.
Thal
350
Moon Latencyy
220
250
150
108
24
Request Rate (gpis!
Requestt Rate (opk)
(a) Mean Latency
(b) Median Latency
Figure 11: Mean and median of end-to-end latency with LLama3-8b-instruct as a serving model distributed across two GPUs using tensor parallelism.
For the second experiment, we tested TRAIL+ (TRAIL with perfect prediction) using Vicuma 13B, distributed across two GPUs using tensor parallelism.
21

---

Published as a conference paper at ICLR 2025
Meain Latercy
04
21 Request (op/5)
66
Request (apis)
(a) Mean Latency
(b) Median Latency
Figure 12: Mean and median end-to-end lateney using Vicuna 13B model, distributed across two GPUs using tensor parallelism. Tested with sampling 10000 prompts the first 1000 prompts of the Alpaca dataset.
H MINISTRAL-BB-NSSIRUCE PREDICTIONS
We evaluated the predictive performance another LLM, Ministral-3b-instruce using 1,000 prompts from the Alpaca dataset. The dataset was split such that 75% of the prompts were used for training and the remaining 25% for evaluation.
These results are preliminary. Figure 13 illustrates the Mean Absolute Error (MAE) for predictions across different layers of the model, while Figure 3 demonstrates the refined predictions for each layer after applying the smoothing method described in Section 3.1 From the results, observe that a specific layer, Layer 10, consistently yields the most accurate predictions.
Prediction w/ per-layer embeddings
100
MAE
50
0
0
1
2
3
4
5
6
7
8
9
13
Layer index
Figure 13: Mean Absolute Error (MAE) per layer for Ministral-3b-inet predictions.
22

---

Published as a conference paper at ICLR 2025
Max: 46.23 (LO)
Mean Absolute Error (MAE)
Refined
45
Mean MAE: 42.50
44
43
40
39
00
L10
11
L12
L13
Layer
Min: 38.99 (L10)
Figure 14: Refined predictions per layer using the smoothing method described in Section 3.1
23

---

Published as a conference paper at ICLR 2025
24


Reviews:
[{"review_id": "reUocCO6ig", "reviewer": "Reviewer_uAxn", "summary": "The paper addresses the challenges of efficient scheduling in interactive Large Language Model (LLM) applications. It uses a combination of length prediction and batching for speedup.  It introduces TRAIL, a method that uses the LLM’s own embeddings to predict the remaining length of requests, enabling a prediction-based Shortest Remaining Process Time (SRPT) variant with limited preemption. This approach aims to reduce memory overhead and optimize resource utilization by allowing preemption early in request execution and restricting it as requests near completion. Experiments show lower mean latency and time to the first token compared to current approaches.", "questions": "1. It seems strange to take the average of prefill embeddings for inference. Is this the only approach you have tried, or did you try other (for e.g. non-linear) ways to aggregate the prefill embeddings?\n\n2. What is 44? (shape of the embedding tensor after the prefill phase is [1,44,4096] as mentioned in Section 3.1 on page 4)\n\n3. What is $\\hat{q}_\\text{prior}$? How is it initialized?\n\n4. What is the LLM forward pass latency per token for the setting(s) in Table 1? Please include that number so we can get a clearer idea of the latency overhead of length prediction.\n\n5. There appears to be an inconsistency in notation between 'C' in Section 3.3 and 'c' in Section 4.2. Please clarify if they are the same or different and fix the notation accordingly. Also do you have any ideas on how one could choose a good value of 'C'/'c' in real world? Can it be learned/adjusted in real-time based on the batch/request characteristics?\n\n6. It is mentioned in Section 4.2 that vLLM-SJF_BERT prioritizes incoming requests over existing running requests. What does that mean? Shouldn't it be scheduling requests just based on the predicted length and not based on arrival time?  \n\n7. I would recommend adding an oracle baseline where requests are scheduled based on their exact decode length. While the true decode length will not be known in practice, this will serve as a useful benchmark and approaches can be ranked based on how close they are to this one.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The use of LLM embeddings for predicting residual request lengths, is novel, interesting, and intuitive. It has a clear potential to improve request scheduling as seen in the evaluation in this paper and can also be useful in any scenario where output length prediction is required and thus can have a broad positive impact.\n\n2. The use of pre-emption based on length prediction has also not been fully explored in other works to the best of my knowledge. The authors propose a clear SRPT-like approach based on predicted request length and also make interesting observations about the extent of pre-emption that is useful (Fig 5) which will be useful for future research in this space.", "weaknesses": "1. The main algorithm is not very clearly explained although the issues are relatively minor and can be fixed by addressing questions 1-4 below.\n\n2. By solely relying on length-based scheduling there is a risk of violating latency SLAs. In general, a real-world scheduler should be able to tune between FCFS and length-based scheduling to handle requests with different SLAs.\n\n3. The evaluation is largely limited to the mean latency and does not adequately capture the tail behavior of the proposed approach or baselines. I would recommend including the corresponding P95 latency plots, along with comments, for Fig 5,6, and 7 to address this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hZYG2EDrml", "reviewer": "Reviewer_KSCH", "summary": "This paper introduces TRAIL, a method for predicting the remaining output length of an LLM based on its internal structural embeddings. Using these predictions, the authors propose a scheduling algorithm, a variant of the shortest remaining processing time, to reduce latency. They derive the theoretical expected response time for a job and implement the scheduling policy within the vLLM serving system for evaluation.", "questions": "Could you explain why the performances of vLLM-FCFS and vLLM-BERT are so close in all metrics? Is it because that BERT prediction does not work at all, as shown in figure 3?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The problem the authors address: the inefficiency of FCFS scheduling in LLM serving, is practical. Overall, I think the proposed solution is useful. Online token length prediction is inherently challenging, and their approach of using bins and classifiers is reasonable based on my experience. The feature design for the predictor is a novel contribution. The paper appropriately tackles key considerations in LLM serving, such as the uncertainty in output token length, the difficulty of making accurate predictions, the drawbacks of adding memory overhead during runtime decoding, and the distinction between the initial token generation (prefill) and the decoding of subsequent tokens.", "weaknesses": "Some terms could be more precise and consistent. For example, in the context of LLM serving, what the paper refers to as a “job” is actually a “request.” The term “job” typically implies a collection of multiple requests or tasks. While the authors note that “job” and “request” are used interchangeably, this distinction isn’t necessary and could lead to confusion. Additionally, the term “TRAIL” is described as a prediction method in the abstract but appears to refer to a scheduling policy in the evaluation section. This inconsistency should be clarified.\n\nIn Figure 6, the main experimental results show that the trends for Trail and Trail-BERT exhibit a faster growth rate compared to vLLM-FCFS and vLLM-SJF_BERT. Could you increase the request rate further to demonstrate whether the latency of Trail and Trail-BERT will remain below that of vLLM-FCFS and vLLM-SJF_BERT as the request rate continues to increase?\n\nIn multi-GPU settings with larger models, the results may differ due to the additional time required for initializing the Ray engine, GPU-to-GPU communication, and other overheads. While I believe the scheduling algorithm remains effective, it would be beneficial to evaluate its performance on larger models in a multi-GPU environment.\n\nI'll raise the score if the authors successfully address the experiment issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dqYO5lYOeU", "reviewer": "Reviewer_DYK6", "summary": "In this work, the authors propose a strategy for predicting output length that builds upon state-of-the-art approaches such as S3. The approach modifies output length prediction by leveraging embeddings from different layers of the LLM and providing them to a trainable, lightweight predictor module. These predictions are then used in conjunction with size-based scheduling approaches such as SRPT, resulting in reduced overall latency and TTFT. The key idea here is that length-aware scheduling reduces memory overhead and, consequently, preemptions.", "questions": "1. Can the generalizability of this approach across different LLMs be evaluated? Is it a generic assumption that for every LLM, a certain layer carries most of the information about the output length, or could this vary depending on the model?\n2. How can this approach be applied in the cases of OpenAI models where the internal architecture is unknown?\n3. What is the output length distribution used in the paper, and how would the performance vary if the buckets are selected based on their size?\n4. Can the paper provide more information on the Alpaca dataset, and what was the distribution of the data used with respect to the root verb of instructions from the Alpaca dataset?\n5. How does the paper ensure that the performance of the approach is consistent across different tasks within the Alpaca dataset, given that some tasks may be more unpredictable than others? Can the authors provide the breakdown of the MAE and heatmap with respect to different root verbs of instructions?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "A major strength of this paper is its output length prediction module, which iteratively refines predictions using embeddings from the LLMs. This approach enables the use of SRPT-based scheduling methods. The work also proposes two approaches for predicting output length and discusses the overhead associated with LLM inference. Additionally, the paper presents closed-form expressions based on queuing theory for mean response time.", "weaknesses": "The paper has a few weaknesses that should be considered. \n1. The output length prediction assumes uniform bucket sizes, which may lead to overfitting if the number of requests in each bucket varies significantly. A comprehensive evaluation of the distribution per bucket is required for a fair evaluation of this approach. \n2. The assumption that the output length predictor needs access to the layers of the LLMs may not hold for many OpenAI models, whose internal architecture is unknown. The generalizability of the approach in such cases needs to be discussed. \n3. Finally, some tasks from the Alpaca dataset are more unpredictable than others. For instance, tasks such as writing a story are far more arbitrary than classifying a sentence. Therefore, remarks on the performance of the approach with respect to different tasks should be added for better generalizability.\n4. It is unclear whether a certain layer carries most of the information about the output length for every LLM, or if this is a generic assumption that may not hold for all models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XhEhQtIuig", "reviewer": "Reviewer_bvWq", "summary": "The authors address the scheduling of LLM inference jobs. In particular, to provide a high-quality experience, LLM inference should be fast enough to allow real-time conversational user interaction. Scheduling jobs according to their length has since long been shown to minimize average waiting time, but it requires knowing the length of jobs in advance. Due to its autoregressive nature, LLM-based text generation can lead to highly varying execution times. To overcome this limit, related work proposes various models to predict the length of LLM inference. The authors here improve on this by, rather than using smaller LLM models such as BERT-like ones, predicting the inference length based on the state of the internal layers of the LLM. Moreover, the authors also address the drawback of job preemption, which requires saving intermediate state, which can be considerable for modern LLMs with consequent issues of either exhausting GPU memory or costly memory transfers, by only allowing preemption at the initial steps of a job. In contrast, jobs close to termination can not be preempted. The authors evaluate both the accuracy of their intermediate layer-based job length prediction and preemption-tuned scheduling against three baselines covering vanilla and state-of-the-art solutions.", "questions": "The rationale for the number and size of bins for predicting the LLM inference length is missing. What is the mean length of responses? New models have significantly larger contexts. Does this influence the choice? How? Also, would it not make more sense to choose a divider of 512 instead of ten (e.g., 8 or 16) since length is intrinsically an integer?\n\nLearning rate decreased to 0 does not seem to make much sense. I guess it is close to 0. Please check.\n\nLength prediction overhead is small in relative terms but could be significant in absolute terms, especially from an energy point of view. It would be nice to comment on this. Also, how are batches accumulated in this scenario since you want to predict after each next token has been generated? What is the prediction time with a batch size of one?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ Significantly lower mean latency compared to the considered baselines.\n+ The model used to predict the LLM output length is way simpler than the ones used by SOTA-related work, which reduces the overall computational burden of LLMs (especially in the context of energy efficiency over the whole sector). \n+ Theoretical proof of the proposed scheduling scheme with limited preemption and dynamic threshold to adapt to the variability of inference lengths.", "weaknesses": "- Layer-based prediction only tested on one LLM, which makes it unclear how well the method generalizes to other LLMs.\n- The choice of the layer used for prediction is model-dependent and requires a preliminary study.\n- A deeper sensitivity analysis of the parameter c would have been nice to analyze the optimal value for c better.\n- Lack of some design rationales (see questions below)\n\nTo give more meaning to the average prediction error presented in Figure 2, it woudl have been nice to provide soem statistics on the inference length.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)