CHAPTER 04

State needs an owner

Chapter 3 ended with an address: a state ID multiplied by a ring length, plus a position modulo that length. Something must hand out those IDs. It must also decide when an ID can be given to another request. A request’s memory need grows one token at a time, and requests finish in an order nobody chose in advance. Two requests may even begin with the same text.

4.1 Hand out memory in blocks

Suppose the engine gave each request one contiguous region sized for its longest possible answer. Most requests stop early, so much of that memory would sit unused. Reserving exactly the current length fails the other way: the next token needs room that may already belong to someone else.

A common answer is to divide memory into equal blocks, each holding a fixed number of token positions. A request owns a list of block IDs, called its block table. Position $p$ lives in block $\mathrm{table}[\lfloor p∕B\rfloor ]$ at offset $pmodB$, where $B$ is the block size. The blocks need not be adjacent. Any free block can serve any request. This paged design was popularised for language-model serving by vLLM [32].

Run the first experiment:

python3 labs/09_blocks.py

It uses a pool of eight blocks with four positions each. Block 0 is never handed out; the number 0 can therefore mean “no block”. Before reading the output, predict when request A, which starts with six tokens, first needs a third block.

admit A(6), B(9)               A=[1, 2] B=[3, 4, 5]         free=[6, 7]
A reaches 8 tokens             A=[1, 2] B=[3, 4, 5]         free=[6, 7]
A reaches 9 tokens             A=[1, 2, 6] B=[3, 4, 5]      free=[7]
A at 11, two drafts reserved   A=[1, 2, 6, 7] B=[3, 4, 5]   free=[]
drafts rejected; A at 12       A=[1, 2, 6] B=[3, 4, 5]      free=[7]
admission check for 4 blocks: LATER
B finishes                     A=[1, 2, 6]                  free=[7, 5, 4, 3]
admission check for 4 blocks: OK
admission check for 8 blocks: NEVER

Six tokens need $\lceil 6∕4\rceil =2$ blocks and nine need three, so A receives block 6 at its ninth token, not at its seventh. The last block in a table is usually only partly full. With block size $B$, at most $B-1$ positions are idle per request, however long the request becomes.

The fourth line reserves room for two tokens that do not yet exist. Chapter 8 will introduce draft tokens: cheap guesses that the model verifies in the next step. Positions for a draft must exist before the step writes it. Here 11 tokens plus two drafts need $\lceil 13∕4\rceil =4$ blocks, so A takes the last free block. Both drafts are then rejected and the verified step adds one token. Twelve tokens need only three blocks, and block 7 returns to the pool. A table can shrink as well as grow.

The admission checks distinguish two kinds of refusal. Four blocks are not free while B runs, so the answer is LATER: the request can wait. Eight blocks exceed the seven that this pool can ever issue, so the answer is NEVER. Waiting would not help.

Each block carries an owner count. The pool allocates from the front of its free list and returns a block to the back when the count reaches zero. When B finishes, its blocks are released in reverse table order. The free list therefore ends as $[7,5,4,3]$: B’s last block will be reused before its first. The check in show confirms after every step that each usable block is either owned or free, never both and never neither.

4.2 Two owners, one block

Owner counts matter once two requests share a block. Suppose a parent request has processed ten prompt tokens and a second request, the child, should continue from the same prefix. Recomputing the prefix would repeat work; copying all its blocks would repeat memory. The child can instead point to the parent’s blocks and add one to each block’s count. We call that pinning a block.

Sharing is safe only for rows that neither owner will write again. Positions 0–7 fill the first two blocks and are complete. Positions 8 and 9 sit in a partly filled third block. The next token of either request would be written there. So the child keeps the two full blocks and receives a private copy of the partial one. Copying only when a write would otherwise affect another owner is called copy-on-write.

python3 labs/10_share.py

The experiment stores token IDs in place of key and value rows, so a wrong block shows up as a wrong number:

