Skip to main content
Inference optimization

3.2 Config, entry points, and one full generate path

Starting from the README and example.py usage, unpack Config / SamplingParams / LLM / LLMEngine and see how the generate main loop drives add_request, step, and postprocess

nano-vllm LLMEngine Inference loop Source reading

📑 Contents

1. Start from a usage snippet

To read an engine, first look at how it is called. The README Quick Start and the root example.py expose only three public pieces: LLM, SamplingParams, and generate. Below is a minimal usage path rewritten in the style of example.py:

import os
from transformers import AutoTokenizer
from nanovllm import LLM, SamplingParams

path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")   # Local model directory

# 1) Build the engine: local path + global switches
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)

# 2) Define the sampling policy
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)

# 3) Chat models: wrap with a chat template before handing to the engine
tokenizer = AutoTokenizer.from_pretrained(path)
questions = ["Introduce PagedAttention in one sentence", "Which primes are under 100"]
prompts = [
    tokenizer.apply_chat_template([{"role": "user", "content": q}],
                                  tokenize=False, add_generation_prompt=True)
    for q in questions
]

# 4) Batch generate: each element is {"text": ..., "token_ids": ...}
outputs = llm.generate(prompts, sampling_params)
print(outputs[0]["text"])

Three details in that snippet:

  • The first argument when constructing LLM is a local directory path, not a HuggingFace repo id. Config asserts os.path.isdir(model); weights load from disk via safetensors. nano-vllm does not download them.
  • generate returns a list of dicts (with "text" and "token_ids"). vLLM returns RequestOutput objects. The README calls this one of the «minor differences» versus the vLLM API.
  • The chat template is the caller’s job. example.py uses apply_chat_template(..., tokenize=False, add_generation_prompt=True) to wrap the question as a chat-format string first; inside the engine only encode runs—no automatic template.

⚠️ Note: the README Quick Start feeds bare text such as "Hello, Nano-vLLM." straight into generate. Fine for showing the API, but Qwen3 is a chat model; without a template, output quality drops clearly. In real use, follow example.py.

The next sections walk those four steps: the fields of Config and SamplingParams, what happens when LLM is constructed, and the loop inside generate.

2. Config: global settings

Config in config.py is a @dataclass(slots=True) that gathers the global constraints the engine needs in one place. Fields and defaults:

@dataclass(slots=True)
class Config:
    model: str                              # Local model directory (required)
    max_num_batched_tokens: int = 16384     # Max tokens processed in one step (Token Budget)
    max_num_seqs: int = 512                 # Max concurrent sequences in one step
    max_model_len: int = 4096               # Max supported context length
    gpu_memory_utilization: float = 0.9     # Allowed GPU memory fraction
    tensor_parallel_size: int = 1           # Number of GPUs for tensor parallelism
    enforce_eager: bool = False             # True disables CUDA Graph
    hf_config: AutoConfig | None = None     # HF config loaded in __post_init__
    eos: int = -1                           # EOS token id (filled from tokenizer after init)
    kvcache_block_size: int = 256           # Tokens per KV block
    num_kvcache_blocks: int = -1            # Total KV blocks (computed from GPU memory at runtime)

After construction, __post_init__ runs three asserts, loads the HF config, and caps context length at the model’s supported maximum:

def __post_init__(self):
    assert os.path.isdir(self.model)                 # Must be an existing local directory
    assert self.kvcache_block_size % 256 == 0        # Block size must be a multiple of 256
    assert 1 <= self.tensor_parallel_size <= 8       # TP degree between 1 and 8
    self.hf_config = AutoConfig.from_pretrained(self.model)
    self.max_model_len = min(self.max_model_len,
                             self.hf_config.max_position_embeddings)

A few fields worth pausing on:

  • 🔑 kvcache_block_size = 256: how many tokens each KV block holds. Larger blocks mean a shorter Block Table and less management overhead, but more internal fragmentation when the last block is only partly filled. The % 256 == 0 assert comes from the underlying operator: the flash-attn flash_attn_with_kvcache docs say that when a block_table (paged KV) is passed, page_block_size must be a multiple of 256. nano-vllm’s decode path and the prefix-hit prefill path both read paged KV through block_table, so the block size can only be values like 256 or 512. 256 is the smallest allowed value, not a compromise picked between fragmentation and overhead.

