Research map / Large language models / Reasoning and chain of thought
DON’T STOP ME NOW: EMBEDDING BASED SCHEDULING FOR LLMS
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TRAIL, which probes LLM intermediate-layer embeddings with a lightweight classifier to predict remaining output length each iteration and uses an SRPT variant with limited preemption in vLLM.
Problem
Size-based scheduling (SRPT) for LLM serving needs output-length predictions, but accurate predictors are costly, and preemption adds KV-cache memory overhead, while FCFS causes head-of-line blocking.
Approach
Recycle the target LLM's own intermediate-layer embeddings (layer 11 of Llama3-8B) with a lightweight MLP probe, refined per token by Bayesian smoothing, to predict remaining length; preemption is allowed only for the first C·r iterations of each request.
Key results
- Refined layer-embedding predictions achieve 2.66x lower MAE than BERT predictions from prompts.
- TRAIL achieves 1.66x to 2.01x lower mean latency than vLLM-FCFS on Alpaca with LLaMA3-8B on one A100.
- TRAIL achieves 1.76x to 24.07x lower mean time to first token compared to the state-of-the-art serving system.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use refined layer-11 embedding predictions instead of BERT prompt predictions | Trail length prediction · Alpaca | mean absolute error: 2.66x lower mean absolute error | Helped |
| Use Trail with embedding predictions and limited-preemption scheduling instead of vLLM | Llama3-8b-instruct serving · Alpaca | mean latency: 1.66x to 2.01x lower mean latency | Helped |
| Use Trail with embedding predictions and limited-preemption scheduling instead of vLLM | Llama3-8b-instruct serving · Alpaca | mean time to first token: 1.76x to 24.07x lower mean time to the first token | Helped |
| Add limited preemption rather than using vLLM scheduling | LLM serving with Trail · Alpaca | mean latency | Helped |
| Replace vLLM-FCFS with SJF based on BERT prompt predictions | vLLM · Alpaca | mean latency | No clear effect |
| Remove the limit on preemption (c=1) (ablation) | Trail · Alpaca, request rate 14 | mean latency | Hurt |
| Remove the limit on preemption (c=1) (ablation) | Trail · Alpaca, request rate 14 | mean TTFT | Hurt |
| Set the preemption threshold to c=0.2 (ablation) | Trail · Alpaca, request rate 14 | mean latency | Hurt |
| Set the preemption threshold to c=0.2 (ablation) | Trail · Alpaca, request rate 14 | mean TTFT | Hurt |
| Set the preemption threshold to c=0.8 rather than c=0.5 (ablation) | Trail · Alpaca, request rate 14 | mean latency | No clear effect |
| Set the preemption threshold to c=0.8 rather than c=0.5 (ablation) | Trail · Alpaca burst workload | mean latency | No clear effect |
Surprising findings
- SJF based on BERT predictions provides minimal benefit over vLLM-FCFS because vLLM prioritizes incoming requests over running ones.
Limitations
- Evaluated only on the Alpaca dataset and a single model/layer choice (layer 11).
- Preemption cap constant c must be tuned and depends on memory and workload.
- No dynamic retraining to handle data drift.
Experimental setup
Datasets: Alpaca
Benchmarks: Alpaca
Metrics: mean absolute error mean latency median latency time to first token
Baselines: vLLM-FCFS vLLM-SJF_BERT TRAIL-BERT FastServe
Builds on: vLLM Orca SPRPT SOAP S3
Closest papers by meaning
- Efficient LLM Scheduling by Learning to Rank (NeurIPS 2024 · similarity 0.90)
- Preble: Efficient Distributed Prompt Scheduling for LLM Serving (ICLR 2025 · similarity 0.89)
- Fast Inference for Augmented Large Language Models (NeurIPS 2025 · similarity 0.88)
- MAPS: Memory-Aware Predictive Scheduling Framework for Large Language Model Serving (ICML 2026 · similarity 0.87)
- Geometry-Aware Online Scheduling for LLM Serving: From Theoretical Bound to System Practice (NeurIPS 2026 · similarity 0.87)
- Beyond Prediction: Tail-Aware Scheduling for LLM Inference (ICML 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.