E4.5 · LESSON 5 OF 6

Preemption

Trace static and iteration-level schedules and explain chunked prefill, preemption and latency tradeoffs.

THE IDEA

When memory runs out: preemption

If memory is short, a scheduler can preempt a request: pause it and take its cache memory for others. To resume later it must either recompute the cache (prefill the tokens again) or swap it to slower memory and back. Both cost time, and the cost grows with the number of cached tokens, so a good policy considers it when picking the victim.

cached tokens of the victim1,000link bandwidth (GB/s)25

option	how	cost (ms)
drop and recompute	prefill the tokens again at 0.1 ms each	100
swap out and back	131.1 MB each way at 25 GB/s	10.49

cheaper here: swap Either way the resumed request pays; eviction is never free, and the policy must include that cost.

With the toy numbers swapping is cheaper. Slow the link and see the choice flip.

Predict first. Is evicting a cache free if the sequence later resumes?

Chapter context

Requests arrive at different times and finish after different output lengths. Scheduling decides which work enters each iteration and how waiting prompts compete with ongoing generation.

TRY IT YOURSELF

Part 1. Do it hands-on

Cost the two ways to resume an evicted sequence, then write the cost functions and a victim rule.

cached tokens of the victim1,000link bandwidth (GB/s)25

option	how	cost (ms)
drop and recompute	prefill the tokens again at 0.1 ms each	100
swap out and back	131.1 MB each way at 25 GB/s	10.49

cheaper here: swap Either way the resumed request pays; eviction is never free, and the policy must include that cost.

recompute 1000 tokens (ms) one-way swap (ms, 3 decimals) cheaper (recompute or swap)

To do: Memory is short, so a sequence with 1000 cached tokens is evicted and its cache dropped. To resume it must recompute those tokens as prefill at 0.1 ms per token. ms?
To do: Alternative: copy its cache to CPU memory and back. 1000 tokens × 131,072 bytes = 131,072,000 bytes over a 25 GB/s link, one way. ms (3 decimals)?
To do: Out and back is two transfers (10.49 ms) against 100 ms of recomputation. With these numbers which is cheaper: recompute or swap? Type the word

Part 2. Do it in code

recompute_ms(tokens, ms_per_token=0.1): the prefill time to rebuild an evicted cache.
swap_ms(tokens, per_token_bytes, bandwidth_bytes_per_s, round_trip=True): time to copy the cache out (and back, if round_trip) in milliseconds. bandwidth must be positive (ValueError).
cheaper(tokens, per_token_bytes, bandwidth_bytes_per_s): "swap" if swap_ms (round trip) is less than recompute_ms, otherwise "recompute".
pick_victim(seqs): seqs maps a name to its cached token count. Evict the sequence that is cheapest to rebuild: the smallest count (the alphabetically first name on ties). ValueError if empty.
eviction_overhead(evictions): the total ms spent on a list of (tokens, strategy) events where strategy is "recompute" or "swap" (swap uses 131072 bytes per token over 25e9 bytes/s, round trip).

def recompute_ms(tokens, ms_per_token=0.1):
    pass

def swap_ms(tokens, per_token_bytes, bandwidth_bytes_per_s, round_trip=True):
    pass

def cheaper(tokens, per_token_bytes, bandwidth_bytes_per_s):
    pass

def pick_victim(seqs):
    pass

def eviction_overhead(evictions):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Is evicting a cache free if the sequence later resumes?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

In a simplified equal-cost decode schedule, utilization = active sequence-iterations / available sequence-iterations. This omits prefill costs and assumes uniform iteration time. A real scheduler needs measured compute, memory and latency constraints.

Worked example

Two slots serve requests A,B,C with output lengths 2,1,2. Static groups run A+B for two iterations, then C alone for two: 5 active slots out of 8 available. Continuous admission runs A+B, then A+C, then C: 5 active slots out of 6. Utilization rises from 62.5% to about 83.33% in this toy schedule; real token costs are not equal.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
static=[['A','B'],['A',None],['C',None],['C',None]]
continuous=[['A','B'],['A','C'],['C',None]]
util=lambda rows: sum(x is not None for row in rows for x in row)/(2*len(rows))
assert util(static) == .625
assert isclose(util(continuous),5/6)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Continuous batching does not guarantee a fixed improvement. Chunked prefill changes scheduling granularity, not the model's context dependencies. Include cache movement or recomputation costs when comparing preemption policies.