⏰ Currency note: 256 is nano-vllm’s current default. vLLM uses its own PagedAttention kernels and commonly defaults to block_size=16; do not mix the two defaults.

  • 🔑 max_num_batched_tokens and max_num_seqs: the former caps how many tokens one step may compute (mainly constrains prefill); the latter caps how many sequences one step may hold (mainly constrains decode). Together they form the scheduler’s Token Budget; 3.5 covers this in detail.
  • 🔑 gpu_memory_utilization = 0.9: the engine does not reserve KV for every request at max length. It loads weights, runs one max-size prefill to measure the activation peak, then carves the remaining quota entirely into KV blocks. That is why num_kvcache_blocks starts at -1 and is only computed at runtime; the formula is in 3.6.
  • tensor_parallel_size ≤ 8: an upper bound for 8 GPUs on one machine. The usable TP degree is also constrained by the model: the Qwen3 implementation requires both the number of attention heads and the number of KV heads to be divisible by the TP degree; see 3.3.

3. SamplingParams: sampling hyperparameters

Compared with vLLM’s dozens of sampling parameters, nano-vllm’s SamplingParams has only three fields:

@dataclass(slots=True)
class SamplingParams:
    temperature: float = 1.0
    max_tokens: int = 64
    ignore_eos: bool = False

    def __post_init__(self):
        assert self.temperature > 1e-10, "greedy sampling is not permitted"
  • temperature: temperature; controls sampling randomness.
  • max_tokens: maximum tokens to generate; default 64 (the example sets it explicitly to 256). It is both a stop condition and a way to fix generation length in benchmarks.
  • ignore_eos: whether to ignore EOS and keep generating. bench.py sets it to True so every sequence runs to max_tokens and throughput numbers stay comparable.

📌 Key point: __post_init__ asserts temperature must be greater than 1e-10, which means nano-vllm does not support greedy sampling (temperature=0). Its Sampler has only one random-sampling path: divide logits by temperature, softmax, then argmax after exponential noise (see 3.3). Temperature 0 would make the division fail; the author did not add a separate argmax branch for that case and instead blocks it at the entry point. For near-deterministic output, set temperature to a very small positive number.

4. LLM and LLMEngine: what construction does

llm.py is only 5 lines. LLM subclasses LLMEngine directly; the class body is pass, and methods such as generate all come from the parent. All initialization work lives in LLMEngine.__init__, in four steps:

class LLMEngine:
    def __init__(self, model, **kwargs):
        # 1) Pick Config-known fields out of kwargs and build global config
        config_fields = {f.name for f in fields(Config)}
        config = Config(model, **{k: v for k, v in kwargs.items()
                                   if k in config_fields})
        # Sync block size onto Sequence (it stores block_size as a class variable)
        Sequence.block_size = config.kvcache_block_size

        # 2) Tensor parallelism: spawn worker child processes for rank>0
        ctx = mp.get_context("spawn")
        self.ps, self.events = [], []
        for rank in range(1, config.tensor_parallel_size):
            event = ctx.Event()
            p = ctx.Process(target=ModelRunner, args=(config, rank, event))
            p.start()
            self.ps.append(p); self.events.append(event)
        # Rank 0's ModelRunner runs in the main process
        self.model_runner = ModelRunner(config, 0, self.events)

        # 3) Tokenizer lives only in the main process; write EOS back into config
        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
        config.eos = self.tokenizer.eos_token_id

        # 4) Scheduler + exit cleanup
        self.scheduler = Scheduler(config)
        atexit.register(self.exit)

A few engineering details here:

  • 🔑 Why spawn instead of Linux’s default fork: the CUDA runtime does not support re-initializing inside a forked child. Once the parent has initialized CUDA, a forked child that then uses the GPU will error. spawn starts a fresh interpreter process; each rank initializes its own GPU.
  • 🔑 Parent and workers talk through shared memory + Event: a rank>0 ModelRunner enters a wait loop at the end of the child’s __init__. The main process pickles the method name and arguments into shared memory, then sets each worker’s Event so it can read. Details in 3.6.
  • Tokenizer only in the main process: encode and decode are CPU work; TP workers only need token ids and do not each load a tokenizer.
  • atexit cleanup: on exit, the same RPC path sends "exit" to the children, then joins them so no processes are left behind.

