server: KV preemption by recompute, a zero host RAM alternative to #184

[danielhanchen]

Summary

A second server-side design for the same problem as #184, built from scratch so the two can be compared on the same hardware: when the unified KV pool cannot fit the next decode, park a slot by dropping its cells and re-prefilling its tokens when room returns, instead of copying its state to host RAM. This is vLLM's default preemption mode. It holds zero bytes while parked; the price is that a resumed answer may differ from the uninterrupted one, which is measured below rather than assumed.

Everything is in tools/server/server-context.cpp, inside server_slot and update_slots(), tagged [TAG_PREEMPT].

Mechanism

• update_preemption() runs in update_slots() immediately before pre_decode(), the same place as server: preempt a slot instead of ending every conversation when the KV pool fills #184: every slot is at a token boundary, prompt.tokens matches the cache, no draft is in flight.
• A victim's cells go with prompt_clear(), which is seq_rm on a common_memory holding both the target and the draft context, so the MTP draft sequence goes with them. spec_draft, spec_i_batch and spec_ckpt are dropped because a draft is a prediction, not a result. Nothing is copied out.
• Resume re-enters prompt processing with a replay list of prompt.tokens + sampled, the token the slot had sampled but not yet decoded. Re-prefilling exactly those puts the logits where the interrupted step was about to read them, and the DONE_PROMPT transition samples the token that step would have sampled. A victim still in its prompt phase needs no replay list; it starts its own prompt over.
• The second pass through prompt processing is guarded: update_prompt_start() only when t_start == 0, no second is_begin or progress chunk into a live stream, no stats.n_gen = 0 and no init_sampler() at DONE_PROMPT, so penalties, grammar and the RNG carry on from where they were.
• The MTP draft context is a pure function of the target prefill, so the re-prefill re-establishes it as the original did; speculation restarts through the existing common_speculative_begin in the DONE_PROMPT to GENERATING transition.
• usage.prompt_tokens stays the client's prompt. timings.prompt_n accumulates the recompute, so prompt_per_second is not meaningful for a recomputed request; /slots reports n_recompute for the clean split.

Policy

Identical to #184 so the comparison is fair: idle finished slots are purged first, the slot with the most tokens is never preempted, the smallest of the rest is parked, a slot parked PREEMPT_N_STARVED (3) times is passed over while another candidate exists, restore is most-preempted first then longest parked with fit-first, slots still processing their prompt are victims, parent/child slots are not. Two additions RECOMPUTE forces: multimodal slots are excluded, since their token list holds placeholders rather than chunks, and a resuming slot is passed over as a victim while another candidate exists. preempt_kv_reserve() charges a resuming slot its whole remaining replay, not one batch, otherwise the next iteration wakes a second slot into cells the first has not claimed yet. A parked sequence that no longer fits the pool at all is ended with the ordinary context-overflow error rather than waiting forever.

--preempt / --no-preempt (env LLAMA_ARG_PREEMPT), on by default, active only with --kv-unified and more than one slot. /metrics gains n_preempt_total, n_resume_total, n_recompute_tokens_total and requests_preempted; /slots entries gain is_preempted, n_preempt and n_recompute.

Exactness

Qwen3.5-4B UD-Q4_K_XL with the embedded MTP head, --parallel 4 --kv-unified --spec-type draft-mtp --spec-draft-n-max 2 --flash-attn on, one request at a time, temperature 0, seed 1234. The unforced path is byte-identical to master on both prompts, twice each. Forced parks (LLAMA_SERVER_PREEMPT_EVERY) are not:

prompt	forced parks	reference chars	recomputed chars	identical	first divergence
0	21, every 200 tokens	13449	17356	no	char 1136
2	34, every 200 tokens	29889	31164	no	char 1034
0	1, at generated token 3002	13449	13404	no	char 11875
2	2, at 3000 and 6001	29889	29877	strict prefix	at the end

The single-park row is the clean one: 11875 characters at about 3.97 characters per token is generated token 2991, against a park at 3002. The divergence lands at the resume. Prefill and decode are different kernels, the recomputed K/V differ in the last bits, and a greedy argmax flips at the first near tie after the resume. A park can also cost nothing, as the strict-prefix row shows. #184 is byte-identical over 160 forced cycles on the 4B and the 35B; RECOMPUTE cannot make that promise and does not.

Cost

Both builds from the same base (e9e0d992), run back to back on the same GPU in one session, three runs per cell, four concurrent chats with roughly 1000-token prompts, no max_tokens. Wall time is not comparable across builds because the answers differ, so aggregate tok/s is the metric.

build	-c	aggregate tok/s, three runs	median	preemptions	recomputed tokens	parked host RAM
this branch	8192	336, 343, 336	336	9, 8, 9	28535, 24940, 28211	0
#184	8192	326, 348, 347	347	9, 7, 8	0	up to 194 MiB per park
this branch	16384	207, 440, 393	393	4, 3, 4	21388, 13537, 20132	0
#184	16384	355, 443, 390	390	1, 1, 4	0	up to 234 MiB per park

12 of 12 completions on both builds, 0 errors, failed to find free space logged zero times. #184 reproduced its own PR row here (17803 generated tokens in 51.2 s, 8 preemptions, against 17803 / 50.5 s / 8 in its body), which validates the harness. RECOMPUTE costs about 3 percent of throughput at 8192 and is at parity at 16384 while re-running through the model more tokens than the whole load generated, because a recompute is batched prefill rather than decode.

Solo cost when preemption never fires, -c 8192, one chat, two runs each: master 264.0 / 236.9 and 289.8 / 312.4 tok/s on prompts 0 and 2; this branch 233.3 / 255.4 and 299.3 / 312.7. Output byte-identical to master.

Tests

tools/server/tests/unit/test_preempt_recompute.py, seven tests on the two-slot unified pool with the stories260K model: forced-park identity (which does hold on that tiny model, so the 4B measurement above is the one that matters), two requests that fit alone but not together both finish, two prompts that do not fit together both finish, a generating slot beside a prefilling one, --no-preempt restores the old failure, /metrics and /slots report the parked state, and usage does not double count the recompute. All pass on this branch; the first two fail on master, the second with Context size has been exceeded on both requests.

When to use which

SWAP (#184) where host RAM allows: it is exact and slightly faster. RECOMPUTE where host RAM is the binding constraint or the parked volume is large: at about 50 KiB per token on this 4B, four 8k chats park about 0.6 GiB, and a 128k-context deployment with 16 slots would pin tens of GiB and pay 300 to 420 ms of copy per restore into a full pool. The right shape is both, SWAP under --preempt-ram and RECOMPUTE as the fallback above it, which is limitation 1 in #184. This PR is opened as a draft so that fallback can be built on it once #184 lands, or so it can stand alone where exactness matters less than memory.