3.6 ModelRunner, CUDA Graph, and benchmarks

📑 Contents

• 1. Where ModelRunner sits
• 2. Turning sequences into tensors: prepare_prefill / prepare_decode
• 3. Global Context: the operators’ “implicit parameters”
• 4. How KV Cache GPU memory is carved out
• 5. CUDA Graph: accelerate decode only
• 6. Process coordination under tensor parallelism
• 7. Benchmarks versus vLLM
• Summary
• Self-check checklist
• References

---

1. Where ModelRunner sits

Earlier sections covered which operators run (3.3), where KV lives (3.4), and who gets scheduled (3.5). ModelRunner (engine/model_runner.py, 257 lines) wires them into one real GPU execution: take the sequences the scheduler picked, build the tensors forward needs, set the global context, call Qwen3, then hand logits to the sampler.

The dashed box is one ModelRunner.run. Inside it, five steps run left to right; the two green ones call the model itself:

run itself is only a few lines:

def run(self, seqs, is_prefill):
    input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
    logits = self.run_model(input_ids, positions, is_prefill)   # eager or CUDA Graph replay
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
    reset_context()
    return token_ids

ModelRunner.__init__ finishes setup in a fixed order at engine start; the later sections unpack each piece:

1. Initialize the process group with NCCL, bind this rank’s GPU, and set the default dtype to the model’s dtype (Qwen3-0.6B is bf16).
2. Build Qwen3ForCausalLM and load weights from safetensors.
3. warmup_model: run one max-size prefill and measure the activation peak (section 4).
4. allocate_kv_cache: carve the remaining GPU memory into KV blocks (section 4).
5. When enforce_eager=False, capture_cudagraph captures decode CUDA Graphs ahead of time (section 5).
6. When TP>1, create shared memory; ranks >0 enter loop and wait for commands (section 6).

2. Turning sequences into tensors: prepare_prefill / prepare_decode

Model forward eats tensors; the scheduler hands over a batch of Sequence objects. ModelRunner converts them with two prepare functions, one per phase, because the two phases need different tensor shapes.

Prefill: pack variable-length sequences end to end

def prepare_prefill(self, seqs):
    input_ids, positions, slot_mapping = [], [], []
    cu_seqlens_q, cu_seqlens_k = [0], [0]
    max_seqlen_q = max_seqlen_k = 0
    block_tables = None
    for seq in seqs:
        start = seq.num_cached_tokens               # tokens already in the Cache
        end = start + seq.num_scheduled_tokens      # how far this round computes
        input_ids.extend(seq[start:end])            # Q only holds tokens computed this round
        positions.extend(range(start, end))         # positions continue from start
        cu_seqlens_q.append(cu_seqlens_q[-1] + end - start)
        cu_seqlens_k.append(cu_seqlens_k[-1] + end) # K covers every token from the start through end
        max_seqlen_q = max(max_seqlen_q, end - start)
        max_seqlen_k = max(max_seqlen_k, end)
        if not seq.block_table:                     # fake warmup sequences have no blocks; do not write Cache
            continue
        # map [start, end) to physical slots (source extends by block segments; shown here token by token)
        for i in range(start, end):
            block_id = seq.block_table[i // self.block_size]
            slot_mapping.append(block_id * self.block_size + i % self.block_size)
    if cu_seqlens_k[-1] > cu_seqlens_q[-1]:         # historical KV present: prefix hit or chunked resume
        block_tables = self.prepare_block_tables(seqs)
    # ...convert to pin_memory tensors, non_blocking copy to GPU, then set_context(True, ...)

• 🔑 cu_seqlens (cumulative sequence lengths): after several variable-length sequences are packed end to end, a cumulative-length array marks each boundary. That is the format flash_attn_varlen_func requires (see 3.3). The Q side only has tokens computed this round; the K side starts from the beginning of the sequence.
• 🔑 When to bring block_tables: with no history, Q and K are the same length, and the kernel uses the k/v freshly computed this round. As soon as any sequence has K longer than Q, the first half of its KV is already in the Cache, so block_tables is prepared and Attention reads the full K/V from the Cache (section 7.2 of 3.3).
• slot_mapping: each token’s physical location in the KV Cache, equal to physical_block_id × block_size + offset_in_block. The Triton store_kvcache in 3.3 writes K/V into the Cache by this map.

The figure below uses block_size = 4 on a 10-token sequence that hits 1 prefix-cache block (gray = tokens already in the Cache; orange = tokens computed this round):

Decode: one token per sequence

def prepare_decode(self, seqs):
    input_ids, positions, slot_mapping, context_lens = [], [], [], []
    for seq in seqs:
        input_ids.append(seq.last_token)          # feed only the token sampled last step
        positions.append(len(seq) - 1)
        context_lens.append(len(seq))             # total length including this new token
        # the new token is the last_block_num_tokens-th token inside the last block
        slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1)
    block_tables = self.prepare_block_tables(seqs)  # pad each sequence's block_table on the right with -1 into a matrix
    # ...convert to GPU tensors, set_context(False, slot_mapping=..., context_lens=..., block_tables=...)

