Nano-vLLM 源码解读 - 1. 推理引擎导论

📢 欢迎关注公众号「大模型推理」

本系列文章首发于公众号「大模型推理」，专注分享大模型推理引擎源码解读、性能优化实战与前沿技术追踪。如果你也对推理加速、KV Cache、Continuous Batching、Tensor Parallel 等话题感兴趣，欢迎扫码关注，第一时间获取后续更新，和我一起把推理引擎的每一个细节都拆开讲透。

公众号二维码

这是「Nano-vLLM 源码解读」第 1 讲。整门课围绕 GeeeekExplorer/nano-vllm 这份约 1200 行的精简 vLLM 实现，把推理引擎的核心机制——KV Cache 块化、Continuous Batching、Tensor Parallel、CUDA Graph——一个一个拆开讲透。这一讲是导论：建立心智模型，看懂全景。

配套源码（路径相对 nanovllm/ 包根目录）：llm.py、engine/llm_engine.py、config.py，以及仓库根目录的 example.py。

学习目标

读完这一讲，你能：

  • 用一句话回答"推理引擎到底在做什么"，说清楚它跟训练框架的根本差异
  • 解释 prefill 与 decode 两个阶段为什么对硬件要求截然相反，以及为什么调度必须分开
  • 在脑中画出 nano-vllm 的三层结构图，说出每一层的职责边界
  • 知道 nano 相对生产 vLLM 砍掉了什么、保留了什么、为什么这套精简还能在小模型场景跑出对等吞吐

1. 推理引擎到底在做什么

把它浓缩成一句话：

给定一组到达时刻不同、长度各异的请求，把 GPU 的算力和显存压榨到极限，同时让每个用户感知到的延迟可控。

四个关键词：

  • 到达时刻不同：请求是流式打进来的，不像训练有个静态 dataset；调度器必须不停做在线决策
  • 长度各异：prompt 从 80 token 到 8000 token 都有，没办法 pad 到一样长（pad 等于直接浪费算力）
  • 算力 + 显存：必须同时优化两者；只有 FLOPS 没用，KV Cache 装不下就要抢占
  • 延迟可控：吞吐再高，单个用户首 token 30 秒也是产品级灾难

训练框架（DeepSpeed、Megatron）解决"我有一个静态 batch，怎么让它快"。
推理引擎解决"我有 N 条用户实时打进来的请求，怎么让总系统快"。
这是两类完全不同的问题，需要不同的抽象。

训练像时间均匀的工厂流水线；推理像不断进出的医院急诊室。

2. Prefill 与 Decode：两个完全不同的阶段

LLM 推理拆成两个连续阶段：

阶段 输入 输出 计算量/token 显存压力 瓶颈
Prefill 整个 prompt（N 个 token） 第 1 个新 token 高（一次算 N 个 q · k） 中 算力（compute-bound）
Decode 1 个新 token（自回归） 下 1 个新 token 低（1 行 q · 全部 k） 高（每步读全部 KV） 显存带宽（memory-bound）

两个阶段对硬件的诉求几乎相反：

  • Prefill 喜欢大 batch、长序列、纯 GEMM。算子是 compute-bound 的，FLOPS 利用率轻松上 70%。
  • Decode 受显存带宽限制，单条序列时算力反而吃不饱。要靠把多条序列的 decode 步聚到一起，让一次 KV 读取能服务 batch 里所有序列——这就是 continuous batching 存在的根本原因。

engine/llm_engine.py:49 的 step() 把这种"二选一"写得很直白：

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)
    self.scheduler.postprocess(seqs, token_ids, is_prefill)

每个 step 要么全 prefill 要么全 decode，由调度器决定走哪条。这种二选一简化了内核分发，代价是 prefill 来了就要打断 decode（vLLM 后来用 chunked prefill 缓解；nano 也支持队首 chunk，详见 L8）。

3. nano-vllm 的三层架构

graph TD A["LLM<br/>(用户 API)"] B["LLMEngine<br/>(编排层)"] C["Scheduler + BlockManager<br/>(资源层)"] D["ModelRunner<br/>(执行层)"] E["Qwen3 模型 + layers/*<br/>(计算层)"] A --> B B --> C B --> D D --> E C -.->|"sequences<br/>block_table"| D

三层各自的"独立性测试"——换掉任何一层，其它两层不需要动：

编排层（LLMEngine）
知道"用户有 prompt"，不知道"模型如何前向"。它的工作就是：把 prompts 转成 Sequence、起 step 循环、把完成的序列汇总给 tokenizer 解码。换底层模型，这一层不动。

资源层（Scheduler + BlockManager）
知道"显存里有 N 个 KV 块"，不知道"模型里有几层 attention"。它决定哪些序列这个 step 跑、哪些块该分配 / 释放 / 复用。换 attention 实现，这一层不动。

