3.4 KV Cache and BlockManager

📑 Contents

• 1. Why KV is split into blocks
• 2. Block: the smallest management unit
• 3. Sequence: a request’s state machine
• 4. BlockManager: allocate, append, and free
• 5. Prefix cache: reuse history with hashes
• Summary
• Self-check checklist
• References

---

1. Why KV is split into blocks

In decode, every new token needs the K/V of every token before it. If you reserve one contiguous chunk of GPU memory at the “max length” for each sequence, most requests never fill it, and the unused part is wasted. That is the internal fragmentation problem from chapter 2.

PagedAttention borrows paging from operating systems: cut the KV Cache into fixed-size blocks, give a sequence as many blocks as it needs, and the blocks do not have to sit next to each other in physical memory. nano-vllm’s block_manager.py is the smallest version of this, about 120 lines. It only keeps the ledger of blocks and never touches matrix math, so it lives in the scheduling layer rather than layers/.

💡 Tip: nano-vllm’s default block size is kvcache_block_size = 256 (256 tokens per block), much larger than the 16 that is common in vLLM. That number is not arbitrary: when flash-attn reads paged KV it requires the block size to be a multiple of 256 (see 3.2), and 256 is already the smallest value you can pick. The cost is more waste when the last block is only partly filled, and coarser matching for the prefix cache (see section 5).

2. Block: the smallest management unit

Each Block corresponds to one physical stretch of KV Cache GPU memory, but the Block object itself is only metadata: how many sequences reference it, and which tokens it holds:

class Block:
    def __init__(self, block_id):
        self.block_id = block_id       # physical block index
        self.ref_count = 0             # reference count: how many sequences share this block
        self.hash = -1                 # content hash (-1 means not full / not registered yet)
        self.token_ids = []            # tokens in the block (used to verify a prefix really matches)

    def update(self, hash, token_ids): # register when the block fills and the hash is ready
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):                   # when taken from the free queue and given to a new sequence
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []

What the three fields are for:

• 🔑 ref_count (reference count): when several sequences share one prefix block, that block’s ref_count is greater than 1, and it can only be freed when the count hits 0. Shared blocks are always already full, so their content no longer changes. That is why nano-vllm does not need copy-on-write: every newly generated token always goes into a sequence’s own last block.
• 🔑 hash: a fingerprint of the block’s content. Only full blocks get a hash and join the prefix cache; a partially filled block keeps hash = -1 and cannot be reused by another sequence, because it will still receive more tokens.
• token_ids: a copy of the tokens in the block. After a hash hit, they are compared one by one so a hash collision cannot reuse the wrong block.

3. Sequence: a request’s state machine

Sequence in sequence.py tracks one request from arrival to finish. An enum covers the lifecycle:

class SequenceStatus(Enum):
    WAITING = auto()     # in the waiting queue, not computed yet
    RUNNING = auto()     # currently scheduled to run
    FINISHED = auto()    # done (hit EOS or reached max_tokens)

The Sequence constructor looks like this:

class Sequence:
    block_size = 256                 # class var: every sequence shares the same block size
    counter = count()                # class var: auto-incrementing ID generator

    def __init__(self, token_ids, sampling_params=SamplingParams()):
        self.seq_id = next(Sequence.counter)
        self.status = SequenceStatus.WAITING
        self.token_ids = copy(token_ids)   # prompt + tokens generated so far (a copy; do not mutate the caller's list)
        self.last_token = token_ids[-1]    # last token; used as model input in decode
        self.num_tokens = len(self.token_ids)
        self.num_prompt_tokens = len(token_ids)
        self.num_cached_tokens = 0         # tokens already computed, with KV written into the cache
        self.num_scheduled_tokens = 0      # tokens scheduled for this step
        self.is_prefill = True             # instance attr: still in the prefill phase?
        self.block_table = []              # logical block → physical block_id
        # sampling hyperparameters flattened onto the instance (no separate sampling_params object)
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos

A few design points:

• 🔑 block_table is like a page table: a Python list whose i-th element is the physical block id for this sequence’s i-th logical block. With it, physical blocks scattered in GPU memory can be read as one logically contiguous sequence.
• 🔑 is_prefill is an instance attribute, not a property: the scheduler sets it to False the first time the sequence enters a decode batch (see the decode branch in 3.5); on preemption it is reset to True, and the next run prefills again.
• num_scheduled_tokens and num_cached_tokens: these support chunked prefill. A long prompt may take several steps to finish. num_scheduled_tokens is how many this step computes; num_cached_tokens is how many have been written into the cache so far. Prefill is done when their sum equals num_tokens. On a prefix-cache hit, allocate sets num_cached_tokens directly to hit_blocks × 256.
• append_token appends a new token to token_ids and updates last_token and num_tokens.

Common derived fields are properties, for example:

@property
def num_blocks(self):                        # how many blocks are needed (ceil)
    return (len(self) + self.block_size - 1) // self.block_size

@property
def last_block_num_tokens(self):             # how many tokens are in the last block
    return len(self) - (self.num_blocks - 1) * self.block_size

📌 Key point: num_blocks and last_block_num_tokens are what BlockManager uses to decide whether a new block is needed. When the sequence length crosses a block boundary, another physical block must be allocated.

3.1 Across processes, only send the fields that matter

When TP>1, every step the main process pickles this batch of Sequences and sends them to the other ranks over shared memory (see 3.6). A sequence’s token_ids can have thousands of elements; sending the full list every step is wasteful. Sequence customizes __getstate__ and packs only what workers need:

def __getstate__(self):
    # prefill needs the full token list; decode only needs the last token
    last_state = self.last_token if not self.is_prefill else self.token_ids
    return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens,
            self.num_scheduled_tokens, self.block_table, last_state)

__setstate__ restores from the type of last_state: if it is a list, restore token_ids; otherwise clear token_ids and keep only last_token. In decode, each sequence a worker receives is a few integers plus a block_table. That is also why prepare_decode reads seq.last_token instead of seq.token_ids[-1]. Sampling parameters are not packed, because sampling only runs on rank 0.

4. BlockManager: allocate, append, and free

BlockManager holds every physical block, a free queue, and a hash table. Outwardly it offers allocate, append, and free. Initialization:

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size
        self.blocks = [Block(i) for i in range(num_blocks)]
        self.hash_to_block_id = {}                        # prefix cache: hash → block id
        self.free_block_ids = deque(range(num_blocks))   # free-block queue (FIFO)
        self.used_block_ids = set()

① Allocate (give enough blocks in one shot before prefill)

The scheduler first calls can_allocate(seq): it walks the hash chain block by block, counts how many prefix-cache hits there are, then checks whether free blocks can cover the rest. If not, it returns -1; otherwise it returns the hit count. After that looks feasible, the scheduler calls allocate(seq, num_cached_blocks) to hang physical blocks on seq.block_table. Matching details are in section 5.

New blocks come from _allocate_block, which takes from the head of the free queue:

def _allocate_block(self):
    block_id = self.free_block_ids.popleft()
    block = self.blocks[block_id]
    assert block.ref_count == 0
    # this block was registered before, and the table still points at it: about to overwrite, drop the entry
    if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
        del self.hash_to_block_id[block.hash]
    block.reset()                  # ref_count = 1, clear hash and token_ids
    self.used_block_ids.add(block_id)
    return block_id

② Append (decode may need one more block each step)

def can_append(self, seq):
    # only when "the last token lands exactly in the first slot of a new block" do we need a free block
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

def may_append(self, seq):
    if len(seq) % self.block_size == 1:
        seq.block_table.append(self._allocate_block())

During decode scheduling, the token sampled in the previous step is already appended to the sequence, and this step must compute its K/V. When len(seq) % block_size == 1 is true, that token is the first token of a new block, so one free block is needed (the boolean True compares as 1); otherwise the current block still has room, and 0 free blocks are enough. may_append takes a block from the free queue and appends it to block_table only when a boundary is crossed.

③ Free (sequence finished or preempted)

def deallocate(self, seq):
    for block_id in reversed(seq.block_table):   # put back from the last block forward
        block = self.blocks[block_id]
        block.ref_count -= 1
        if block.ref_count == 0:                 # only return to the free-queue tail when nobody uses it
            self._deallocate_block(block_id)
    seq.num_cached_tokens = 0
    seq.block_table.clear()