38def fork(pool, parent, processed, copies):
39    child = list(parent)
40    pool.pin(child)
41    keep = processed // BLOCK_SIZE
42    unpin = child[keep:]
43    child[keep:] = []
44    if processed % BLOCK_SIZE:
45        tail = pool.allocate()
46        copies.append((unpin[0], tail))
47        child.append(tail)
48    return child, unpin

In both runs, the fork gives blocks 1, 2 and 3 two owners, as Figure 4.1 shows. The parent then finishes. Its release lowers every count by one, but blocks 1 and 2 still belong to the child and do not enter the free list. The script asserts this. Ending one owner’s lifetime has not ended the data’s lifetime.

Release the tail pin before the queued copy
  fork: parent=[1, 2, 3] child=[1, 2, 4] refs 1:2 2:2 3:2 4:1 copy=[(3, 4)]
  parent finishes: free=[]; shared [1, 2] kept
  writer receives block 3 before the copy runs
  child reads [100, 101]...[900, 901]: correct=False

Release the tail pin after the queued copy
  fork: parent=[1, 2, 3] child=[1, 2, 4] refs 1:2 2:2 3:2 4:1 copy=[(3, 4)]
  parent finishes: free=[]; shared [1, 2] kept
  writer waits; receives block 3 after the copy
  child reads [100, 101]...[108, 109]: correct=True

The difference lies in the tail. The copy from block 3 to block 4 is queued, standing for work submitted to a device that runs later. In the first run, the host releases the child’s pin on block 3 as soon as it has queued the copy. The parent has already finished, so the block’s count reaches zero, so a new writer receives it and stores its own rows. When the queued copy finally runs, it copies the writer’s data. The child reads 900 and 901 where it expected 108 and 109.

In the second run, the host keeps the pin until the copy has run. The writer must wait for a block, then receives block 3 safely. The safe policy costs capacity for a short time: a block that no request will use again stays unavailable until queued work has finished reading it.

The writer in this lab does not wait for the copy queue. If every later write were queued behind the copy in the same order, the early release would not corrupt this example. A host-side allocator usually cannot assume that about every future writer. The rule “release a source only after work that reads it has completed” avoids depending on that assumption.

4.3 Follow DeepInfer’s ownership path

DeepInfer implements both experiments’ ideas, with one addition: there are two pools. Figure 4.2 places them on the request path from chapter 1. We will follow one request from admission to release, then a fork of its prefix.

4.3.1 Two pools, one allocator class

The engine creates both pools when it starts:

265self.block_manager = BlockManager(cache_config.num_gpu_blocks)
266self.state_manager = BlockManager(cache_config.num_state_seqs) if cache_config.bytes_of_state_per_seqs > 0 else None

The first manager issues paged KV blocks. For the V4 model, the block size is 128 token positions (model_metadata_utils.py, lines 134–135). The cache engine’s field descriptions place the compressed-attention cache and the index cache in these blocks, each block holding $128∕\text{ratio}$ rows for its compression ratio (cache_engine.py, lines 206–219). The compression and indexing appendix (Appendix A) explains those rows. Here they are simply rows assigned to token ranges.

The second manager issues state slots, one per running request. A slot holds the SWA ring from chapter 3 and the compression state. The pool exists only if the model has such state. Its size is $\lfloor 1.2\times \mathtt{\text{max\_num\_seqs}}\rfloor +1$ by default (cache_engine.py, lines 201–202; environ.py, line 162). Because slot 0 is never issued, this provides $\lfloor 1.2\times \mathtt{\text{max\_num\_seqs}}\rfloor$ usable slots [23, 30].

The order of sizing is revealing. After a profiling run, the worker subtracts the fixed state reservation from available device memory, doubled for what its comment calls the state store and a copy buffer. Only the remainder becomes paged blocks (deepinfer/worker/base.py, lines 237–246). The number of state slots follows from the request limit; the number of blocks follows from whatever memory is left [36].

The block table retains the paged-addressing idea introduced at the start of this chapter. The additional allocation boundary follows the model: compressed KV and index rows belong to token blocks, while the SWA ring and rolling compression state belong to a request. DeepInfer gives both regions the same allocator machinery but different capacities and copy lifetimes. This is a comparison between a general paging mechanism and its use for this model, not a claim that every other engine has only one kind of state.