• Read last_token, not token_ids[-1]: when TP>1, the decode sequences a worker receives only have last_token, not token_ids (section 3.1 of 3.4).
• Where the slot sits: the scheduler already allocated a new block in may_append for tokens that cross a boundary, so the new token always lands in block_table[-1].
• 🔑 context_lens: tells flash_attn_with_kvcache how many historical K/V entries each sequence should read. It counts the new token too, because Attention first runs store_kvcache to write the new token’s K/V into the Cache, then computes attention.

⚠️ Note: the -1 padding on the right of block_tables is only for matrix alignment. The kernel reads only the valid leading blocks according to context_lens and never touches the padding. That is not the same as the slot == -1 that store_kvcache skips; those come from padding the batch during CUDA Graph replay (section 5).

3. Global Context: the operators’ “implicit parameters”

Tensors such as cu_seqlens, slot_mapping, block_tables, and is_prefill are needed by Attention, LMHead, and other operators. Passing them as function arguments layer by layer would add seven or eight parameters to Qwen3ForCausalLM, every DecoderLayer, and every Attention.forward. nano-vllm puts them in a global Context (utils/context.py); operators fetch them with get_context():

@dataclass(slots=True)
class Context:
    is_prefill: bool = False
    cu_seqlens_q: torch.Tensor | None = None
    cu_seqlens_k: torch.Tensor | None = None
    max_seqlen_q: int = 0
    max_seqlen_k: int = 0
    slot_mapping: torch.Tensor | None = None
    context_lens: torch.Tensor | None = None
    block_tables: torch.Tensor | None = None

_CONTEXT = Context()

def set_context(is_prefill, cu_seqlens_q=None, ...):   # called at the end of prepare_*
    global _CONTEXT
    _CONTEXT = Context(is_prefill, cu_seqlens_q, ...)

def reset_context():                                     # called at the end of run
    global _CONTEXT
    _CONTEXT = Context()

max_seqlen_q / max_seqlen_k are the longest sequence lengths in the batch that flash_attn_varlen_func needs, used together with cu_seqlens. Decode only sets slot_mapping, context_lens, and block_tables.

💡 Tip: this trades a global variable for cleaner call sites. Qwen3.forward only takes input_ids and positions; the cost is that operators implicitly depend on global state, so every forward must be preceded by set_context. nano-vllm’s inference loop is single-threaded and serial, one process per rank, so that premise holds. At the end of run, reset_context drops the references so this step’s tensors are not reused by mistake on the next step.

4. How KV Cache GPU memory is carved out

The physical GPU memory that section 2’s slot_mapping points at is allocated once at engine start. ModelRunner does not reserve max length per request. It loads weights first, measures the activation peak with warmup, then cuts all remaining GPU memory into KV blocks. That is why num_kvcache_blocks starts as -1 in 3.2 and is only computed at runtime.

Warmup runs one max-size prefill on a batch of fake sequences:

def warmup_model(self):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()              # peak stats restart from here
    seq_len = min(max_num_batched_tokens, max_model_len)
    num_seqs = min(max_num_batched_tokens // seq_len, max_num_seqs)
    seqs = [Sequence([0] * seq_len) for _ in range(num_seqs)]   # fake sequences with no block_table
    for seq in seqs:
        seq.num_scheduled_tokens = seq_len
    self.run(seqs, True)
    torch.cuda.empty_cache()

Under the default config that is 16384 // 4096 = 4 sequences of 4096 tokens, which exactly fills one step’s token budget. The KV Cache is not allocated yet, k_cache.numel() is 0, Attention’s write-Cache branch in 3.3 is skipped, and the measured peak only covers weights and activations. Then allocate_kv_cache cuts blocks from the remainder:

def allocate_kv_cache(self):
    free, total = torch.cuda.mem_get_info()
    used = total - free                                                # GPU memory already occupied on the whole card
    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]      # allocation peak during warmup
    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
    num_kv_heads = hf_config.num_key_value_heads // self.world_size    # KV heads this rank owns
    block_bytes = 2 * num_layers * block_size * num_kv_heads * head_dim * dtype.itemsize
    config.num_kvcache_blocks = int(total * gpu_memory_utilization - used - peak + current) // block_bytes
    assert config.num_kvcache_blocks > 0
    # one 6-D tensor holds K and V for every layer and every block
    self.kv_cache = torch.empty(2, num_layers, num_kvcache_blocks, block_size, num_kv_heads, head_dim)
    layer_id = 0
    for module in self.model.modules():                # bind in order to each layer's Attention
        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

