About

Inside nano-vLLM: Anatomy of a LLM Inference System

nano-vLLM

nano-vLLM is based on vLLM, which is known by paged attention, it is a simple implementation but contained the core theory of paged attention and flash attention.

1. Lifecycle of Request

1.1 Public API

The external entrypoint is LLM.generate(), it receives prompts and sampling parameters, then calls add_request() for each prompt, which will do two things:

  1. Converts string prompt into token ids through HF tokenizer.(BPE)
  2. Wraps token ids into a Sequence and pushes it into the scheduler.

At this point, we will no longer see the prompt as an identity of a request but a runtime object called Sequence instead.

1.2 Sequence

Sequence is the central request state in nano-vLLM. It stores tokens and kv-cache metadata.

Fields:

  • token_ids: prompt tokens plus generated tokens.
  • num_prompt_tokens: length of the original prompt.
  • num_cached_tokens: how many tokens have already been computed and written into KV cache.
  • num_scheduled_tokens: how many tokens will be executed in the current step.
  • block_table: mapping from logical blocks to physical KV-cache blocks.
  • status: WAITING, RUNNING, or FINISHED.

This design supports cyclic decoding, enabling fair scheduling during multi-user requests, the core idea is round-robin-like decode scheduling.

1.3 Scheduler

Scheduler owns two queues:

  • waiting: sequences that still need prefill work, prefill queue.
  • running: sequences that are ready for decode, decode queue.

schedule() first tries to schedule prefill work from waiting. If there is no prefill work, it schedules decode work from running.

Note: modern vLLM schedulers can mix prefill and decode, often prioritizing decode to reduce inter-token latency.

1.4 One Engine Step

One LLMEngine.step() is the basic execution unit.

  1. Scheduler selects a batch.
  2. ModelRunner executes the batch and returns sampled token ids.
  3. Scheduler postprocesses the result.

1.5 Postprocess

After model execution, postprocess() updates each scheduled sequence.

For both prefill and decode:

  • Hash completed KV blocks for prefix cache.
  • Increase num_cached_tokens.
  • Clear num_scheduled_tokens.

For prefill:

  • If prefill is complete, the sampled token becomes the first generated token, otherwise do nothing.

For decode:

  • Append the sampled token to token_ids.
  • If EOS or max_tokens is reached, mark the sequence as finished.
  • Release KV block through BlockManager.

2. KV Block Management

Instead of storing each sequence in one contiguous kv tensor region, paged attention splits kv cache into fixed-size physical blocks. Each sequence owns a block_table, which maps its logical blocks to physical blocks.

2.1 Blocks

A Sequence is a logical token stream, it is split into logical blocks by block_size, e.g. block_size = 4

token_ids = [t0, t1, t2, ..., t9]
->
logical block 0: [t0, t1, t2, t3]
logical block 1: [t4, t5, t6, t7]
logical block 2: [t8, t9]

The Sequence.block_table maps these logical blocks to physical KV blocks:

seq.block_table = [7, 3, 12]
logical block 0 -> physical block 7
logical block 1 -> physical block 3
logical block 2 -> physical block 12

2.2 Block Metadata

class Block:
    block_id  # physical block id.
    ref_count # how many sequences are currently using this block.
    hash      # prefix-cache hash for this block.
    token_ids # original token ids used to verify hash correctness.

hash is not enough by itself because hash collision is theoretically possible.

2.3 BlockManager State

It maintains the four core data structures:

blocks: list[Block]              # all physical block metadata.
hash_to_block_id: dict[int, int]  # prefix-cache index from block hash to physical block id.
free_block_ids: deque[int] # blocks that can be allocated.
used_block_ids: set[int]   # blocks currently referenced by running sequences.

2.4 Allocation

When a new sequence enters prefill, scheduler calls:

num_cached_blocks = block_manager.can_allocate(seq)
block_manager.allocate(seq, num_cached_blocks)

can_allocate() is a dry-run check. It scans complete logical blocks of the sequence and tries to find cached prefix blocks.

If a cached block is found, it increases num_cached_blocks.

If the cached block is still used by another sequence, it does not consume a new free block, because it can be shared through ref_count.

If the cached block is free, it still needs to be removed from free_block_ids when reused, so it still counts as consuming one free block.

allocate() then mutates state:

  • cached used block: ref_count += 1
  • cached free block: move from free to used, set ref_count = 1
  • uncached block: allocate a new physical block
  • append all physical block ids to seq.block_table

After allocation:

seq.num_cached_tokens = num_cached_blocks * block_size

This tells ModelRunner how many prefix tokens can be skipped.

2.5 Prefix cache hash