4.3.2 An owner count decides when a block is free

Both managers use the same small structure, Blocks. Its allocate method, lines 82–86, takes the head of an intrusive free list, a linked list threaded through the block records themselves, and sets the block’s owner count to one. Its comment says, in the author’s translation, that removal from the list has already checked for a zero count. The other two operations change an existing count:

88void pin_block(int block_id) {
89    TORCH_CHECK(block_id > 0 && block_id < static_cast<int>(_blocks.size()));
90    auto &block = _blocks[block_id];
91    TORCH_CHECK(block.ref_count > 0, "try to pin free block");
92    block.ref_count += 1;
93}
94
95void free(int block_id) {
96    TORCH_CHECK(block_id > 0 && block_id < static_cast<int>(_blocks.size()));
97    auto &block = _blocks[block_id];
98    block.ref_count -= 1;
99    TORCH_CHECK(block.ref_count >= 0, "ref_count corrupted: ref_count < 0");
100    if (block.ref_count == 0) {
101        _push_back(block_id);
102    }
103}

pin_block adds an owner but refuses a free block: sharing must start from a live one. free removes one owner and returns the block to the back of the list only at zero. The Python binding exposes free_blocks under the name unpin_blocks (pybind_scheduler.cpp, line 68). Releasing a request and releasing a pin are the same operation [34, 35].

The structure’s constructor explains the lab’s reserved block. Its comment at line 27 says, in the author’s translation, “GPU block 0 is unused, so 0 here means empty”. Block 0 terminates the free list; later we will see it mark a copy with no destination.

can_allocate_nolock, lines 128–136, returns the lab’s three answers. A request for at least num_blocks blocks is NEVER, because block 0 can never be issued. The scheduler combines this answer with the state pool: if blocks are available but no state slot is, the result becomes LATER (scheduler.hpp, lines 194–201). A request needs both kinds of storage to run [33].

4.3.3 Reserve positions that the history does not show

How many blocks should a request own? DeepInfer’s answer includes positions that are not yet tokens:

291int get_num_blocks() const {
292    return (get_len() + schedule.total_inflight_tokens + get_nextn_len() + cache.block_size - 1) / cache.block_size;
293}

The first term is the token history from chapter 1. The second counts positions of steps that have been scheduled but whose results have not yet reached the host. Under ZOS, the scheduler plans step $N+1$ while step $N$ is still running, so the history is behind the work already submitted. The third term counts draft tokens awaiting verification. Memory must cover all three, or a step could write beyond its table.

When a request is first admitted for prompt work, allocate_seq takes this many blocks and one state slot (scheduler.hpp, lines 518–519; block_manager.hpp, lines 207–214). The count uses the whole current prompt. A long prompt may be processed in several chunks, but its blocks are reserved at the first one. Chunking divides the computation, not the memory claim [33, 34].

4.3.4 Each decode step adjusts the table

Before each decode step, the scheduler brings every running table back to the required length:

817void _append_slot(Sequence &seq) {
818    int num_blocks = seq.get_num_blocks();
819    if (static_cast<decltype(num_blocks)>(seq.cache.block_table.size()) < num_blocks) {
820        seq.cache.block_table.emplace_back(block_manager->allocate_single_nolock());
821    } else if (static_cast<decltype(num_blocks)>(seq.cache.block_table.size()) > num_blocks) {
822        std::vector<int> redundant_blocks(seq.cache.block_table.begin() + num_blocks, seq.cache.block_table.end());
823        block_manager->free_blocks_nolock(redundant_blocks);
824        seq.cache.block_table.erase(seq.cache.block_table.begin() + num_blocks, seq.cache.block_table.end());
825    }
826    TORCH_CHECK(static_cast<decltype(num_blocks)>(seq.cache.block_table.size()) == num_blocks,
827        "len(block_table) = ", seq.cache.block_table.size(), ", but should be ", num_blocks, ". seq_len = ", seq.get_len());
828}

