Inference Systems · Accelerator Scheduling
Preemption by Swap Versus Recomputation
When the cache pool cannot serve a running sequence, the scheduler must take memory back; the two recovery strategies have different costs that depend on sequence length.
Preemption is the scheduler’s admission of defeat. A sequence has been admitted, has consumed device memory, has made progress, and now the pool cannot satisfy its next allocation. Something must give: either the sequence stops holding memory, or the system stops admitting new work indefinitely, which is not a stable state when the queue is non-empty. The two ways to make a sequence stop holding memory are to move its cache somewhere else or to throw it away and rebuild it later. The costs of those two are not close, and which one is cheaper depends almost entirely on how long the sequence is.
The two recovery paths
Swap. The sequence’s key-value blocks are copied from device memory to host memory, and the block table is invalidated so the blocks return to the pool. When the sequence is resumed, the blocks are copied back and the table is restored. Generation resumes exactly where it stopped, because the cache is intact; nothing needs recomputing.
Recompute. The sequence’s blocks are simply returned to the pool. When the sequence is resumed, it re-enters the prefill path with its prompt plus every token generated so far, and the cache is rebuilt by running the model over that whole prefix. Generation resumes from the same point semantically, but at the cost of a prefill over $T$ tokens.
A third variant exists and is a hybrid: swap only the prefix that would be expensive to recompute, and discard the rest. It appears in systems that want to bound both the transfer volume and the recompute volume, and it is more complex than either pure strategy.
The cost curves
Swapping costs bytes moved. For a sequence of $T$ tokens at $c$ bytes per token, the eviction moves $Tc$ bytes out and the resumption moves $Tc$ bytes back, so the round trip is $2Tc$ bytes over a link whose bandwidth is $W$. The time is $2Tc/W$, linear in length. The PCIe link is roughly three orders of magnitude slower than device memory bandwidth, so even a modest sequence takes a measurable time to move.
Recomputation costs a prefill. Prefill over $T$ tokens does $T$ tokens of work at high arithmetic intensity, so its time is approximately $2NT$ operations divided by the device’s throughput, plus the attention term that grows with $T^2$ within the batch. That is also linear in $T$ for the weight-bound part, but the constant is different and, importantly, the work is done on the device rather than on a link.
The comparison therefore turns on the ratio between link bandwidth and device throughput, and on the sequence length. Two regimes exist. For short sequences the transfer is small in absolute terms and the recompute is cheap; both are tolerable. For long sequences the transfer volume becomes large and the recompute becomes a long prefill. The break-even sits where $2Tc/W$ equals the prefill time for $T$ tokens, and since both are linear in $T$, the comparison is largely independent of length in the linear regime. What changes with length is the quadratic attention term, which grows faster for recompute and shifts the balance toward swapping as sequences get longer.
Why recompute is often chosen anyway
Recomputation has properties that make it attractive beyond the raw time comparison.
It frees memory instantly. Swapping requires the blocks to be readable and the copy to complete before the pool can reuse them, unless the copy is asynchronous and the blocks are held until it lands. Recomputation frees blocks in one operation, which matters when the pool is critically low and the scheduler needs room in the current step.
It needs no host-side staging area. A swap policy requires host memory to hold evicted caches, and host memory is finite. With a large pool and a large batch, the evicted set can be hundreds of gigabytes, which forces the system to also decide which evicted sequences to drop entirely. Recompute has no such second-order capacity problem.
It interacts correctly with prefix sharing. A recomputed sequence can re-derive its prefix from a shared prefix cache if one exists, which makes the rebuild cheaper than a full prefill and, in the limit, close to free for the shared portion. A swapped sequence gains nothing from sharing.
The costs are real too. Recomputation consumes device time that could have generated tokens for other sequences, so a recompute wave is visible as a latency spike across the whole batch. It also invalidates the assumption that a sequence’s token count monotonically increases: the recomputed prefill may produce a slightly different cache if the model is nondeterministic, and any sampling state must be preserved separately.
Selecting victims
Given that some sequence must be preempted, the choice of victim is a policy decision with measurable consequences. The objective is to minimise the total cost of the preemption plus the lost progress, and the candidates differ on both.
Evicting the longest sequence frees the most memory per eviction, which is what a critically full pool needs. It also has the highest per-eviction cost, since the transfer or recompute scales with length. Evicting the shortest sequence is cheap per eviction and frees little, so it may need to be done repeatedly.
A policy that evicts the most recently admitted sequence exploits locality: the newest sequence has the least progress to lose and, if recomputation is used, the shortest prefix to rebuild. This is effectively last-in-first-out, and it has the property of preserving the sequences that have been running longest, which are the ones closest to finishing.
A policy that evicts the sequence with the furthest remaining distance to its maximum length minimises wasted work, since a sequence that was going to stop soon anyway is a poor victim.
The choice also has a fairness dimension. Evicting the same tenant’s sequences repeatedly under contention is a policy decision, not an accident, and the scheduler should make it deliberately. Per-tenant limits, when they exist, are what prevent one tenant’s long sequences from being evicted every time another tenant’s burst arrives.
Sequencing the preemption
The order of operations matters when the pool is nearly empty. A preemption triggered by a failed allocation in the middle of a step has to complete before the step’s kernels launch, or the sequence will be included in the batch with an invalid table. That makes the eviction synchronous with respect to the step, which is why recomputation is often preferred: the copy-out for a swap is the expensive part, and it sits directly on the critical path.
A scheduler can avoid this by preempting ahead of need, using the watermark from admission control and the memory watermark as the trigger. If the pool’s free count falls below the watermark plus the next step’s growth, a victim is chosen and preempted while the current step is still running on the device, so the recovery overlaps with useful work. This is the same structure as a write-back cache evicting before it is forced to.
What resumption costs the scheduler
A preempted sequence is not simply returned to the waiting queue at its original priority. It has a claim on the memory it lost, and if it is placed at the tail it may be starved while later arrivals proceed. Placing it at the head gives it priority but risks evicting it again immediately if the pool is still tight, which produces a livelock: the sequence is admitted, evicted, admitted, evicted, and never advances.
The standard guard against livelock is to admit a preempted sequence only when its full reservation can be satisfied, not merely its next block. A sequence that cannot get its reservation waits without consuming device time, and the pool is left to make progress on sequences that fit. The cost is that a large preempted sequence may wait a long time, which is why some systems instead shrink the reservation by truncating the sequence’s context, discarding the oldest tokens. Truncation changes the model’s output, so it is only acceptable where the application allows it.
The edge cases
Zero-length resume. A sequence preempted before generating any token has no generated prefix to recompute beyond its prompt. Its recompute cost is exactly a prompt prefill, which is what it would have paid anyway, so recomputation is free in the sense that no progress is lost.
Swapping a shared prefix. If two sequences share blocks, a swap of one must not move blocks the other still references. Reference counting handles the bookkeeping, but the copy is per-sequence logical view, so the transfer volume may be larger than the sequence’s private residency.
Preemption during chunked prefill. A sequence partway through prompt processing holds a partially written cache. Recomputing it restarts the prompt from the beginning, which wastes the chunks already processed; swapping it preserves partial progress but requires the table to record how many positions are valid.
Device-to-device preemption. In a tensor-parallel deployment, a sequence’s cache is spread across every device in the group. Swapping means every device copies its shard out, and the eviction is only complete when the slowest one finishes. Recomputation likewise runs across all devices. The cost model is unchanged in shape but the constant includes the group’s slowest member, which is why tensor-parallel residency and interconnect pressure changes the arithmetic.
The choice between the two strategies is ultimately a statement about which resource the system is short of. Swapping spends interconnect bandwidth and host memory to preserve device work. Recompute spends device work to preserve nothing. A system with idle link capacity and a full pool should swap; a system with a saturated link and idle matrix units should recompute. Neither is universally right, and a scheduler that hard-codes one has decided in advance which resource will be scarce.
References
- Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP 2023
- Operating Systems: Three Easy Pieces — Arpaci-Dusseau & Arpaci-Dusseau, Chapter on Swapping
- Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al., OSDI 2022