⚠️ Note: on free, decrement the reference count first, and only return the block to the free queue when it hits 0. A prefix block shared by several sequences is not freed just because one of them finished. Also, _deallocate_block only moves the block id back into the free queue; it does not clear the block’s hash and token_ids, and the hash-table entry stays. Until _allocate_block actually takes the block and overwrites it, later requests can still hit it.

Two ordering details keep the cache around longer:

• The free queue is FIFO: _allocate_block takes from the head, free puts at the tail, so the most recently freed blocks are overwritten last.
• deallocate frees in reverse order, so a sequence’s later blocks enter the queue first and are overwritten first, while the early blocks sit further back. Early blocks are the ones most likely to be shared prefixes across requests, so they stay the longest.

5. Prefix cache: reuse history with hashes

Many requests share a prefix (the same system prompt, the same few-shot examples). Sharing that prefix’s KV skips repeated prefill work. nano-vllm does this with block-level hashes: fingerprint every full block, and let a new request look up blocks that were already computed.

5.1 Hash chain

Each block’s hash mixes in the previous block’s hash, forming a hash chain:

@classmethod
def compute_hash(cls, token_ids, prefix=-1):
    # feed "prefix hash + this block's tokens" into xxhash together
    h = xxhash.xxh64()
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))
    h.update(np.array(token_ids).tobytes())
    return h.intdigest()

🔑 Why mix in the prefix hash: even if two blocks have the same content, they are not the same context if the history before them differs, so they must not be reused. Chaining the prefix hash in means each block’s hash stands for “the full prefix from the start of the sequence through this block.” That is the same idea as a Merkle chain.

5.2 Register: call hash_blocks after a block fills

A block only gets a hash once it is full. Registration happens in the scheduler’s postprocess (see 3.5): after each model step, before updating num_cached_tokens, call hash_blocks once for the sequence:

def hash_blocks(self, seq):
    start = seq.num_cached_tokens // self.block_size                               # blocks that were already full before this step
    end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size    # blocks that are full after this step
    if start == end: return                                                        # no newly filled block this step
    h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1          # previous link in the chain
    for i in range(start, end):
        block = self.blocks[seq.block_table[i]]
        token_ids = seq.block(i)
        h = self.compute_hash(token_ids, h)
        block.update(h, token_ids)
        self.hash_to_block_id[h] = block.block_id                                  # overwrite any older entry

[start, end) is the set of blocks newly filled this step. One prefill step may fill several blocks; decode computes 1 token per step, so only the step that exactly fills a block registers 1 block. When a chunked-prefill chunk stops in the middle of a block, that block is not hashed until a later chunk fills it. The previous link is read from the hash already stored on block_table[start - 1]; there is no need to recompute from the start.

hash_to_block_id[h] overwrites directly. When two blocks have the same content (for example two requests with the same prefix prefill at the same time and both miss), the table only keeps the one registered last. That is why _allocate_block checks hash_to_block_id.get(block.hash) == block_id before deleting an entry: the table may already point at a different block.

5.3 Match: can_allocate and allocate

When a new request arrives, the scheduler first calls can_allocate and walks the hash chain block by block:

def can_allocate(self, seq):
    h = -1
    num_cached_blocks = 0
    num_new_blocks = seq.num_blocks              # worst case: every block needs a free block
    for i in range(seq.num_blocks - 1):          # note: the last block is never looked up
        token_ids = seq.block(i)
        h = self.compute_hash(token_ids, h)
        block_id = self.hash_to_block_id.get(h, -1)
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
            break                                # stop at the first miss
        num_cached_blocks += 1
        if block_id in self.used_block_ids:      # hit an in-use block: does not consume a free block
            num_new_blocks -= 1
    if len(self.free_block_ids) < num_new_blocks:
        return -1
    return num_cached_blocks

Once that fits, allocate hangs the hit blocks on the same hash chain and takes the rest from the free queue:

def allocate(self, seq, num_cached_blocks):
    h = -1
    for i in range(num_cached_blocks):
        h = self.compute_hash(seq.block(i), h)
        block_id = self.hash_to_block_id[h]
        block = self.blocks[block_id]
        if block_id in self.used_block_ids:      # another sequence is using it: share
            block.ref_count += 1
        else:                                    # freed but not yet overwritten: pull it back out of the free queue
            block.ref_count = 1
            self.free_block_ids.remove(block_id)
            self.used_block_ids.add(block_id)
        seq.block_table.append(block_id)
    for i in range(num_cached_blocks, seq.num_blocks):
        seq.block_table.append(self._allocate_block())
    seq.num_cached_tokens = num_cached_blocks * self.block_size

The full matching flow is below. Diamonds are the two decisions, blue boxes are what happens on a hit, and red boxes are taking a new block from the free queue:

A few details in this code:

• 🔑 The last block is never looked up: the loop only goes to num_blocks - 1. Even if the whole prompt is in the cache, the last block is still recomputed. Prefill needs the logits of the last token to sample the next one; if every token hit, this step would send no token into the model and there would be no logits.
• Prefix hits must be contiguous: the first miss breaks. If one block in the middle is missing, later blocks still in the table are useless, because the earlier KV they depend on is no longer in the cache.
• Two kinds of hit blocks: an in-use block just does ref_count += 1 and does not consume a free block; a freed block that has not been overwritten must be pulled out of the free queue and reactivated, and still counts as one free slot. So can_allocate only decrements num_new_blocks when the hit is an in-use block. A preempted sequence that prefills again may hit its own recently freed blocks this way.
• token_ids double-check: after the hashes match, tokens are compared one by one so an xxhash collision cannot reuse the wrong block.

The figure below uses block_size = 4 on a 12-token request: L0 hits #1 that another request is using, L1 hits #3 in the free queue, and the last block is not looked up; it is taken from the free-queue head as #6 (warm colors are this request’s blocks; hatched blocks are held by other requests):

📌 Key point: prefix-cache granularity is a whole block, and with the “never look up the last block” rule, reuse is truncated when block_size = 256. Two requests that share a 600-token prefix can only reuse the first 512 tokens (2 blocks); the remaining 88 must be recomputed. Even a 512-token prompt that matches the cache exactly can only reuse block 1; block 2 is recomputed. Smaller blocks match more finely, but the block table is longer and hashing runs more often.

📝 Summary

• nano-vllm cuts the KV Cache into fixed-size physical blocks (default 256 tokens/block, the smallest flash-attn paged KV allows), and BlockManager keeps the ledger without touching matrix math.
• Block’s ref_count / hash / token_ids support shared free, prefix caching, and collision checks. Shared blocks are already full and no longer change, so copy-on-write is not needed.
• Sequence is the request state machine: block_table acts as a page table; is_prefill becomes False on the first decode entry and resets on preemption; __getstate__ only sends workers last_token in decode.
• can_append uses len(seq) % block_size == 1 to detect a block crossing; deallocate decrements the reference count first and only returns a block to the free-queue tail at 0, keeping the hash so the block can still be hit until it is overwritten.
• Prefix caching builds a hash chain with block-level xxhash that mixes in the prefix: hash_blocks registers newly filled blocks in postprocess; can_allocate matches contiguously and stops on the first miss; the last block is always recomputed.

🎯 Self-check checklist

• Can explain what fragmentation problem splitting KV into blocks solves versus “reserve max-length contiguous GPU memory”
• Can name what Block’s ref_count, hash, and token_ids each do, and why nano-vllm does not need copy-on-write
• Can explain why Sequence.block_table is like an OS page table, and what __getstate__ skips in decode
• Can read why len(seq) % block_size == 1 in can_append detects “needs a new block”
• Can say why deallocate decrements the reference count before freeing, and why a freed block can still be hit
• Can explain why prefix caching chains the prefix hash into the current block, and when hash_blocks registers which blocks
• Can say why can_allocate skips the last block, and how hitting an in-use block versus a free block differs for free-block demand
• Can compute how many tokens of a 600-token shared prefix are actually reused when block_size = 256

📚 References

• nano-vllm GitHub repository
• Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)
• vLLM docs - Automatic Prefix Caching
• xxHash official repository