The growth branch takes one block; the lab’s general resize can take more. With 128-position blocks and, by default, at most six new positions per request in a step, one block is enough; the final check would catch a shortfall. The shrink branch is the lab’s rejected-draft case: when fewer positions are reserved, surplus blocks go back immediately. The decode loop calls this function only after checking that a block is free. If none is, it preempts another request (lines 629–640). Chapter 5 examines that choice [33].

4.3.5 A written row is not yet a valid row

A table says which blocks a request owns. It does not say which rows inside them are meaningful. That is the job of the request’s processed-token count:

216void advance_optimistically(int token_count) {
217    previous_num_tokens = num_tokens;
218    num_tokens += token_count;
219}
220
221void commit_processed_tokens(int token_count) { num_tokens = previous_num_tokens + token_count; }

When a step’s output is processed, the scheduler first advances num_tokens by every position the step evaluated (scheduler.hpp, line 927). For decode, it then commits the true count. The comment at line 998 says, in the author’s translation: “decode: correct cache.num_tokens from the optimistically advanced value to the number actually accepted (rejected drafts do not occupy KV)” [33, 6].

“Do not occupy” describes meaning, not bytes. The verification step wrote rows for its drafts into the request’s blocks. After rejection those bytes remain until a later token overwrites them, but they lie beyond num_tokens and no reader should use them. Validity is decided by the count, not by whether bytes are present. Chapter 3 drew the same distinction between leaving the window and being overwritten. Here it applies to rows that were never accepted.

4.3.6 Fork a prefix by pinning it

Prefix sharing enters through Scheduler::fork_seq (lines 1061–1122). The released benchmark uses it to give a batch of decode requests one cached prompt, in fork_request_batch (examples/benchmark.py, lines 331–353). Production prefix reuse also involves the external KV cache system, which the release does not include [1, 15]. The fork itself is a few lines:

1110block_manager->pin_blocks(fork_seq->cache.block_table);
1111if (state_manager) {
1112    fork_seq->cache.state_id = seq->cache.state_id;
1113    state_manager->pin_block(seq->cache.state_id);
1114}

At this point the child shares everything, including the parent’s state slot. That cannot last. A state slot holds a ring that every decode step writes, as chapter 3 showed. Two requests cannot advance one ring.

The benchmark’s forks share a fully processed prompt, so they take the scheduler’s path for forks entering decode. It separates what must be private:

602if (forked) {
603    if ((state_manager && !state_manager->can_allocate_single_nolock()) || !block_manager->can_allocate_single_nolock()) {
604        break;
605    }
606    _clone_state(*seq);
607    int new_block = block_manager->allocate_single_nolock();
608    int num_blocks = (seq->cache.num_tokens + seq->cache.block_size - 1) / seq->cache.block_size;
609    TORCH_CHECK_LE(num_blocks, static_cast<int>(seq->cache.block_table.size()));
610    clone_blocks.emplace_back(seq->cache.block_table[num_blocks - 1]);
611    clone_blocks.emplace_back(new_block);
612    for (size_t i = num_blocks; i < seq->cache.block_table.size(); ++i) {
613        clone_blocks.emplace_back(seq->cache.block_table[i]);
614        clone_blocks.emplace_back(0);
615    }
616    seq->cache.block_table[num_blocks - 1] = new_block;
617    seq->cache.block_table.erase(seq->cache.block_table.begin() + num_blocks, seq->cache.block_table.end());
618}

The fork waits unless one state slot and one block are free. _clone_state allocates a new state slot and records an (old, new) copy (lines 203–212). num_blocks counts the blocks that hold processed rows. The last of them is copied into the new block; this path copies it even when it happens to be full. Every later parent block is recorded with destination 0: release without copying. The table then points to the shared complete blocks followed by the private copy.

Forks whose prompt continues beyond the shared prefix are admitted through the prompt path instead (lines 412–437). There the shared part is the num_reserve_blocks complete blocks, and the tail is copied only when the processed count ends partway through a block. The comment at line 427 says “has a tail, needs copy” (author’s translation). That is the rule in our lab [33].

This is the lab’s rule with one more kind of state. Complete prefix blocks are read but not written, so they are shared. The partial tail and the state ring will be written, so each owner gets its own copy.

