Nano-vLLM, line by line

Step 00Orientation

Before reading any function, get three things straight: which files exist, how they stack, and the five ideas from LLM serving that every one of them assumes.

0.1Repository map

Twenty-one Python files. The package is nanovllm/; example.py and bench.py are the two programs that use it. Every file name below links to the step that explains it. Line counts include blank lines.

The .gitignore is the stock GitHub Python template, LICENSE is MIT, and assets/logo.png is the README logo. None of them affect behaviour.

0.2How the pieces stack

Read the picture as a call stack. Your script calls LLM.generate. The engine decides, on the CPU, which sequences run in the next step and where their keys and values live in GPU memory. The runner turns that decision into tensors, runs the model, and returns one sampled token per sequence. The model is an ordinary PyTorch module tree whose leaves are the layer files.

0.3Five ideas the code assumes

Each idea below names where it lives in this repository, so you can jump ahead if you already know it.

Prefill and decode

A request starts with a prefill: the whole prompt goes through the model in one forward pass, which writes every prompt token's keys and values (K/V) into a cache and produces the first output token. After that, each decode step feeds exactly one new token per sequence. Prefill is compute-bound and has many tokens per sequence; decode is memory-bound and has one. nano-vllm never mixes the two inside one step (Step 06).

KV cache

Attention for token t needs the K/V of every earlier token, at every layer. Keeping them on the GPU means a decode step only computes K/V for the one new token. For Qwen3-0.6B that cache costs 112 KiB per token (Step 07.4).

Paged KV

The cache is cut into fixed-size blocks of 256 token slots. A sequence owns a list of block ids, its block table, which do not have to be adjacent. That removes fragmentation and lets two sequences point at the same block when they share a prompt prefix (Step 05).

Continuous batching

The batch is rebuilt every step. Finished sequences leave and waiting ones join between steps, so the GPU stays busy without waiting for the slowest request in a fixed batch (Step 06).

Tensor parallelism

With tensor_parallel_size = N, N processes each hold 1/N of most weight matrices and 1/N of the attention heads, and combine partial results with NCCL collectives (Step 10.1).

0.4The route through this guide

1. Steps 01–03 cover the surface: how the package is used, every configuration knob, and the engine loop that drives everything.
2. Steps 04–06 are the CPU-side bookkeeping: the Sequence object, the paged block allocator with prefix caching, and the scheduler. Step 06 ends with a simulator that runs a JavaScript port of those three files.
3. Steps 07–08 cross to the GPU: the per-rank ModelRunner that sizes the KV cache, builds the batch tensors and replays CUDA graphs, and the Context side channel it fills.
4. Steps 09–11 are the model: Qwen3's module tree, each layer file, and the weight loader that fills sharded parameters.
5. Steps 12–15 put it back together: one step traced with real shapes, an index of where each optimisation lives, sharp edges found while reading, and a glossary.

Code blocks show the exact lines from commit e86692c with their real line numbers; highlighted lines are the ones the text discusses. Use Copy for the code without line numbers, or GitHub ↗ to open the same range in the repository. The “Mark this step as read” buttons tick the route in the sidebar; that progress is kept only in this browser.

Step 01The front door

What a user of the package touches: two exported names, two example programs, and the packaging metadata.

1.1The public surface: LLM and SamplingParams

The package exports exactly two names. LLM is an empty subclass of LLMEngine; it exists so that from nanovllm import LLM reads the same as vLLM's from vllm import LLM. All behaviour is in LLMEngine (Step 03).

1from nanovllm.llm import LLM2from nanovllm.sampling_params import SamplingParams

1from nanovllm.engine.llm_engine import LLMEngine2 3 4class LLM(LLMEngine):5    pass

1.2example.py: the smallest real use

• The model path must be a local directory. Config asserts os.path.isdir(model), so a Hugging Face hub id such as "Qwen/Qwen3-0.6B" will not work. The README's huggingface-cli download … --local-dir ~/huggingface/Qwen3-0.6B/ command produces the directory this script expects.
• The script loads its own tokenizer only to apply Qwen3's chat template with tokenize=False, so the prompts it passes are strings. The engine tokenizes them again inside add_request.
• enforce_eager=True skips CUDA graph capture. Start-up is faster; decode steps are slower (Step 07.9).
• One SamplingParams object is shared by both prompts; generate broadcasts it.
• Each output is a dict with "text" and "token_ids", in the same order as the prompts.

1import os2from nanovllm import LLM, SamplingParams3from transformers import AutoTokenizer4 5 6def main():7    path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")8    tokenizer = AutoTokenizer.from_pretrained(path)9    llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)10 11    sampling_params = SamplingParams(temperature=0.6, max_tokens=256)12    prompts = [13        "introduce yourself",14        "list all prime numbers within 100",15    ]16    prompts = [17        tokenizer.apply_chat_template(18            [{"role": "user", "content": prompt}],19            tokenize=False,20            add_generation_prompt=True,21        )22        for prompt in prompts23    ]24    outputs = llm.generate(prompts, sampling_params)25 26    for prompt, output in zip(prompts, outputs):27        print("\n")28        print(f"Prompt: {prompt!r}")29        print(f"Completion: {output['text']!r}")30 31 32if __name__ == "__main__":33    main()

1.3bench.py: the throughput benchmark

