vLLM 官方仓库数十万行代码，劝退了很多想搞懂 LLM 推理引擎的人。DeepSeek 工程师俞星凯从零写出的 nano-vllm 把核心机制——连续批处理、PagedAttention 分页 KV Cache、前缀缓存、CUDA Graph、张量并行——浓缩进约 1200 行 Python，且吞吐不输 vLLM。本文逐文件精读这份代码，回答一个问题：一个现代 LLM 推理引擎最少需要哪些零件，它们如何咬合在一起。
直接用 HuggingFace model.generate() 逐条跑请求，功能上完全正确，但在服务场景下会浪费掉大部分 GPU 算力。nano-vllm 要解决的所有问题，都可以追溯到三个结构性矛盾。
自回归生成分两个阶段。Prefill 一次性处理整个 prompt，$n$ 个 token 并行过模型，计算量 $\sim 2Pn$（$P$ 为参数量），是 compute-bound 的；Decode 每步只处理 1 个新 token，却要把全部权重和整条 KV Cache 从显存搬一遍，算术强度极低，是 memory-bandwidth-bound 的。把两个阶段混在一个静态 batch 里跑，两边的效率都会被拖累。nano-vllm 的调度器每个 iteration 只跑一种阶段（schedule() 返回 is_prefill 标志），正是对这个二分法的直接回应。
Decode 阶段为了不重复计算历史 token 的注意力，必须缓存每层每头的 K/V。每个 token 的 KV 占用为：
以 nano-vllm 默认支持的 Qwen3-0.6B 为例（$L{=}28$、$n_{kv}{=}8$、$d_{head}{=}128$、bf16），每 token 约 $2{\times}28{\times}8{\times}128{\times}2 = 114{,}688$ 字节 ≈ 112 KB；一条 4096 token 的序列就要约 448 MB——已经赶上 0.6B 模型 bf16 权重本身（约 1.2 GB）的三分之一。换成 70B 级模型，单条长序列的 KV Cache 轻松超过 10 GB。KV Cache 才是推理显存管理的主角，这正是 nano-vllm 用整块 BlockManager 来管它的原因。
传统 serving 把一批请求拼成大 tensor 一起跑、一起收尾，浪费来自两处：
两笔浪费各自对应一篇奠基论文和一个 nano-vllm 的核心模块：
Scheduler.schedule() 92 行实现了后者。推理引擎的本质 = 调度器（让 GPU 每步都满载）× 显存管理器（让每 MB 显存都装真 token）。nano-vllm 的全部代码都是围绕这两件事展开的。
nano-vllm 只支持 Qwen3 一个模型族，换取了极致的可读性。下表逐文件列出 nanovllm 包（19 个文件、1,385 行）的职责划分，另有 bench.py / example.py 共 65 行：
| 文件 | 职责 | 行数 | 对应 vLLM 子系统 |
|---|---|---|---|
engine/llm_engine.py | 顶层引擎：启动 worker、主循环 step()、进度条与吞吐统计 | 90 | LLMEngine |
engine/scheduler.py | 连续批处理调度：prefill 优先、chunked prefill、抢占 | 92 | Scheduler（v1） |
engine/block_manager.py | 分页 KV Cache 分配 / 释放 / 前缀缓存哈希 | 120 | BlockManager + PrefixCache |
engine/sequence.py | 请求的状态对象：token 流、block_table、状态机 | 83 | Sequence / SequenceGroup |
engine/model_runner.py | 模型执行：显存测量、输入拼装、CUDA Graph、多卡指令分发 | 257 | ModelRunner + Worker |
layers/attention.py | Triton KV 写入 kernel + FlashAttention 调用 | 75 | Attention backend |
layers/linear.py | 并行 Linear：Replicated / Column / Row 三种基础形态，外加 QKV / Merged 两种合并列并行 | 156 | tensor_parallel layers |
layers/{rotary,layernorm,sampler,embed_head,activation}.py | RoPE、RMSNorm（融合残差）、采样、Embedding/LMHead、SiLU | 198 | 对应同名 layers |
models/qwen3.py | Qwen3 模型搭建与权重打包映射 | 216 | model registry 的一项 |
utils/{context,loader}.py | 全局 attention 上下文；safetensors 权重加载 | 55 | — |
config.py · sampling_params.py · llm.py · __init__.py | 全局配置（块大小 / 批上限 / TP 等）、采样参数、LLM 类与包入口 | 43 | — |
bench.py / example.py 的 65 行，全仓 Python 共 1,450 行（wc -l）。
LLMEngine.generate() 把所有 prompt 包成 Sequence 塞进调度器后，就是一个朴素的 while 循环。每个 iteration 固定三段：
没有异步流水线、没有流式输出、没有 API server——这是离线批推理的最小闭环，也因此每一步都能在 90 行的 llm_engine.py 里被一眼看穿。
Sequence 是贯穿引擎的「单据」：token_ids（prompt + 已生成）、block_table（逻辑块 → 物理块号）、num_cached_tokens / num_scheduled_tokens 两个游标，外加 WAITING → RUNNING → FINISHED 状态机。它有一个容易被忽略但非常实用的设计：__getstate__ / __setstate__ 让序列在被 pickle 发给 worker 进程时，decode 阶段只携带 last_token 而非整条 token 流——多卡通信的带宽就这样被省下来了（详见第 7 章）。
从 example.py → llm_engine.py（主循环）→ sequence.py（数据模型）入手，先建立「每个 step 每条序列出一个 token」的心智模型，再下钻 scheduler / block_manager / model_runner。
调度器维护两个队列：waiting（新请求 / 被抢占的请求）和 running（正在生成）。schedule() 每次被调用时做出一个全局决策：这一步是 prefill 还是 decode，跑哪些序列，各跑多少个 token。两个硬约束来自配置：max_num_seqs = 512（一步最多多少条序列）和 max_num_batched_tokens = 16384（一步最多多少个 token）。
只要 waiting 非空，本轮就是 prefill 轮：调度器从队首开始依次接纳新请求，直到序列数或 token 预算耗尽。这里藏着一个 vLLM v1 同款的重要机制——chunked prefill：
remaining 个 token，剩下的下一个 iteration 接着算（seq.num_scheduled_tokens = min(num_tokens, remaining)）；if remaining < num_tokens and scheduled_seqs: break——只有队首第一条允许被切块，后续的要么整进要么等下轮。这避免了 batch 里塞满「半截 prompt」导致频繁的小步 prefill，是个简单有效的折中。切块 prefill 的意义在于：一条 8k token 的长 prompt 不会独占 GPU 好几步，decode 中的老请求最多只被推迟一个 iteration，TTFT（首 token 延迟）与 TPOT（逐 token 延迟）之间取得平衡。代价是切块边界上 KV 还没写全，注意力要靠 cu_seqlens 正确掩蔽（第 6 章）。
当没有 prefill 可做时进入 decode 轮：每条 running 序列各生成 1 个 token。问题是有时候显存不够了——某条序列写满当前块、需要追加新块，而空闲块已耗尽。nano-vllm 的处理是教科书式的 preemption（抢占）：
can_append(seq) 检查「序列长度 mod 256 == 1」（即待处理 token 是新块的第一个位置）时是否还有空闲块可分；preempt(self.running.pop())——从队尾（最新加入的序列）开始牺牲，直到挤出空间；若队列里只剩当前序列，则抢占它自身、下轮重来；deallocate 全部物理块、状态打回 WAITING、插回 waiting 队首，下次以 prefill 重算（recompute 模式，不做 swap 到 CPU）。抢占选队尾是 FCFS 下的自然选择：队首序列跑得更久、沉没成本更高。而「重算而非换出」在小模型 + 高带宽场景下往往比 PCIe swap 更划算，也让 BlockManager 少了一整套 CPU 池。
被抢占序列释放块时，块只是引用计数归零回到空闲队列，哈希映射仍然保留（见第 4 章）。于是它重新进入 can_allocate() 时，只要这些块还没被别人复用，就会逐块命中、直接「复活」——名义上是 recompute，实际上重算被前缀缓存免掉了；只有已被复用的块才需要真正重算。抢占机制和 APC 本来各自独立，在这里形成了一层免费的容错。
模型跑完后，postprocess() 做三件账：给新填满的块计算哈希（喂给前缀缓存，见第 4 章）、推进 num_cached_tokens 游标、把新 token 追加进序列并判断结束条件（token == eos 或达到 max_tokens）。结束的序列当场释放全部物理块——注意这些块只是引用计数归零回到空闲队列，哈希映射仍然保留，这正是前缀缓存能命中「已经结束的历史请求」的关键。
nano-vllm 的调度是 FCFS，没有优先级、没有 max_num_partial_prefills 之类的细粒度控制，也不支持 prefill/decode 混合 batch（vLLM v1 已支持同一步混跑）。作为教学实现这是有意的删减，但高并发在线服务场景下这些正是拉开延迟差距的地方。
92 行调度器装下了连续批处理的全部要义：iteration 级重组、prefill 优先 + 队首切块、显存不足时队尾抢占重算。读懂它，vLLM v1 那几千行调度代码就只剩下工程复杂度，没有思想上的新东西了。
PagedAttention 的核心类比是虚拟内存：每条序列看到的「逻辑块」是连续的（第 0 块装 token 0–255，第 1 块装 256–511……），但物理块可以散落在显存任意位置，靠 block_table（逻辑块号 → 物理块号）做地址翻译，attention kernel 拿着这张表去真正读数据。nano-vllm 里块大小固定 256 token：
may_append），不预留；ref_count，多条序列的 block_table 可以指向同一个物理块——这是前缀缓存的地基。nano-vllm 实现了 vLLM 的 Automatic Prefix Caching（APC）的精髓。做法是给每个写满的块算一个 链式哈希：
「链式」意味着块的哈希不仅取决于自身内容，还取决于之前所有块——两个序列只有从头开始逐块相同才算命中，避免了「中间碰巧相同」的误命中。hash_to_block_id 字典把哈希映射到物理块，新请求进来时 can_allocate() 逐块比对：
used_block_ids 中）：ref_count += 1，零拷贝共享；注意比对范围只到 range(seq.num_blocks - 1)——最后一个没写满的块不参与缓存，它的内容还会增长；缓存的最小单位永远是一个写满的 256 token 块。
命中的收益直接体现在 prefill 输入上：num_tokens = seq.num_tokens - num_cached_blocks * block_size，缓存部分根本不进模型；attention 时再通过 block_table 从 KV Cache 读回（第 6 章）。多轮对话、共享 system prompt 的批量任务因此获得数倍加速。
postprocess → hash_blocks() 计算哈希——没写满的块内容还会变，不能入缓存。_allocate_block 里的 del hash_to_block_id[...]），期间它还能被「复活」。xxhash.xxh64 是快速非加密哈希，对 token id 序列的字节视图直接计算，避免 Python 对象开销；64 位输出配合「链式 + 逐块内容比对」（blocks[block_id].token_ids != token_ids 的二次校验）把碰撞概率压到可以忽略。
block_table + ref_count + 链式哈希 三件套，120 行同时实现了 PagedAttention 的按需分配和 APC 的跨请求共享。这是全仓库「行数 / 思想密度」比最高的文件。
调度器管「跑谁」，ModelRunner 管「怎么跑」。它做了四件事：加载模型、测量显存并划分 KV Cache 池、把调度结果翻译成 attention kernel 需要的扁平输入、为 decode 路径提前录制 CUDA Graph。
初始化顺序非常讲究：warmup_model() 先用假数据跑一遍最大 batch 的前向，让 PyTorch 缓存分配器达到峰值；然后 allocate_kv_cache() 读取真实显存占用，把剩余预算全部划给 KV Cache：
分子是 gpu_memory_utilization（默认 0.9）预算下减去权重和预热峰值后的富余；分母正是第 1 章的「每 token 字节数」乘上 256。减 peak 加回 current 是在补偿 PyTorch 分配器的保留显存——这种「先压测再切蛋糕」的做法与 vLLM 的 profile_run 如出一辙。KV Cache 本体是一块形状为 [2, L, num_blocks, 256, n_kv, d] 的巨型张量，随后按层切片，把视图直接挂到每个 Attention 模块的 k_cache / v_cache 属性上——层与缓存池之间没有拷贝，只有指针。
nano-vllm 不做 padding。prefill 时所有序列的待算 token 被拼接成一维长向量，序列边界用累积长度数组 cu_seqlens_q / cu_seqlens_k 标记（FlashAttention 的 varlen 接口约定）。同时构建 slot_mapping：第 $i$ 个输入 token 的 KV 应该写到物理池的哪个槽位——逐块把 block_table[i] * 256 + 块内偏移 展开。若存在前缀缓存命中（cu_seqlens_k > cu_seqlens_q），还要额外传 block_tables，让 kernel 去缓存里读历史 KV。所有张量都走 pin_memory + cuda(non_blocking=True)，主机到设备的拷贝与计算重叠。
Decode 每步每条序列只有 1 个 token，计算量极小，此时 Python 侧逐层 launch kernel 的 CPU 开销会超过 GPU 计算本身。解法是 CUDA Graph：把整次 decode 前向录制成一张图，之后每步只改输入缓冲区、一次 graph.replay() 整体重放。nano-vllm 的实现要点：
[1,2,4,8,16,32,...,512] 分桶，各录一张，运行时向上取整到最近的桶，不足部分用 slot_mapping=-1 的哑元填充（attention 层写缓存时跳过 -1 槽位）；graph_pool 显存池，避免每张图各留一份激活值；enforce_eager=True 时回退 eager——健壮性兜底。小模型 + 小 batch 的 decode 场景，CUDA Graph 往往能带来数倍端到端加速——nano-vllm 能在 0.6B 模型上反超 vLLM，很大一部分功劳属于这套「录制一次、零开销重放」的机制。调试期记得开 enforce_eager=True，否则报错栈会被图重放吞掉。
attention 前向需要一堆与层无关的元数据（cu_seqlens、slot_mapping、block_tables……）。vLLM 用精心设计的 attention metadata 对象层层透传；nano-vllm 则直接设了一个模块级全局变量 _CONTEXT：ModelRunner 在每次 run() 前 set_context()，Attention 层里 get_context() 自取，跑完 reset_context()。这是教科书会批评的写法，但它让 75 行代码承载了 vLLM 里分散在多个抽象层中的逻辑——对学习者而言，数据流向反而更显眼。
store_kvcache 把本步算出的 K/V 按 slot_mapping 散写进物理池：每个 program 处理一个 token，读出其目标槽位 slot，把 num_heads × head_dim 的向量连续写入 k_cache[slot] 和 v_cache[slot]。slot == -1 的哑元直接返回——这就是上一章 CUDA Graph 填充位不产生副作用的原因。分页管理的全部复杂性在调度侧已经消化完，kernel 看到的只是一张「逻辑位置 → 物理槽位」的映射表。
| 场景 | kernel | K/V 从哪来 |
|---|---|---|
| prefill（无前缀命中） | flash_attn_varlen_func | 本步新算的 K/V，因果掩蔽由 causal=True 配合 varlen 边界保证 |
| prefill（有前缀命中） | flash_attn_varlen_func + block_table | k, v 直接换成 k_cache, v_cache：已缓存前缀从物理池读，新 token 刚写入同一池，一次 kernel 调用同时看到新旧 KV |
| decode | flash_attn_with_kvcache | q 只有 1 个 token，全部历史 KV 按 block_table + cache_seqlens 从物理池分页读入 |
最妙的是第二行：前缀命中时 prefill 并不需要单独的「读缓存再拼接」逻辑，而是把 query 限定为新 token、把 K/V 整个指向缓存池——FlashAttention 的 block_table 接口天然支持「一部分 KV 是旧的」，PagedAttention 论文里 kernel 层的分页读取在这里落地为一次函数调用。
nano-vllm 没有自己实现分页注意力 kernel，而是把调度侧维护好的 block_table / slot_mapping / cu_seqlens 精确喂给 FlashAttention 的三个接口。推理引擎的核心竞争力在「元数据的正确性」，kernel 交给最专业的实现——这也是 vLLM 后端架构的真实思想。
nano-vllm 的 TP 用 torch.multiprocessing spawn 为 rank 1..N-1 各起一个进程，rank 0 留在主进程，全组用 NCCL 通信（tcp://localhost:2333）。调度器和 BlockManager 只在 rank 0 存在——它每步调完，需要把「跑哪些序列」广播给 worker。这里的通信设计非常轻巧：
"nanovllm" 的 1 MB 共享内存 + 每个 worker 一个 Event：rank 0 把 (方法名, 参数) pickle 进共享内存并 set 事件；worker 的 loop() 阻塞在 event.wait()，醒来后读缓冲区、本地执行同名方法；model_runner.call("run", seqs, is_prefill) 这一行，在 rank 0 是「广播 + 本机执行」，在 worker 是「收到指令后执行」——所有 rank 跑完全对称的代码，NCCL 集合通信要求的天性被顺势满足；last_token，prefill 才带整条 token 流。只有 rank 0 做采样（prepare_sample 和 sampler(...) 都有 rank == 0 分支），结果也只需 rank 0 知道——下一步调度本来就只发生在 rank 0。
权重加载侧的精妙在 weight_loader 协议：每个参数挂上自己的装载函数，loader.py 遍历 safetensors 时按 packed_modules_mapping 把 HuggingFace 的 q_proj/k_proj/v_proj 三块权重塞进合并后的 qkv_proj 的对应切片，每块再按 rank 取自己的 shard——「合并权重」与「TP 切分」两件麻烦事在同一个函数里完成。细节亮点：
RowParallelLinear 的 bias 只在 rank 0 加（否则 all_reduce 后 bias 被加了 N 遍）；VocabParallelEmbedding 用 mask 把不属于本 rank 的 token id 归零、查表后用 mask 清零再 all_reduce——等价于「谁查到谁贡献」；ParallelLMHead 各 rank 算自己那半词表的 logits，dist.gather 到 rank 0 拼接成完整词表再采样；prefill 时先用 cu_seqlens_q[1:] - 1 只取每条序列最后一个位置的隐状态——中间位置的 logits 根本不算，这是 prefill 阶段最大的单次省算力点；n_kv_heads 同样除以 world_size，每张卡的缓存池只存自己那几头。共享内存方案简单但有边界：缓冲区固定 1 MB、单次只广播一条指令、无流水线。vLLM 用专门的 worker 消息队列和更复杂的集合通信调度。另外 tcp://localhost:2333 硬编码，单机多实例会撞端口。
models/qwen3.py 216 行，结构是标准 decoder-only：嵌入 → N 层（RMSNorm → Attention → RMSNorm → MLP）→ RMSNorm → LMHead，注意力采用 GQA（16 个 Q 头、8 个 KV 头）。没有 registry、没有 mapping 表——想加新模型就照抄这个文件。两个 Qwen3 细节被忠实保留：当 qkv_bias=False 时对 q/k 按头做 q_norm / k_norm（head 维度的 RMSNorm），这是 Qwen3 稳定注意力的关键设计；0.6B 开启 tie_word_embeddings，构造时让 lm_head.weight 与 embedding 权重指向同一份数据，不重复占显存、也不重复加载。
常规写法每层是 h = h + attn(norm(h)); h = h + mlp(norm(h))，产生多次独立的 add 与 norm kernel。nano-vllm 采用 vLLM 同款融合：add_rms_forward(x, residual) 一次完成「加残差 + 归一化」，残差以参数形式在层内显式传递（Qwen3DecoderLayer.forward 里的 hidden_states, residual 双返回值）。再叠加 @torch.compile，add、rsqrt、乘权重被融合成更少的 kernel。
所有层共享同一个 RoPE 模块——get_rope 上挂了 @lru_cache(1)，全模型只有一份 cos_sin_cache 表（形状 [max_position, 1, head_dim]）。前向时按 positions 索引取 cos/sin 做旋转，整个函数同样被 @torch.compile 接管。
Sampler 核心逻辑只有寥寥数行，却藏着一个经典技巧——Gumbel-max 采样的指数分布变体：
除以指数随机数再取 argmax，数学上等价于按 $p$ 做分类采样（取 log 后即 Gumbel-max 形式），全程无循环、可向量化，还能被 @torch.compile 编译进 CUDA Graph 里。温度除法在 softmax 前的 logits 上完成。值得注意的限制：SamplingParams 断言 temperature > 1e-10，即不允许贪心解码（temperature=0 会除零）——想要 greedy 得自己改两行，这算是教学项目的任性之处。
权重加载的 weight_loader 协议（参数自带装载函数）、KV Cache 池按层切片挂属性、RoPE 单例、Gumbel-max 采样——这些 10 行内的巧思都可以直接搬进自己的项目。
RMSNorm 与 RoPE 内部都先升到 float32 计算再转回原 dtype；Sampler 同样先 logits.float()。低精度存储、高精度计算是 LLM 推理的标准卫生。
项目自带的 benchmark（RTX 4070 Laptop 8GB、Qwen3-0.6B、256 条随机长度请求）：
| 推理引擎 | 输出 tokens | 耗时 (s) | 吞吐 (tokens/s) |
|---|---|---|---|
| vLLM | 133,966 | 98.37 | 1,361.84 |
| nano-vllm | 133,966 | 93.41 | 1,434.13 |
nano-vllm 略胜 5%。这不意味着它「比 vLLM 强」——0.6B 小模型、单机离线批处理恰好是开销敏感而功能需求少的场景，vLLM 为通用性付出的抽象成本在这里体现为固定开销。但它证明了：连续批处理 + 分页 KV + CUDA Graph 这套组合拳，就是 vLLM 性能的主体；其余几十万行是为生产环境的广度付出的代价。
口径说明：bench.py 先跑一次预热请求再计时，吞吐 = Σmax_tokens ÷ 墙钟时间；因设置了 ignore_eos=True，两个引擎都恰好输出 133,966 个 token，对比的是相同工作量下的端到端耗时。
| 能力 | vLLM | nano-vllm | 影响 |
|---|---|---|---|
| 模型支持 | 数百种 | 仅 Qwen3 | 去掉模型 registry，模型文件即文档 |
| 在线服务 | OpenAI 兼容 server、流式输出 | 无 | 纯离线批推理；generate 是一次性返回 |
| 采样 | top-p / top-k / greedy / beam 等 | 仅温度采样 | Sampler 12 行；不支持 temperature=0 |
| 调度 | 优先级、混部 prefill+decode、多LoRA | FCFS + 队首切块 + 队尾抢占 | 在线高并发延迟表现会弱于 vLLM |
| 抢占恢复 | recompute / swap 双模式 | 仅 recompute | 无 CPU 缓存池，实现减半 |
| 并行 | TP / PP / EP / DP、跨机 | 单机 TP ≤ 8 卡 | 共享内存指令分发，无流水线并行 |
| 量化 / 投机解码 / 多模态 | 支持 | 无 | 这些正是 vLLM 代码量的主要去向 |
Sequence.block_size 是类属性，由 LLMEngine.__init__ 在运行时改写——grep 赋值处找不到别慌；seq.token_ids 在 worker 进程里是空列表（pickle 瘦身），只有 last_token 有效，别在 worker 侧调试时打印全长；can_append 里的 len(seq) % block_size == 1 表示「待处理的这个 token 是新块的第一个位置」，布尔值在比较时自动按 0/1 计——即此刻需要 1 个空闲块；k, v 会被整个替换成 k_cache, v_cache（attention.py 第 66 行），第一次读容易以为是笔误，实为分页读取的关键；__post_init__ 里 kvcache_block_size % 256 == 0 的断言意味着你不能把它改成 16——flash-attn 的 block_table 接口要求 256 对齐，这是性能与通用性的取舍。nano-vllm 是「推理引擎的最小完备集」：它证明 1200 行足以装下现代 LLM serving 的全部核心思想。把它读通之后再去读 vLLM，你会看到同样的骨架外面包着生产的血肉——而不是一堵陌生的墙。
三条循序渐进的改造练习：① 给 Sampler 加 top-p / top-k；② 给调度器加 prefill+decode 混合 batch；③ 实现 beam search（提示：需要 fork block_table 并管理 ref_count——PagedAttention 论文的并行采样场景）。每一个都会让你对对应模块的理解深入一层。