block manager only hash completed blocks, the hash is chained:

hash(block_0)
hash(hash(block_0), block_1)
hash(hash(block_1), block_2)

After ModelRunner finishes executing scheduled tokens, postprocess() calls block_manager.hash_blocks(seq) to save completed blocks.

2.6 Decode

Most decode steps write into the current last block, but when the new token crosses a block boundary, BlockManager must allocate a new physical block.

can_append(seq)
may_append(seq)

judge condition:

len(seq) % block_size == 1

2.7 Deallocation

When a sequence finishes or is preempted, BlockManager deallocates its blocks:

for block_id in reversed(seq.block_table):
    block.ref_count -= 1
    if block.ref_count == 0:
        move block from used to free

Deallocation does not clear hash and token_ids

At this point, BlockManager only decides which physical blocks a sequence owns. It does not compute GPU tensor offsets directly. ModelRunner later converts block_table into slot_mapping and block_tables, which are consumed by the attention kernel.

3. Model Runner

ModelRunner converts scheduled sequences into model inputs, runs the model, samples the next token, and returns token ids.

3.1 Initialization

load model weights
    -> warmup_model()
    -> allocate_kv_cache()
    -> bind KV cache to each attention layer

warmup_model() estimates temporary peak memory, and allocate_kv_cache() uses that estimate to decide how many KV blocks can be allocated.

After the KV cache tensor is allocated, ModelRunner injects each layer’s KV cache slice into the corresponding attention module:

for module in self.model.modules():
    if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
        module.k_cache = self.kv_cache[0, layer_id]
        module.v_cache = self.kv_cache[1, layer_id]
        layer_id += 1

3.2 Run

It receives two parameters: seqs and is_prefill

The important idea is that prefill and decode use different input construction strategies.

Prefill may execute many prompt tokens for each sequence. Decode executes only one token per sequence.

3.3 Prefill Preparation

For each scheduled sequence, prepare_prefill() uses:

start = seq.num_cached_tokens
end = start + seq.num_scheduled_tokens
input_ids = seq[start:end]
positions = range(start, end)

Cached token are skiped and do not appera in input_ids.

attention metadata:

- input_ids: token ids to execute in this batch.
- positions: position ids for RoPE.
- cu_seqlens_q: query sequence boundaries in the flattened batch.
- cu_seqlens_k: full context boundaries, including cached prefix.
- max_seqlen_q: maximum number of query tokens in the batch.
- max_seqlen_k: maximum context length in the batch.
- slot_mapping: physical KV slots where current tokens should write K/V.
- block_tables: physical block tables used when historical KV needs to be read.

3.4 Decode Preparation

Each running sequence only contributes one input token: its last token.

input_ids = [seq.last_token for seq in seqs]
positions = [len(seq) - 1 for seq in seqs]

Even though decode only computes one new token per sequence, the new token must attend to all previous tokens. Therefore decode always needs:

context_lens: current full length of each sequence.
block_tables: where historical KV blocks are located.
slot_mapping: where the new token's K/V should be written.

current token -> compute Q/K/V
historical tokens -> read K/V from KV cache through block_tables
new K/V -> write into KV cache through slot_mapping

3.5 Runtime Context

Instead of passing all metadata through every model layer manually, ModelRunner sets a global context before model forward. The attention layer then reads this context during its forward pass.

ModelRunner.set_context(...)
    -> Qwen model forward
    -> Attention.forward()
    -> get_context()
    -> use parameters

this keeps the model forward signature simple while still allowing attention kernels to access runtime metadata.

3.6 Attention and KV Cache

The attention layer has two responsibilities related to KV cache:

  1. store current token K/V into KV cache.
  2. read historical K/V when needed. Current K/V is written by slot_mapping.
token index in current batch
  -> slot_mapping
  -> physical KV cache slot

Historical K/V is read through block_tables.

sequence logical block
  -> block_table
  -> physical KV block
  -> KV cache tensor

Sequence.block_table
  -> ModelRunner.block_tables
  -> Attention kernel reads historical KV

Sequence.block_table + token position
  -> ModelRunner.slot_mapping
  -> Attention kernel writes current KV

3.7 Sampling

The sampler applies temperature and samples the next token from the probability distribution. The returned token_ids have the same order as the scheduled seqs.

seqs      = [seq0, seq1, seq2]
token_ids = [tok0, tok1, tok2]

Then scheduler update each sequence.

for seq, token_id in zip(seqs, token_ids):
    ...

In general, scheduler decides which sequences and tokens should run, BlockManager decides which physical KV blocks each sequence owns, ModelRunner converts this decision into GPU-facing tensors, and Attention uses those tensors to read/write KV cache.