• 256 requests with prompts of 100–1024 random token ids drawn from [0, 10000], and max_tokens drawn from 100–1024. seed(0) makes the workload repeatable.
• ignore_eos=True makes every request produce exactly max_tokens tokens, so the total output token count is known in advance.
• enforce_eager=False turns CUDA graphs on; max_model_len=4096 covers the longest request (1024 + 1024).
• A one-prompt generate(["Benchmark: "], SamplingParams()) runs before the timer starts, so first-call costs stay out of the measurement.
• Throughput is output tokens divided by wall time, and the wall time includes all prefill work.
• The commented lines switch the same script to vLLM, which takes token-id prompts as dict(prompt_token_ids=p).

1import os2import time3from random import randint, seed4from nanovllm import LLM, SamplingParams5# from vllm import LLM, SamplingParams6 7 8def main():9    seed(0)10    num_seqs = 25611    max_input_len = 102412    max_ouput_len = 102413 14    path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")15    llm = LLM(path, enforce_eager=False, max_model_len=4096)16 17    prompt_token_ids = [[randint(0, 10000) for _ in range(randint(100, max_input_len))] for _ in range(num_seqs)]18    sampling_params = [SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=randint(100, max_ouput_len)) for _ in range(num_seqs)]19    # uncomment the following line for vllm20    # prompt_token_ids = [dict(prompt_token_ids=p) for p in prompt_token_ids]21 22    llm.generate(["Benchmark: "], SamplingParams())23    t = time.time()24    llm.generate(prompt_token_ids, sampling_params, use_tqdm=False)25    t = (time.time() - t)26    total_tokens = sum(sp.max_tokens for sp in sampling_params)27    throughput = total_tokens / t28    print(f"Total: {total_tokens}tok, Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s")29 30 31if __name__ == "__main__":32    main()

The README reports this run on an RTX 4070 Laptop GPU (8 GB) with Qwen3-0.6B: vLLM produced 133,966 tokens in 98.37 s (1361.84 tok/s), nano-vllm the same tokens in 93.41 s (1434.13 tok/s).

1.4Packaging and dependencies

Python 3.10–3.12. Five direct dependencies, each with one job in the code:

safetensors, numpy and tqdm are imported too; they arrive transitively through transformers.

1[build-system]2requires = ["setuptools>=61"]3build-backend = "setuptools.build_meta"4 5[project]6name = "nano-vllm"7version = "0.2.0"8authors = [{ name = "Xingkai Yu" }]9license = "MIT"10license-files = ["LICENSE"]11readme = "README.md"12description = "a lightweight vLLM implementation built from scratch"13requires-python = ">=3.10,<3.13"14dependencies = [15    "torch>=2.4.0",16    "triton>=3.0.0",17    "transformers>=4.51.0",18    "flash-attn",19    "xxhash",20]21 22[project.urls]23Homepage="https://github.com/GeeeekExplorer/nano-vllm"24 25[tool.setuptools.packages.find]26where = ["."]27include = ["nanovllm*"]

Step 02Configuration

Two small dataclasses carry every setting. Config is per engine; SamplingParams is per request.

2.1Config: every engine knob

Config is a @dataclass(slots=True). slots=True fixes the attribute set and makes attribute access a little cheaper. Three fields start as placeholders and are filled later by other objects: hf_config in __post_init__, eos by LLMEngine, and num_kvcache_blocks by ModelRunner.allocate_kv_cache.

1import os2from dataclasses import dataclass3from transformers import AutoConfig4 5 6@dataclass(slots=True)7class Config:8    model: str9    max_num_batched_tokens: int = 1638410    max_num_seqs: int = 51211    max_model_len: int = 409612    gpu_memory_utilization: float = 0.913    tensor_parallel_size: int = 114    enforce_eager: bool = False15    hf_config: AutoConfig | None = None16    eos: int = -117    kvcache_block_size: int = 25618    num_kvcache_blocks: int = -119 20    def __post_init__(self):21        assert os.path.isdir(self.model)22        assert self.kvcache_block_size % 256 == 023        assert 1 <= self.tensor_parallel_size <= 824        self.hf_config = AutoConfig.from_pretrained(self.model)25        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)

2.2SamplingParams: per-request decoding

Three fields. The assertion forbids temperature 0: the sampler divides logits by the temperature (Step 10.7), so greedy decoding is not available. There is no top-k, top-p or repetition penalty. Sequence.__init__ copies the three values onto the sequence, so the params object is not referenced afterwards.

1from dataclasses import dataclass2 3 4@dataclass(slots=True)5class SamplingParams:6    temperature: float = 1.07    max_tokens: int = 648    ignore_eos: bool = False9 10    def __post_init__(self):11        assert self.temperature > 1e-10, "greedy sampling is not permitted"

Step 03The engine loop: LLMEngine

LLMEngine owns the tokenizer, the scheduler and rank 0's model runner, and spawns the other ranks. Its generate method is a loop around step(), and each step is one schedule, one forward, one round of bookkeeping.

3.1Construction, and why the order matters

