31. Engine v2: one core for every request

Five good pieces that don’t fit together

Look at what Parts IV-VI gave you. Chapter 19’s FastDecoder removes launch overhead, but serves one request. Chapter 24’s ContinuousBatchingEngine serves many, but gives each request a fixed-size row of a static cache and prefills one request’s chunk per forward call. Chapter 25’s PagedKVCache packs caches without waste and shares prefixes, but nothing schedules requests onto it. Chapter 26’s speculative decoding handles one request with a contiguous cache. Each piece passed its tests. None of them can be combined with the others without rewriting the seams between them.

That’s the situation every serving engine reaches. vLLM’s first version grew one feature at a time and was rewritten in 2025 as “V1” around a simpler core; SGLang, TensorRT-LLM and LMDeploy converged on similar designs. This chapter builds that core. The next twelve chapters plug the rest of a production server into it: fast kernels, CUDA graphs, a full sampler, constrained decoding, a tokenizer, an HTTP API, speculative decoding, real checkpoint formats, other hardware, several GPUs and more model families.

The design rests on three decisions:

1. One kind of work. A request’s state is its token list (prompt, then output) and one number, num_computed_tokens: how many of those tokens already have keys and values in the cache. The scheduler’s only output is a list of (request, n): compute the next n tokens of this request.
2. One batch layout. All scheduled tokens of all requests are concatenated into one flat sequence of $N$ tokens. Matmuls see [N, D]; attention gets a small metadata object that says where each request’s rows start and where its blocks live.
3. One memory pool. Every layer’s K and V live in a pool of fixed-size blocks shared by all requests, with reference counts and a prefix cache keyed by hash chains.

One kind of work

With num_computed_tokens as the only state, every situation an engine meets is a value of n:

The invariant that makes this work: the newest token is always uncomputed. After a decode step samples token $t$, it’s appended to the list, so num_tokens = num_computed_tokens + 1 and the next step feeds exactly that token. A request needs a sample from the model exactly when a step’s tokens reach the end of its list, which is how the batch builder knows which rows to send to the LM head.

@dataclass(eq=False)
class Request:
    request_id: str
    prompt_token_ids: list
    params: SamplingParams = field(default_factory=SamplingParams)
    eos_token_id: int | None = None
    priority: int = 0                 # lower runs first under the "priority" policy (vLLM's convention)
    arrival_time: float = field(default_factory=time.monotonic)
    cache_key: tuple = ()             # anything else that changes activations: LoRA, images, tenant salt
    status: Status = Status.WAITING
    num_computed_tokens: int = 0      # tokens whose K/V are in the cache
    num_cached_tokens: int = 0        # of those, how many came from the prefix cache at admission
    spec_token_ids: list = field(default_factory=list)   # draft tokens to verify this step (Chapter 37)
    swapped_out: bool = False         # KV lives in host memory (swap preemption)
    num_preemptions: int = 0
    first_token_time: float | None = None
    finish_time: float | None = None
    stop_reason: int | str | None = None
    extra: dict = field(default_factory=dict)            # per-feature state: guided FSM, LoRA slot, images

    def __post_init__(self):
        if not self.prompt_token_ids:
            raise ValueError(f"{self.request_id}: empty prompt")
        self.token_ids = list(self.prompt_token_ids)     # prompt then output, one list
        self.num_prompt_tokens = len(self.token_ids)
        self.block_hashes = []                           # filled by the block manager
        self.num_placeholders = 0                        # trailing PLACEHOLDER tokens (async scheduling)

    @property
    def output_token_ids(self):
        return self.token_ids[self.num_prompt_tokens:]

    @property
    def num_tokens(self):
        return len(self.token_ids)

    @property
    def num_output_tokens(self):
        return len(self.token_ids) - self.num_prompt_tokens

    @property
    def num_tokens_with_spec(self):
        return len(self.token_ids) + len(self.spec_token_ids)

    def append(self, token_id):
        self.token_ids.append(int(token_id))

num_cached_tokens and num_preemptions are statistics; cache_key and extra are hooks for later chapters (LoRA adapters, images and per-tenant salts change activations, so they must be part of the prefix-cache key). Status distinguishes waiting, running, preempted and the three ways to finish.

Flattened ragged batches

Chapter 24’s engine ran one forward call per prefill chunk, plus one for all decodes. With twenty requests in flight that’s up to twenty small forward calls per step, each reading every weight from memory. The fix is to give the model all of a step’s tokens at once, without padding:

request:          a (chunk of 3)    b (decode)   c (decode)   d (new, 4 tokens)
input_ids:        a6 a7 a8          b41          c9           d0 d1 d2 d3
positions:        6  7  8           41           9            0  1  2  3
query_start_loc:  0                 3            4            5             9
seq_lens:         9                 42           10           4

The linear layers, norms and MLPs don’t care which request a row belongs to: they process [9, D] as one matrix, so the weights are read once for the whole step and the matmul is as large as the batch allows. Only attention needs the structure, and it gets it from the metadata:

@dataclass
class BatchMeta:
    query_start_loc: torch.Tensor
    seq_lens: torch.Tensor
    block_table: torch.Tensor
    slot_mapping: torch.Tensor
    block_size: int
    query_start_loc_cpu: list            # host copies: kernels launch without a device sync
    seq_lens_cpu: list
    max_query_len: int
    max_seq_len: int

    @property
    def num_reqs(self):
        return len(self.seq_lens_cpu)

    @property
    def decode_only(self):
        return self.max_query_len == 1

