DON’T STOP ME NOW: EMBEDDING BASED SCHEDULING FOR LLMS

Thank you for recognizing the strengths of our work, including the gains in mean latency, the simplicity of our prediction model compared to SOTA methods, and the theoretical proof of our proposed scheduling scheme.

• “Layer-based prediction only tested on one LLM, which makes it unclear how well the method generalizes to other LLMs.”

We agree that generalization is an important concern. Due to resource and time constraints during the rebuttal period, we were unable to evaluate additional LLMs. However, we plan to include experiments with another LLM in the camera-ready version.

• “A deeper sensitivity analysis of the parameter c would have been nice.”

The value of 'c' primarily depends on the available memory for the KV cache, which is influenced by model size, batch sizes, and incoming request sizes. For instance, preempting long requests could monopolize memory and reduce throughput during periods with long requests followed by short ones before finishing. As shown in Appendix D (Figure 8), workload simulations can help determine suitable 'c' values, which can also be adjusted in real time without further system changes. This adjustment directly impacts the rank of jobs (see Equation 1 in Appendix C), connecting the 'c' value to the request ranking.

• “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?”

We followed the settings of $S^3$[1], which also uses the Alpaca dataset. For the camera-ready version, we will explore non-uniform bin sizes (e.g., bins of geometrically increasing size such as 2, 4, 8, …) to address this. In the revised paper, we have added the mean and distribution of response lengths (Appendix F) to provide additional context. For future models with larger contexts, we agree that adjusting bin sizes, including exploring non-uniform bins, may better capture the distribution. We will explore alternatives in the future.

[1] $S^3$: Increasing GPU Utilization during Generative Inference for Higher Throughput, Jin, Yunho, Neurips 2023.

• ”Learning rate decreased to 0 does not seem to make much sense. I guess it is close to 0. Please check.”

We use the minimal learning rate default parameter from the PyTorch implementation of the learning rate decay schedule (https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html), which indeed decays the learning rate to zero at the end of training. This is a fairly common choice, which is suggested, for example, in [2].

[2] SGDR: STOCHASTIC GRADIENT DESCENT WITH WARM RESTARTS, Ilya Loshchilov & Frank Hutter, ICLR 2017

• “Length prediction overhead is small in relative terms but could be significant in absolute terms, especially from an energy point of view.”

We acknowledge the importance of minimizing prediction overhead, particularly from an energy perspective. As noted in the limitations section of the paper, one potential optimization is to compute embedding predictions at specific intervals rather than every iteration, which could significantly reduce computational costs. Additionally, we have added the LLM forward pass latency per token to Table 1 in the revised paper to clarify the computational overhead associated with predictions in our tested model.

Thank you for your thoughtful feedback and for engaging with our responses.

• We appreciate your understanding regarding another LLM evaluation. As promised, we have now added preliminary results for the new model, Ministral-3b-instruct, to Appendix H of the revised paper. We will continue this experiment with the full dataset and include the results in the camera-ready version.

• Regarding the parameter c, we agree with your suggestion to include a summary of our observations. This has been addressed in the revised version (Appendix D).

• To address your concern about prediction time, we have extended Table 1 in the revised paper to include inference time for batch size 1 on both CPU and GPU.

Thank you for supporting our approach and recognizing the contributions of the output length prediction module and the theoretical closed-form equation for the mean response time of the proposed scheduling policy.

• Q1: 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?

A1: We appreciate your comment regarding another LLM evaluation. We have now added preliminary results for the new model, Ministral-3b-instruct, to Appendix H of the revised paper. We will continue this experiment with the full dataset and include the results in the camera-ready version (due to time and resource constraints during the rebuttal period). We see that the assumption that certain layers in LLMs carry output length information also holds for the newly evaluated model, as these architectures share foundational design principles.

• Q2: How can this approach be applied in the cases of OpenAI models where the internal architecture is unknown?

A2: While our work assumes access to the model's internal layers, the proposed scheduling policy can still work with black-box models, provided output size predictions are available. For OpenAI models, if the model can return the remaining token count alongside each generated token, our approach, TRAIL, could be directly applied. Even without this, our scheduling policy, which limits preemption based on external output predictions, can still be effective if some other external prediction is available. Analyzing scheduling for black-box models is an interesting direction for future work.

• Q3, Q4, Q5: More information on the Alpaca dataset, output length distribution, and consistency of performance across tasks with varying unpredictability.

We appreciate the reviewer’s concerns. In the revised paper, we have added the output length distribution and bins histogram (Appendix F). We acknowledge that uniform bin sizes may lead to biases, and we followed the same settings used in $S^3$[1], which also uses the Alpaca dataset. For the camera-ready version, we will explore non-uniform bin sizes (e.g., bins of geometrically increasing size such as 2, 4, 8, …) to address this. Regarding the Alpaca dataset, while we considered breaking the data based on root verbs, we believe comparing results across datasets would provide a clearer evaluation. Due to time and resource constraints, we could not perform this comparison during the rebuttal but aim to include an additional dataset in the camera-ready version.

[1] $S^3$: Increasing GPU Utilization during Generative Inference for Higher Throughput, Jin, Yunho, Neurips 2023.

We thank the reviewer for the thoughtful feedback and valuable suggestions, which have helped us improve the clarity and quality of our work.