1. Filter kwargs into Config. dataclasses.fields(Config) gives the allowed names; everything else is dropped.
2. Set Sequence.block_size. It is a class attribute, so every Sequence in this process uses the configured block size for num_blocks and friends. The assignment happens only in the main process (see Step 14 for what that means for workers).
3. Spawn ranks 1…N−1. mp.get_context("spawn") starts fresh interpreters, which CUDA requires; a forked child cannot reuse the parent's CUDA context. The process target is the ModelRunner class itself, so constructing it is the whole life of the worker: its __init__ ends in self.loop(). Each worker gets one Event; rank 0 later receives the list of all of them.
4. Build rank 0's ModelRunner in this process. It must come after the spawns because init_process_group blocks until all ranks have joined. This call loads the model, measures memory, allocates the KV cache and captures CUDA graphs (Step 07).
5. Load the tokenizer and record EOS. config.eos is set after the workers were spawned, so their copies of config keep −1. They never need it.
6. Build the Scheduler last. Its BlockManager needs config.num_kvcache_blocks, which only exists after step 4. Spawned workers received a pickled copy of config earlier, so this value is rank 0's own measurement.
7. Register exit with atexit, so workers are told to stop when the interpreter exits.

1import atexit2from dataclasses import fields3from time import perf_counter4from tqdm.auto import tqdm5from transformers import AutoTokenizer6import torch.multiprocessing as mp7 8from nanovllm.config import Config9from nanovllm.sampling_params import SamplingParams10from nanovllm.engine.sequence import Sequence11from nanovllm.engine.scheduler import Scheduler12from nanovllm.engine.model_runner import ModelRunner13 14 15class LLMEngine:16 17    def __init__(self, model, **kwargs):18        config_fields = {field.name for field in fields(Config)}19        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}20        config = Config(model, **config_kwargs)21        Sequence.block_size = config.kvcache_block_size22        self.ps = []23        self.events = []24        ctx = mp.get_context("spawn")25        for i in range(1, config.tensor_parallel_size):26            event = ctx.Event()27            process = ctx.Process(target=ModelRunner, args=(config, i, event))28            process.start()29            self.ps.append(process)30            self.events.append(event)31        self.model_runner = ModelRunner(config, 0, self.events)32        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)33        config.eos = self.tokenizer.eos_token_id34        self.scheduler = Scheduler(config)35        atexit.register(self.exit)

3.2add_request: text or token ids in, a Sequence out

Strings go through tokenizer.encode; lists of ints are used as they are, which is how bench.py sends random prompts. The new Sequence is appended to scheduler.waiting; nothing else happens until the next step().

43    def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):44        if isinstance(prompt, str):45            prompt = self.tokenizer.encode(prompt)46        seq = Sequence(prompt, sampling_params)47        self.scheduler.add(seq)

3.3step: one iteration of the serving loop

1. scheduler.schedule() returns the batch and whether it is a prefill batch.
2. num_tokens is measured before the forward, because postprocess resets every num_scheduled_tokens to 0. The sign carries the phase: positive means “prefill tokens processed”, negative means “number of sequences decoded”. generate uses the sign to pick which throughput figure to update.
3. model_runner.call("run", seqs, is_prefill) runs the forward on every rank and returns one token id per sequence on rank 0.
4. scheduler.postprocess appends tokens, registers newly full blocks for prefix caching, and retires finished sequences.
5. The return value lists (seq_id, completion_token_ids) only for sequences that finished in this step.

49    def step(self):50        seqs, is_prefill = self.scheduler.schedule()51        num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs)52        token_ids = self.model_runner.call("run", seqs, is_prefill)53        self.scheduler.postprocess(seqs, token_ids, is_prefill)54        outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]55        return outputs, num_tokens56 57    def is_finished(self):58        return self.scheduler.is_finished()

3.4generate: batch in, ordered results out

• A single SamplingParams is repeated for every prompt; a list is used element-wise.
• All requests are enqueued before the first step, so the scheduler sees the whole workload at once.
• The progress bar counts finished requests. Its postfix shows the most recent prefill rate (tokens / step time) and decode rate (sequences / step time, one token each).
• Results are collected in a dict keyed by seq_id and then sorted. seq_id comes from a process-wide itertools.count(), so sorting restores submission order even though requests finish in any order, and it keeps working across repeated generate calls.
• tokenizer.decode(token_ids) does not skip special tokens. A request that stopped on EOS includes the EOS token in both token_ids and text (for Qwen3 chat models, <|im_end|>).
• The return annotation says list[str]; the function actually returns a list of dicts.

57    def is_finished(self):58        return self.scheduler.is_finished()59 60    def generate(61        self,62        prompts: list[str] | list[list[int]],63        sampling_params: SamplingParams | list[SamplingParams],64        use_tqdm: bool = True,65    ) -> list[str]:66        pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True, disable=not use_tqdm)67        if not isinstance(sampling_params, list):68            sampling_params = [sampling_params] * len(prompts)69        for prompt, sp in zip(prompts, sampling_params):70            self.add_request(prompt, sp)71        outputs = {}72        prefill_throughput = decode_throughput = 0.73        while not self.is_finished():74            t = perf_counter()75            output, num_tokens = self.step()76            if num_tokens > 0:77                prefill_throughput = num_tokens / (perf_counter() - t)78            else:79                decode_throughput = -num_tokens / (perf_counter() - t)80            pbar.set_postfix({81                "Prefill": f"{int(prefill_throughput)}tok/s",82                "Decode": f"{int(decode_throughput)}tok/s",83            })84            for seq_id, token_ids in output:85                outputs[seq_id] = token_ids86                pbar.update(1)87        pbar.close()88        outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())]89        outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]90        return outputs