Two pieces connect tokens to the pool. The block table of request $r$ lists its physical blocks in logical order, exactly as in Chapter 25. The slot mapping gives, for each of the $N$ new tokens, the flat pool slot where its K and V go: position $p$ of a request lives in block table[p // block_size] at offset p % block_size, so its slot is table[p // block_size] * block_size + p % block_size. Writing a step’s keys is then one scatter over $N$ slots, whatever mix of requests the step contains.

def build_batch(scheduled, block_tables, block_size, device="cpu", pad_to=None):
    """Lay out [(request, n), ...] as one flattened batch.  (Your engine: Chapter 31)

    block_tables[request_id] lists the request's physical blocks in logical order. Request r's
    n new tokens are token_ids[c : c + n] (plus its draft tokens, if any), at positions
    c .. c + n - 1, where c is its num_computed_tokens. Position p lives in slot
    table[p // block_size] * block_size + p % block_size.

    pad_to (CUDA graphs, Chapter 33) appends dummy decode rows with slot -1 and length 0.
    """
    ids, positions, slots, starts, lengths, tables, logits_rows, counts = [], [], [], [0], [], [], [], []
    prompt_rows, prompt_spans = [], []
    for request, n in scheduled:
        c = request.num_computed_tokens
        tokens = (request.token_ids + request.spec_token_ids)[c:c + n]
        if len(tokens) != n:
            raise ValueError(f"{request.request_id}: scheduled {n} tokens but only {len(tokens)} exist")
        table = block_tables[request.request_id]
        ids.extend(tokens)
        positions.extend(range(c, c + n))
        slots.extend(table[p // block_size] * block_size + p % block_size for p in range(c, c + n))
        starts.append(starts[-1] + n)
        lengths.append(c + n)
        tables.append(table)
        # Rows from the last real token onwards produce samples: one for a finished prefill or a
        # decode, 1 + (draft tokens scheduled) when verifying. A mid-prompt chunk produces none.
        k = min(n, max(0, c + n - (request.num_tokens - 1)))
        logits_rows.extend(range(starts[-1] - k, starts[-1]))
        counts.append(k)
        if request.params.prompt_logprobs is not None:      # rows whose next token is a prompt token
            last = min(c + n, request.num_prompt_tokens - 1)
            if last > c:
                prompt_rows.extend(range(starts[-2], starts[-2] + last - c))
                prompt_spans.append((request, c, request.token_ids[c + 1:last + 1]))
    if pad_to is not None:
        for _ in range(pad_to - len(lengths)):
            ids.append(0), positions.append(0), slots.append(-1)
            starts.append(starts[-1] + 1)
            lengths.append(0)
            tables.append([])
    width = max(1, max(len(t) for t in tables))
    table_tensor = torch.zeros((len(tables), width), dtype=torch.int32)
    for r, table in enumerate(tables):
        table_tensor[r, :len(table)] = torch.tensor(table, dtype=torch.int32)
    meta = BatchMeta(
        query_start_loc=torch.tensor(starts, dtype=torch.int32, device=device),
        seq_lens=torch.tensor(lengths, dtype=torch.int32, device=device),
        block_table=table_tensor.to(device),
        slot_mapping=torch.tensor(slots, dtype=torch.int64, device=device),
        block_size=block_size, query_start_loc_cpu=starts, seq_lens_cpu=lengths,
        max_query_len=max(b - a for a, b in zip(starts, starts[1:])), max_seq_len=max(lengths))
    return Batch(torch.tensor(ids, dtype=torch.long, device=device),
                 torch.tensor(positions, dtype=torch.long, device=device), meta,
                 torch.tensor(logits_rows + prompt_rows, dtype=torch.long, device=device), counts,
                 [r.request_id for r, _ in scheduled], {}, prompt_spans)

logits_indices selects the rows whose logits the sampler needs. A decode contributes its one row; a prefill chunk that ends at the prompt’s last token contributes that row; a mid-prompt chunk contributes nothing. This is a large saving: Qwen3-8B’s LM head is a 4,096 × 151,936 matmul, the single largest in the model. A 2,000-token prefill that sent every row through it would spend more time there than in the 36 layers’ attention. Production engines all compute logits only where they’re sampled.

The pad_to argument, used in Chapter 33, adds dummy rows with slot -1 and length 0, so a decode batch of 13 requests can run in a CUDA graph captured for 16.

Attention over the pool

The pool stores, per layer, k_cache and v_cache of shape [num_blocks, block_size, kv_heads, head_dim]. That’s a different axis order from Chapter 25’s [num_blocks, kv_heads, block_size, head_dim]. Slot-major rows make writing one index_copy_ over flat slot numbers:

def write_kv(k_cache, v_cache, k, v, slot_mapping):
    """Scatter k, v [N, Hkv, D] into the pools at flat slots.  (Your engine: Chapter 31)

    Slot -1 marks a padding row (Chapter 33's CUDA graphs pad batches). Its write goes to the
    pool's last slot, a scratch block the runner allocates but the block manager never hands
    out, instead of being filtered out: filtering would need the host to know how many rows
    are real, a sync that a CUDA graph can't contain.
    """
    scratch = k_cache.shape[0] * k_cache.shape[1] - 1
    slots = torch.where(slot_mapping >= 0, slot_mapping, scratch)
    k_cache.view(-1, *k_cache.shape[2:]).index_copy_(0, slots, k.to(k_cache.dtype))
    v_cache.view(-1, *v_cache.shape[2:]).index_copy_(0, slots, v.to(v_cache.dtype))

The reference backend gathers each request’s context through its block table and calls the causal_attention you wrote in Chapter 5. The positions do the work, as they have since Chapter 16: request $r$’s $t$ queries are the last $t$ positions of its seq_len-token context, so query $i$ is at position seq_len - t + i.

class ReferenceBackend:
    """Gathers each request's context and calls Chapter 5's causal_attention. Slow and obviously
    correct: every faster backend is tested against it."""
    name = "reference"

    def allocate(self, shape, dtype, device):
        return torch.zeros(shape, dtype=dtype, device=device), torch.zeros(shape, dtype=dtype, device=device)

    def write(self, cache, k, v, meta):
        write_kv(cache[0], cache[1], k, v, meta.slot_mapping)

    def forward(self, q, cache, meta, scale=None, window=0):
        """(Your engine: Chapter 31)

        Request r's queries are rows query_start_loc[r]:query_start_loc[r+1]; they are the
        LAST t positions of its seq_len-token context, so query i sits at position
        seq_len - t + i and may see keys 0 .. that position (and, with a sliding window of
        w > 0, only the last w of them: Chapter 42).
        """
        k_cache, v_cache = cache[0], cache[1]
        out = torch.zeros_like(q)
        starts, lengths = meta.query_start_loc_cpu, meta.seq_lens_cpu
        for r, length in enumerate(lengths):
            begin, end = starts[r], starts[r + 1]
            if length == 0 or begin == end:                  # padding row
                continue
            t = end - begin
            k = gather_context(k_cache, meta.block_table[r], length)
            v = gather_context(v_cache, meta.block_table[r], length)
            q_pos, k_pos = torch.arange(length - t, length, device=q.device), torch.arange(length, device=q.device)
            allowed = (k_pos[None, :] > q_pos[:, None] - window) if window > 0 else None
            y = causal_attention(q[begin:end].transpose(0, 1)[None], self.dequantize(k, cache, 0, meta, r, length),
                                 self.dequantize(v, cache, 1, meta, r, length), q_pos, k_pos, allowed, scale)
            out[begin:end] = y[0].transpose(0, 1)
        return out

    @staticmethod
    def dequantize(x, cache, which, meta, r, length):
        """[length, Hkv, D] -> [1, Hkv, length, D] in float; quantized pools carry scales at cache[2:]."""
        x = x.float()
        if len(cache) > 2:
            x = x * gather_context(cache[2 + which], meta.block_table[r], length)[..., None]
        return x.transpose(0, 1)[None]

It loops over requests in Python, which is slow, and it’s meant to be: every faster backend in Chapter 32 is tested against it. Chapter 32’s Triton kernel does the same computation in one launch, reading K/V tiles straight from the blocks.

The flat model

The model needs one change: attention reads and writes the pool through a backend, and every tensor is [N, ...] instead of [B, T, ...]. Rather than writing a new model, FlatModel reuses the modules and weights of the Qwen3 (or Qwen3Moe) you already loaded and tested. Nothing is copied or renamed, so every parity test of Chapters 17 and 27 still covers the weights.

class FlatModel(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model.eval()
        self.cfg = model.cfg
        self.backbone = model.model
        self.layers = model.model.layers
        self.rope = getattr(model, "rope_tables", None) or self.default_rope
        self.fused = False

    def default_rope(self, positions, dtype):
        cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta, dtype)
        return cos[0, 0][:, None], sin[0, 0][:, None]          # [N, 1, rotary_dim]: broadcast over heads

    def kv_spec(self):
        """(layers, kv_heads, head_dim) of the pool every attention layer writes."""
        custom = getattr(self.model, "kv_spec", None)            # latent caches (Chapter 42)
        if custom is not None:
            return custom()
        return self.cfg.num_hidden_layers, self.cfg.num_key_value_heads, self.cfg.head_dim

    def attention(self, layer, attn, h, positions, rope, kv, meta, backend):
        """One attention sublayer on [N, D] rows.  (Your engine: Chapter 31)

        The same steps as Qwen3Attention.forward, but heads are the middle axis of [N, H, D]
        and the cache is the pool: write this step's keys (after RoPE), then attend.
        """
        paged = getattr(attn, "paged_forward", None)          # attention of another shape: MLA (Chapter 42)
        if paged is not None:
            return paged(h, positions, rope, kv, meta, backend)
        n = h.shape[0]
        c = self.cfg
        q, k, v = self.project_qkv(attn, h)
        q = q.view(n, c.num_attention_heads, c.head_dim)
        k = k.view(n, c.num_key_value_heads, c.head_dim)
        v = v.view(n, c.num_key_value_heads, c.head_dim)
        if hasattr(attn, "q_norm"):
            q, k = attn.q_norm(q), attn.k_norm(k)
        cos, sin = rope
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        backend.write(kv, k, v, meta)
        out = backend.forward(q, kv, meta, window=getattr(attn, "sliding_window", 0))
        return attn.o_proj(out.reshape(n, -1))

    def forward(self, input_ids, positions, kv_caches, meta, backend, embeds=None, **features):
        """input_ids [N] -> final hidden states [N, D] (before the LM head).  (Your engine: Chapter 31)

        add_norm(norm, x, delta) returns (norm(x + delta), x + delta): the residual add and the
        next norm, which Chapter 33 fuses into one kernel. mlp(module, h) runs the MLP or MoE.
        """
        bank = getattr(self, "adapters", None)
        with bank.activate(features.get("lora_ids"), meta.decode_only) if bank is not None else nullcontext():
            return self._forward(input_ids, positions, kv_caches, meta, backend, embeds, features)

    def _forward(self, input_ids, positions, kv_caches, meta, backend, embeds, features):
        x = self.backbone.embed_tokens(input_ids) if embeds is None else embeds
        if "mrope_positions" in features:
            from ..multimodal import mrope_cos_sin
            rope = mrope_cos_sin(features["mrope_positions"], self.cfg.head_dim, self.cfg.rope_theta,
                                 features["mrope_sections"], x.dtype)
        else:
            rope = self.rope(positions, x.dtype)
        delta = None
        for i, layer in enumerate(self.layers):
            h, x = self.add_norm(layer.input_layernorm, x, delta)
            delta = self.attention(i, layer.self_attn, h, positions, rope, kv_caches[i], meta, backend)
            h, x = self.add_norm(layer.post_attention_layernorm, x, delta)
            delta = self.mlp(layer.mlp, h)
        h, _ = self.add_norm(self.backbone.norm, x, delta)
        return h

    def compute_logits(self, hidden):
        """Only the rows the sampler needs reach the LM head: for a 2,000-token prefill
        that is 1 row instead of 2,000, the largest matmul in the model skipped."""
        return self.model.lm_head(hidden).float()

RoPE tables are computed once per step for all $N$ positions and broadcast over heads. The MoE block of Chapter 27 already flattens its input to [N, D], so Qwen3Moe runs through the same code unchanged. Chapter 42 turns this class into a registry that covers Llama, Mistral, Qwen2, DeepSeek’s latent attention and more.

Prefix caching with hash chains

Chapter 25’s PrefixCache keyed each cached block by the entire token prefix up to its end. That’s correct but costly: the key for block 1,000 of a long document is a 16,000-token tuple, and comparing keys is proportional to their length. Production engines use a hash chain instead:

$$ h_0 = H(\text{root},\ \text{tokens}0,\ \text{extra}), \qquad h_i = H(h{i-1},\ \text{tokens}_i,\ \text{extra}) $$

Each key has a fixed size, depends on the whole history through $h_{i-1}$, and is computed once per block as the request grows. extra holds anything else that changes the keys and values: the LoRA adapter (Chapter 43), hashes of images in the prompt (Chapter 43), and a per-tenant salt if users must never share cache entries.

def hash_block(parent, tokens, extra=()):
    """Key for one full block: depends on the previous block's key, this block's tokens and
    any extra keys (LoRA adapter, image hashes, tenant salt).  (Your engine: Chapter 31)

    A cryptographic hash, not Python's hash(): a crafted collision would let one user's
    request read another user's cached keys and values.
    """
    h = hashlib.blake2b(digest_size=16)
    h.update(parent if parent is not None else b"\x00root")
    h.update(array("q", tokens).tobytes())
    if extra:
        h.update(repr(tuple(extra)).encode())
    return h.digest()

Warning

Use a cryptographic hash. With Python’s hash() or a weak 64-bit hash, an attacker who can submit prompts can search for a token sequence whose block hash collides with someone else’s, and their request would then attend to the victim’s cached keys and values, which leaks information about the victim’s prompt. vLLM offers SHA-256 block hashing (--prefix-caching-hash-algo) for this reason. BLAKE2b with a 128-bit digest is as safe for this purpose and fast in Python.

Blocks that outlive their requests

The second improvement over Chapter 25 is what happens when a request finishes. Chapter 25’s prefix cache held an extra reference on every published block, so cached blocks were never free, and you had to evict them explicitly. vLLM V1’s design is better: a freed block keeps its contents and its hash, and goes onto the free queue. It’s still findable by hash and can be reused for free, but it’s also the next candidate for allocation. Only when the allocator actually hands it out is it evicted from the cache. A block is in one of three states:

The free queue is an LRU list: allocation takes from the front, freeing appends to the back. One detail matters for hit rates: a finished request’s blocks are appended tail first, so its last blocks (unique to that request) are evicted before its first blocks (the system prompt and shared history, most likely to be reused).

class BlockManager:
    """Reference-counted block pool with an LRU free queue and a hash -> block map."""

    def __init__(self, num_blocks, block_size, enable_prefix_caching=True):
        if num_blocks < 1 or block_size < 1:
            raise ValueError("Need at least one block of at least one token")
        self.num_blocks, self.block_size = num_blocks, block_size
        self.enable_prefix_caching = enable_prefix_caching
        self.ref_count = [0] * num_blocks
        self.block_hash = [None] * num_blocks
        self.free_queue = OrderedDict((b, None) for b in range(num_blocks))   # front = evicted first
        self.cached = {}                 # hash -> block id
        self.req_blocks = {}             # request id -> [block ids], logical order
        self.num_registered = {}         # request id -> how many of its blocks are published by hash
        self.stats = {"queries": 0, "hits": 0}

    @property
    def num_free(self):
        return len(self.free_queue)

    @property
    def usage(self):
        return 1.0 - self.num_free / self.num_blocks

    def blocks_for(self, num_tokens):
        return -(-num_tokens // self.block_size)

    def update_hashes(self, request):
        """Extend request.block_hashes to cover every full block of its known tokens (not the
        placeholders of tokens still being computed, Chapter 33)."""
        size, hashes = self.block_size, request.block_hashes
        known = request.num_tokens - request.num_placeholders
        while (len(hashes) + 1) * size <= known:
            start = len(hashes) * size
            parent = hashes[-1] if hashes else None
            hashes.append(hash_block(parent, request.token_ids[start:start + size], request.cache_key))
        return hashes

    def find_cached_prefix(self, request):
        """Longest run of cached full blocks at the start of the request.  (Your engine: Chapter 31)

        Stops at the first miss. Leaves at least one token to compute, because the last
        prompt position's logits are what choose the first output token.
        """
        if not self.enable_prefix_caching or request.params.prompt_logprobs is not None:
            return []                       # prompt logprobs need every prompt position computed
        hashes = self.update_hashes(request)
        limit = (request.num_tokens - 1) // self.block_size
        found = []
        for h in hashes[:limit]:
            block = self.cached.get(h)
            if block is None:
                break
            found.append(block)
        self.stats["queries"] += request.num_tokens
        self.stats["hits"] += len(found) * self.block_size
        return found

    def allocate_slots(self, request, num_new_tokens, cached_blocks=()):
        """Make room for num_new_tokens more tokens of request, after its computed tokens and
        any cached_blocks it is about to adopt. Returns False, changing nothing, if the pool is
        too small; the scheduler then preempts someone or waits.  (Your engine: Chapter 31)
        """
        blocks = self.req_blocks.setdefault(request.request_id, [])
        total = request.num_computed_tokens + len(cached_blocks) * self.block_size + num_new_tokens
        needed = self.blocks_for(total) - len(blocks) - len(cached_blocks)
        # Adopting a cached-but-free block takes it out of the free queue too.
        reclaimed = sum(1 for b in cached_blocks if self.ref_count[b] == 0)
        if needed > self.num_free - reclaimed:
            return False
        for b in cached_blocks:
            if self.ref_count[b] == 0:
                del self.free_queue[b]
            self.ref_count[b] += 1
            blocks.append(b)
        for _ in range(max(needed, 0)):
            blocks.append(self._pop_free())
        if cached_blocks:
            self.num_registered[request.request_id] = len(cached_blocks)
        return True

    def _pop_free(self):
        block, _ = self.free_queue.popitem(last=False)          # least recently freed
        h = self.block_hash[block]
        if h is not None:                                       # evict it from the prefix cache
            if self.cached.get(h) == block:
                del self.cached[h]
            self.block_hash[block] = None
        self.ref_count[block] = 1
        return block

    def cache_full_blocks(self, request, num_tokens=None):
        """Publish request's full blocks among its first num_tokens tokens (default: the computed
        ones) under their hashes.  (Your engine: Chapter 31)

        The scheduler publishes blocks as soon as they are allocated for this step's tokens,
        before they are computed. A request admitted later in the same step may then adopt
        them: in a flattened batch every layer writes all K/V before any attention reads, so
        it reads them after they are written.
        """
        if not self.enable_prefix_caching:
            return
        blocks = self.req_blocks.get(request.request_id, [])
        hashes = self.update_hashes(request)
        num_tokens = request.num_computed_tokens if num_tokens is None else num_tokens
        full = min(num_tokens // self.block_size, len(hashes), len(blocks))
        start = self.num_registered.get(request.request_id, 0)
        for i in range(start, full):
            block, h = blocks[i], hashes[i]
            if self.block_hash[block] is None:
                self.block_hash[block] = h
                self.cached.setdefault(h, block)               # first writer wins; duplicates stay private
        self.num_registered[request.request_id] = max(start, full)

    def free(self, request):
        """Drop request's references. Blocks are queued tail-first, so a sequence's last blocks
        are evicted before its prefix, which is the part most likely to be shared.  (Your engine: Chapter 31)"""
        for block in reversed(self.req_blocks.pop(request.request_id, [])):
            self.ref_count[block] -= 1
            if self.ref_count[block] == 0:
                self.free_queue[block] = None
            elif self.ref_count[block] < 0:
                raise RuntimeError(f"Block {block} freed more often than referenced")
        self.num_registered.pop(request.request_id, None)

    def trim(self, request, num_tokens):
        """Release blocks past num_tokens (rejected speculative tokens, Chapter 37)."""
        blocks = self.req_blocks.get(request.request_id, [])
        keep = self.blocks_for(num_tokens)
        for block in reversed(blocks[keep:]):
            self.ref_count[block] -= 1
            if self.ref_count[block] == 0:
                h = self.block_hash[block]
                if h is not None and self.cached.get(h) == block:
                    del self.cached[h]
                self.block_hash[block] = None
                self.free_queue[block] = None
        del blocks[keep:]

    def reset_prefix_cache(self):
        """Forget every cached-free block (for example after loading new weights)."""
        for block in self.free_queue:
            self.block_hash[block] = None
        self.cached = {h: b for h, b in self.cached.items() if self.ref_count[b] > 0}

    @property
    def hit_rate(self):
        return self.stats["hits"] / self.stats["queries"] if self.stats["queries"] else 0.0

Three rules in this code are easy to get wrong:

1. Leave one token to compute. find_cached_prefix never matches the block that contains the prompt’s last token, even if it’s cached. The first output token is sampled from the logits at the last prompt position, so that position must run through the model. (A request with a 32-token prompt and 16-token blocks reuses at most the first block.)
2. Adopting a cached-free block costs a free slot. allocate_slots counts how many of the adopted blocks are currently in the free queue, because taking them out of it shrinks what’s left for new allocations. Getting this wrong over-commits the pool.
3. Allocation is all or nothing. If the request doesn’t fit, allocate_slots returns False before changing anything. The scheduler can then preempt someone and retry, without undoing a half-finished allocation.

Only full blocks are published, and they’re published as soon as they’re allocated for this step’s tokens, before the forward pass computes them. That looks premature, but it’s safe and it matters: four requests that share a long prompt and arrive together (or the four samples of an n=4 request, Chapter 34) are admitted in the same step, and only the first would compute the prompt if the others could adopt its blocks right away. In a flattened batch every layer writes all $N$ tokens’ keys and values before any attention reads them, so a request that adopted a block in this step reads it after it was written. The engine also publishes after each step, which covers generated tokens as they fill blocks. That’s what makes multi-turn chat cheap: turn 2’s prompt is turn 1’s prompt plus turn 1’s answer plus the new message, and most of it is already cached.

The scheduler

With one kind of work, the scheduler is short:

class Scheduler:
    def __init__(self, config, block_manager):
        self.config, self.blocks = config, block_manager
        self.waiting, self.running = deque(), []
        self.requests = {}
        self.on_finish = None             # called before a finished request's blocks are freed (Chapter 41)

    def add(self, request):
        capacity = self.blocks.num_blocks * self.blocks.block_size
        if request.num_tokens + request.params.max_tokens > capacity:
            raise ValueError(f"{request.request_id}: needs more KV slots than the whole pool holds ({capacity})")
        if not self.config.enable_chunked_prefill and request.num_tokens > self.config.max_num_batched_tokens:
            raise ValueError(f"{request.request_id}: prompt exceeds the token budget and chunked prefill is off")
        self.requests[request.request_id] = request
        self._enqueue(request)

    def _enqueue(self, request, front=False):
        if self.config.policy == "priority":
            key = (request.priority, request.arrival_time)
            index = next((i for i, r in enumerate(self.waiting) if (r.priority, r.arrival_time) > key), len(self.waiting))
            self.waiting.insert(index, request)
        elif front:
            self.waiting.appendleft(request)
        else:
            self.waiting.append(request)

    def _chunk(self, n, budget):
        limit = self.config.long_prefill_token_threshold
        return min(n, budget, limit) if limit else min(n, budget)

    def _preempt(self, request, out):
        """Take request out of the running set and give its blocks back.  (Your engine: Chapter 31)

        recompute: forget its KV; when readmitted it prefills prompt + output again (and
        probably hits its own blocks in the prefix cache). swap: its blocks are copied to host
        memory first, and num_computed_tokens is kept.
        """
        self.running.remove(request)
        request.spec_token_ids = []                      # drafts are only valid for the very next step
        if self.config.preemption_mode == "swap":
            out.swap_out.append((request, list(self.blocks.req_blocks[request.request_id])))
            request.swapped_out = True
        else:
            request.num_computed_tokens = 0
        self.blocks.free(request)
        request.status = Status.PREEMPTED
        request.num_preemptions += 1
        out.preempted.append(request)
        self._enqueue(request, front=True)

    def schedule(self):
        """Choose (request, num_new_tokens) pairs for one step.  (Your engine: Chapter 31)"""
        out = SchedulerOutput()
        budget = self.config.max_num_batched_tokens
        if self.config.policy == "priority":
            self.running.sort(key=lambda r: (r.priority, r.arrival_time))
        # 1. Running requests: decodes, unfinished prefill chunks, verification blocks.
        index = 0
        while index < len(self.running) and budget > 0:
            request = self.running[index]
            n = self._chunk(request.num_tokens_with_spec - request.num_computed_tokens, budget)
            if request.num_output_tokens >= request.params.max_tokens:
                n = 0                                        # its last token is in flight (Chapter 33)
            if n <= 0:
                index += 1
                continue
            while not self.blocks.allocate_slots(request, n):
                victim = self.running[-1]                    # lowest priority: last in order
                self._preempt(victim, out)
                if victim is request:
                    break
            if request.status is Status.PREEMPTED:
                break                                        # it was the last one; nothing left to try
            out.scheduled.append((request, n))
            self.blocks.cache_full_blocks(request, request.num_computed_tokens + n)
            budget -= n
            index += 1
        # 2. Waiting requests, unless memory was just so tight that we preempted.
        while self.waiting and budget > 0 and len(self.running) < self.config.max_num_seqs and not out.preempted:
            request = self.waiting[0]
            cached = [] if request.num_computed_tokens else self.blocks.find_cached_prefix(request)
            remaining = request.num_tokens - request.num_computed_tokens - len(cached) * self.blocks.block_size
            if not self.config.enable_chunked_prefill and remaining > budget:
                break
            n = self._chunk(remaining, budget)
            if not self.blocks.allocate_slots(request, n, cached):
                break
            self.waiting.popleft()
            if request.swapped_out:
                out.swap_in.append((request, list(self.blocks.req_blocks[request.request_id])))
                request.swapped_out = False
            request.num_computed_tokens += len(cached) * self.blocks.block_size
            if request.status is Status.WAITING:
                request.num_cached_tokens = len(cached) * self.blocks.block_size
            request.status = Status.RUNNING
            self.running.append(request)
            out.scheduled.append((request, n))
            self.blocks.cache_full_blocks(request, request.num_computed_tokens + n)   # visible to later admissions
            budget -= n
        return out

Running requests come first. They hold blocks and users are watching their streams; a new request can wait one more step. Within the running list, order is arrival (FCFS) or (priority, arrival). Each running request asks for everything it still needs: one token if it’s decoding, the rest of its prompt if it’s mid-prefill (capped by long_prefill_token_threshold, the chunk size), or its history if it was just resumed.

Admission spends what’s left. A waiting request starts after its longest cached prefix, so a request whose 2,000-token system prompt is cached costs only its own question. If the pool can’t hold it, admission stops: the scheduler never preempts a running request to admit a new one.

Chunked prefill falls out of the budget: a 5,000-token prompt with a 2,048-token budget runs as 2,048 + 2,048 + 904 over three steps, sharing each step with every running request’s decode token. Chapter 24 explained why: decode latency stays steady while long prompts arrive. With enable_chunked_prefill=False a prompt must fit in one step’s budget, which some engines prefer for simplicity at low load.

When memory runs out, the scheduler preempts the lowest-priority running request (the last in order) and retries. If the victim is the request it was trying to schedule, it stops: nothing with lower priority is left. After any preemption it also stops admitting, because the pool is evidently full.

Preemption: recompute or swap

A preempted request gives back its blocks. There are two ways to bring it back later:

• Recompute (preemption_mode="recompute"): forget its K/V and set num_computed_tokens = 0. When readmitted, it prefills prompt plus output again, in chunks like any prompt. Its blocks were just freed into the queue with their hashes, so if they haven’t been reused yet, the prefix cache gives most of them back for free.
• Swap (preemption_mode="swap"): copy its blocks to host memory before giving them back, and keep num_computed_tokens. When readmitted, allocate fresh blocks and copy the K/V back. The runner does both copies before the step’s forward pass, so the freed blocks can be reused by other requests in the same step.

Which costs less? For Qwen3-8B in BF16, one token’s K/V across all 36 layers is $2 \times 36 \times 8 \times 128 \times 2 = 147{,}456$ bytes. Swapping it out and back over PCIe 5.0 at about 50 GB/s takes $2 \times 147{,}456 / 50 \times 10^9 \approx 5.9$ µs. Recomputing it costs about $2 \times 8 \times 10^9 = 16$ GFLOPs, roughly 27 µs at a realistic 600 TFLOP/s on an H100. So swapping is cheaper per token, and the copy can overlap compute. vLLM V1 nevertheless supports only recompute: preemption should be rare in a well-sized engine, recompute needs no host-memory management, and with prefix caching the “recompute” is often a cache hit. Both modes are here so you can measure the trade-off on your own workload (stretch exercise 2). For models with a tiny KV cache per token, such as DeepSeek’s latent attention (Chapter 42), swapping is cheaper still.

The engine loop

EngineCore.step ties the parts together: schedule, perform swaps, build one batch, run it, sample, update.

class EngineCore:
    def __init__(self, model, config=None, eos_token_id=None, sampler=None, vocab_bytes=None):
        self.config = config = config or EngineConfig()
        p = next(model.parameters())
        self.device = p.device
        flat = model if isinstance(model, FlatModel) else FlatModel(model)
        options = {"kv_cache_dtype": config.kv_cache_dtype} if config.kv_cache_dtype != "auto" else {}
        backend = get_backend(config.attention_backend, **options)
        layers, kv_heads, head_dim = flat.kv_spec()
        num_blocks = config.num_blocks
        if num_blocks is None:
            if self.device.type == "cuda":
                num_blocks = profile_num_blocks(flat, backend, config.block_size, p.dtype, self.device,
                                                config.max_num_batched_tokens, config.gpu_memory_utilization)
            else:
                num_blocks = num_blocks_for_memory(config.kv_cache_bytes, layers, kv_heads, head_dim,
                                                   config.block_size, p.dtype, getattr(flat.model, "kv_pools", 2))
        self.runner = ModelRunner(flat, num_blocks, config.block_size, backend)
        self.blocks = BlockManager(num_blocks, config.block_size, config.enable_prefix_caching)
        self.scheduler = Scheduler(config.scheduler_config(), self.blocks)
        self.sampler = sampler or (Sampler() if config.sampler == "full" else SimpleSampler())
        vocab_size = flat.model.cfg.vocab_size
        self.guides = GuideFactory(vocab_bytes, eos_token_id, vocab_size) if vocab_bytes is not None else None
        self.eos_token_id = eos_token_id
        self.max_model_len = config.max_model_len or getattr(flat.model, "context_limit", 1 << 30)
        self.generators = {}
        self.engine_generator = torch.Generator(device=self.device).manual_seed(config.seed)
        self.steps = self.preemptions = 0
        self.drafter = None                      # speculative decoding (Chapter 37)
        self.inflight = None                     # the launched, unread step (async scheduling)
        self.mrope_sections = getattr(flat.cfg, "mrope_sections", None)
        if config.cuda_graphs != "off" and (config.cuda_graphs != "auto" or self.device.type == "cuda"):
            mode = {"static": "off"}.get(config.cuda_graphs, config.cuda_graphs)
            self.runner.enable_graphs(config.max_num_seqs, self.max_model_len, mode)

    def add_request(self, request_id, prompt_token_ids, params=None, priority=0, cache_key=(), arrival_time=None,
                    lora=None, features=None):
        params = params or SamplingParams()
        if params.n > 1:                       # n samples = n requests; the prefix cache shares the prompt
            return [self.add_request(f"{request_id}:{i}", prompt_token_ids, params.child(i), priority, cache_key,
                                     arrival_time, lora, features) for i in range(params.n)]
        if request_id in self.scheduler.requests:
            raise ValueError(f"Duplicate request id {request_id!r}")
        if (lora is not None or features) and self.drafter is not None:
            raise ValueError("This drafter does not support image or adapter requests")
        guided = params.guided_regex is not None or params.guided_json is not None
        if guided and (self.guides is None or self.config.async_scheduling or self.config.sampler != "full"):
            raise ValueError("Guided decoding needs vocab_bytes, sampler='full' and synchronous scheduling")
        if (params.logprobs is not None or params.prompt_logprobs is not None or params.needs_penalties
                or params.logit_bias or params.allowed_token_ids is not None or params.min_tokens) \
                and self.config.sampler != "full":
            raise ValueError("This request needs EngineConfig(sampler='full')")
        if len(prompt_token_ids) + params.max_tokens > self.max_model_len:
            raise ValueError(f"{request_id}: prompt + max_tokens exceed the model's {self.max_model_len} positions")
        request = Request(request_id, list(prompt_token_ids), params, self.eos_token_id, priority,
                          cache_key=tuple(cache_key))
        if arrival_time is not None:
            request.arrival_time = arrival_time
        if guided:
            request.extra["guide"] = self.guides(params)
        if features:
            from ..multimodal import validate_features
            extra, key = validate_features(features, len(prompt_token_ids), self.runner.flat.cfg.hidden_size,
                                           self.runner.flat.cfg.head_dim)
            if "mrope_sections" in extra:
                if getattr(self.runner.flat.cfg, "rope_scaling", None):
                    raise ValueError("This M-RoPE path does not implement scaled rotary frequencies")
                if self.mrope_sections is not None and self.mrope_sections != extra["mrope_sections"]:
                    raise ValueError("M-RoPE sections must match the model")
                self.mrope_sections = extra["mrope_sections"]
            request.extra.update(extra)
            request.cache_key += (key,) if key else ()
        bank = getattr(self.runner.flat, "adapters", None)
        if lora is not None:
            if bank is None:
                raise ValueError("No adapter bank installed")
            slot, key = bank.acquire(lora)
            request.extra["lora_slot"] = slot
            request.cache_key += (key,)
        try:
            self.scheduler.add(request)
        except Exception:
            if lora is not None:
                bank.release(slot)
            raise
        if params.seed is not None:
            self.generators[request_id] = torch.Generator(device=self.device).manual_seed(params.seed)
        return request

    def abort_request(self, request_id):
        request = self.scheduler.abort(request_id)
        if request is not None:
            self.release_adapter(request)
        self.generators.pop(request_id, None)
        self.runner.swapped.pop(request_id, None)
        return request

    @property
    def has_unfinished(self):
        return self.scheduler.has_unfinished or self.inflight is not None

    def check_stop(self, request, index=-1):
        """Finish reason after the token at `index` (the newest by default), or None.  (Your engine: Chapter 31)"""
        p, token = request.params, request.token_ids[index]
        produced = (index % request.num_tokens) + 1 - request.num_prompt_tokens     # output tokens up to it
        if produced >= p.min_tokens:
            if token in p.stop_token_ids:
                request.stop_reason = token
                return Status.FINISHED_STOPPED
            if not p.ignore_eos and token == request.eos_token_id:
                return Status.FINISHED_STOPPED
        if produced >= p.max_tokens or produced + request.num_prompt_tokens >= self.max_model_len:
            return Status.FINISHED_LENGTH
        guide = request.extra.get("guide")
        if guide is not None and guide.finished:          # the pattern admits no further byte
            request.stop_reason = "guided"
            return Status.FINISHED_STOPPED
        return None

    def step(self):
        """One iteration: schedule, run, sample, update. Returns a RequestOutput per request that
        produced tokens or finished.  (Your engine: Chapter 31)"""
        if self.config.async_scheduling:
            return self.step_async()
        plan = self.scheduler.schedule()
        self.preemptions += len(plan.preempted)
        for request, blocks in plan.swap_out:
            self.runner.swap_out(request, blocks)
        for request, blocks in plan.swap_in:
            self.runner.swap_in(request, blocks)
        if plan.empty:
            return []
        if self.drafter is not None:
            return self.drafter.step(self, plan)            # Chapter 37 replaces sample-and-update
        batch = self.runner.prepare(plan.scheduled, self.blocks.req_blocks)
        logits = self.runner.execute(batch)
        if batch.prompt_spans:
            self.record_prompt_logprobs(batch, logits[batch.num_sample_rows:])
            logits = logits[:batch.num_sample_rows]
        sampling = [r for (r, _), k in zip(plan.scheduled, batch.sample_counts) if k]
        tokens, logprobs = self.sampler(logits, sampling, self.generators) if sampling else ([], None)
        sampled = dict(zip((r.request_id for r in sampling), tokens))
        lp = dict(zip((r.request_id for r in sampling), logprobs)) if logprobs else {}
        outputs, now = [], time.monotonic()
        for request, n in plan.scheduled:
            request.num_computed_tokens += n
            new = sampled.get(request.request_id, [])
            for token in new:
                request.append(token)
            self.blocks.cache_full_blocks(request)
            if not new:
                continue                                      # mid-prompt chunk: nothing to report
            request.first_token_time = request.first_token_time or now
            status = self.check_stop(request)
            if status is not None:
                self.scheduler.finish(request, status)
                self.generators.pop(request.request_id, None)
            outputs.append(self.make_output(request, new, lp.get(request.request_id)))
        self.steps += 1
        return outputs

After the forward pass, every scheduled request advances num_computed_tokens by its n. Requests whose rows produced logits get a new token, which is appended to the list (it becomes the next step’s input) and checked against the stop conditions: a stop token, EOS unless ignore_eos, max_tokens, or the model’s context limit. min_tokens suppresses the stop checks until enough tokens exist; Chapter 34’s sampler also masks those tokens so they can’t be chosen too early. Finished requests free their blocks at once, which is why the scheduler can admit a new request in the same step that an old one finishes.

Notice what the loop doesn’t do: there’s no padding, no per-request forward, no slot bookkeeping, and no copying of caches. The sampler here is SimpleSampler, which takes greedy rows with one argmax and calls Chapter 8’s sample for each sampled row. Chapter 34 replaces it.

Sizing the pool

On a GPU, the engine gives the KV pool everything that’s left after the weights and the activations of the largest possible step:

def kv_bytes_per_block(layers, kv_heads, head_dim, block_size, dtype, pools=2):
    """K and V (pools=2; a latent cache has one), every layer, one block."""
    return pools * layers * block_size * kv_heads * head_dim * torch.empty((), dtype=dtype).element_size()


def num_blocks_for_memory(free_bytes, layers, kv_heads, head_dim, block_size, dtype, pools=2):
    return int(free_bytes // kv_bytes_per_block(layers, kv_heads, head_dim, block_size, dtype, pools))


@torch.inference_mode()
def profile_num_blocks(flat, backend, block_size, dtype, device, max_tokens, utilization=0.9):
    """CUDA: weights are loaded; run the largest batch the scheduler may build with a dummy
    single-block pool, measure peak memory, and give everything else (up to `utilization` of
    the GPU) to the KV pool. That is how an engine fills a GPU without running out mid-traffic."""
    from .batch import build_batch as _build
    from .request import Request
    layers, kv_heads, head_dim = flat.kv_spec()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    dummy = Request("profile", [0] * max_tokens)
    probe = backend.allocate((-(-max_tokens // block_size), block_size, kv_heads, head_dim), dtype, device)
    batch = _build([(dummy, max_tokens)], {"profile": list(range(probe[0].shape[0]))}, block_size, device)
    per_block = sum(t[0].numel() * t.element_size() for t in probe) * layers
    hidden = flat(batch.input_ids, batch.positions, [probe] * layers, batch.meta, backend)
    flat.compute_logits(hidden[batch.logits_indices])
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated()
    total = torch.cuda.get_device_properties(device).total_memory
    del probe, hidden
    return int((total * utilization - peak) // per_block)

profile_num_blocks runs one forward pass of max_num_batched_tokens tokens with a dummy one-request pool, records the peak, and divides what remains of gpu_memory_utilization × total by the bytes per block. For Qwen3-8B in BF16 on an 80 GB H100 at 90% utilization: 72 GB available, minus 16.4 GB of weights and about 2 GB of activations for a 2,048-token step, leaves about 53 GB, or $53 \times 10^9 / 147{,}456 \approx 360{,}000$ tokens of cache. At 2,000 tokens per request that’s 180 concurrent requests’ worth of memory, against the 22 that a static 16,384-token slot per request would allow.

The runner owns the pools and does the swaps:

class ModelRunner:
    def __init__(self, model, num_blocks, block_size, backend="reference", device=None, dtype=None):
        self.flat = model if isinstance(model, FlatModel) else FlatModel(model)
        p = next(self.flat.parameters())
        self.device, self.dtype = torch.device(device or p.device), dtype or p.dtype
        self.backend = get_backend(backend) if isinstance(backend, str) else backend
        self.block_size, self.num_blocks = block_size, num_blocks
        layers, kv_heads, head_dim = self.flat.kv_spec()
        shape = (num_blocks + 1, block_size, kv_heads, head_dim)      # + 1: the scratch block for padding rows
        where = getattr(self.flat, "layer_device", lambda i: self.device)    # Chapter 40: layers on several devices
        allocate = getattr(self.flat.model, "allocate_kv", None) or self.backend.allocate     # MLA: one latent pool
        self.kv_caches = [allocate(shape, self.dtype, where(i)) for i in range(layers)]
        self.swapped = {}                 # request id -> [(k, v) per layer] in host memory
        self.graphs = None                # GraphRunner for decode batches (Chapter 33)
        self.last_hidden = None

    def swap_out(self, request, blocks):
        """Copy a preempted request's blocks to (pinned, on CUDA) host memory."""
        index = torch.tensor(blocks, device=self.device)
        host = (lambda t: t.cpu().pin_memory()) if self.device.type == "cuda" else (lambda t: t.cpu())
        self.swapped[request.request_id] = [tuple(host(_bytes(t)[index]) for t in layer) for layer in self.kv_caches]

    def swap_in(self, request, blocks):
        """Copy a resumed request's saved KV into its newly allocated blocks."""
        saved = self.swapped.pop(request.request_id)
        count = saved[0][0].shape[0]
        index = torch.tensor(blocks[:count], device=self.device)
        for layer, host in zip(self.kv_caches, saved):     # K, V and, when quantized, their scales
            for tensor, copy in zip(layer, host):
                _bytes(tensor).index_copy_(0, index, copy.to(self.device, non_blocking=True))

    def prepare(self, scheduled, block_tables, pad_to=None, previous=None):
        """Build the batch. With async scheduling, a request's newest input token may still be a
        placeholder: its value is in `previous.tokens` on the device, so copy it there, on the
        device, without reading it back (Chapter 33)."""
        batch = build_batch(scheduled, block_tables, self.block_size, self.device, pad_to)
        if previous is not None:
            rows, sources = [], []
            for (request, n), start in zip(scheduled, batch.meta.query_start_loc_cpu):
                if request.token_ids[request.num_computed_tokens] == PLACEHOLDER:
                    rows.append(start)
                    sources.append(previous.row_of[request.request_id])
            if rows:
                index = torch.tensor(rows, device=self.device)
                batch.input_ids[index] = previous.tokens[torch.tensor(sources, device=self.device)]
        self.prepare_features(batch, scheduled)
        return batch

    @torch.inference_mode()
    def prepare_features(self, batch, scheduled):
        """Map absolute feature positions and adapter slots to this step's ragged rows.  (Your engine: Chapter 43)
        A prefill chunk can cut through an image; recomputation uses the same saved rows.
        """
        extras = [r.extra for r, _ in scheduled]
        if any("lora_slot" in e for e in extras):
            slots = torch.full_like(batch.input_ids, -1)
            for (request, n), start in zip(scheduled, batch.meta.query_start_loc_cpu):
                slots[start:start + n] = request.extra.get("lora_slot", -1)
            batch.extra["lora_ids"] = slots
        mrope = [e for e in extras if "mrope_positions" in e]
        if mrope:
            sections = mrope[0]["mrope_sections"]
            if any(e["mrope_sections"] != sections for e in mrope):
                raise ValueError("All requests must use the model's M-RoPE sections")
            axes = batch.positions[None].expand(3, -1).clone()
            for (request, n), start in zip(scheduled, batch.meta.query_start_loc_cpu):
                if "mrope_positions" not in request.extra:
                    continue
                for offset in range(n):
                    pos = request.num_computed_tokens + offset
                    if pos < request.num_prompt_tokens:
                        axes[:, start + offset] = request.extra["mrope_positions"][:, pos].to(self.device)
                    else:
                        axes[:, start + offset] = pos + request.extra["mrope_delta"]
            batch.extra.update(mrope_positions=axes, mrope_sections=sections)
        replacements = []
        for (request, n), start in zip(scheduled, batch.meta.query_start_loc_cpu):
            for pos, row in request.extra.get("image_rows", {}).items():
                offset = pos - request.num_computed_tokens
                if 0 <= offset < n:
                    replacements.append((start + offset, row))
        if replacements:
            embeds = self.flat.backbone.embed_tokens(batch.input_ids)
            for row, value in replacements:
                embeds[row] = value.to(embeds)
            batch.extra["embeds"] = embeds

    def enable_graphs(self, max_batch, max_model_len, mode="auto"):
        from .graphs import GraphRunner
        blocks = -(-max_model_len // self.block_size)
        self.graphs = GraphRunner(self, max_batch, blocks, use_graphs=None if mode == "auto" else mode == "on")
        self.graphs.capture()

    @torch.inference_mode()
    def execute(self, batch):
        """Run one flattened batch; returns logits [L, vocab] for batch.logits_indices."""
        if self.graphs is not None and not batch.extra and self.graphs.can_run(batch):
            return self.graphs.run(batch)
        hidden = self.flat(batch.input_ids, batch.positions, self.kv_caches, batch.meta, self.backend,
                           **(batch.extra or {}))
        if hidden is None:                                       # a pipeline stage other than the first (Chapter 41)
            return None
        if batch.logits_indices.numel() == 0:
            return torch.empty((0, 0), device=self.device)
        self.last_hidden = hidden[batch.logits_indices]          # what Medusa / MTP heads read (Chapter 37)
        return self.flat.compute_logits(self.last_hidden)

The correctness test

As in Chapters 16, 24 and 25, the optimizations must not change any output. The milestone runs five requests of different lengths through the engine in four configurations: small chunks and a tight budget; a 12-block pool that forces preemption by recompute; the same pool with swapping; and prefix caching off with only two sequences at a time. One request shares its first two blocks with another, so the prefix cache is exercised too. In every configuration, each request’s greedy output must equal sampling.generate run on that prompt alone, and when the engine is idle every block must be free again.

The test catches the bugs that this design invites: a slot mapping off by one at a block boundary, a request resumed with a stale num_computed_tokens, a cached block adopted without leaving a token to compute, a swapped request restored into the wrong blocks, logits taken from a mid-prompt row, and blocks leaked by aborted or preempted requests.

Run it

python run.py core --requests 16

Sixteen requests that share a 192-token system prompt, each with its own 8-48-token question and 8-40 output tokens, through the 2-layer test model on a laptop CPU. First Chapter 24’s engine with 8 slots of 320 tokens, then the engine core with an ample pool and with a 24-block (384-token) pool:

{"engine": "ch24 slots", "kv_slots": 2560, "output_tokens": 322, "tok_s": 516.5, "steps": 55}
{"engine": "ch31 core, ample pool", "kv_slots": 65536, "output_tokens": 322, "tok_s": 1108.1, "steps": 52, "prefix_hit_rate": 0.814, "preemptions": 0}
{"engine": "ch31 core, 24-block pool", "kv_slots": 384, "output_tokens": 322, "tok_s": 804.9, "steps": 96, "prefix_hit_rate": 0.866, "preemptions": 8}

The core is about twice as fast as Chapter 24’s engine, even on a CPU (timings vary by 10-20% from run to run on a laptop), for two reasons: 81% of prompt tokens come from the prefix cache, and a step’s prefill chunks run in one forward pass instead of one per request. With a pool of only 384 token slots, 15% of what the slot engine reserved, it still serves every request correctly: it preempts 8 times and takes more steps, and the prefix hit rate rises because resumed requests find their own blocks still cached.

Build it

pytest tests/test_ch31_engine_v2.py
python run.py core --impl engine

Stretch exercises

1. ★ Add a prefix_cache_stats endpoint-style method that reports, per step, how many tokens were cached, computed and evicted. Run a multi-turn chat workload (each turn’s prompt is the previous prompt + answer + a new message) and plot the hit rate against the pool size. Where: add a reporting method to EngineCore in engine/serve/engine.py, reading BlockManager counters in engine/serve/blocks.py.
2. ★★ Measure recompute against swap on a GPU with Qwen3-0.6B: force preemption with a small pool and long outputs, and compare total time and p99 time per output token. Then repeat with prefix caching off. When does swapping win? Where: experiments/ch31.py (create it), configuring engine.serve.engine.EngineCore and adapting run.py’s cmd_core.
3. ★★ Implement SGLang’s cache-aware scheduling: when several requests are waiting, admit first the one with the longest cached prefix (longest-prefix-match), instead of FCFS. Show the hit-rate improvement on a workload with several different system prompts, and the starvation risk it introduces. Where: waiting-request selection in Scheduler.schedule in engine/serve/scheduler.py.
4. ★★★ Replace the flat hash map with a radix tree over token blocks (SGLang’s RadixAttention). Support partial-block matches for the last block by copying it, and evict leaves in LRU order. Compare hit rates with the hash chain on a tree-shaped workload (many branches from shared prefixes). Where: prefix lookup/storage/eviction in BlockManager in engine/serve/blocks.py.

Check your understanding

1. Why is a prefill chunk, a decode step and a resumed request the same kind of work for this scheduler?
2. In a flattened batch, which operations need the metadata and which don’t? Why does that make the matmuls faster than Chapter 24’s per-request prefill calls?
3. Why can’t the prefix cache match the block containing the prompt’s last token?
4. A finished request’s blocks go to the back of the free queue tail-first. What would happen to hit rates if they went head-first?
5. Why must the block hash include the LoRA adapter and image hashes, and why must it be cryptographic?
6. Why does the scheduler stop admitting new requests in a step where it preempted one?
7. For a model with 4 KiB of K/V per token and 3 B parameters, which preemption mode would you expect to cost less, and why?

Going deeper

• vLLM V1: vllm/v1/core/sched/scheduler.py (the unified token-budget scheduler), vllm/v1/core/kv_cache_manager.py and vllm/v1/core/block_pool.py (hash-chained prefix caching with a free-queue LRU), and the design notes vLLM V1: A Major Upgrade to vLLM’s Core Architecture (vLLM blog, 2025) and Automatic Prefix Caching in the vLLM docs.
• Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs (2024), for RadixAttention and cache-aware scheduling; python/sglang/srt/mem_cache/radix_cache.py.
• Agrawal et al., Sarathi-Serve (OSDI 2024), for chunked prefill and the token budget; Kwon et al., PagedAttention (SOSP 2023), §4.5 on swapping versus recomputation.
• GPU Mode L35 (SGLang performance optimization) and L40 (FlashInfer) for how the flattened batch’s metadata reaches the kernels.