The block-count formula:

$$\text{num\_blocks} = \left\lfloor \frac{\text{total} \times \text{util} - \text{used} - (\text{peak} - \text{current})}{\text{block\_bytes}} \right\rfloor$$

used is what the driver sees as whole-card occupancy: weights, CUDA context, the PyTorch caching allocator, and anything else on the same GPU. peak - current is how much more activation occupied during warmup than now; that spike will return in real runs, so it must be reserved ahead of time. Bytes per block:

$$\text{block\_bytes} = 2 \times L \times \text{block\_size} \times \frac{H_{kv}}{\text{tp}} \times d_{\text{head}} \times \text{itemsize}$$

Take Qwen3-0.6B as an example (28 layers, 8 KV heads, head_dim = 128, bf16, TP=1): $2 \times 28 \times 256 \times 8 \times 128 \times 2 = 29{,}360{,}128$ bytes, i.e. 28 MiB per block and 112 KiB per token. About 36 blocks fit in 1 GiB of GPU memory, roughly 9200 tokens.

📌 Key point: the KV Cache is one 6-D tensor (2, layers, blocks, block_size, num_kv_heads, head_dim), with dim 0 separating K and V. After allocation, each layer’s slice is bound to that Attention module’s k_cache / v_cache attributes, and store_kvcache plus flash_attn_with_kvcache in 3.3 read and write self.k_cache directly. Remaining GPU memory is cut entirely into KV blocks rather than reserved per request; how many sequences fit at once depends on their actual lengths.

⚠️ Note: the KV Cache dtype matches the model’s default dtype (set to hf_config.dtype in __init__). nano-vllm does not implement KV Cache quantization.

5. CUDA Graph: accelerate decode only

In decode each step computes one token per sequence. A single kernel does very little work, so CPU kernel-launch overhead becomes a larger share of the time: the GPU finishes a small kernel and waits for the CPU to launch the next one. A CUDA Graph records a chain of kernel launches into one graph; later a single replay() replays the whole chain and skips per-kernel CPU launch cost.

What a CUDA Graph records is the kernels and the fixed GPU addresses they read and write. On replay you can only swap data, not addresses or shapes. So nano-vllm first allocates a set of static buffers at the max batch size, points Context at slices of those buffers, and captures one graph per batch bucket:

@torch.inference_mode()
def capture_cudagraph(self):
    max_bs = min(self.config.max_num_seqs, 512)
    max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
    input_ids = torch.zeros(max_bs, dtype=torch.int64)       # static buffers shared by every graph
    positions = torch.zeros(max_bs, dtype=torch.int64)
    slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
    context_lens = torch.zeros(max_bs, dtype=torch.int32)
    block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
    outputs = torch.zeros(max_bs, hf_config.hidden_size)
    self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
    self.graphs, self.graph_pool = {}, None
    for bs in reversed(self.graph_bs):                       # capture large to small
        graph = torch.cuda.CUDAGraph()
        set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs],
                    block_tables=block_tables[:bs])
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # run once normally first
        with torch.cuda.graph(graph, self.graph_pool):
            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # then capture
        if self.graph_pool is None:
            self.graph_pool = graph.pool()                   # the largest graph's memory pool is shared by later graphs
        self.graphs[bs] = graph
        torch.cuda.synchronize()
        reset_context()
    self.graph_vars = dict(input_ids=input_ids, positions=positions, slot_mapping=slot_mapping,
                           context_lens=context_lens, block_tables=block_tables, outputs=outputs)

• 🔑 Buckets: with the default max_num_seqs = 512, the buckets are 1, 2, 4, 8, then every 16 from 16 to 512, 36 graphs in total.
• Capture large to small: the first (largest) graph’s memory pool is passed via graph.pool() to the later smaller graphs, so 36 graphs do not each hold their own intermediate activations.
• Run once before capture: operators decorated with @torch.compile may recompile on a new input shape, and Triton kernels JIT on first call; none of that can happen during capture.
• The graph stops at the hidden state; compute_logits and sampling run outside the graph (section 8 of 3.3).