3.5exit: shutting down the ranks

call("exit") writes "exit" into shared memory so every worker leaves its loop and runs its own ModelRunner.exit, then runs rank 0's exit locally. Deleting the runner and joining the processes finishes the shutdown.

37    def exit(self):38        self.model_runner.call("exit")39        del self.model_runner40        for p in self.ps:41            p.join()

3.6The life of one request

Keep this picture in mind for the next four steps. Steps 04–06 explain the three left-hand lifelines. Steps 07–11 explain the two on the right.

Step 04The unit of work: Sequence

A Sequence is one request: its token ids, four counters that say how far its KV cache has progressed, its block table, and its sampling settings. Almost every decision in the scheduler and runner is a comparison between these counters.

4.1Identity, counters and block table

1from copy import copy2from enum import Enum, auto3from itertools import count4 5from nanovllm.sampling_params import SamplingParams6 7 8class SequenceStatus(Enum):9    WAITING = auto()10    RUNNING = auto()11    FINISHED = auto()12 13 14class Sequence:15    block_size = 25616    counter = count()17 18    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):19        self.seq_id = next(Sequence.counter)20        self.status = SequenceStatus.WAITING21        self.token_ids = copy(token_ids)22        self.last_token = token_ids[-1]23        self.num_tokens = len(self.token_ids)24        self.num_prompt_tokens = len(token_ids)25        self.num_cached_tokens = 026        self.num_scheduled_tokens = 027        self.is_prefill = True28        self.block_table = []29        self.temperature = sampling_params.temperature30        self.max_tokens = sampling_params.max_tokens31        self.ignore_eos = sampling_params.ignore_eos

seq_id

From the class-level itertools.count(), shared by every sequence in the process. Used only to order outputs.

token_ids

Prompt plus generated tokens. The constructor copies the caller's list so append_token never mutates it.

num_tokens

len(token_ids), kept as a counter because worker processes receive a Sequence without its token list during decode (4.4).

num_prompt_tokens

Fixed at creation. Splits token_ids into prompt and completion.

num_cached_tokens

How many leading tokens already have their K/V in the cache: prefix-cache hits, earlier prefill chunks and every decoded token. Advanced only in postprocess; reset to 0 by deallocate.

num_scheduled_tokens

How many tokens this step will compute, starting at position num_cached_tokens. Set by the scheduler, read by prepare_* and hash_blocks, reset to 0 in postprocess.

is_prefill

True until the sequence is first scheduled for decode, and again after a preemption. It only decides how the sequence is pickled.

block_table

Physical KV block ids, one per logical block, filled by BlockManager.

4.2Block arithmetic and list behaviour

• __len__ and __getitem__ let a sequence behave like its token list, so the runner writes len(seq) and seq[start:end].
• num_blocks is ceil(num_tokens / block_size), written with integer arithmetic.
• last_block_num_tokens is how many tokens occupy the last block, from 1 to block_size. Decode uses it to find the slot of the newest token.
• block(i) returns the token ids of logical block i; the last block may be partial. Prefix caching hashes these lists.
• append_token updates token_ids, last_token and num_tokens together.

33    def __len__(self):34        return self.num_tokens35 36    def __getitem__(self, key):37        return self.token_ids[key]38 39    @property40    def is_finished(self):41        return self.status == SequenceStatus.FINISHED42 43    @property44    def num_completion_tokens(self):45        return self.num_tokens - self.num_prompt_tokens46 47    @property48    def prompt_token_ids(self):49        return self.token_ids[:self.num_prompt_tokens]50 51    @property52    def completion_token_ids(self):53        return self.token_ids[self.num_prompt_tokens:]54 55    @property56    def num_blocks(self):57        return (self.num_tokens + self.block_size - 1) // self.block_size58 59    @property60    def last_block_num_tokens(self):61        return self.num_tokens - (self.num_blocks - 1) * self.block_size62 63    def block(self, i):64        assert 0 <= i < self.num_blocks65        return self.token_ids[i*self.block_size: (i+1)*self.block_size]66 67    def append_token(self, token_id: int):68        self.token_ids.append(token_id)69        self.last_token = token_id70        self.num_tokens += 1

4.3Status transitions

8class SequenceStatus(Enum):9    WAITING = auto()10    RUNNING = auto()11    FINISHED = auto()

4.4Compact pickling for worker ranks

With tensor parallelism, rank 0 sends the batch's sequences to the other ranks through shared memory every step (Step 07.2). __getstate__ keeps that message small by sending a 6-tuple instead of the object's __dict__:

• the four counters and the block table, which the slot and position computations need;
• the whole token list while the sequence is in prefill, because prepare_prefill slices seq[start:end];
• only last_token in decode, because prepare_decode needs nothing else. A 4,000-token sequence costs one int instead of 4,000.

Nothing else crosses: no seq_id, status, temperature, max_tokens or ignore_eos. Workers never sample and never schedule, so they never need them.

