Nano-vLLM源码学习(2):配置、入口与一次 generate 的全流程

📑 目录

• 1. 从一段示例代码进入
• 2. Config：全局配置
• 3. SamplingParams：采样超参
• 4. LLM 与 LLMEngine：构造时做了什么
• 5. 主循环：add_request → step → postprocess
• 总结

1. 从一段示例代码进入

读引擎先看它怎么被调用。README 的 Quick Start 和根目录的 example.py用到的对外 API 只有三样：LLM、SamplingParams和 generate。下面按 example.py的写法整理了一段最简用法：

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

path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")   # 本地模型目录

# 1) 构造引擎：本地路径 + 全局开关
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)

# 2) 定义采样策略
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)

# 3) Chat 模型先套对话模板，再交给引擎
tokenizer = AutoTokenizer.from_pretrained(path)
questions = ["用一句话介绍 PagedAttention", "100 以内有哪些质数"]
prompts = [
    tokenizer.apply_chat_template([{"role": "user", "content": q}],
                                  tokenize=False, add_generation_prompt=True)
for q in questions
]

# 4) 批量生成：每个元素是 {"text": ..., "token_ids": ...}
outputs = llm.generate(prompts, sampling_params)
print(outputs[0]["text"])

这段代码里有三个细节：

• 构造 LLM时第一个参数是本地目录路径，不是 HuggingFace 仓库名。Config会断言 os.path.isdir(model)，权重由 safetensors从本地加载，nano-vllm 不负责下载。
• generate返回 dict 列表（含 "text"和 "token_ids"），vLLM 返回的是 RequestOutput对象。README 把这称为与 vLLM 接口的「minor differences」。
• 对话模板由调用方处理。example.py用 apply_chat_template(..., tokenize=False, add_generation_prompt=True)先把问题包成对话格式的字符串；引擎内部只做 encode，不会自动套模板。

⚠️ 注意：README 的 Quick Start 直接把 "Hello, Nano-vLLM."这样的裸文本传给 generate，用来演示 API 可以，但 Qwen3 是对话模型，不套模板时输出质量会明显变差。实际使用按 example.py的写法。

下面按这四步依次看 Config和 SamplingParams的字段、LLM构造时做的事，以及 generate内部的循环。

2. Config：全局配置

config.py里的 Config是一个 @dataclass(slots=True)，把引擎运行需要的全局约束集中在一处。字段与默认值如下：

●●●Python@dataclass(slots=True)
class Config:
    model: str                              # 本地模型目录（必填）
    max_num_batched_tokens: int = 16384# 单步最多处理的 token 总数（Token Budget）
    max_num_seqs: int = 512# 单步最多并发的序列数
    max_model_len: int = 4096# 支持的最大上下文长度
    gpu_memory_utilization: float = 0.9# 允许占用的显存比例
    tensor_parallel_size: int = 1# 张量并行 GPU 数
    enforce_eager: bool = False# True 则禁用 CUDA Graph
    hf_config: AutoConfig | None = None# 在 __post_init__ 里加载的 HF 配置
    eos: int = -1# EOS token id（初始化后由 tokenizer 填入）
    kvcache_block_size: int = 256# 每个 KV 块容纳的 token 数
    num_kvcache_blocks: int = -1# KV 块总数（运行时按显存推算）

构造后 __post_init__做三条断言，加载 HF 配置，并把上下文长度截到模型支持的上限：

●●●Pythondef__post_init__(self):
assert os.path.isdir(self.model)                 # 必须是存在的本地目录
assertself.kvcache_block_size % 256 == 0# 块大小需 256 的整数倍
assert1 <= self.tensor_parallel_size <= 8# TP 度在 1~8 之间
self.hf_config = AutoConfig.from_pretrained(self.model)
self.max_model_len = min(self.max_model_len,
self.hf_config.max_position_embeddings)

几个值得注意的字段：

• 🔑 kvcache_block_size = 256：每个 KV 块装多少个 token。块越大，块表（Block Table）越短、管理开销越小，但最后一块没填满造成的内部碎片也越大。% 256 == 0这条断言来自底层算子的要求：flash-attn 的 flash_attn_with_kvcache文档写明，传入 block_table（分页 KV）时 page_block_size必须是 256 的倍数。nano-vllm 的 decode 路径和前缀命中的 prefill 路径都通过 block_table读分页 KV，所以块大小只能取 256、512 这类值，256 是允许的最小值，不是在碎片和开销之间挑出来的折中。