4.3.7 Release a source only after its copy

The copy lists reach the device during LLMEngine.step (llm_engine.py, lines 889–905). The C++ helper queues a device copy for each pair with a nonzero destination and adds every source to a release list, whether or not it was copied (scheduler.hpp, lines 1288–1300). The engine does not release those sources at once:

95def collect_clone_unpins(self, states_to_unpin: np.ndarray[np.int32], blocks_to_unpin: np.ndarray[np.int32]) -> None:
96    assert self.clone_unpins_to_release is None
97    self.clone_done_evt.record()
98    self.clone_unpins_to_release = (states_to_unpin, blocks_to_unpin)

An event marks a point in the device’s work queue; the host can later wait until the device has passed it. The release happens in DeferredRelease.release. The engine calls it from pre_sampler_sync later in the same step, as the comment at line 904 says. By then the worker has launched the forward pass and sampling (worker.py, lines 113–164):

109if self.clone_unpins_to_release is not None:
110    states_to_unpin, blocks_to_unpin = self.clone_unpins_to_release
111    self.clone_unpins_to_release = None
112    self.clone_done_evt.synchronize()
113    if states_to_unpin.size > 0:
114        self.state_manager.unpin_blocks(states_to_unpin)
115    if blocks_to_unpin.size > 0:
116        self.block_manager.unpin_blocks(blocks_to_unpin)

This is the second run of the sharing lab. The host waits for the copy event before the sources can return to either free list. The copies were queued before the forward pass, so the host’s wait does not hold back the forward and sampling work already queued behind them [7]. The class docstring states the general purpose: collect resources to release, then release them at the next pre_sampler_sync, overlapping with device computation (lines 68–71). The same object delays the output-pool release from chapter 1 and the release of per-step metadata [5].

4.3.8 Finishing returns every reference

A finished or preempted request gives back what it owns:

738void free_seq_nolock(Sequence &seq) {
739    release_mtp_info_slot(seq);
740    if (seq.cache.state_id != -1) {
741        state_manager->free_blocks_nolock({seq.cache.state_id});
742    }
743    block_manager->free_blocks_nolock(seq.cache.block_table);
744    seq.cache.block_table.clear();
745}

release_mtp_info_slot returns a third resource, a slot for draft information that chapter 8 will use. Then the state slot and every block lose one owner. A block shared with a fork survives; an unshared block returns to its free list. free_blocks_nolock walks the table in reverse (block_manager.hpp, lines 147–149), which is why the lab’s free list ended with B’s blocks in reverse order. The source does not explain that choice.

4.3.9 Output storage has a different owner

Chapter 1’s result records belong to RequestOutputPool. The caller must clone any record it keeps before the pool releases it. That ownership change differs from pinning a KV block: cloning creates a new object, while pinning adds an owner of the same block. The binding’s clone constructs a new object from the one at the supplied address:

16py::class_<RequestOutput>(engine, "RequestOutput")
17    .def_static(
18        "clone", [](uint64_t ptr) { return std::make_unique<RequestOutput>(*reinterpret_cast<RequestOutput *>(ptr)); }, py::arg("ptr"))

Here reinterpret_cast treats the integer as an object address; make_unique allocates a separate object managed by the binding. In contrast, the pool’s release loop destroys the original and makes its slot available again:

72void release(py::array_t<uint64_t> ptrs_py) {
73    const uint64_t *ptrs = ptrs_py.data();
74    int size = ptrs_py.size();
75    for (int i = 0; i < size; ++i) {
76        auto p = reinterpret_cast<RequestOutput *>(ptrs[i]);
77        std::destroy_at(p);
78        free_.push_back(p);
79    }
80}

The deferred-release helper invokes this operation at a later pre_sampler_sync. Merely appending raw_output to a Python list would leave an integer referring to destroyed or reused storage. Cloning during collection preserves the result across that boundary. The pointer-based interface allows explicit reuse of pool slots; the cost is that callers must follow a lifetime contract that ordinary Python references do not express [9, 10, 5].

4.4 A state slot has regions with different lifetimes