72    def __getstate__(self):73        last_state = self.last_token if not self.is_prefill else self.token_ids74        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state)75 76    def __setstate__(self, state):77        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state = state78        if isinstance(last_state, list):79            self.token_ids = last_state80            self.last_token = self.token_ids[-1]81        else:82            self.token_ids = []83            self.last_token = last_state

Step 05Paged KV memory: BlockManager

BlockManager is a CPU-side allocator for GPU memory it never touches. It hands out block ids, counts how many sequences share each block, and remembers which token prefixes each block holds so later requests can reuse them.

5.1Blocks, block tables and slots

The physical pool is one big tensor allocated by the runner (Step 07.4). The block manager only deals in integer ids from 0 to num_kvcache_blocks − 1. Two numbers follow from a block table and a position: the block block_table[pos // block_size] and the slot block_id × block_size + pos % block_size. The runner computes slots; the block manager decides which ids appear in the table.

5.2Block and the four bookkeeping structures

A Block records how many sequences use it (ref_count) and, once it is full and computed, the hash of the prefix it completes and its token ids. reset() sets ref_count to 1, not 0, because it is only called on a block that is being handed to a sequence.

1from collections import deque2import xxhash3import numpy as np4 5from nanovllm.engine.sequence import Sequence6 7 8class Block:9 10    def __init__(self, block_id):11        self.block_id = block_id12        self.ref_count = 013        self.hash = -114        self.token_ids = []15 16    def update(self, hash: int, token_ids: list[int]):17        self.hash = hash18        self.token_ids = token_ids19 20    def reset(self):21        self.ref_count = 122        self.hash = -123        self.token_ids = []24 25 26class BlockManager:27 28    def __init__(self, num_blocks: int, block_size: int):29        self.block_size = block_size30        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]31        self.hash_to_block_id: dict[int, int] = dict()32        self.free_block_ids: deque[int] = deque(range(num_blocks))33        self.used_block_ids: set[int] = set()

5.3Chained block hashes

compute_hash feeds xxh64 the previous block's hash (8 little-endian bytes) followed by the block's token ids as a NumPy array. Chaining makes a block's hash identify the entire prefix up to and including that block. That is the right key: in every layer after the first, a token's K and V depend on every earlier token, so the same four tokens after a different prefix have different K/V.

35    @classmethod36    def compute_hash(cls, token_ids: list[int], prefix: int = -1):37        h = xxhash.xxh64()38        if prefix != -1:39            h.update(prefix.to_bytes(8, "little"))40        h.update(np.array(token_ids).tobytes())41        return h.intdigest()

5.4can_allocate: count cache hits, check capacity

The scheduler calls this before admitting a waiting sequence that has no blocks yet. It changes nothing; it answers how many leading blocks are cache hits, or −1 if the sequence cannot fit.

1. Walk the logical blocks, recomputing the hash chain. Stop at the first block whose hash is unknown or whose stored token ids differ. The token comparison protects against a 64-bit hash collision.
2. The loop runs to num_blocks − 1, so the last block is never matched, even when the prompt is block-aligned and fully cached. At least one block's worth of tokens is always computed, and its final position produces the logits for the next token.
3. num_new_blocks starts at the full block count. A hit on a block that is in use by another sequence costs no new memory, so it is subtracted. A hit on a block sitting in the free deque still has to be taken out of the deque, so it still counts.
4. Return −1 if the free deque is shorter than num_new_blocks, otherwise the number of cached blocks.

58    def can_allocate(self, seq: Sequence) -> int:59        h = -160        num_cached_blocks = 061        num_new_blocks = seq.num_blocks62        for i in range(seq.num_blocks - 1):63            token_ids = seq.block(i)64            h = self.compute_hash(token_ids, h)65            block_id = self.hash_to_block_id.get(h, -1)66            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:67                break68            num_cached_blocks += 169            if block_id in self.used_block_ids:70                num_new_blocks -= 171        if len(self.free_block_ids) < num_new_blocks:72            return -173        return num_cached_blocks

5.5allocate: build the block table

• For each cached block, recompute the hash and look it up. If the block is in use, increment its ref_count. If it is free, bring it back: ref_count = 1, remove it from the free deque (an O(n) deque.remove), add it to the used set. Its contents are untouched, so no recomputation is needed.
• Every remaining logical block gets a fresh block from _allocate_block. The whole sequence is allocated at once, even if chunked prefill will fill it over several steps.
• num_cached_tokens becomes num_cached_blocks × block_size; prefill will start computing at that position.

75    def allocate(self, seq: Sequence, num_cached_blocks: int):76        assert not seq.block_table77        h = -178        for i in range(num_cached_blocks):79            token_ids = seq.block(i)80            h = self.compute_hash(token_ids, h)81            block_id = self.hash_to_block_id[h]82            block = self.blocks[block_id]83            if block_id in self.used_block_ids:84                block.ref_count += 185            else:86                block.ref_count = 187                self.free_block_ids.remove(block_id)88                self.used_block_ids.add(block_id)89            seq.block_table.append(block_id)90        for i in range(num_cached_blocks, seq.num_blocks):91            seq.block_table.append(self._allocate_block())92        seq.num_cached_tokens = num_cached_blocks * self.block_size

5.6_allocate_block: take a block, evict its old identity