⏰ 信息时效：256是 nano-vllm 当前默认值。vLLM 使用自己的 PagedAttention 内核，常见默认 block_size=16，两者的默认值不要混为一谈。

• 🔑 max_num_batched_tokens与 max_num_seqs：前者限制一步最多算多少 token（主要约束 prefill），后者限制一步最多有多少条序列（主要约束 decode）。两者共同构成调度器的 Token Budget，3.5 会详细讲。
• 🔑 gpu_memory_utilization = 0.9：引擎不按最大长度给每条请求预留 KV，而是先加载权重、用一次最大规模的 prefill 测出激活峰值，再把剩余配额全部切成 KV 块。num_kvcache_blocks因此初始为 -1，运行时才算出来，公式见 3.6。
• tensor_parallel_size ≤ 8：这是单机 8 卡的上限。实际可用的 TP 度还受模型约束，Qwen3 的实现要求注意力头数和 KV 头数都能被 TP 度整除，见 3.3。

3. SamplingParams：采样超参

相比 vLLM 几十个采样参数，nano-vllm 的 SamplingParams只有三个字段：

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

def__post_init__(self):
assertself.temperature > 1e-10, "greedy sampling is not permitted"

• temperature：温度，控制采样随机性。
• max_tokens：最多生成多少个 token，默认 64（示例里显式设成了 256）。它既是停止条件，也用于 benchmark 时固定生成量。
• ignore_eos：是否忽略 EOS 继续生成。bench.py把它设为 True，保证每条序列都生成满 max_tokens，吞吐数据才可比。

📌 关键点：__post_init__断言温度必须大于 1e-10，也就是说 nano-vllm 不支持贪心采样（temperature=0）。它的 Sampler只有一条随机采样路径：logits 除以温度后做 softmax，再用指数噪声取 argmax（见 3.3）。温度为 0 时除法会出错，作者没有为这种情况单独写 argmax 分支，而是在入口处挡掉。想要接近确定性的输出，可以把温度设成一个很小的正数。

4. LLM 与 LLMEngine：构造时做了什么

llm.py只有 5 行，LLM直接继承 LLMEngine，类体是 pass，generate等方法全部来自父类。所有初始化工作都在 LLMEngine.__init__里，分四步：

●●●Pythonclass LLMEngine:
def__init__(self, model, **kwargs):
# 1) 从 kwargs 里挑出 Config 认识的字段，构造全局配置
        config_fields = {f.name for f infields(Config)}
        config = Config(model, **{k: v for k, v in kwargs.items()
if k in config_fields})
# 把块大小同步给 Sequence（它用类变量记 block_size）
        Sequence.block_size = config.kvcache_block_size

# 2) 张量并行：用 spawn 启动 rank>0 的 worker 子进程
        ctx = mp.get_context("spawn")
