# Don't Stop Me Now: Embedding Based Scheduling for LLMs

arXiv: https://arxiv.org/abs/2410.01035 · HTML: https://mlsanity.com/paper/2410.01035/

## Key facts

- Cited by 35 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 21 references that are in the map.
- First posted on arXiv 1 Oct 2024.

## Builds on (most cited first, 13 of 13)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [Designing and Interpreting Probes with Control Tasks](https://mlsanity.com/paper/1909.03368/) (2019)
- [Fast Distributed Inference Serving for Large Language Models](https://mlsanity.com/paper/2305.05920/) (2023)
- [DynamoLLM: Designing LLM Inference Clusters for Performance and Energy Efficiency](https://mlsanity.com/paper/2408.00741/) (2024)
- [Response Length Perception and Sequence Scheduling: An LLM-Empowered LLM Inference Pipeline](https://mlsanity.com/paper/2305.13144/) (2023)
- [Efficient LLM Scheduling by Learning to Rank](https://mlsanity.com/paper/2408.15792/) (2024)
- [S3: Increasing GPU Utilization during Generative Inference for Higher Throughput](https://mlsanity.com/paper/2306.06000/) (2023)
- [Efficient Interactive LLM Serving with Proxy Model-based Sequence Length Prediction](https://mlsanity.com/paper/2404.08509/) (2024)
- [Scheduling with Predictions and the Price of Misprediction](https://mlsanity.com/paper/1902.00732/) (2019)
- [Enabling Efficient Batch Serving for LMaaS via Generation Length Prediction](https://mlsanity.com/paper/2406.04785/) (2024)
- [SkipPredict: When to Invest in Predictions for Scheduling](https://mlsanity.com/paper/2402.03564/) (2024)

## Built on by (most cited first, 30 of 35)

- [Online Scheduling for LLM Inference with KV Cache Constraints](https://mlsanity.com/paper/2502.07115/) (2025)
- [Taming the Titans: A Survey of Efficient LLM Inference Serving](https://mlsanity.com/paper/2504.19720/) (2025)
- [Fast Inference for Augmented Large Language Models](https://mlsanity.com/paper/2410.18248/) (2024)
- [Adaptively Robust LLM Inference Optimization under Prediction Uncertainty](https://mlsanity.com/paper/2508.14544/) (2025)
- [Predicting LLM Output Length via Entropy-Guided Representations](https://mlsanity.com/paper/2602.11812/) (2026)
- [Scheduling LLM Inference with Uncertainty-Aware Output Length Predictions](https://mlsanity.com/paper/2604.00499/) (2026)
- [LLM Serving Optimization with Variable Prefill and Decode Lengths](https://mlsanity.com/paper/2508.06133/) (2025)
- [Queueing, Predictions, and LLMs: Challenges and Open Problems](https://mlsanity.com/paper/2503.07545/) (2025)
- [Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management](https://mlsanity.com/paper/2605.06472/) (2026)
- [Libra: Efficient Resource Management for Agentic RL Post-Training](https://mlsanity.com/paper/2606.03077/) (2026)
- [FlowPrefill: Decoupling Preemption from Prefill Scheduling Granularity to Mitigate Head-of-Line Blocking in LLM Serving](https://mlsanity.com/paper/2602.16603/) (2026)
- [Robust Length Prediction: A Perspective from Heavy-Tailed Prompt-Conditioned Distributions](https://mlsanity.com/paper/2604.07931/) (2026)
- [Orla: A Library for Serving LLM-Based Multi-Agent Systems](https://mlsanity.com/paper/2603.13605/) (2026)
- [SMetric: Rethink LLM Scheduling for Serving Agents with Balanced Session-centric Scheduling](https://mlsanity.com/paper/2607.08565/) (2026)
- [TimelyLLM: Segmented LLM Serving System for Time-sensitive Robotic Applications](https://mlsanity.com/paper/2412.18695/) (2024)
- [CASTILLO: Characterizing Response Length Distributions of Large Language Models](https://mlsanity.com/paper/2505.16881/) (2025)
- [Efficient Serving of LLM Applications with Probabilistic Demand Modeling](https://mlsanity.com/paper/2506.14851/) (2025)
- [Intra-request branch orchestration for efficient LLM reasoning](https://mlsanity.com/paper/2509.24957/) (2025)
- [Position: LLM Serving Needs Mathematical Optimization and Algorithmic Foundations, Not Just Heuristics](https://mlsanity.com/paper/2605.01280/) (2026)
- [PALS: Power-Aware LLM Serving for Mixture-of-Experts Models](https://mlsanity.com/paper/2605.21427/) (2026)
- [Mind the Delay: Latency-aware Query Routing for Heterogeneous LLM pools](https://mlsanity.com/paper/2607.18253/) (2026)
- [OUTLETS: Output-Length Prediction from Speculative Decoding Backbones](https://mlsanity.com/paper/2609.01068/) (2026)
- [Do Stop Me Now: Detecting Boilerplate Responses with a Single Iteration](https://mlsanity.com/paper/2510.22679/) (2025)
- [SageSched: Efficient LLM Scheduling Confronting Demand Uncertainty and Hybridity](https://mlsanity.com/paper/2603.07917/) (2026)
- [Tackling the Data-Parallel Load Balancing Bottleneck in LLM Serving: Practical Online Routing at Scale](https://mlsanity.com/paper/2605.06113/) (2026)
- [Beyond Prediction: Tail-Aware Scheduling for LLM Inference](https://mlsanity.com/paper/2606.18431/) (2026)
- [Geometry-Aware Online Scheduling for LLM Serving: From Theoretical Bound to System Practice](https://mlsanity.com/paper/2606.22327/) (2026)
- [Online Linear Programming for Multi-Objective Routing in LLM Serving](https://mlsanity.com/paper/2607.03948/) (2026)
- [NELSSA: A GPU-PNM Heterogeneous System for Mixed-Length LLM Serving via Length-based Request Placement](https://mlsanity.com/paper/2607.26633/) (2026)
- [PowerAtlas: Towards Electricity-Computing Co-Scheduling for Power Systems](https://mlsanity.com/paper/2607.26710/) (2026)

## Abstract (arXiv)

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.