Pop the leftmost free id. If the block still carries a hash and the cache map still points at this block, delete that map entry: the block is about to be overwritten. The == block_id check matters because one hash can be registered twice. Two identical prompts admitted in the same prefill step both get fresh blocks (neither is hashed yet), both are hashed in postprocess, and the map ends up pointing at the second. Recycling the first block later must not delete the entry that belongs to the second.

43    def _allocate_block(self) -> int:44        block_id = self.free_block_ids.popleft()45        block = self.blocks[block_id]46        assert block.ref_count == 047        if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:48            del self.hash_to_block_id[block.hash]49        block.reset()50        self.used_block_ids.add(block_id)51        return block_id

5.7deallocate: release a sequence's blocks

Used when a sequence finishes and when it is preempted. Each block's ref_count drops by one; a block that reaches zero moves to the right end of the free deque, keeping its hash. The table is walked in reverse so a sequence's tail blocks enter the deque before its prefix blocks and are recycled first. Prefix blocks, the ones most likely to be shared again, survive longest.

53    def _deallocate_block(self, block_id: int):54        assert self.blocks[block_id].ref_count == 055        self.used_block_ids.remove(block_id)56        self.free_block_ids.append(block_id)

94    def deallocate(self, seq: Sequence):95        for block_id in reversed(seq.block_table):96            block = self.blocks[block_id]97            block.ref_count -= 198            if block.ref_count == 0:99                self._deallocate_block(block_id)100        seq.num_cached_tokens = 0101        seq.block_table.clear()

5.8can_append and may_append: growing during decode

When a decode step is scheduled, the token sampled in the previous step is already in token_ids but its K/V is not stored yet. This step stores it at position len(seq) − 1. That position starts a new logical block exactly when len(seq) % block_size == 1, and only then is a new block needed. can_append turns the comparison into 0 or 1 by letting Python treat the boolean as an int.

103    def can_append(self, seq: Sequence) -> bool:104        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)105 106    def may_append(self, seq: Sequence):107        if len(seq) % self.block_size == 1:108            seq.block_table.append(self._allocate_block())

5.9hash_blocks: publish newly full blocks

Called from postprocess for every sequence in the batch, after the forward and before num_cached_tokens advances. start is the first block that was not yet full before this step; end is the number of full blocks after it. Every block in between is hashed, using the previous block's stored hash as the chain prefix, and registered in hash_to_block_id.

• It runs after decode steps as well as prefill. Blocks filled by generated tokens become cacheable, which is what makes recomputation after preemption cheap (Step 06.4).
• A partial last block is never hashed, so a block's hash always covers exactly block_size tokens.
• Blocks are registered only after their K/V has really been written, so a later lookup can never hit a block whose contents are still pending.

110    def hash_blocks(self, seq: Sequence):111        start = seq.num_cached_tokens // self.block_size112        end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size113        if start == end: return114        h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1115        for i in range(start, end):116            block = self.blocks[seq.block_table[i]]117            token_ids = seq.block(i)118            h = self.compute_hash(token_ids, h)119            block.update(h, token_ids)120            self.hash_to_block_id[h] = block.block_id

Step 06Choosing each batch: Scheduler

The scheduler is 92 lines and holds two queues. schedule() picks the next batch, preempt() frees memory under pressure, and postprocess() applies the sampled tokens. Read it with Step 05 open: almost every branch is a call into the block manager.

6.1Queues and limits

Two deques. waiting holds new requests, preempted ones and a sequence whose prefill is part-way through its chunks. running holds sequences whose prompt is fully in the cache and that are generating. The limits and EOS id are copied from config, and the block manager is created here with rank 0's measured block count.

1from collections import deque2 3from nanovllm.config import Config4from nanovllm.engine.sequence import Sequence, SequenceStatus5from nanovllm.engine.block_manager import BlockManager6 7 8class Scheduler:9 10    def __init__(self, config: Config):11        self.max_num_seqs = config.max_num_seqs12        self.max_num_batched_tokens = config.max_num_batched_tokens13        self.eos = config.eos14        self.block_size = config.kvcache_block_size15        self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)16        self.waiting: deque[Sequence] = deque()17        self.running: deque[Sequence] = deque()18 19    def is_finished(self):20        return not self.waiting and not self.running21 22    def add(self, seq: Sequence):23        self.waiting.append(seq)

6.2Prefill phase and chunked prefill

Prefill always goes first. If anything at all is scheduled in this loop, the step is a prefill step and no sequence decodes in it.

1. First come, first served. The loop only looks at waiting[0]. When the head cannot be served it breaks; it never skips ahead to a smaller request.
2. Budget. remaining is what is left of max_num_batched_tokens. A batch that has used it all stops.
3. Fresh or continuing. A sequence without blocks is new or was preempted: can_allocate reports cache hits or −1, and the work is everything after the cached blocks. A sequence that already has blocks is mid-way through a chunked prefill, and its work is everything after num_cached_tokens.
4. The chunking rule. If the work does not fit in remaining and the batch already has a sequence, stop. Only the first sequence of a batch can be cut. That sequence always sees the full budget, so in practice a prompt is chunked only when its uncached part is longer than max_num_batched_tokens. With the defaults (16,384 tokens of budget, max_model_len 4,096) that never happens; lower the budget or raise max_model_len to see it.
5. Commit. Blocks are allocated only now that the sequence is definitely in the batch. num_scheduled_tokens is the work, capped by the budget.
6. Done or not. If this chunk reaches num_tokens, the sequence becomes RUNNING and moves to running. Otherwise it stays at waiting[0]; it used the whole remaining budget, so the next iteration hits remaining == 0 and the loop ends.

