LEAP: Efficient LLM Inference Scheduling with Generalizable Length Prediction

Teams running LLM inference should read this, because LEAP (the paper's length prediction system) replaces the most expensive and most fragile part of length-aware scheduling, the length predictor, with a component that costs almost nothing and corrects itself after deployment. Running short requests first clearly lowers average latency, but only if you know roughly how long each request will generate, and existing length predictors rank worse when the domain changes while costing GPU time of their own. LEAP finds that the hidden states the serving model computes during prefill already carry an output-length signal, and a single linear layer on one layer's hidden state can read it out, with no separate encoder. The linear layer runs on CPU, in parallel with the rest of prefill on the GPU, so it takes no GPU time. After each request finishes, LEAP uses recursive least squares to correct the layer's weights from that request's true length with a closed-form update, with no gradients and no replay of old data. In single-GPU experiments on ShareGPT traffic, average per-token latency is 15.9% to 40.8% lower than with several existing predictors. A length predictor should be judged by latency after scheduling and by its own cost.

What problem is it actually solving?

An inference engine can only advance a limited number of requests per decode step. In the paper's setup, up to 512 requests sit in memory but only 64 decode slots exist per step. When requests outnumber slots, something must decide who goes first. The trouble with first-come-first-served (FCFS) is that one 2,000-token answer holds a slot while several short answers wait behind it.

An illustrative example: three requests arrive together and one advances per step.

Short-first cuts the average finish time from 1067 to 450 (illustrative). But the scheduler does not know A will take 1000 steps; it has to predict. Two requirements follow. The prediction must separate 50 from 500, while guessing 53 instead of 50 hardly matters; the paper builds on this point, since scheduling needs ranking. And the prediction itself must be cheap, or the predictor eats the GPU time it was meant to save.

Existing predictors fall short on both ranking quality and cost. TIE (an uncertainty-aware length predictor) and SSJF (a small proxy-model predictor) run a separate text encoder over the prompt; LTR trains a separate ranking model that outputs only ordering scores between requests, not lengths. These three have 110 to 185 million parameters and compete with the serving model for the GPU. TRAIL (a hidden-state MLP predictor) and EGTP (which pools hidden states weighted by each token's entropy, then predicts length) reuse the serving model's hidden states but attach MLP heads with millions of parameters. Worse, a predictor trained on one dataset ranks noticeably worse on another task or language, and recovering means new data, retraining and redeployment.

The pipeline

LEAP has three core components: the linear readout, online correction and asynchronous CPU execution. The table follows the order a request meets them; countdown scheduling is the scheduler that consumes the predictions:

The division of labor: the serving model understands the request, the linear layer only converts that understanding into a length, and RLS (recursive least squares) recalibrates the conversion to the new traffic after deployment, and asynchronous CPU execution keeps all of this from slowing the GPU.

Why one hidden state and one linear layer are enough

During prefill the serving model has already read the whole prompt, and it is the very model that will write the answer. Its hidden state holds both the prompt content and information about how this model is about to respond. Running a second encoder over the prompt repeats that work.

LEAP takes the last prompt token's hidden state at layer 30 of Qwen3-8B (late but not final), standardizes it with training-set means and scales, and fits a Ridge regression. The whole predictor has 4,097 parameters, 16.39 KB in FP32, small enough to stay in one CPU core's private cache.

The target is the square root of the length, not the length itself, because the target decides which errors matter. Fitting raw length lets a few very long answers dominate the loss. Fitting log length penalizes a 100-token error on a long request so little that long requests blur together. Square root sits in between. For the same 100-token miss (illustrative):

On the log target the miss on the long request is barely penalized (0.10 against 0.69); on the square-root target it is still 1.5 against 4.1. The ablation in Table 6 shows where each target is strong. On short requests of at most 256 tokens, the log target has the lowest mean error (50.7) and raw length the highest (105.4). On long requests over 1,024 tokens, raw length is best (266.3), square root is close (279.3) and log is worst (310.4). Square root does well on both ends and has the lowest overall error (149.6), so it is the default. At prediction time the output is clipped to between 1 and √M before squaring (M is the output budget, 2,048 in the experiments, so the upper bound is about 45.3), so a very negative value cannot turn into a huge length.

Is one layer enough? In Table 7, concatenating all layers lowers mean error by only 3.3%, while the data sent to the CPU per request grows from 16 KiB to 592 KiB. So adding other layers beyond layer 30 gains little and costs much more to move. Layer 30 is the paper's pre-existing default, not selected on the test sets.

When the domain shifts: correcting from finished requests

Trained on LMSYS conversations and moved unchanged to Alpaca instructions, MATH problems and multilingual Aya, the readout still ranks well on Alpaca and MATH, but Spearman correlation on Aya falls to 0.498. When the language changes, the mapping from hidden state to length changes too.

The correction signal is free: when a request finishes, its true length is known. RLS is the incremental form of linear regression. It keeps a matrix $P$ that records how much uncertainty remains in each direction of the feature space: small where old data covers it well, large where it has rarely been seen. for each new batch it computes the correction in closed form, without gradients or old samples. The core update is one line:
$${\theta }^{+}=\theta +G(y-U\theta )$$
$\theta$ is the current readout weight, $U$ the features of the finished batch, $y$ the square roots of their true lengths, so $y-U\theta$ is the batch's prediction error, and $G$ sets how much that error moves the weights: compute $V=P{U}^{\top }$, then $G=V(I+UV{)}^{-1}$, and afterwards $P^{+}=P-G{V}^{\top }$ for the next batch. $P$ starts as $(U_s^{\top }{U}_s+\lambda D{)}^{-1}$ from the LMSYS training features $U_s$. Directions well covered by old data get a small correction and barely move; directions that only appear in new traffic get a large correction and change more.

RLS starts from the LMSYS weights, not from zero. With little feedback the model stays close to the original; with more it shifts toward the new traffic. The paper proves the batch-by-batch result equals solving the same anchored objective once with all feedback at hand, so batch size only changes timing. On Aya, RLS raises Spearman correlation from 0.498 to 0.820, a 64.7% relative gain, touching only the readout.

Does running on CPU slow the scheduler?

The risk is that the GPU is ready to decode while the length is still being computed. LEAP copies the hidden state to the CPU as soon as prefill reaches layer 30; the GPU continues with the remaining layers while the CPU does one dot product and sets the initial order. A GPU probe before the final RMSNorm found the result ready in all 144 asynchronous trials and in none of the 72 controls that started only after prefill.

Throughput is enough too: one CPU core keeps up with 10,000 predictions per second, saturates near 58,000, and keeps P95 delay under 3 ms. RLS runs in another thread, updating after every 8 finished requests, while the scheduler keeps using the last published weights.

Is serving actually faster?

Table 3 replays the same ShareGPT requests and generation trajectories at 2.65 requests per second under one scheduling rule, changing only the predictor and counting all its cost. PTL (per-token latency) is each request's latency from planned arrival to completion (including prediction, queueing and execution) divided by its actual output tokens, then averaged over requests; it is not inter-token latency.

LEAP's mean PTL is 15.9% below the strongest baseline, LTR, and 40.8% below TRAIL; throughput is 1.9% above LTR, and TTFT is nearly the same as FCFS. That is consistent with keeping prediction on the CPU and off the GPU, though Table 3 is a single run.

In the load sweep of Figure 6, all three predictors stay far below FCFS, so most of the gain over FCFS comes from length ordering itself. LEAP's remaining margin mixes more accurate ranking with keeping prediction off the GPU, and this single-run experiment does not separate the two.

Scheduling after domain transfer is simulated (Table 4, prediction cost excluded, all methods sharing the same waiting protection). With arrivals at 1.5 times the capacity measured in an FCFS simulation, RLS helps most on Aya: mean response time falls from 180.59 s with frozen LEAP to 161.18 s. On MATH at the same load, it falls from 118.34 s with TIE to 100.94 s with LEAP and 94.36 s with RLS. The gap widens with load, because longer queues make order matter more.

Better prediction does not always mean faster serving, and Table 2 gives a counterexample: refreshing predictions during decode lowers error but raises LEAP's mean response time by 8.37% to 9.42% and TRAIL's by 9.16% to 12.89%, because repeated prediction costs time. The paper also notes that the two prediction heads use different training targets, so this comparison does not isolate feature refresh alone. That is why LEAP captures features once at prefill and counts down; only when RLS publishes new weights does it rescore live requests from the cached features.

Limitations I see

First, long requests pay. Length-aware scheduling favors short work. In Table 4 every method already uses the same waiting protection (TIE's exponential aging rule, which raises a request's priority the longer it waits), yet on MATH RLS still raises P95, and LEAP, LEAP with RLS and TIE all have higher P95 than FCFS. The live engine adds a separate 128-step waiting cap. Services with SLAs on long jobs should measure tail latency separately and tune how strong this protection is.

Second, the serving setup is simple: single-GPU, eager-mode vLLM 0.8.5, no prefix caching, no chunked prefill, no multi-GPU. With chunked prefill, when the layer-30 state becomes available is less clear, and whether the CPU still finishes in time must be re-measured.

Third, the readout ignores decoding settings. Temperature and top-k change output length for the same prompt, and a prefill-only feature cannot see them. Table 3 is a single run with no variance.

How a search or recommendation team could try it

Teams serving generative search summaries or conversational recommendation usually mix short and long answers:

1. Validate the signal offline: take logged requests, run prefill with your serving model, grab a late-layer hidden state, fit a Ridge regression on square-root length, and check Spearman correlation. For reference, frozen LEAP reaches only 0.498 on Aya in the paper and still schedules better than TIE.
2. Change only the scheduler in replay: use predicted remaining length to assign decode slots, compare mean and P90 latency with FCFS, and watch long-request P95 with a waiting cap.
3. Add RLS: test transfer to traffic a week later or from another product line and see how much online correction recovers. At 4096 dims, RLS state is about 64 MiB.
4. Go live last: confirm CPU prediction finishes within your prefill time, then run a small A/B.

Verdict

The key lesson: length prediction serves ordering, so judge it by post-scheduling latency and its own cost; judging by prediction error alone picks the wrong design.