At runtime, pick the smallest graph that is at least the actual batch size, copy data into the static buffers, and replay:

@torch.inference_mode()
def run_model(self, input_ids, positions, is_prefill):
    if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
        return self.model.compute_logits(self.model(input_ids, positions))  # eager
    bs = input_ids.size(0)
    context = get_context()
    graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]   # round up to a bucket
    graph_vars = self.graph_vars
    graph_vars["input_ids"][:bs] = input_ids
    graph_vars["positions"][:bs] = positions
    graph_vars["slot_mapping"].fill_(-1)                  # padded slots: slot = -1, do not write Cache
    graph_vars["slot_mapping"][:bs] = context.slot_mapping
    graph_vars["context_lens"].zero_()                    # padded slots: length 0, do not read Cache
    graph_vars["context_lens"][:bs] = context.context_lens
    graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
    graph.replay()
    return self.model.compute_logits(graph_vars["outputs"][:bs])

The figure below uses 37 sequences as an example: there is no bucket 37, so graph 48 is replayed, and the extra 11 positions are padding (green = real sequences; gray = padding):

• 🔑 Capture decode only: prefill’s total token count changes every step, so a fixed-shape graph does not work; decode is always 1 token per sequence and only the batch size changes, so a few dozen pre-captured buckets cover it.
• 🔑 How padding is handled: slot_mapping is filled with -1, and store_kvcache in 3.3 returns immediately on -1 without writing the Cache; context_lens is filled with 0 so attention reads no history. Padded input_ids / positions keep whatever was left from last time; the computed results are dropped when slicing outputs[:bs]. block_tables only overwrites the first bs rows and some leading columns; leftover old block ids are never read because context_lens limits the read.
• enforce_eager=True: run_model always takes the eager path, and graphs are not captured at start. Turn it on when debugging so every kernel really launches and problems are easier to locate; turn it off when measuring throughput.
• Batch above 512 falls back to eager: buckets go up to 512; if max_num_seqs is set higher, anything beyond runs normally.

⚠️ Note: when max_num_seqs is less than 512 and not a multiple of 16 (and greater than 8), the largest bucket is smaller than max_num_seqs. For example, set to 100, the largest bucket is 96, and for batch 97–100 next(...) finds no bucket and raises StopIteration. The default 512 does not hit this.

6. Process coordination under tensor parallelism

Section 3.2 noted that LLMEngine uses spawn to start worker child processes for rank>0, and those children also run ModelRunner. Tensor gathering goes through NCCL (init_process_group("nccl", "tcp://localhost:2333", ...)). “What this step should do” is issued by rank 0 over a shared-memory region, a tiny RPC:

def write_shm(self, method_name, *args):        # rank 0
    data = pickle.dumps([method_name, *args])
    n = len(data)
    self.shm.buf[0:4] = n.to_bytes(4, "little") # first 4 bytes store the length
    self.shm.buf[4:n+4] = data
    for event in self.event:                    # rank 0 holds every worker's Event
        event.set()

def read_shm(self):                             # rank > 0
    self.event.wait()
    n = int.from_bytes(self.shm.buf[0:4], "little")
    method_name, *args = pickle.loads(self.shm.buf[4:n+4])
    self.event.clear()
    return method_name, args

def call(self, method_name, *args):
    if self.world_size > 1 and self.rank == 0:
        self.write_shm(method_name, *args)      # broadcast to workers first
    return getattr(self, method_name)(*args)    # then run locally

def loop(self):                                 # rank > 0 enters at the end of __init__ and never returns
    while True:
        method_name, args = self.read_shm()
        self.call(method_name, *args)
        if method_name == "exit":
            break

• Shared memory is named "nanovllm", size 2**20 bytes (1 MiB), created by rank 0; after dist.barrier() workers open it by name.
• Each step writes a pickled method name and Sequence list. Sequence.__getstate__ in decode only packs a few integers, block_table, and last_token (section 3.1 of 3.4), so the per-step message stays small.
• The NCCL init address and the shared-memory name are hard-coded. Even with TP=1 the process group is initialized on port 2333, so two nano-vllm instances on the same machine collide on the port; with TP>1 the shared-memory name collides too.

📌 Key point: rank 0 owns scheduling and sampling and holds the tokenizer; ranks 1 through N-1 only run forward on their shard of the weights and join results with rank 0 through the all_reduce / gather from 3.3. The single call("run", seqs, is_prefill) entry that step uses in 3.2 makes every rank run the same step.

7. Benchmarks versus vLLM