25    def schedule(self) -> tuple[list[Sequence], bool]:26        scheduled_seqs = []27        num_batched_tokens = 028 29        # prefill30        while self.waiting and len(scheduled_seqs) < self.max_num_seqs:31            seq = self.waiting[0]32            remaining = self.max_num_batched_tokens - num_batched_tokens33            if remaining == 0:34                break35            if not seq.block_table:36                num_cached_blocks = self.block_manager.can_allocate(seq)37                if num_cached_blocks == -1:38                    break39                num_tokens = seq.num_tokens - num_cached_blocks * self.block_size40            else:41                num_tokens = seq.num_tokens - seq.num_cached_tokens42            if remaining < num_tokens and scheduled_seqs:  # only allow chunked prefill for the first seq43                break44            if not seq.block_table:45                self.block_manager.allocate(seq, num_cached_blocks)46            seq.num_scheduled_tokens = min(num_tokens, remaining)47            num_batched_tokens += seq.num_scheduled_tokens48            if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:49                seq.status = SequenceStatus.RUNNING50                self.waiting.popleft()51                self.running.append(seq)52            scheduled_seqs.append(seq)53 54        if scheduled_seqs:55            return scheduled_seqs, True

6.3Decode phase and preemption

1. Only reached when nothing could be prefilled.
2. Take sequences from the front of running, up to max_num_seqs.
3. Make sure the sequence can store one more token (can_append). While it cannot, preempt the sequence at the back of running, the most recently admitted one, which frees its blocks. If no other sequence is left, preempt this one and move on.
4. Python's while … else: the else block runs when the loop condition turns false, meaning there is room, and is skipped after the break that follows self-preemption. In the else branch the sequence is set to compute one token, marked is_prefill = False so it pickles compactly, and given a new block if it needs one.
5. assert scheduled_seqs fires only if a single sequence needs more blocks than the whole pool, because then even preempting everything else frees nothing usable.
6. running.extendleft(reversed(scheduled_seqs)) puts the batch back at the front of running in its original order.

57        # decode58        while self.running and len(scheduled_seqs) < self.max_num_seqs:59            seq = self.running.popleft()60            while not self.block_manager.can_append(seq):61                if self.running:62                    self.preempt(self.running.pop())63                else:64                    self.preempt(seq)65                    break66            else:67                seq.num_scheduled_tokens = 168                seq.is_prefill = False69                self.block_manager.may_append(seq)70                scheduled_seqs.append(seq)71        assert scheduled_seqs72        self.running.extendleft(reversed(scheduled_seqs))73        return scheduled_seqs, False

6.4preempt: give memory back, recompute later

The victim goes back to WAITING, is marked as prefill again, loses all its blocks and is put at the front of waiting, so it is the first sequence re-admitted. There is no swapping to CPU memory; the K/V will be recomputed. Two details make that cheaper than it sounds:

• deallocate keeps hashes, and decode steps hashed every block they filled. When the sequence is re-admitted, can_allocate finds all of its full blocks that have not been recycled yet, and only the rest is recomputed.
• The last token of a preempted sequence was sampled but never fed back. The re-prefill computes it, so the token sampled at the end of that prefill is a correct next token and postprocess appends it.

75    def preempt(self, seq: Sequence):76        seq.status = SequenceStatus.WAITING77        seq.is_prefill = True78        self.block_manager.deallocate(seq)79        self.waiting.appendleft(seq)

6.5postprocess: apply the sampled tokens

1. hash_blocks first, while num_cached_tokens and num_scheduled_tokens still describe what was just computed.
2. Advance num_cached_tokens, clear num_scheduled_tokens.
3. A prefill chunk that did not reach the end of the sequence produced a token for a position in the middle of the prompt. It is discarded with continue.
4. Otherwise append the token. Finish when the token is EOS (unless ignore_eos) or the completion has reached max_tokens. The token is appended before the check, so a finishing EOS is part of the output.
5. A finished sequence frees its blocks immediately and leaves running (an O(n) deque.remove).

81    def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):82        for seq, token_id in zip(seqs, token_ids):83            self.block_manager.hash_blocks(seq)84            seq.num_cached_tokens += seq.num_scheduled_tokens85            seq.num_scheduled_tokens = 086            if is_prefill and seq.num_cached_tokens < seq.num_tokens:87                continue88            seq.append_token(token_id)89            if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:90                seq.status = SequenceStatus.FINISHED91                self.block_manager.deallocate(seq)92                self.running.remove(seq)

6.6Simulator: run the scheduler yourself

This runs a line-by-line JavaScript port of sequence.py, block_manager.py and scheduler.py, plus the slot arithmetic of prepare_prefill and prepare_decode. Block size is 4 instead of 256 so a whole sequence fits on screen, a 32-bit FNV hash stands in for xxh64, and the model is replaced by random tokens 1–9 (no EOS, so sequences finish at max_tokens). Everything else follows the Python.

