In on-policy distillation (OPD), a small student writes its own answers and a larger teacher scores every token of them. The student learns from those scores. On the serving side, almost all of the teacher's work is one kind of request: here is a context, here is a response, tell me what you would have predicted at each response position.
Those requests overlap a lot. A training step usually samples several responses per prompt, and agent or document tasks carry thousands of tokens of shared context. vLLM has a prefix cache for exactly this kind of overlap, yet scoring requests didn't use it. I ran into that while wiring Vime's distillation path to vLLM's per-position scoring, and the reason turned out to be a good one.
The request sends the full input, context followed by the student's response, and asks for scores starting at the last context position:
SamplingParams(
max_tokens=1,
prompt_logprob_start=P - 1, # P = number of context tokens
prompt_logprob_token_ids=[[t] for t in response_token_ids],
skip_reading_prefix_cache=False,
)Call that start position S and the input length N. Row j of the result comes from the model's output at position S + j. Position S is the last context token. Its output distribution predicts the first response token, so it is the first position whose output we need.
A transformer forward pass produces two different things at each position: the keys and values that later positions attend to, and the hidden state that the LM head turns into logits. The prefix cache stores only the first. That's all generation needs, because generation only ever asks for the logits of the last position.
Scoring breaks that assumption. If the cache covers position S or anything after it, the engine skips computing those positions, the KV for them is restored from cache, and their logits are simply never produced. The request comes back with rows missing. So scoring requests skipped cache reads by default, and turning reads on could return an incomplete result without an error.
Everything before S is still safe to reuse. Positions S and later attend to the cached prefix through its KV exactly as they would to freshly computed KV, so their outputs don't change. The fix is a cap on the cache lookup:
max_cache_hit_length = min(S, N - 1)
cached_prefix H <= STwo details fall out of this. The cache works in blocks, so the usable hit rounds down to a block boundary: with 16-token blocks and S = 2047, the hit is 2032 tokens and the last 15 context tokens are recomputed. That costs almost nothing. And the cap holds even when an earlier request with the same response is fully cached, response included. Those response positions still have to run, because their KV is cached but their logits are not.
The rule is one line. Keeping it true is the harder part, because a request doesn't run once from start to finish. Chunked prefill splits a long prompt across engine steps, and preemption can evict a half-finished request and bring it back later, at which point it looks up the cache again.
So the cap is resolved once, at admission, and stored on the request without touching the caller's sampling parameters. The KV cache manager applies it on every lookup, including the one after preemption, and the worker assembles scores by absolute position and rejects an incomplete buffer instead of returning it. Block allocation, reference counting and scheduling order don't change. Tensor parallelism needs no special handling: the scheduler makes the cache decision and every rank executes the same plan.
I kept the first version narrow on purpose: text inputs, one full-attention KV group, unquantized KV, standard tensor parallelism. Everything else falls back to full recomputation, which is exactly the old behavior. It's easier to widen a correct rule than to fix a fast wrong one.
Before measuring, it's worth predicting. Say each prompt has P context tokens and K responses of R tokens each, all scored by the same teacher. Without reuse, every request prefills the whole thing. With perfect reuse, the context is computed once and only the responses are computed per request. Counting prefilled tokens, and ignoring that attention cost grows with position:
speedup <= K (P + R) / (P + K R)The two workloads I tested sit at opposite ends of this formula. Long documents had about 3,000 to 4,000 context tokens and short responses of 13 to 58 tokens. With K = 4 and typical values, the bound is about 3.9×. Math prompts were the reverse: 140 to 210 context tokens and responses of 437 to 1,024 tokens. There the bound is about 1.17×. Reuse can only remove the part that is shared, and in math the shared part is small.
Measured with Qwen3-32B at TP=4 on 4×A100 80GB, four responses per prompt, each group starting from an empty cache:
| Workload | Token budget | Uncached | Bounded reuse | Speedup | Simple bound |
|---|---|---|---|---|---|
| Long documents | 512 | 6.13 s | 1.60 s | 3.84× | ≈3.9× |
| Long documents | 2048 | 5.16 s | 1.39 s | 3.71× | |
| Short math prompts | 512 | 1.88 s | 1.66 s | 1.13× | ≈1.17× |
| Short math prompts | 2048 | 1.61 s | 1.42 s | 1.14× |
The measurements land just under the bound on both workloads. That is the result I care about more than the 3.8× itself: it says the implementation is close to removing all of the shared work there is to remove. Most of the remaining gap is likely the fixed cost every request still pays, including the one generated token that the current request lifecycle requires, plus the extra attention cost of response tokens reading a long cached context.
Reuse changes the shapes that kernels run on, and in BF16 that changes results slightly. On the 32B runs the largest logprob difference from uncached scoring was 0.50, on the same order as the 0.37 between scoring one request alone and in a batch with no cache at all. In batch-invariant mode, where kernels are forced to reduce in a fixed order, uncached, reused and preempted paths matched exactly on a small model at TP=2 and TP=4. So the difference comes from execution shape, not from the reuse logic. Whether noise of that size matters for training is a question for loss and gradient comparisons, not something to assume either way.
The teachers people most want to distill from are long-context models, and their caches don't look like the simple case. A sliding-window layer doesn't keep the whole prefix, so a cache hit has to be consistent across layers that remember different spans. Hybrid models carry several KV groups that all have to agree on one hit length. Quantized KV raises the numerics question again. The boundary rule stays the same in all of them. What changes is how much of the cache machinery has to understand it. After that comes reusing the scores themselves, not just the context, which is a different kind of cache: one keyed on outputs rather than on attention state.