• We appreciate the suggestion to extend the evaluation:

  • In the revised paper, we have considered increased request rates and updated Figure 6 accordingly. As predicted, Trail continued to exhibit benefits at higher request rates due to the mitigation of head-of-line blocking, which is consistent with queueing theory predictions. We also added to Figure 6 Trail+, a baseline oracle where requests are scheduled based on their exact length.

  • For testing multi-GPU settings, we used a machine with dual AMD EPYC 7313 CPUs (16 cores per CPU, totaling 64 threads), 503 GB of RAM, and two NVIDIA A100 GPUs with 80 GB memory each connected via NVLink. We evaluated our approach in multi-GPU settings using two methods:

    1- Distributing the existing tested model across two GPUs. The results in Figure 11, Appendix G, show that TRAIL continues to perform effectively.

    2- Since training our classifier for larger models takes resources and time (While we aim to profile and train our classifier for these larger models for the camera-ready version), we tested TRAIL with a larger model size (vicuna 13B), using "perfect" predictions (denoted as Trail+) using sampling 10000 prompts from the first 1000 prompts from Alpaca dataset. In the camera-ready, we will complete all baselines and the full dataset. Trail+ serves as an upper-bound benchmark for this scenario. These results appear in Figure 12 Appendix G.

• We acknowledge the reviewer’s concern regarding using both the terms request and job. We clarified the description of TRAIL throughout the paper and revised the paper to consistently use "request" in the context of LLM to avoid confusion. We do note that in queueing “job” is commonly used (as in the Shortest Job First policy), but we believe this change provides helpful clarification.

• Q: “Could you explain why the performances of vLLM-FCFS and vLLM-BERT are so close in all metrics?”

A: vLLM-SJF_BERT sticks to vLLM's implementation, where incoming requests are prioritized. For the remaining slots, SJF is applied using BERT predictions. This approach adjusts the ordering policy but does not modify the underlying scheduling mechanism of vLLM. Comparing BERT predictions with our approach, we compare TRAIL-BERT with TRAIL in Figure 6.

Dear Reviewer KSCH,

Thank you once again for your thoughtful feedback. We appreciate the time and effort you spend reviewing our work.

As the ICLR public discussion phase draws to a close, we wanted to confirm whether our responses properly addressed your concerns. If there are any remaining questions or points requiring further clarification, please let us know—we would be glad to provide additional details.

We thank the reviewer for supporting our paper. We appreciate the thoughtful questions that clarify essential details of our approach, as well as your pointing out typos. In the revised paper, we have added to Figure 6 a new baseline Trail+, which shows results for an oracle that provides exact lengths and schedules accordingly. We have also extended Figure 6 with increased request rates. As predicted, Trail continued to exhibit benefits at higher request rates due to the mitigation of head-of-line blocking, which is consistent with queueing theory predictions. We also added LLM forward pass latency per token in Table 1.

• Q1: “Is this the only approach you have tried, or did you try other (for e.g. non-linear) ways to aggregate the prefill embeddings?”

We used averaging as a simple and efficient baseline that does not depend on input size to demonstrate the feasibility of using embeddings for length prediction. While we did not explore non-linear methods in this work, we agree they could improve accuracy and plan to investigate them in future research. We aimed to establish the foundation and demonstrate the utility of embedded-informed scheduling.

• Q2: “What is 44?”

Thank you for pointing this out. The value "44" was a mistake in the original text and represents the number of input tokens. The correct parameter list should be (1, [input tokens], 4096), where "44" is just an example of the number of input tokens. This clarification has been added to the revised paper.

• Q3: “What is $q_{\mathrm{prior}}$​? How is it initialized?”

$q_{\mathrm{prior}}$ represents the prior probability estimate of refined predictions at iteration t during Bayesian inference. At t=0, it​ is initialized to the initial prediction p(0). We have fixed the notation in the revised paper.

• Q4: “What is the LLM forward pass latency per token for the setting(s) in Table 1?”

We added LLM forward pass latency per token to Table 1 to clarify the prediction overhead in our tested model.

• Q5: “There appears to be an inconsistency in notation between 'C' in Section 3.3 and 'c' in Section 4.2.”

This is a typo; thank you for pointing to this. Both notations refer to the same parameter. We have corrected this in the revised version.

• “How to choose a good value of 'c' in real-world scenarios? Can it be learned/adjusted in real-time based on batch/request characteristics?”

The value of 'c' primarily depends on the available memory for the KV cache, which is influenced by model size, batch sizes, and incoming request sizes. For instance, preempting long requests could monopolize memory and reduce throughput during periods with long requests followed by short ones before finishing. As shown in Appendix D (Figure 8), workload simulations can help determine suitable 'c' values, which can also be adjusted in real time without further system changes. This adjustment directly impacts the rank of jobs (see Equation 1 in Appendix C), connecting the 'c' value to the request ranking. This discussion was added to Appendix D as well.

• Q6: “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?”

vLLM-SJF_BERT sticks to vLLM's implementation, where incoming requests are prioritized. For the remaining slots, SJF is applied using BERT predictions. This approach adjusts the ordering policy but does not modify the underlying scheduling mechanism of vLLM.

• Q7 “I would recommend adding an oracle baseline where requests are scheduled based on their exact decode length.”

We agree and added this Oracle baseline (TRAIL+) to Figure 6 in the revised version.

• We understand the concern about the risk of violating latency SLAs. Our intuition from previous work on queueing is the benefits of prioritizing jobs by length (or predicted length) provide such significant gains in expected latency that the system can be modified to handle tail latency issues while maintaining much of these gains. In particular, one way to address this is by adding a starvation prevention mechanism. We will add this to the camera ready and evaluate the tail latency under different parameters.