With the default settings, stepping through shows each mechanism in turn: step 1 admits A alone, step 2 admits B with a two-block prefix hit on A's blocks, steps 3–4 chunk C's 22-token prompt across the 16-token budget, step 7 preempts C when the pool runs dry, and step 9 re-admits C with six cached blocks and recomputes one token.

Things to try: set num_kvcache_blocks to 8 and max_tokens to 12, then run until idle; C alone outgrows the pool and the scheduler's assertion fires, the same failure Step 14 describes. After everything finishes, add A again: its blocks are still in the free list with their hashes, so it is admitted with a prefix hit.

Step 07From sequences to tensors: ModelRunner

One ModelRunner exists per GPU. It builds the model, measures memory, owns the KV cache, captures CUDA graphs, converts a list of Sequences into the flat tensors FlashAttention wants, and, on rank 0, samples. On other ranks it is a loop that waits for instructions.

7.1Boot sequence

1. Join the process group. NCCL, rendezvous at tcp://localhost:2333, so all ranks must be on one machine. torch.cuda.set_device(rank) pins rank i to the i-th visible GPU.
2. Build the model directly on the GPU. The default dtype becomes the checkpoint's (hf_config.dtype, bfloat16 for Qwen3) and the default device becomes CUDA, so every torch.empty inside the layer constructors allocates GPU memory in model precision. Nothing is initialised on the CPU first. Older transformers releases call this attribute torch_dtype; the code switched to dtype in commit 8d63a98, and an install whose config objects lack dtype fails on this line.
3. Load weights with load_model (Step 11) and create the Sampler.
4. Warm up before allocating the KV cache, so the measured activation peak does not include it (7.3).
5. Allocate the KV cache (7.4), then capture CUDA graphs unless enforce_eager (7.9). Capture has to come after the cache exists, because the captured kernels hold the cache's addresses.
6. Restore the defaults to CPU and the previous dtype. From here on, tensors built from Python lists in prepare_* start on the CPU and are copied over explicitly.
7. Set up shared memory when world_size > 1. Rank 0 creates a 1 MiB segment named "nanovllm" and then waits at the barrier; workers wait at the barrier first, so the segment exists before they attach, then enter loop() and stay there.

1import pickle2import torch3import torch.distributed as dist4from multiprocessing.synchronize import Event5from multiprocessing.shared_memory import SharedMemory6 7from nanovllm.config import Config8from nanovllm.engine.sequence import Sequence9from nanovllm.models.qwen3 import Qwen3ForCausalLM10from nanovllm.layers.sampler import Sampler11from nanovllm.utils.context import set_context, get_context, reset_context12from nanovllm.utils.loader import load_model13 14 15class ModelRunner:16 17    def __init__(self, config: Config, rank: int, event: Event | list[Event]):18        self.config = config19        hf_config = config.hf_config20        self.block_size = config.kvcache_block_size21        self.enforce_eager = config.enforce_eager22        self.world_size = config.tensor_parallel_size23        self.rank = rank24        self.event = event25 26        dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)27        torch.cuda.set_device(rank)28        default_dtype = torch.get_default_dtype()29        torch.set_default_dtype(hf_config.dtype)30        torch.set_default_device("cuda")31        self.model = Qwen3ForCausalLM(hf_config)32        load_model(self.model, config.model)33        self.sampler = Sampler()34        self.warmup_model()35        self.allocate_kv_cache()36        if not self.enforce_eager:37            self.capture_cudagraph()38        torch.set_default_device("cpu")39        torch.set_default_dtype(default_dtype)40 41        if self.world_size > 1:42            if rank == 0:43                self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)44                dist.barrier()45            else:46                dist.barrier()47                self.shm = SharedMemory(name="nanovllm")48                self.loop()

7.2The control plane: call, write_shm, read_shm, loop

• call(name, *args) is the only entry point the engine uses. On rank 0 with several ranks it first broadcasts the call, then runs the method locally by name through getattr. In practice only "run" and "exit" go through it.
• The message format is a 4-byte little-endian length followed by pickle.dumps([name, *args]). The segment is fixed at 2²⁰ bytes, and the payload has to fit.
• Each worker has its own Event. A worker waits on it, reads, clears it, then runs the method.
• No acknowledgement flows back. Ordering still holds: every forward pass contains NCCL collectives, so rank 0 cannot finish a run until every worker has entered the same run, which they only do after reading and clearing. Rank 0 therefore never overwrites a message a worker has not read yet.

61    def loop(self):62        while True:63            method_name, args = self.read_shm()64            self.call(method_name, *args)65            if method_name == "exit":66                break67 68    def read_shm(self):69        assert self.world_size > 1 and self.rank > 070        self.event.wait()71        n = int.from_bytes(self.shm.buf[0:4], "little")72        method_name, *args = pickle.loads(self.shm.buf[4:n+4])73        self.event.clear()74        return method_name, args75 76    def write_shm(self, method_name, *args):77        assert self.world_size > 1 and self.rank == 078        data = pickle.dumps([method_name, *args])79        n = len(data)80        self.shm.buf[0:4] = n.to_bytes(4, "little")81        self.shm.buf[4:n+4] = data82        for event in self.event:83            event.set()84 85    def call(self, method_name, *args):86        if self.world_size > 1 and self.rank == 0:87            self.write_shm(method_name, *args)88        method = getattr(self, method_name, None)89        return method(*args)