nano-vllm 源码剖析：千行 Python 手写离线推理引擎的调度、显存与并行时序

1. 概览

nano-vllm 是一个用大约一千多行 Python 从零写的离线推理库，接口接近 vLLM 的 LLM.generate。它只在本进程里把一批提示跑完，不提供 HTTP 服务。模型实现目前只有 Qwen3，权重从本地目录的 *.safetensors 加载。

2. 特性

调度：请求先进入等待队列，每一轮只做 prefill，或者只做 decode。Prefill 有 token 预算，一条长提示可以拆成多轮算，也就是 chunked prefill。Decode 时如果 KV 块不够，会抢占别的正在生成的序列，把它们的缓存释放掉，退回等待队列后重新 prefill。
显存：KV 按页管理，块大小必须是 256 的倍数，默认 256。块数按 GPU 剩余显存和 gpu_memory_utilization 估算，每张卡只保存自己负责的 KV 头。前缀缓存用链式 xxhash 认完整块，命中后再比对 token 内容，命中的块直接复用，不把 KV 拷回 CPU。
并行：张量并行支持 1 到 8，每个进程绑一张 GPU。词表、QKV、MLP 按卡切开，行并行层用 NCCL all_reduce，lm_head 把 logits gather 到 rank 0。多卡时主进程用一块名为 nanovllm 的共享内存把本轮命令发给工作进程。单卡不创建这块内存。
计算：Prefill 用 flash_attn_varlen_func，decode 用 flash_attn_with_kvcache，K/V 由 Triton 按 slot 写入缓存。RMSNorm 和采样器用 torch.compile。未设置 enforce_eager 时，decode 且 batch 不超过 512 会回放 CUDA Graph；图里只有模型前向，lm_head 和采样在图外面。
采样和停止 ：只支持温度、max_tokens 和 ignore_eos。温度必须大于 0，不能用 0 做贪心。没有 top-k、top-p。遇到 EOS，或者生成 token 数达到 max_tokens，请求结束。
GQA（grouped query attention）：多个 query head 共用一组 K/V head，从而减少 KV cache，节省显存空间。

3. 各模块职责

4. 时序

4.1 权重初始化时序

每个进程在构造函数里自己跑。ModelRunner.__init__ 里，所有 rank 都会执行到这两处，此时工作进程还没进 loop()：

warmup_model() 里的 self.run(seqs, True)，每个进程跑一次最大 prefill。
没设 enforce_eager 时，capture_cudagraph() 里对每个 decode 批次大小直接调用 self.model(...)，先热身再捕获 CUDA Graph。

这两次都不经过共享内存。多卡时各进程在 forward 内部的 all_reduce 上会合，所以会一起进入、一起离开。

构造函数后半段才分叉：rank 0 创建共享内存后返回；其他 rank 打开共享内存，停在 loop() 的 event.wait()。

4.2 一次对话的时序

LLMEngine 只活在 rank 0。generate 的 while 每轮调用一次 step，调度器选出本批序列后：

python 复制代码

def step(self):
    seqs, is_prefill = self.scheduler.schedule()
    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)

rank 0 的 call("run", ...) 先把 "run" 和这批 seqs 写入共享内存并 event.set()，然后自己执行 run()，进入 run_model，从而执行本卡的 forward。

工作进程从 wait() 返回，读到 "run"，也执行自己的 run() → run_model() → 本卡的 forward。prefill 每步一次；decode 在能用 CUDA Graph 时，各卡回放自己捕获的图，相当于再跑一遍已录好的 forward。

各进程的 Python 调用不是同一条语句，rank 0 写完共享内存就先进入 run()，工作进程要等事件。它们在 forward 里第一次 all_reduce 时对齐：embedding、每个 Decoder 层的输出投影和 MLP 都会等所有卡都执行到同一次集合通信。