执行层（ModelRunner + Qwen3* + layers/*）
知道"hidden_size、num_heads、TP shard 怎么切"，不知道"调度策略"。它接受"一组 sequences、是 prefill 还是 decode"，吐出 token id。换调度策略，这一层不动。

这套划分是 vLLM 的核心设计精髓。nano 1200 行能在小模型上跟上 vLLM 几万行的吞吐，前提就是它没在边界上偷懒——每一层都对它不该知道的东西保持无知。

4. 一条 prompt 的端到端调用链

example.py：

from nanovllm import LLM, SamplingParams

llm = LLM("/YOUR/MODEL/PATH", enforce_eager=True, tensor_parallel_size=1)
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
prompts = ["Hello, Nano-vLLM."]
outputs = llm.generate(prompts, sampling_params)

generate() 内部的骨架（engine/llm_engine.py:60）：

def generate(self, prompts, sampling_params, use_tqdm=True):
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp)        # 1. 入队
    while not self.is_finished():           # 2. step 直到全部完成
        output, num_tokens = self.step()
        for seq_id, token_ids in output:
            outputs[seq_id] = token_ids
    return [tokenizer.decode(...) for ...]  # 3. 解码

把这 11 步画成时序图：

用户 LLMEngine Tokenizer Scheduler + BlockManager ModelRunner (Qwen3 + Sampler) loop [while not is_finished()] prepare_* → Context forward → logits Sampler → token_id 更新状态 / append_token 检查 EOS / max_tokens generate(prompts, sampling_params) encode(prompt) token_ids add_request → waiting 入队 Sequence schedule() (seqs, is_prefill) run(seqs, is_prefill) token_ids postprocess(seqs, token_ids) decode(token_ids) text outputs 1 2 3 4 5 6 7 8 9 10 11 12

对照源码逐步说明：

  1. LLM.generate(prompts, sampling_params) 收外部输入
  2. tokenizer 把 prompt 编码成 token_ids
  3. 包成 Sequence 对象，放进 Scheduler.waiting 队列
  4. 进入 while not is_finished() 主循环
  5. Scheduler.schedule() 返回这个 step 要跑的 seqs + is_prefill 标记
  6. ModelRunner.run(seqs, is_prefill) 准备 batch tensor + 设置 Context
  7. Qwen3ForCausalLM.forward 跑前向，得到 logits
  8. Sampler 采样下一个 token id
  9. Scheduler.postprocess 更新序列状态、追加 token、检查 EOS
  10. 完成的序列汇总到 outputs 字典
  11. tokenizer 解码回字符串，返回给用户

这 11 步是这门课反复展开的脚手架——每一讲都会盯着其中某一两步深入。

5. nano-vllm vs 生产 vLLM：取舍清单

功能 vLLM nano 备注
Continuous Batching ✓ ✓ nano 的双队列实现更直接
PagedAttention + Prefix Cache ✓ ✓ nano 用 xxhash + dict，没有 Trie
Tensor Parallelism ✓ ✓ nano 走 mp.spawn + SHM，不依赖 ray
CUDA Graph ✓ ✓ nano 只对 decode 路径捕获
Chunked Prefill ✓ 部分 nano 仅允许队首 chunk
Speculative Decoding ✓ ✗ 砍掉
LoRA / 多模型 ✓ ✗ 砍掉
多模态 ✓ ✗ 砍掉，仅 Qwen3
分布式调度（ray / zmq） ✓ ✗ 单机多卡足够
Streaming（SSE）/ OpenAI Server ✓ ✗ nano 是离线推理库

nano 的"刻意省略"才是它的教学价值。 每个被砍掉的功能都对应一个工程复杂度跳跃。当你看清"少了它 1200 行就够"，你也就理解了"加上它为什么要 5 万行"。

需要给一个量化对照——README 的 benchmark（RTX 4070 Laptop 8GB / Qwen3-0.6B / 256 sequences / 长度 100–1024 随机）：

引擎 输出 token 时间 吞吐
vLLM 133,966 98.37s 1361.84 t/s
nano 133,966 93.41s 1434.13 t/s

注意这是小模型 + 单卡场景。换成 70B + 8 卡 TP，nano 大概率不再领先（生产 vLLM 在大模型上的众多工程优化是 nano 砍掉的）。这门课的目标不是论证"nano 比 vLLM 快"，而是借 nano 的简洁把"vLLM 是怎么工作的"讲透。

6. 这门课接下来讲什么

整门课分 6 个模块、17 讲，都围绕这一讲铺开的"三层架构 + 11 步链路"展开：

模块 主题 解决什么问题
一 全景与抽象 建立心智模型；看清 Sequence 这个核心数据结构在状态机中的流转
二 KV Cache 与内存管理 PagedAttention 块化布局、BlockManager 引用计数、Prefix Cache、显存预算反推
三 调度与批处理 Continuous Batching 双队列、chunked prefill、抢占机制
四 模型前向与算子 Context 元数据、FlashAttention 双 API、Triton store_kvcache、Qwen3 模型搭建
五 张量并行与权重加载 Column/Row/QKV/Merged ParallelLinear、weight_loader 协议、词表并行
六 系统级加速与扫尾 CUDA Graph 捕获策略、多进程 TP 编排、Sampler、性能闭环复盘

下一讲：Sequence 状态机与请求生命周期。 我们会盯着 Sequence 这个类，把它从入队到出队的所有字段变化跟一遍，看它如何同时承担"用户请求"、"调度单位"、"KV 块持有者"三重角色——这是理解后面 KV Cache 和调度章节的前提。

posted @ 2026-05-09 17:25  大模型推理  阅读(23)  评论(0)    收藏  举报