💡 Tip: the line Sequence.block_size = config.kvcache_block_size is easy to miss. Sequence stores block size as a class variable shared by every sequence; properties such as num_blocks and last_block_num_tokens all depend on it.

5. Main loop: add_request → step → postprocess

The skeleton of generate (tqdm progress bar and throughput stats removed):

def generate(self, prompts, sampling_params, use_tqdm=True):
    # A single SamplingParams is copied to every prompt; or pass an equal-length list
    if not isinstance(sampling_params, list):
        sampling_params = [sampling_params] * len(prompts)
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp)

    outputs = {}
    while not self.is_finished():
        finished, num_tokens = self.step()   # Run one step; collect sequences finished this step
        for seq_id, token_ids in finished:
            outputs[seq_id] = token_ids

    # Finish order differs from submission order; sort by seq_id to restore it
    outputs = [outputs[i] for i in sorted(outputs)]
    return [{"text": self.tokenizer.decode(ids), "token_ids": ids}
            for ids in outputs]

seq_id is assigned by a class-level counter on Sequence, so sorting by seq_id is the same as sorting by submission order. Short prompts often finish first; without this sort, return order would not match the input.

Two key methods:

add_request: turn one prompt into a Sequence and hand it to the scheduler. The prompt may be a string (encoded internally first) or an existing list of token ids:

def add_request(self, prompt, sampling_params):
    if isinstance(prompt, str):
        prompt = self.tokenizer.encode(prompt)
    self.scheduler.add(Sequence(prompt, sampling_params))

step: one full «schedule → forward → write back» cycle:

def step(self):
    seqs, is_prefill = self.scheduler.schedule()                  # Pick the batch
    # Must count before postprocess: postprocess clears num_scheduled_tokens
    num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs)
    token_ids = self.model_runner.call("run", seqs, is_prefill)  # Run model + sample
    self.scheduler.postprocess(seqs, token_ids, is_prefill)      # Write back, decide finished
    outputs = [(seq.seq_id, seq.completion_token_ids)
               for seq in seqs if seq.is_finished]
    return outputs, num_tokens

📌 Key point: num_tokens uses the sign to distinguish phases. Prefill returns a positive number—the actual token count computed this step (with chunked prefill that may be only part of the prompt). Decode returns the negative batch size -len(seqs). generate uses the sign to compute Prefill and Decode throughput (tok/s) separately for the progress bar. This line cannot move after postprocess, or num_scheduled_tokens has already been cleared and Prefill throughput would always show 0.

Division of labor in the loop: model_runner.call("run", ...) decides how to compute; scheduler decides who to compute and what to do afterward. With TP=1, call is an ordinary method call; with TP>1 it also broadcasts the same call to the other ranks. The next three sections expand the execution layer (3.3), KV management (3.4), and the scheduling layer (3.5); 3.6 returns to ModelRunner and ties them together.

📝 Summary

  • The public API used by the README and example.py is just LLM, SamplingParams, and generate, returning a list of dicts with "text" and "token_ids"; the chat template is applied by the caller.
  • Config is a @dataclass(slots=True) that centralizes GPU memory fraction, Token Budget, TP degree, and KV block size. kvcache_block_size % 256 == 0 comes from flash-attn’s paged-KV requirement.
  • SamplingParams has only three fields and asserts temperature > 1e-10; greedy sampling is unsupported.
  • LLM is an empty subclass of LLMEngine. LLMEngine.__init__ uses spawn to start TP workers, keeps the tokenizer only in the main process, and registers atexit cleanup.
  • The main loop is generate → add_request → step → postprocess. step counts num_tokens before postprocess and uses the sign to mark prefill vs decode; generate finally sorts by seq_id to restore submission order.

🎯 Self-check checklist

  • Can write a minimal inference snippet with LLM and SamplingParams that applies a chat template, and pull outputs[0]["text"]
  • Can say what Config’s max_num_batched_tokens, max_num_seqs, gpu_memory_utilization, and kvcache_block_size each constrain
  • Can explain why kvcache_block_size must be a multiple of 256
  • Can explain why nano-vllm does not support greedy sampling
  • Can explain why LLMEngine starts tensor-parallel processes with spawn rather than fork
  • Can retell the order inside step(), and why num_tokens must be computed before postprocess
  • Can explain why generate sorts by seq_id before returning

📚 References