nano-vllm’s README gives a throughput comparison with vLLM (numbers from the README committed on 2026-04-26):

Before reading those numbers, look at how bench.py measures:

• seed(0) fixes the RNG; the prompt is random token ids from randint(0, 10000), length randint(100, 1024); each sequence’s max_tokens is randint(100, 1024).
• ignore_eos=True, temperature=0.6, so every sequence generates a full max_tokens and both engines produce the same output total.
• LLM(path, enforce_eager=False, max_model_len=4096) turns CUDA Graph on; a short request warms up before formal timing.
• Throughput = sum(max_tokens) / total_time. The numerator only counts output tokens, but the denominator includes all prefill time. That is end-to-end output throughput, not decode throughput.
• The README does not name the vLLM version used in the comparison.

💡 Tip: these numbers say that on a single GPU, a 0.6B model, and offline batch generation, nano-vllm’s throughput sits at the same level as vLLM, about 5% higher. They do not say nano-vllm is also faster on large models, long context, online serving, or multi-GPU setups. In those settings vLLM’s mixed batching, more Attention backends, and quantization support all matter, and this benchmark does not cover them. A fairer reading: nano-vllm implements PagedAttention, prefix caching, continuous batching, and CUDA Graph in about 1200 lines, and in this test it reaches throughput comparable to vLLM.

Tie the chapter’s six sections together and one step’s full call chain looks like this. The blue label at the top-left of each box is the section that covers it:

📝 Summary

• ModelRunner.run goes: tensorize sequences and set_context → run_model (eager or CUDA Graph replay) → compute_logits → rank 0 samples → reset_context.
• Prefill packs variable-length sequences end to end; the Q side only holds tokens computed this round, with boundaries marked by cu_seqlens_q / cu_seqlens_k. When K is longer than Q (prefix hit or chunked resume), block_tables is included. Decode feeds only last_token per sequence and reads history from the Cache with context_lens and block_tables. In both phases slot_mapping writes new K/V to physical_block_id × block_size + offset_in_block.
• A global Context lets operators fetch those tensors via get_context(), so the forward signature is only input_ids and positions; the cost is that every step must set_context first.
• At start, one max-size prefill (default 4 × 4096 tokens) measures the activation peak, then remaining GPU memory is cut into KV blocks by total × util − used − (peak − current). The KV Cache is a 6-D tensor (2, layers, blocks, block_size, KV heads, head_dim), sliced per layer onto each Attention. For Qwen3-0.6B each block is 28 MiB and each token 112 KiB.
• CUDA Graph captures decode only: default 36 batch buckets captured large-to-small sharing one memory pool; at runtime round up to a bucket, pad with slot_mapping=-1 and context_lens=0. When max_num_seqs is not a multiple of 16, batches above the largest bucket raise StopIteration.
• Under tensor parallelism, rank 0 writes a pickled method name and Sequences into 1 MiB of shared memory and signals workers with Events to run the same method; tensor gathering goes through NCCL. Port 2333 and the shared-memory name are hard-coded.
• In the README benchmark (RTX 4070 Laptop, Qwen3-0.6B, 256 random requests) nano-vllm’s end-to-end output throughput is about 5% above vLLM. That only shows the two are comparable for single-GPU small-model offline batch generation.

🎯 Self-check checklist

• Can name the five steps of ModelRunner.run, and which ones run only on rank 0
• Given block_table, num_cached_tokens, and this round’s token count, can hand-compute prefill slot_mapping, cu_seqlens_q, cu_seqlens_k, and decide whether block_tables is needed
• Can explain why decode context_lens includes the new token, and how the -1 padding on block_tables differs from slot_mapping’s -1
• Can state the design tradeoff of the global Context (clean signature vs implicit dependency) and the premise it rests on
• Can say why warmup runs before KV Cache allocation, and compute Qwen3-0.6B’s bytes per KV block from the formula
• Can explain why CUDA Graph only captures decode, why capture is large-to-small, and why padding does not corrupt the Cache
• Can give an example max_num_seqs value that makes run_model raise StopIteration
• Can describe how rank 0 uses shared memory and Events so workers run the same step under tensor parallelism
• Can point out how the README benchmark is measured (denominator includes prefill; vLLM version unnamed) and what it can and cannot conclude

📚 References

• nano-vllm GitHub repository
• nano-vllm bench.py
• PyTorch blog - Accelerating PyTorch with CUDA Graphs
• PyTorch docs - CUDA Graphs
• CUDA C++ Programming Guide - CUDA Graphs
• Python docs - multiprocessing.shared_memory
• Qwen3-0.6B model card
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
• Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)