Ownership is not only about who may write. It also determines how long data must survive and where it may go. The StateBlock docstring in cache_engine.py, lines 221–241, describes one request’s slot. In the author’s translation of its central rule: state that must be stored goes in a front stored region; state that need not be stored, but is still needed while running, goes in a rear dropped region. Taking the first state_size_to_disk bytes of the inference layout therefore excludes the dropped region by truncation. The next sentences of the docstring add a condition (lines 226–227). If extra prefill state is enabled, each slot holds the longer kvstore ring from chapter 3, which must first be reordered into the inference window; only when it is disabled can the original buffer be cut directly. GPUCacheInterface.select_state_slice_to_store implements both branches: a direct slice when DEEPINFER_PREFILL_STORE_EXTRA_STATE is off, and a per-view reordering otherwise (lines 763–817) [23].

The stored region holds the SWA rings of the layers whose window must survive across requests. The dropped region holds the SWA layers that the deployment Q&A says are recovered by prefill recomputation, and the compression state. Chapter 9 explains why those layers can be recomputed [2].

The code sets state_size_to_disk to the padded total for the stored SWA layers, using the inference ring length (line 306). Chapter 3 derived that ring as 134 positions of 528 bytes:

That is already a multiple of 32, so the 32-byte padding adds nothing. The Q&A states that V4.1 stores its first 20 layers, giving

the figure it reports [2]. The layer count comes from the model configuration, which this book does not inspect; the arithmetic shows that the Q&A’s number follows from the inspected layout.

The released engine prepares this layout but leaves its use to the missing storage system. LLMEngine.should_copy_state returns False, and do_save_kvcache and related methods are empty hooks; their docstrings call them virtual functions (llm_engine.py, lines 811–822). The Q&A describes the production behaviour: copying the window at 128-token boundaries, keeping SWA state in the producing machine’s host memory and fetching it across machines on a hit. That is documented service design, not released engine code [5, 2].

4.5 What a cache entry means

We can now answer chapter 1’s third question, which memory holds the state needed to continue, with some precision. Two different questions hide in it.

The first asks which rows form a committed prefix: rows that later steps, a fork or a store may reuse. Such a row must satisfy four conditions together:

1. The request owns the block or state slot, alone or through a pin.
2. The row’s position is below the request’s processed-token count, num_tokens.
3. The row has not been replaced by a later write to the same physical slot, as in chapter 3’s ring.
4. Any queued copy that fills it has completed before it is read.

The second asks which rows the current step’s computation may read. That range is larger. When the scheduler builds a step’s metadata, it sets the key length to num_tokens plus the positions of steps still in flight plus this step’s query positions (scheduler.hpp, line 697). Attention in this step reads rows that the step has just written for its own queries and drafts, before the host has advanced num_tokens. Those rows are visible because the step’s metadata and indices say so. They join the committed prefix only if the later correction accepts them [33].

Ownership is recorded in owner counts, the committed prefix in num_tokens, the current step’s visibility in its metadata and indices, physical replacement in address arithmetic, and copy completion in device events. No single data structure answers the whole question. The next chapter asks what happens when the free lists run short and the scheduler must decide who waits.

4.6 Experiments to try next

1. In labs/09_blocks.py, reserve five drafts instead of two at length 11. How many blocks must resize add at once? DeepInfer’s _append_slot adds at most one. State a condition on block size and positions added per step under which one block is always enough.
2. In labs/10_share.py, set PROCESSED to 8 instead of 10. Predict the copy list, the owner counts and whether releasing the pin early still corrupts the child. The script’s final check expects corruption only when a tail copy was queued; confirm that it passes.
3. Change release in either lab to walk blocks in forward order. Which block does the next allocation receive? Suggest a reason an engine might prefer one order, and describe an experiment that could test it. The source gives no reason, so label your answer as a hypothesis.
4. The state pool provides $\lfloor 1.2\times \mathtt{\text{max\_num\_seqs}}\rfloor$ usable slots. List situations in which more than max_num_seqs slots can be held at once. Check each against the fork and release code before relying on it.