server: one preemption path, swap under --preempt-ram and recompute above it

[danielhanchen]

Summary

#184 parks a slot by copying its sequence to host RAM and putting it back; #185 parks it by dropping its cells and re-prefilling them later. This makes them one mechanism with one policy, and lets the host RAM budget choose per victim: a slot that fits under --preempt-ram is copied out and copied back, which is exact to the byte; a slot that does not fit has its cells dropped and its tokens re-prefilled when room returns, which holds nothing at all. That closes limitation 1 of #184, where above the budget nothing more was parked and the KV-full ladder ran as before.

Each slot records which way it was parked, so the restore takes the matching path: llama_state_seq_set_data_ext for a swapped slot, a replay through prompt processing for a recomputed one. Everything else is shared: one update_preemption(), one victim choice, one restore order, one set of metrics.

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

Policy

Unchanged from #184, with one change and one addition.

• Victim: never the leader (the slot with the most tokens), then the smallest of the rest, a slot parked PREEMPT_N_STARVED (3) times passed over while another candidate exists, prompt-phase slots are victims, parent/child (n_cmpl > 1) are not. Idle finished slots are purged first through try_clear_idle_slots().
• Changed: a slot that does not fit the RAM budget is no longer skipped, it is parked by recompute. The only slot still ruled out is one that neither mechanism can take: a multimodal slot over the budget, whose token list holds media placeholders that cannot be re-prefilled from the list alone. Under the budget a multimodal slot is swapped exactly as in server: preempt a slot instead of ending every conversation when the KV pool fills #184.
• Added (from server: KV preemption by recompute, a zero host RAM alternative to #184 #185, and required once recompute exists): a slot part-way through a recompute is passed over as a victim while another candidate exists, and preempt_kv_reserve() charges a resuming slot its whole remaining replay rather than one batch, so the next iteration cannot wake a second slot into cells the first has not claimed yet.
• Restore order is unchanged: most-preempted first, then longest parked, fit-first. A swapped slot's need is its cells plus its next step; a recomputed slot's need is its whole replay plus its next step.
• A parked sequence that no longer fits the pool at all is ended with the ordinary context-overflow error rather than waiting forever. A copy that fails (host RAM refusing the allocation) falls through to recompute instead of back to the KV-full ladder.

Flags

• --preempt-ram N (env LLAMA_ARG_PREEMPT_RAM) keeps its name and its 8192 MiB default, and gains a meaning at each end. -1 is no limit, so nothing is ever recomputed and this is server: preempt a slot instead of ending every conversation when the KV pool fills #184 exactly. N is the host RAM parked sequences may hold; over it, parks fall through to recompute. 0 no longer disables preemption, it means never copy, always recompute, which is server: KV preemption by recompute, a zero host RAM alternative to #184 #185 exactly.
• --preempt / --no-preempt (env LLAMA_ARG_PREEMPT), on by default, is the switch that --preempt-ram 0 used to be. It is the one incompatible change here and the README says so next to the flag.
• LLAMA_SERVER_PREEMPT_EVERY=N is kept, and parks through the same mode decision, which is what makes the exactness check below measure the hybrid rather than one half of it.

Reporting

/metrics keeps n_preempt_total, n_resume_total, requests_preempted and preempt_ram_bytes, and gains n_preempt_swap_total, n_preempt_recompute_total (which sum to n_preempt_total) and n_recompute_tokens_total. Each /slots entry keeps is_preempted and n_preempt and gains n_recompute and preempt_mode, which is "swap", "recompute", or null before the task has been parked at all.

Results

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, four concurrent streaming chats with roughly 1000-token prompts, no max_tokens, temperature 0, seed 1234. Three runs per cell. The GPU is shared, so wall time varies between identical runs and aggregate tok/s is the metric.

--preempt-ram	-c	completed	parks (swap / recompute)	recomputed tokens	aggregate tok/s	median
default (8192 MiB)	8192	12 of 12	25 (25 / 0)	0	311, 311, 302	311
150 MiB	8192	12 of 12	66 (39 / 27)	30225, 29345, 30344	308, 299, 297	299
64 MiB	8192	12 of 12	81 (36 / 45)	42388, 42795, 42407	307, 306, 310	307
0	8192	12 of 12	74 (0 / 74)	42795, 42335, 31910	317, 284, 303	303
default (8192 MiB)	16384	12 of 12	10 (10 / 0)	0	369, 407, 456	407
150 MiB	16384	12 of 12	26 (9 / 17)	32642, 32642, 12078	413, 411, 401	411
64 MiB	16384	12 of 12	22 (8 / 14)	32642, 32630, 0	393, 398, 456	398
0	16384	12 of 12	32 (0 / 32)	32642, 32642, 32740	411, 400, 377	400

12 of 12 requests completed in every cell, 96 in all, 0 errors, failed to find free space in the KV cache and Context size has been exceeded logged zero times anywhere. The four budgets are within the run-to-run spread of each other; the choice is about host RAM and exactness, not throughput.

The split is worth reading against the parked sizes rather than the counts. At -c 8192 a park that holds a real sequence is 118 to 194 MiB of host RAM for 1947 to 4083 cells, target state plus draft state. So at 150 MiB the boundary falls inside the distribution: of the 39 swaps, 15 carried a sequence of 1928 to 2185 cells at 118 to 127 MiB and the other 24 were slots parked before they had taken any cells at all, while all 27 recomputes were sequences of 2660 cells and up. At 64 MiB, and at either budget on -c 16384, every sequence is over the line and the only things copied are the empty slots. A budget only mixes where it lands inside the distribution of park sizes, and preempt_ram_bytes on /metrics is what an operator sizes it against.

A recompute is also cheaper to park than a copy, 6 to 9 ms against 120 to 350 ms, and that shows in the counts: the tighter the budget the more parks a run does, because parking is nearly free and the pool is rebalanced more often. It costs re-prefill instead, which is batched work and the reason the throughput columns barely move.

Exactness

One request at a time, LLAMA_SERVER_PREEMPT_EVERY=200, against the same request without the knob. With one request on an idle server the batch has the same shape at every step, so the pause is the only difference.

--preempt-ram	prompt	parks	mode	reference chars	got	result
default	0	50 total over both prompts	swap	13449	13449	byte-identical
default	2	(same server)	swap	29889	29889	byte-identical
0	0	55 total over both prompts	recompute	13449	17356	differs, first divergence at char 1136
0	2	(same server)	recompute	29889	31164	differs, first divergence at char 1034

So the hybrid keeps #184's guarantee wherever the budget holds and reproduces #185's divergence, to the same character, wherever it does not. 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. That is the trade the budget buys, and it is why the copy is preferred whenever there is room for it.

Cost when it does not fire

update_preemption() is a loop over the slots that returns as soon as the pool has room. One chat at a time on the same flags, -c 8192, two runs each, master (e9e0d992) against this branch:

prompt	gen tokens	master tok/s	this branch tok/s	output
0	3387	253.5, 204.9	272.0, 260.7	byte-identical to master
2	6917	263.0, 250.8	312.1, 311.0	byte-identical to master

The GPU is shared with another tenant, so the spread within a build is wider than the difference between builds; the point of the row is the last column.

Tests

tools/server/tests/unit/test_preempt.py from #184 runs unchanged except for its one test that asserted --preempt-ram 0 restores the old failure. That is precisely the semantic this change alters, so the test now uses --no-preempt and keeps its assertions. The other five are untouched.

tools/server/tests/unit/test_preempt_hybrid.py is new, four tests on the stories260K model:

• under the default budget the forced-knob run is byte-identical to the unforced one and /metrics shows n_preempt_swap_total == n_preempt_total with no recomputes and no parked RAM;
• with --preempt-ram 0 two requests that overflow the pool together both finish with their full token count, /metrics shows only recomputes, and no request is left parked. On server: preempt a slot instead of ending every conversation when the KV pool fills #184 this configuration is the old failure, both requests dying with Context size has been exceeded;
• with --preempt-ram 1 on four slots, where 1 MiB is inside the range of park sizes for that model, both counters come back non-zero, they sum to n_preempt_total, and every request finishes;
• a recomputed request's usage.prompt_tokens is still the prompt the client sent, timings.prompt_n is still its prompt and timings.prompt_ms is still its prefill, while n_recompute_tokens_total shows the work. That last pair is not a formality: on server: KV preemption by recompute, a zero host RAM alternative to #184 #185 the same request reports prompt_n 384 for a 20-token prompt and a prompt_ms larger than its whole generation time.

Ten tests, three runs on GPU and one on CPU, all pass.

Limitations

1. A recompute is not exact, and the budget is what decides how often one happens. An operator who needs server: preempt a slot instead of ending every conversation when the KV pool fills #184's guarantee has to set --preempt-ram -1 and accept the host RAM, or size the budget above the largest park they expect. preempt_ram_bytes and the two new counters are there to size it.
2. The recomputed work is deliberately kept out of everything the client reads: usage.prompt_tokens, timings.prompt_n and timings.prompt_ms describe the request as it was made, and the work is reported on /slots as n_recompute and on /metrics as n_recompute_tokens_total. A request that was recomputed therefore looks, in its own timings, like one that was not, apart from the wall clock.
3. Multimodal slots can only be swapped. Over the budget they are not candidates at all, so a pool full of multimodal requests over the budget still fails the old way.
4. A parked slot still cannot evict anyone to make room for itself, and a single chat that outgrows the whole pool on its own still gets Context size has been exceeded. Both are inherited from server: preempt a slot instead of ending every conversation when the KV pool fills #184. A sequence that grew to within PREEMPT_N_MARGIN + 1 + n_spec cells of the whole pool while parked is ended with that same error rather than left waiting; it was within a dozen tokens of the limit in any case.
5. Measured on one model, a dense 4B with MTP drafting. SWA, recurrent and hybrid memory take different branches inside the state save and restore, and their interaction with context checkpoints is unverified on either path.

Relation to #184 and #185

This is #184 with #185 built into it as the path above the budget, on #184's branch. --preempt-ram -1 is #184 exactly and --preempt-ram 0 is #185 exactly, both reproduced here: byte-identical output over 50 forced parks at one end, the same two divergence points as #185 measured at the other. #185 was opened as a draft proposing exactly this, so it can be closed in favour of this change.