self.ps, self.events = [], []
for rank inrange(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 的 ModelRunner 跑在主进程
self.model_runner = ModelRunner(config, 0, self.events)

# 3) 分词器只在主进程持有，并把 EOS 回填到 config
self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
        config.eos = self.tokenizer.eos_token_id

# 4) 调度器 + 退出清理
self.scheduler = Scheduler(config)
        atexit.register(self.exit)

这里有几处工程细节：

• 🔑 为什么用 spawn而不是 Linux 默认的 fork：CUDA 运行时不支持在 fork出的子进程里重新初始化。父进程一旦初始化过 CUDA，fork出的子进程再用 GPU 就会报错。spawn启动的是全新的解释器进程，每个 rank 各自初始化自己的 GPU。
• 🔑 主从进程靠共享内存 + Event通信：rank>0的 ModelRunner在子进程的 __init__末尾进入循环，等待任务。主进程把方法名和参数 pickle 后写进共享内存，再 set每个 worker 的 Event通知它读取。细节见 3.6。
• 分词器只在主进程：编码和解码都是 CPU 任务，TP worker 只需要 token id，不需要各自加载一份分词器。
• atexit清理：退出时通过同一套 RPC 给子进程发 "exit"，再 join回收，避免残留进程。

💡 提示：Sequence.block_size = config.kvcache_block_size这一行容易被忽略。Sequence用类变量记录块大小，所有序列共享同一个值，num_blocks、last_block_num_tokens等属性都靠它计算。

5. 主循环：add_request → step → postprocess

generate的骨架如下（去掉了 tqdm 进度条与吞吐统计）：

●●●Pythondefgenerate(self, prompts, sampling_params, use_tqdm=True):
# 单个 SamplingParams 复制给每条 prompt；也可以传一个等长列表
ifnotisinstance(sampling_params, list):
        sampling_params = [sampling_params] * len(prompts)
for prompt, sp inzip(prompts, sampling_params):
self.add_request(prompt, sp)

    outputs = {}
whilenotself.is_finished():
        finished, num_tokens = self.step()   # 跑一步，拿回本步完成的序列
for seq_id, token_ids in finished:
            outputs[seq_id] = token_ids

# 序列完成的先后顺序和提交顺序不同，按 seq_id 排序还原
    outputs = [outputs[i] for i insorted(outputs)]
return [{"text": self.tokenizer.decode(ids), "token_ids": ids}
for ids in outputs]

seq_id由 Sequence的类级计数器递增分配，所以按 seq_id排序就等于按提交顺序排序。短 prompt 往往先结束，没有这一步排序，返回顺序会和输入对不上。

两个关键方法：

add_request：把一条 prompt 变成一个 Sequence放进调度器。prompt 可以是字符串（内部先 encode）也可以是现成的 token id 列表：

●●●Pythondefadd_request(self, prompt, sampling_params):
ifisinstance(prompt, str):
        prompt = self.tokenizer.encode(prompt)
self.scheduler.add(Sequence(prompt, sampling_params))

step：一次完整的「调度 → 前向 → 写回」：

●●●Pythondefstep(self):
    seqs, is_prefill = self.scheduler.schedule()                  # 选 batch
# 必须在 postprocess 之前统计：postprocess 会把 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)  # 跑模型 + 采样
self.scheduler.postprocess(seqs, token_ids, is_prefill)      # 写回、判断结束
    outputs = [(seq.seq_id, seq.completion_token_ids)
for seq in seqs if seq.is_finished]
return outputs, num_tokens

📌 关键点：num_tokens用正负号区分阶段。prefill 时返回正数，即这一步实际计算的 token 总量（开启 chunked prefill 后可能只是 prompt 的一部分）；decode 时返回负的 batch 大小 -len(seqs)。generate根据符号分别计算 Prefill 和 Decode 吞吐（tok/s）显示在进度条上。这一行的位置不能挪到 postprocess之后，否则 num_scheduled_tokens已经被清零，prefill 吞吐会一直显示为 0。

循环里的分工是：model_runner.call("run", ...)负责怎么算，scheduler负责算谁、算完怎么办。call在 TP=1 时就是普通的方法调用，TP>1 时还会把同一个调用广播给其他 rank。后面三节分别展开执行层（3.3）、KV 管理（3.4）和调度层（3.5），3.6 再回到 ModelRunner把它们串起来。

📝 总结

• README 和 example.py用到的对外 API 只有 LLM、SamplingParams、generate，返回含 "text"和 "token_ids"的 dict 列表；对话模板由调用方自己套。
• Config是 @dataclass(slots=True)，集中管理显存比例、Token Budget、TP 度、KV 块大小。kvcache_block_size % 256 == 0来自 flash-attn 分页 KV 的要求。
• SamplingParams只有三个字段，且断言温度 > 1e-10，不支持贪心采样。
• LLM是 LLMEngine的空子类。LLMEngine.__init__用 spawn启动 TP worker、只在主进程持有分词器、注册 atexit清理。
• 主循环是 generate → add_request → step → postprocess。step在 postprocess之前统计 num_tokens，用正负号区分 prefill 和 decode；generate最后按 seq_id还原提交顺序。

📚 参考资料

• nano-vllm GitHub 仓库
• FlashAttention 接口文档（flash_attn_interface.py）
• Python multiprocessing — Contexts and start methods
• PyTorch - Multiprocessing best practices（CUDA in multiprocessing）
• HuggingFace Transformers - Chat templates