大模型推理引擎入门：从 Continuous Batching 到 Scheduler，对着日志读懂 nano-vllm

茜布
0 阅读36分钟

致谢与声明 本文基于 nano-vllm 项目（MIT License，作者 Xingkai Yu）进行学习和分析。文中引用的代码片段均来自该项目。日志中 TEACHING/... 格式的输出是笔者为了学习自行添加的，非项目原始功能。感谢 nano-vllm 作者提供了这个优秀的教学级推理引擎实现。

前言

为什么要写这个

我当初第一次想看懂推理引擎的代码，打开 vLLM 仓库，几十万行，目录翻了半天都没找到主循环在哪，直接劝退。

后来慢慢摸出路子了。马斯克那句话我觉得说得挺对：再难的事也能拆成十几个能执行的步骤。难就难在怎么拆——从哪一层开始看、哪些东西先当黑盒、调用链按什么顺序跟。

这篇笔记就是想把我自己拆解的过程记下来：

  • 给同样是程序员背景、想入门推理引擎但不知道从哪下手的人，省点时间
  • 也是费曼学习法——能讲清楚给别人听，自己才算真搞懂了

如果你之前有过打开一个大项目"不知道从哪读起"的体验，这篇应该对你有帮助。

我怎么读这类代码（这篇也会按这个路子走）

  1. 黑盒 + 接口：这次不讲的东西（比如 Paged Attention 内部 hash 怎么算、TP 多进程怎么同步），就当它有明确的输入输出，像 DDD 里一个 bounded context。先把整体模型搭起来，下次再拆黑盒。
  2. 日志对着代码看：分析任何流程的时候，要阅读真实跑出来的日志。光看代码容易脑补，日志不会骗你。
  3. 读函数前先搞清楚输入输出：这个函数的输入是什么、输出是什么、为什么要存在。如果牵涉前置概念（比如 continuous batching、cu_seqlens），先把概念讲透再进代码。
  4. 看不懂就举例子跑一遍：如果一段代码逻辑不好理解，我们就构造一个具体的输入例子，用这个例子过一遍函数的执行流程，比抽象地读代码有效得多。

导读

这篇文章会讲什么

  • 单 GPU 前提下，把一次 llm.generate() 从用户调用到返回结果，在 LLMEngine → Scheduler → ModelRunner 这条链路上走通
  • 搞懂 Continuous Batching 在 nano-vllm 里是怎么落地的：什么时候 prefill、什么时候 decode、一个 step 里到底干了什么
  • 对着真实日志看调度决策、看 prefill/decode 切换、看两条请求怎么从入队跑到结束
  • 把 Sequence 当"一条请求的状态机"读一遍，把 Scheduler 的 schedule() 和 ModelRunner 的 prepare_* 读细
  • 理解引擎初始化：权重怎么加载、KV cache 块数怎么算、CUDA Graph 什么时候录——这些是运行时的基础设施

这篇文章不讲什么（留到后面）

  • Paged Attention / BlockManager 内部实现（prefix hash、块复用算法）——下篇单独讲
  • Tensor Parallel 多进程（SharedMemory 广播、NCCL 同步）——另开一篇
  • Transformer 各层内部 batch 怎么展开（和 FlashAttention kernel 绑太紧）——单独讲
  • CUDA Graph 原理、各种 MLSys 调优——这篇只点到"日志里你会看到什么"

读之前最好有的知识背景

  1. Decoder-only Transformer 一次 forward 大概怎么走：Q/K/V 矩阵乘、KV cache 是什么意思
  2. Token 是什么：prompt 和 completion 怎么拼成一条序列
  3. GPU 的基本特性：适合并行算一大堆同类运算；推理引擎本质上在抢"显存"和"吞吐"之间的平衡

推荐看过的参考材料：

1. 项目背景和概念

1.1 为什么选择 nano-vllm 来入门

如果你直接去读 vLLM 的源码，大概率会遇到这些问题：

  • 代码量巨大，光 engine 相关的模块就散落在十几个文件里，调用链很深
  • 为了兼容各种硬件、各种模型、各种部署模式，有大量条件分支和抽象层
  • online serving、流式输出、多租户这些工程细节和核心推理逻辑混在一起，初学者很难分清哪些是"必须理解的"、哪些是"先跳过的"

nano-vllm 做的事情，是把工业推理引擎里几个核心机制用最小实现摆出来：

包含的核心机制说明
Continuous Batching动态调度 prefill / decode，不等齐、不浪费
Paged KV CacheKV cache 按块管理，避免显存碎片
Prefix Caching相同前缀的请求复用已有的 KV 块
CUDA Graph（可选）录制 decode 路径减少 kernel launch 开销
Tensor Parallel（可选）多卡切分模型参数并行推理

它没打算做的事情：

  • Online serving（HTTP API、流式输出）
  • 多租户隔离、各种企业级策略
  • 多种模型架构支持（目前只支持 Qwen 系列）

整个项目核心代码大约 1600 行，engine 相关的几个文件加起来不超过 700 行。目标是"先搞懂调度 + 执行主循环"，nano-vllm 比直接扎进 vLLM 仓库友好得多。

1.2 Continuous Batching：在讲什么，解决什么问题

在深入代码之前，我们需要先理解一个核心概念：Continuous Batching（连续批处理）。后面读 Scheduler 和 ModelRunner 的时候，所有调度决策都建立在这个机制之上。

先理解一个前提：LLM 生成文本是迭代式的

这一点非常关键，是理解 Continuous Batching 的基础：

Decoder-only 的语言模型（GPT、LLaMA、Qwen 等）不是一次推理就能生成全部文本的。它的工作方式是：

  1. 第一步（Prefill）：把整段 prompt 一次性过模型，计算所有 token 的 K/V 写入 cache，同时产出第一个续写 token
  2. 之后每一步（Decode）：把上一步生成的 token 作为输入，过一遍模型，生成下一个 token
  3. 重复步骤 2，直到生成了 EOS（结束符）或者达到最大长度限制

也就是说，生成 N 个 token，模型需要被调用 N 次（1 次 prefill + N-1 次 decode）。每两次 GPU 推理之间，存在一个天然的调度窗口——这个观察是 Continuous Batching 的基础。

演进过程：从不 batch 到 static batching 再到 continuous batching

阶段一：不做 batch（逐请求处理）

最简单的做法：一个请求一个请求地处理，处理完一个再接下一个。问题很明显——decode 阶段每步只算 1 个 token，GPU 利用率极低。

阶段二：Static Batching（静态批处理）

改进思路：把多个请求攒成一批，一起送进模型。同一个 batch 里的请求共享一次 GPU 调用，利用率上来了。

但 static batching 有一个根本性的约束：batch 的组成在整个生成周期内是固定的。一旦一批请求开始跑，中途不能有人退出，也不能有新人加入，必须等这批里所有请求全部跑完，才能开始下一批。

这带来两个严重的问题：

问题一：Early-Finished（早结束浪费）

假设一个 batch 里有 4 个请求，其中请求 A 只需要生成 20 个 token 就遇到 EOS 结束了，但请求 D 要生成 256 个 token。在 static batching 下，请求 A 结束后的 236 步里，它占着 batch 里的一个位置，但既不产出 token，也不释放显存。GPU 在为它做无用功。

请求 A: [===]..................   (生成 20 token 后结束，但必须等到 batch 结束)
请求 B: [==========]...........  (生成 80 token 后结束)
请求 C: [==================]...  (生成 200 token 后结束)
请求 D: [======================] (生成 256 token，整个 batch 才能结束)
         ↑                      ↑
       batch 开始            batch 结束
       (. 表示已结束但仍占位的浪费)

问题二：Late-Joining（晚加入延迟）

一次完整的文本生成可能需要几秒到几十秒。如果在当前 batch 执行过程中有新请求到达，它只能在队列里等，直到当前 batch 里所有请求全部结束。这意味着新请求的等待时间取决于当前 batch 里最慢那个请求的剩余生成时间。

核心洞察：Iteration-Level Scheduling（迭代级调度）

2022 年 OSDI 上发表的 Orca 论文提出了关键洞察：

既然 LLM 生成文本是一次一个 token 地迭代调用 GPU，那么在每次 GPU 推理的间隙，就是一个天然的调度点。我们完全可以在这个间隙插入调度操作，实现 batch 成员的动态增删。

这就是 Iteration-Level Scheduling——把调度粒度从"batch 级别"（整个生成周期）细化到"iteration 级别"（每生成一个 token）。不同的框架对这个概念有不同的叫法：

  • vLLM / HuggingFace TGI 叫 Continuous Batching
  • TensorRT-LLM 叫 In-Flight Batching
  • Orca 论文原文叫 Iteration Batching

本质是同一件事。

Continuous Batching 具体怎么工作

在 Continuous Batching 下，每一步（每一次 GPU 推理之前）调度器都会重新决定"这一步算哪些请求"：

  1. 新请求进来，先放到 waiting 队列
  2. 每一步调度时，Scheduler 先检查 waiting 里有没有可以做 prefill 的请求
    • 有的话，拉出来做 prefill：一次把它们的 prompt 算完，写入 KV cache，产出第一个 token
    • 这些请求随后进入 running 队列
  3. 如果没有新请求要 prefill，就从 running 队列里拉出所有正在生成的序列，做一步 decode——每个序列往前走 1 个 token
  4. 某个序列生成完了（遇到 EOS 或达到 max_tokens），立刻从 running 中移除，释放它占用的 KV cache 块
  5. 释放的资源马上可以给 waiting 里的下一个新请求使用

对比 static batching，这两个问题都解决了：

  • Early-Finished：序列结束后立刻移除，不再占位浪费
  • Late-Joining：新请求不需要等当前 batch 跑完，下一步调度时就可能被拉进来做 prefill

用同样的 4 个请求举例，Continuous Batching 下的时间线大致是这样的：

请求 A: [===]                           (生成 20 token 后结束，立刻释放)
请求 B: [==========]                    (生成 80 token 后结束)
请求 C: [==================]            (生成 200 token 后结束)
请求 D: [======================]        (生成 256 token 后结束)
请求 E:      [================]         (A 结束后立刻加入，不用等)
请求 F:               [========]        (B 结束后立刻加入)

在 nano-vllm 里，这个逻辑落在哪

在 nano-vllm 中，Continuous Batching 的核心逻辑在 Scheduler.schedule() 方法中。它的返回值是：

(scheduled_seqs: list[Sequence], is_prefill: bool)

即"这一步要算的序列列表"和"这一步是 prefill 还是 decode"。后面第 6 节我们会详细读这段代码，看它是怎么做调度决策的。

参考材料（建议读完概念再看）：

2. 阅读 nano-vllm 代码

2.1 先把代码跑起来

开始学习任何一个项目，我们需要有 hands-on 代码的习惯。要做的第一件事情，就是把代码跑起来。环境配好、代码跑通之后，才可以开始通过增加日志的方式来学习项目代码的执行流程。

环境准备

# 克隆项目
git clone https://github.com/niconielsen32/nano-vllm.git
cd nano-vllm

# 安装依赖（需要 Python 3.10~3.12，CUDA 环境）
pip install -e .

# 下载模型（这里用 Qwen3-0.6B，体积小，适合学习）
# 模型放到 ~/huggingface/Qwen3-0.6B/ 或你自己的路径

跑 example.py

项目自带了一个 example.py，用法很简单：

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

path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")
tokenizer = AutoTokenizer.from_pretrained(path)
llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)

sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
prompts = ["introduce yourself", "list all prime numbers within 100"]

# 套用 chat template
prompts = [
    tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=False, add_generation_prompt=True,
    )
    for prompt in prompts
]

outputs = llm.generate(prompts, sampling_params)
for prompt, output in zip(prompts, outputs):
    print(f"Prompt: {prompt!r}")
    print(f"Completion: {output['text']!r}\n")

跑通之后，建议自己在代码的关键位置加打印日志，方便对着日志学习执行流程。本文后续贴的那些 TEACHING/... 格式的日志，就是我为了学习特意在 nano-vllm 源码里加的，不是项目原始的输出。你跑出来的 seq_id、耗时可能不同，但只要在对应的位置加了日志，阶段和形态应该一致。

2.2 整体架构鸟瞰

在读具体代码之前，我们先从整体上了解一下 nano-vllm 的模块设计和它们之间的关系。

目录结构

nanovllm/
├── __init__.py              # 对外暴露 LLM、SamplingParams
├── llm.py                   # LLM 类（LLMEngine 的别名）
├── config.py                # 全局配置
├── sampling_params.py       # 单条请求的采样参数
├── engine/
│   ├── llm_engine.py        # 引擎主循环：generate / step / add_request
│   ├── scheduler.py         # 调度器：Continuous Batching 的核心
│   ├── model_runner.py      # 模型执行：加载权重、拼 batch、forward、采样
│   ├── sequence.py          # 单条请求的状态管理
│   └── block_manager.py     # KV cache 块分配与回收
├── models/
│   └── qwen3.py             # Qwen3 模型结构定义
├── layers/                  # 各算子层：attention、linear、layernorm 等
└── utils/
    ├── context.py           # forward 时的元数据上下文
    ├── loader.py            # 权重加载
    └── trace_log.py         # 教学日志工具

各模块职责

模块一句话说清楚它干什么
LLMEngine (llm_engine.py)对外入口。generate() 把请求入队，然后循环调 step() 直到全部完成
Scheduler (scheduler.py)调度核心。管理 waiting / running 两个队列，schedule() 决定本步是 prefill 还是 decode，postprocess() 更新 token、处理结束的序列
ModelRunner (model_runner.py)执行层。加载模型、管理 KV cache，run() 里拼 tensor → forward → 采样，返回生成的 token
Sequence (sequence.py)一条请求的状态。token_ids、block_table、状态机（WAITING → RUNNING → FINISHED）
BlockManager (block_manager.py)KV cache 块的分配与回收（本篇当黑盒，只看接口）
Config (config.py)全局配置：max_num_seqs、max_num_batched_tokens、gpu_memory_utilization 等
SamplingParams (sampling_params.py)单条请求的采样参数：temperature、max_tokens、ignore_eos

2.3 找到入口：从 Python 包到引擎主循环

对于一个新项目，我个人的习惯是先梳理清楚代码从头到尾的执行流程，这样方便在读任何代码的时候，能够找到代码执行的上游，弄明白为什么这里的代码要做这件事情，上游调用这里的需求是什么。

从刚才运行 example.py 的代码可以看到，用户使用 nano-vllm 的方式是：

from nanovllm import LLM, SamplingParams

llm = LLM(path, enforce_eager=False, max_model_len=4096)
outputs = llm.generate(prompts, sampling_params)

我们追一下这个调用链：

  1. nanovllm/__init__.py 导出了 LLM 和 SamplingParams
  2. LLM 定义在 nanovllm/llm.py，打开来看：
from nanovllm.engine.llm_engine import LLMEngine

class LLM(LLMEngine):
    pass

LLM 直接继承自 LLMEngine，没有任何额外逻辑。这是为了保持和 vLLM 项目命名的一致性——用户习惯写 LLM(...)，但所有真正的逻辑都在 LLMEngine 里。

所以我们的入口就是 nanovllm/engine/llm_engine.py。后面所有的代码阅读，都以这个文件为起点往下追。

2.4 一个请求从进来到出去：完整流程

这一节是全文的核心。我们先从函数调用链的角度搞清楚"一次 generate 里发生了什么"，然后用时序图画出来，最后对着真实日志从头到尾走一遍。

这里只讲**"发生了什么"**，不深入每个模块内部的逻辑——那是后面几节的事。

2.4.1 调用链（函数粒度）

打开 llm_engine.py，从 generate() 开始追，核心流程可以概括为：

  1. generate(prompts, sampling_params)

    • 对每个 prompt 调用 add_request()，把请求入队
    • 然后进入 while not is_finished(): step() 循环，直到 scheduler 里 waiting 和 running 都空了
    • 最后把结果按 seq_id 排序，decode 成文本返回

  2. add_request(prompt, sampling_params)

    • 如果 prompt 是字符串，先 tokenize 成 list[int]
    • 用 token_ids 和 sampling_params 构造一个 Sequence 对象
    • 调用 scheduler.add(seq) 放入 waiting 队列队尾

  3. step()（每一轮引擎步，也就是一次"调度 + 推理 + 后处理"）

    • 阶段一：调度 — seqs, is_prefill = scheduler.schedule()
    • 阶段二：执行 — token_ids = model_runner.call("run", seqs, is_prefill)
    • 阶段三：后处理 — scheduler.postprocess(seqs, token_ids)
    • 收集本步刚结束的序列，计算 num_tokens 统计值（prefill 为正，decode 为负），返回

对应到代码（llm_engine.py:185-251）的关键几行：

def step(self):
    self._global_step += 1
    # 阶段一：调度
    seqs, is_prefill = self.scheduler.schedule()
    # 阶段二：执行
    token_ids = self.model_runner.call("run", seqs, is_prefill)
    # 阶段三：后处理
    self.scheduler.postprocess(seqs, token_ids)
    # 收集结果
    outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]
    num_tokens = sum(len(seq) for seq in seqs) if is_prefill else -len(seqs)
    return outputs, num_tokens

这里 num_tokens 的设计值得说一下：

  • Prefill 时为正数：sum(len(seq) for seq in seqs) — 注意这是 postprocess 之后算的，所以每个 seq 的长度已经包含了刚生成的第一个 token
  • Decode 时为负数：-len(seqs) — 表示本步有多少个序列各 decode 了 1 个 token，用负数是为了和 prefill 区分，方便外层算吞吐

2.4.2 时序图

sequenceDiagram
    participant G as generate()
    participant S as Scheduler
    participant M as ModelRunner

    G->>G: add_request × N (入队到 waiting)

    loop 直到 waiting 和 running 都空
        G->>S: schedule()
        S-->>G: (seqs, is_prefill)
        G->>M: run(seqs, is_prefill)
        M-->>G: token_ids
        G->>S: postprocess(seqs, token_ids)
        Note over G: 收集本步结束的序列
    end

    G->>G: 按 seq_id 排序，decode 成文本返回

2.4.3 端到端 trace：两条请求 + 真实日志

下面用本地真实运行的日志来走一遍完整流程。运行条件：

  • 模型：Qwen3-0.6B，单卡
  • 两条 prompt："introduce yourself"（经 chat template 后 11 token）、"list all prime numbers within 100"（17 token）
  • SamplingParams(temperature=0.6, max_tokens=256)
（1）入队
generate：入队 2 条请求，循环 step 直至 waiting/running 皆空

TEACHING/GENERATE_START | phase=ingress | reason=批量请求进入引擎
  | stats={n_prompts=2}

TEACHING/REQUEST_ENQUEUE | phase=ingress | reason=新请求进入 waiting 队列
  | stats={prompt_len=11, seq_id=4, max_tokens=256}

TEACHING/REQUEST_ENQUEUE | phase=ingress | reason=新请求进入 waiting 队列
  | stats={prompt_len=17, seq_id=5, max_tokens=256}

两条请求依次被 add_request 加入 waiting 队列。seq_id 从 4 开始是因为初始化阶段 warmup 已经用掉了 0~3。

（2）Step #1：Prefill — 两条请求一起进入 running
TEACHING/SCHEDULE_PREFILL | phase=prefill | reason=waiting 队列有可准入请求
  | stats={engine_step=1, scheduled=2, waiting_len=0, running_len=2}

Scheduler 发现 waiting 里有 2 条请求，总 token 数 11+17=28 远小于 max_num_batched_tokens=16384，显存也够分配 KV 块，于是把两条都拉进来做 prefill。

step#1 prefill done in 772.24ms num_tokens_stat=30 finished_this_step=[]

这里 num_tokens_stat=30 而不是 28——因为 num_tokens 是在 postprocess 之后算的，此时每条序列已经 append_token 了第一个生成 token，所以 len(seq) 分别变成了 12 和 18，加起来 30。

Prefill 耗时 772ms（包含了两条 prompt 的完整 forward + 采样）。没有序列结束。

（3）Step #2 起：Decode — waiting 空了，从 running 拉人
TEACHING/SCHEDULE_DECODE | phase=decode | reason=无可 prefill 请求，转入 decode 连续批
  | stats={decode_step_index=0, engine_step=2, scheduled=2, running_len=2, waiting_len=0}

waiting 已经空了，Scheduler 转入 decode 分支，把 running 里的 2 条序列都拉出来，每人 decode 1 个 token。

step#2 decode done in 3.02ms num_tokens_stat=-2 finished_this_step=[]

num_tokens_stat=-2 表示本步有 2 条序列各 decode 了 1 个 token。耗时 3ms——相比 prefill 的 772ms 快了两个数量级，因为 decode 每步只算 1 个 token 的 Q，KV 从 cache 里读。

后续 step#3 ~ step#147 形态相同：每步都是 decode，scheduled=2，耗时稳定在 2~3ms。

（4）Step #148：seq_id=4 结束，batch size 从 2 变 1
TEACHING/SCHEDULE_DECODE | phase=decode
  | stats={decode_step_index=146, engine_step=148, scheduled=2, running_len=2}

TEACHING/STATE_TRANSITION | phase=postprocess | reason=命中 eos 或达到 max_tokens
  | result=RUNNING -> FINISHED
  | stats={seq_id=4, token_id=151645, num_completion_tokens=148, max_tokens=256}

step#148 decode done in 2.61ms num_tokens_stat=-2 finished_this_step=[4]

seq_id=4 在第 148 步生成了 token_id=151645（这是 Qwen3 的 EOS token），postprocess 中判断命中 EOS，标记为 FINISHED，释放 KV 块，从 running 中移除。

注意 num_completion_tokens=148——这条请求的 prompt 是 11 token，一共生成了 148 个 token 就遇到了 EOS，没有用满 max_tokens=256。

（5）Step #149 起：只剩 seq_id=5，batch size = 1
TEACHING/SCHEDULE_DECODE | phase=decode
  | stats={decode_step_index=147, engine_step=149, scheduled=1, running_len=1}

step#149 decode done in 73.78ms num_tokens_stat=-1 finished_this_step=[]

注意这里 step#149 耗时突然跳到 74ms——这是因为 batch size 从 2 变成了 1，CUDA Graph 需要切换到另一个预录制的图（bs=1 的图），第一次 replay 有额外的初始化开销。之后 step#150 起又恢复到 ~2.7ms。

（6）Step #256：seq_id=5 结束，全部完成
TEACHING/SCHEDULE_DECODE | phase=decode
  | stats={decode_step_index=254, engine_step=256, scheduled=1, running_len=1}

TEACHING/STATE_TRANSITION | phase=postprocess | reason=命中 eos 或达到 max_tokens
  | result=RUNNING -> FINISHED
  | stats={seq_id=5, token_id=773, num_completion_tokens=256, max_tokens=256}

step#256 decode done in 2.85ms num_tokens_stat=-1 finished_this_step=[5]

seq_id=5 生成了 256 个 token 后达到 max_tokens 限制被结束（注意 token_id=773 不是 EOS，是正常 token，但因为达到长度上限所以被截断）。

TEACHING/GENERATE_FINISH | phase=egress | reason=waiting/running 均为空
  | stats={total_steps=256, prefill_steps=1, decode_steps=255, n_outputs=2}

waiting 和 running 都空了，generate() 退出循环，返回结果。

小结
engine_stepphasescheduled关键事件
1prefill2两条 prompt 一起算，KV 写入 cache，各产出第一个 token
2 ~ 147decode2两条序列每步各生成 1 token，耗时稳定 2~3ms
148decode2seq_id=4 命中 EOS 结束，从 running 移除
149decode1batch size 变 1，CUDA Graph 切换导致首次 74ms
150 ~ 255decode1只剩 seq_id=5，每步 ~2.7ms
256decode1seq_id=5 达到 max_tokens 结束，generate 返回

总计：1 步 prefill + 255 步 decode = 256 步。这就是 Continuous Batching 的真实运行形态——prefill 和 decode 交替调度，序列结束后立刻释放，batch size 动态变化。

最终输出

generate() 返回后，两条请求的生成结果如下：

请求 1（seq_id=4，prompt: "introduce yourself"，148 tokens，命中 EOS 结束）：

<think>
Okay, the user asked me to introduce myself. I need to be friendly and
straightforward. Let me start by saying I'm a helpful AI assistant. Then,
I should mention my purpose, like providing information and support. It's
important to keep the tone positive and not sound too technical. I should
also add something about being a language model to make it more personal.
Let me check if I'm covering all bases: introduction, purpose, and
personal touch. Yeah, that should work.
</think>

Hello! I'm an AI language model designed to assist you with questions,
provide information, and help you in whatever way you need. I'm here to
help! Let me know how I can be of assistance! 😊<|im_end|>

请求 2（seq_id=5，prompt: "list all prime numbers within 100"，256 tokens，达到 max_tokens 截断）：

<think>
Okay, so I need to list all the prime numbers between 100. Let me think.
First, I remember that a prime number is a number greater than 1 that has
no positive divisors other than 1 and itself. So, I need to check numbers
starting from 100 upwards and see if they meet that condition.

Starting with 100. Let me check if it's prime. Well, 100 is even, so it's
divisible by 2. So, 100 is not prime. Next, 101. Hmm, 101. Let me think.
What's the square root of 101? Approximately 10.05. So numbers up to 10.
Let's check divisibility. Starting from 2: 2 doesn't divide 101 because
101 is odd. 3? 1+0+1=2, which isn't divisible by 3. 5? Ends with 1, so
no. 7? Let me do 7*14=98, 7*15=105, so 101-98=3, so
</think>

（因达到 max_tokens=256 被截断，模型还在 <think> 里逐个检查质数，没来得及输出最终回答）

可以看到：请求 1 的回答比较短，模型自己生成了 EOS 就结束了；请求 2 的回答因为涉及逐个数字检查，模型还没说完就被 max_tokens 截断了。这也解释了为什么两条请求的结束时间差了 100 多步——这正是 Continuous Batching 能发挥优势的典型场景。

2.5 引擎是怎么初始化的

在 2.4 的 trace 里，我们看到了请求从入队到结束的完整流程。但有个问题：Scheduler 的 block_manager 里有 3052 个 KV 块——这个数字哪来的？ModelRunner 是怎么加载模型、怎么分配 KV cache 的？CUDA Graph 什么时候录的？

这一节就来回答这些问题。先搞清楚初始化流程，后面深入 Scheduler 和 ModelRunner 的运行时逻辑时就不会断。

2.5.1 LLMEngine.init 的整体流程

打开 llm_engine.py:72，__init__ 的主线可以概括为 6 步：

class LLMEngine:
    def __init__(self, model, **kwargs):
        # 1. 构建 Config
        config = Config(model, **config_kwargs)

        # 2. 若 TP > 1，spawn 子进程（本篇不展开）
        for i in range(1, config.tensor_parallel_size):
            process = ctx.Process(target=ModelRunner, args=(config, i, event))
            process.start()

        # 3. 创建 ModelRunner（rank 0）—— 加载模型、warmup、分配 KV cache、录 CUDA Graph
        self.model_runner = ModelRunner(config, 0, self.events)

        # 4. 加载 tokenizer，拿到 eos_token_id
        self.tokenizer = AutoTokenizer.from_pretrained(config.model)
        config.eos = self.tokenizer.eos_token_id

        # 5. 创建 Scheduler —— 注意：必须在 ModelRunner 之后
        self.scheduler = Scheduler(config)

        # 6. 注册退出清理
        atexit.register(self.exit)

这里有一个关键的顺序依赖：Scheduler 必须在 ModelRunner 之后创建。原因是 Scheduler 内部的 BlockManager 需要知道 config.num_kvcache_blocks（能分多少个 KV 块），而这个值是在 ModelRunner.__init__ 里通过显存估算算出来的。如果先创建 Scheduler，num_kvcache_blocks 还是初始值 -1，BlockManager 无法正确工作。

对应日志：

TEACHING/ENGINE_INIT | phase=bootstrap
  | stats={max_num_batched_tokens=16384, max_num_seqs=512,
           model=/home/.../Qwen3-0.6B/, tp_size=1}

2.5.2 ModelRunner.init：四步走

ModelRunner 的初始化是整个引擎启动中最重的部分，它完成了从"裸 GPU"到"可以跑推理"的全部准备工作。

打开 model_runner.py:69，四步：

第一步：初始化通信 + 加载模型权重

dist.init_process_group("nccl", ...)    # 即使单卡也会走，建立进程组
torch.cuda.set_device(rank)              # 绑定 GPU
self.model = Qwen3ForCausalLM(hf_config) # 创建模型结构（空权重）
load_model(self.model, config.model)     # 从 safetensors 加载权重到 GPU
self.sampler = Sampler()                 # 采样器

对应日志：

TEACHING/TP_RUNTIME_INIT | phase=init
  | stats={device=cuda:0, rank=0, world_size=1}

TEACHING/MODEL_WEIGHTS_LOADED | phase=init
  | stats={num_files=1, num_keys=311}

Qwen3-0.6B 只有 1 个 safetensors 文件，311 个权重 key。

第二步：Warmup — 用假数据顶出显存峰值

def warmup_model(self):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    # 构造最大规模的假 prefill：max_num_batched_tokens / max_model_len 条序列，每条 max_model_len 长
    num_seqs = min(max_num_batched_tokens // max_model_len, max_num_seqs)
    seqs = [Sequence([0] * max_model_len) for _ in range(num_seqs)]
    self.run(seqs, True)  # 跑一次 prefill
    torch.cuda.empty_cache()

为什么要 warmup？因为 GPU 上的显存占用分三块：

  1. 模型权重：加载后一直在
  2. 前向时的临时占用：激活、cuBLAS workspace 等，forward 时冲到峰值，结束后部分释放
  3. KV cache：之后要分配的大块常驻显存

如果不 warmup 直接算 KV cache 能分多少，会高估可用显存——因为你还没见过真实 forward 的峰值有多大。后续跑真实请求时，forward 峰值 + KV cache 同时在，就可能 OOM。

所以流程是：先 warmup 让 peak 暴露出来 → 再按 (总显存 × 利用率 - 已用 - peak + current) 算 KV 预算。

对应日志：

warmup_model：用假数据跑一轮最大规模 prefill，使 cuBLAS 等分配 workspace 并记录 peak

第三步：allocate_kv_cache — 算块数，分配显存，绑定到各层

这一步的核心公式：

available_bytes = total * gpu_memory_utilization - used - peak + current
num_kvcache_blocks = available_bytes // block_bytes

其中 block_bytes = 2（K+V）× 层数 × block_size × kv_heads × head_dim × dtype_size。

算出块数后，分配一个形状为 (2, num_layers, num_blocks, block_size, num_kv_heads, head_dim) 的大 tensor，然后遍历模型的所有 Attention 层，把对应层的切片绑上去：

self.kv_cache = torch.empty(2, num_layers, num_blocks, block_size, num_kv_heads, head_dim)

layer_id = 0
for module in self.model.modules():
    if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
        module.k_cache = self.kv_cache[0, layer_id]  # 本层的 K cache
        module.v_cache = self.kv_cache[1, layer_id]  # 本层的 V cache
        layer_id += 1

这样每层 Attention 在 forward 的时候直接往自己的 k_cache / v_cache 写就行了。

对应日志：

TEACHING/KV_BUDGET_ESTIMATE | phase=init
  | stats={total_gib=95.085, free_gib=93.426, peak_gib=1.599,
           block_bytes=29360128, num_blocks=3052}

allocate_kv_cache：total_mem=95.09GiB free=93.43GiB used=1.66GiB
  peak_alloc=1.60GiB curr_alloc=1.16GiB gpu_memory_utilization=0.9
  -> num_kvcache_blocks=3052 block_size=256

算一下：95.09 GiB × 0.9 ≈ 85.58 GiB 是预算上限。减去已用和峰值后，可用的显存除以每块 ~28 MiB，得到 3052 个块。每块能存 256 个 token 的 KV，总共能缓存约 78 万 token 的 KV cache。

第四步：capture_cudagraph — 录制 decode 的 CUDA Graph

CUDA Graph 是把一段 GPU 执行序列录制下来，之后重放时跳过 CPU 侧的 kernel launch 开销。对 decode 来说特别有效——因为 decode 每步计算量小，kernel launch 的开销占比反而很高。

if not self.enforce_eager:
    self.capture_cudagraph()

nano-vllm 会为一系列 batch size 各录一张图：

CUDA Graph 录制完成：decode batch 档位
  [1, 2, 4, 8, 16, 32, 48, 64, ..., 496, 512]

运行时根据实际 batch size 选最近的档位重放。这也解释了 2.4 节里 step#149 突然变慢（73ms）的原因——batch size 从 2 切到 1，第一次 replay 新图时有额外开销。

录完后，引擎就绪：

引擎就绪：num_kvcache_blocks=3052 eos_token_id=151645

2.5.3 初始化流程小结

flowchart TD
    A["Config 构建"] --> B["ModelRunner.__init__"]
    B --> B1["1. init_process_group + 加载模型权重"]
    B1 --> B2["2. warmup_model()：假数据跑最大 prefill，暴露显存峰值"]
    B2 --> B3["3. allocate_kv_cache()：算块数，分配 KV tensor，绑定到各层"]
    B3 --> B4["4. capture_cudagraph()：录制各 batch size 的 decode 图"]
    B4 --> C["加载 tokenizer，拿 eos_token_id"]
    C --> D["Scheduler(config)：用 num_kvcache_blocks 创建 BlockManager"]
    D --> E["引擎就绪"]
阶段做了什么对应日志标记
Config解析参数，从 HF config 读模型结构ENGINE_INIT
加载权重safetensors → GPU，绑定到模型结构MODEL_WEIGHTS_LOADED
Warmup假数据最大 prefill，暴露 cuBLAS peakwarmup_model
KV cache算块数，分配大 tensor，按层绑定KV_BUDGET_ESTIMATE
CUDA Graph为各 batch size 录制 decode 图CUDA Graph 录制完成
Scheduler用 num_kvcache_blocks 创建 BlockManager引擎就绪

到这里，引擎初始化完成。generate() 的第一个 step() 开始时，模型权重在 GPU 上，KV cache 池子分好了，CUDA Graph 录好了，Scheduler 的 BlockManager 知道自己有 3052 个块可以分配——一切就绪。

2.6 Scheduler：怎么用队列做 Continuous Batching

这一节我们深入 Scheduler 的内部逻辑。在 2.4 的 trace 里我们已经看到了 Scheduler 的外在表现——什么时候 prefill、什么时候 decode、什么时候移除序列。现在来看它是怎么做这些决策的。

2.6.0 前置：Sequence — 一条请求的状态机

在读 Scheduler 之前，先熟悉它操作的核心对象。打开 sequence.py，一个 Sequence 就是一条请求在推理引擎里的完整生命周期状态。

状态机：

WAITING ──(schedule: prefill)──> RUNNING ──(postprocess: EOS/max_tokens)──> FINISHED
   ↑                                 |
   └──(preempt: 显存不够)─────────────┘

关键字段：

字段说明
seq_id全局唯一 ID，自增分配
token_ids当前完整序列（prompt + 已生成的 token）
last_token最后一个 token（decode 时只需要算这一个的 Q）
num_tokens当前总长度
num_prompt_tokensprompt 部分的长度（不变）
num_cached_tokens前缀中命中 prefix cache 的 token 数（本篇不展开）
block_table该序列占用的物理 KV 块 ID 列表
statusWAITING / RUNNING / FINISHED
temperature / max_tokens / ignore_eos从 SamplingParams 拷贝过来的采样参数

核心方法：

  • append_token(token_id)：decode 一步后追加一个 token，更新 last_token 和 num_tokens
  • num_completion_tokens（属性）：num_tokens - num_prompt_tokens，已生成了多少 token
  • num_blocks（属性）：当前长度需要多少个 KV 块（向上取整）

举个例子：我们 trace 里的 seq_id=4，prompt 有 11 个 token。

  • 创建时：num_tokens=11, num_prompt_tokens=11, status=WAITING
  • Prefill 后 append 第一个 token：num_tokens=12, num_completion_tokens=1, status=RUNNING
  • Decode 到第 148 步，append 到 EOS：num_tokens=159, num_completion_tokens=148, status=FINISHED

2.6.1 BlockManager 当黑盒：只认接口

Scheduler 里会频繁调用 BlockManager——问它"有没有块给这条序列"、"decode 下一步要不要新块"。Paged Attention 的内部实现（hash 怎么算、prefix 怎么匹配）留到下一篇，这里只需要知道每个接口的语义：

方法问题谁调用
can_allocate(seq)空闲块够不够给这条 seq 当前长度分配？prefill 准入判断
allocate(seq)给 seq 填 block_table，可能更新 num_cached_tokens从 waiting 拉出来时
can_append(seq)下一步 decode 会不会跨块、需要新块？够不够？decode 准入判断
may_append(seq)如果跨块了，分配新块；如果当前块刚满，固化 hashdecode 分支里拉入前
deallocate(seq)释放 seq 占用的所有块（ref_count--，归零则回收）序列结束或被 preempt 时

读 schedule() 的时候，手里拿着这张表，代码会顺很多。

2.6.2 schedule()：两阶段调度

输入：无显式参数（读 self.waiting / self.running）

输出：(scheduled_seqs: list[Sequence], is_prefill: bool)

这是 Continuous Batching 的核心。逻辑分两阶段：

阶段一：尝试 Prefill

从 waiting 队列头部开始，尽可能多地拉序列进来做 prefill：

while self.waiting and num_seqs < self.max_num_seqs:
    seq = self.waiting[0]
    # 三个退出条件
    over_token_budget = num_batched_tokens + len(seq) > self.max_num_batched_tokens
    cannot_alloc = not self.block_manager.can_allocate(seq)
    if over_token_budget or cannot_alloc:
        break
    # 通过检查，拉出来
    self.block_manager.allocate(seq)
    num_batched_tokens += len(seq) - seq.num_cached_tokens
    seq.status = SequenceStatus.RUNNING
    self.waiting.popleft()
    self.running.append(seq)
    scheduled_seqs.append(seq)

if scheduled_seqs:
    return scheduled_seqs, True  # 有人被 prefill，返回

三个停止条件：

  1. Token 预算超了：再加这条 seq 的 token 数会超过 max_num_batched_tokens（16384）
  2. KV 块不够：can_allocate 返回 False，空闲块不够给这条 seq 分配
  3. 序列数到上限：达到 max_num_seqs（512）

我们 trace 里的场景：waiting 里有 2 条（11 + 17 = 28 token），远小于 16384，块也够，所以两条都被拉进来做 prefill。

阶段二：没有可 Prefill 的，做 Decode

如果 waiting 空了（或者所有 waiting 里的都因为块不够/预算超了而进不来），就从 running 里拉序列做 decode：

while self.running and num_seqs < self.max_num_seqs:
    seq = self.running.popleft()
    # 块不够时 preempt 腾位
    while not self.block_manager.can_append(seq):
        if self.running:
            victim = self.running.pop()  # 从队尾踢，FIFO 下后来的先让出
            self.preempt(victim)
        else:
            self.preempt(seq)  # 自己也进不了，回 waiting
            break
    else:
        self.block_manager.may_append(seq)
        scheduled_seqs.append(seq)

# 把拉出来的序列放回 running 队头，保持 FIFO
self.running.extendleft(reversed(scheduled_seqs))
return scheduled_seqs, False

这里有个 preempt（抢占） 机制：如果某条序列的下一步 decode 需要新块，但空闲块不够了，Scheduler 会把 running 队尾的序列"踢出去"——释放它的块，把它放回 waiting 队头（下次 prefill 时优先处理）。这样"牺牲"一条序列来保证其他序列能继续 decode。

在我们的 trace 里，只有 2 条序列、3052 个块，完全不会触发 preempt。但在高并发场景下，这个机制是保证系统不死锁的关键。

对照日志：

TEACHING/SCHEDULE_PREFILL | phase=prefill
  | stats={engine_step=1, scheduled=2, waiting_len=0, running_len=2}

TEACHING/SCHEDULE_DECODE | phase=decode
  | stats={engine_step=2, decode_step_index=0, scheduled=2, running_len=2, waiting_len=0}

2.6.3 postprocess()：更新 token，判断结束

每步 model_runner.run() 返回 token_ids 后，Scheduler 的 postprocess 负责收尾：

def postprocess(self, seqs, token_ids):
    for seq, token_id in zip(seqs, token_ids):
        seq.append_token(token_id)
        if (not seq.ignore_eos and token_id == self.eos) or \
           seq.num_completion_tokens == seq.max_tokens:
            seq.status = SequenceStatus.FINISHED
            self.block_manager.deallocate(seq)
            self.running.remove(seq)

逻辑很清晰：

  1. 给每条序列 append_token
  2. 检查两个结束条件：
    • 命中 EOS 且没有设置 ignore_eos
    • 已生成的 token 数达到 max_tokens
  3. 如果结束了：标记 FINISHED，释放 KV 块，从 running 移除

对照 trace 里的两次结束：

# seq_id=4：命中 EOS
TEACHING/STATE_TRANSITION | reason=命中 eos 或达到 max_tokens
  | stats={seq_id=4, token_id=151645, num_completion_tokens=148}

# seq_id=5：达到 max_tokens
TEACHING/STATE_TRANSITION | reason=命中 eos 或达到 max_tokens
  | stats={seq_id=5, token_id=773, num_completion_tokens=256}

seq_id=4 的 token_id=151645 是 Qwen3 的 EOS token，seq_id=5 的 token_id=773 是普通 token 但 num_completion_tokens 达到了 256。

2.6.4 Scheduler 整体小结

flowchart TD
    S["schedule() 被调用"]
    S --> P{"waiting 非空？"}
    P -- "是" --> PF["尝试 Prefill：<br/>从 waiting 拉序列<br/>直到预算/块/数量到上限"]
    PF --> PF_OK{"拉到了？"}
    PF_OK -- "是" --> RET_P["返回 (seqs, True)"]
    PF_OK -- "否" --> P
    P -- "否" --> D["Decode：<br/>从 running 拉序列<br/>块不够则 preempt"]
    D --> RET_D["返回 (seqs, False)"]

一句话总结：有新请求就 prefill，没有就 decode；块不够就 preempt。

2.7 ModelRunner：怎么把一批 Sequence 变成 GPU 上的 batch

Scheduler 决定了"这一步算谁"，ModelRunner 负责"怎么算"。这一节我们来看 run() 的内部流程，以及 prefill 和 decode 两种模式下输入是怎么准备的。

2.7.0 前置：Context — 为什么需要一个全局上下文

在读 prepare_prefill / prepare_decode 之前，先理解一个设计决策。

Prefill 和 decode 的输入形态差异很大：prefill 要传 cu_seqlens_q/k（变长序列边界），decode 要传 context_lens + block_tables。模型的每一层（特别是 Attention）都需要知道"当前是 prefill 还是 decode"以及这些元数据，来决定调用哪个 kernel。

如果把这些信息从函数参数一路传下去，每一层的 forward 签名都会很臃肿。nano-vllm 的做法是用一个进程内全局的 Context 对象（utils/context.py）：

  • prepare_* 末尾调用 set_context(...) 把本步的元数据写进去
  • 各层 forward 里调用 get_context() 取出来用
  • run 结束前 reset_context() 清空

就是一个典型的"thread-local 避免参数透传"模式，后端开发里很常见。

2.7.1 run() 整体流程

打开 model_runner.py:586，run() 的逻辑很清晰：

def run(self, seqs, is_prefill):
    # 1. 准备输入：根据 prefill/decode 拼不同形态的 tensor + set_context
    input_ids, positions = self.prepare_prefill(seqs) if is_prefill \
                           else self.prepare_decode(seqs)

    # 2. 准备采样温度（仅 rank 0 需要）
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None

    # 3. 跑模型：得到 logits
    logits = self.run_model(input_ids, positions, is_prefill)

    # 4. 采样：logits + temperature → token_ids（仅 rank 0）
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None

    # 5. 清理上下文
    reset_context()
    return token_ids

2.7.2 prepare_prefill：把变长序列拼成一维

Prefill 的难点在于：不同请求的 prompt 长度不同，但要拼成一个 batch 送进 GPU。

nano-vllm 的做法是把所有序列的 token 展平成一维，然后用 cu_seqlens（cumulative sequence lengths，累积长度）来标记每条序列的边界。这是 FlashAttention 变长输入的标准格式。

我们用 trace 里的真实数据来举例：两条序列，seq_id=4（11 token）、seq_id=5（17 token），num_cached_tokens 都是 0。

# 展平 input_ids：所有 token 拼成一维
input_ids = [t0,t1,...,t10, t0,t1,...,t16]  # 长度 28

# positions：每个 token 的绝对位置（RoPE 用）
positions = [0,1,...,10, 0,1,...,16]  # 长度 28

# cu_seqlens_q：累积长度，标记序列边界
cu_seqlens_q = [0, 11, 28]  # seq_id=4 占 [0,11)，seq_id=5 占 [11,28)

# cu_seqlens_k：K 侧累积长度（无 prefix cache 时和 Q 相同）
cu_seqlens_k = [0, 11, 28]

# slot_mapping：每个 token 写入 KV cache 的物理位置
# seq_id=4 的 block_table=[0]，所以 slot 是 0*256+0, 0*256+1, ..., 0*256+10
# seq_id=5 的 block_table=[1]，所以 slot 是 1*256+0, 1*256+1, ..., 1*256+16
slot_mapping = [0,1,...,10, 256,257,...,272]  # 长度 28

最后调用 set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=17, max_seqlen_k=17, slot_mapping, ...)。

对应日志：

TEACHING/PREPARE_PREFILL_TENSORS | phase=prefill
  | stats={flat_tokens=28, n_seqs=2, max_seqlen_q=17, max_seqlen_k=17,
           has_block_tables=False}

has_block_tables=False 表示没有 prefix cache 命中（两条都是全新 prompt），Attention 层会用当前 forward 产出的 K/V 而不是从 cache 里读。

2.7.3 prepare_decode：每人一个 token

Decode 比 prefill 简单得多——每条序列只需要算最后一个 token 的 Q，K/V 全从 cache 里读。

同样用 trace 数据，step#2 时两条序列各 decode 1 个 token（prefill 后各有 12 和 18 个 token）：

# 每人一个 last_token
input_ids = [token_4_last, token_5_last]  # shape (2,)

# 位置 = 当前序列长度 - 1
positions = [11, 17]  # shape (2,)

# 每序列的当前总长度
context_lens = [12, 18]  # shape (2,)

# 每人只写一个 slot：最后一个 token 在 cache 中的位置
# seq_id=4: block_table[-1]=0, last_block_num_tokens=12, slot = 0*256+12-1 = 11
# seq_id=5: block_table[-1]=1, last_block_num_tokens=18, slot = 1*256+18-1 = 273
slot_mapping = [11, 273]  # shape (2,)

# block_tables：始终需要（decode 的 K/V 全从 cache 读）
block_tables = [[0], [1]]  # shape (2, 1)

对应日志：

TEACHING/PREPARE_DECODE_TENSORS | phase=decode
  | stats={batch_size=2, context_lens_shape=(2,),
           slot_mapping_shape=(2,), block_tables_shape=(2, 1)}

decode 详单#1/3：bs=2 input_ids.shape=(2,) context_lens=[12, 18]
  slot_mapping=[11, 273] block_tables.shape=(2, 1)

日志里的 context_lens=[12, 18] 和 slot_mapping=[11, 273] 和我们算的完全一致。

2.7.4 run_model：Eager vs CUDA Graph

run_model 根据条件选择两种执行路径：

def run_model(self, input_ids, positions, is_prefill):
    if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
        # Eager：直接跑 model.forward + compute_logits
        return self.model.compute_logits(self.model(input_ids, positions))
    else:
        # CUDA Graph：选档位 → 拷数据到静态 buffer → replay → compute_logits
        bs = input_ids.size(0)
        selected_bs = next(x for x in self.graph_bs if x >= bs)
        # 把本步数据拷到图的静态 buffer 前 bs 行
        graph_vars["input_ids"][:bs] = input_ids
        graph_vars["positions"][:bs] = positions
        graph_vars["slot_mapping"][:bs] = context.slot_mapping
        # ... context_lens, block_tables
        graph.replay()
        return self.model.compute_logits(graph_vars["outputs"][:bs])

什么时候走 Eager：

  • Prefill（变长，shape 不固定）
  • enforce_eager=True（调试用）
  • batch size > 512（太大了录图没意义）

什么时候走 CUDA Graph：

  • Decode 且 batch size ≤ 512
  • 从预录的档位 [1, 2, 4, 8, ..., 512] 中选一个 ≥ 当前 bs 的最小值
  • 把数据拷到静态 buffer → graph.replay() → 从静态 buffer 读 output

对应日志：

# Prefill 走 eager
run_model：eager（prefill=True enforce_eager=False bs=28>512=False）

# Decode 走 CUDA Graph
TEACHING/CUDAGRAPH_REPLAY | phase=decode
  | stats={batch_size=2, graph_batch_size=2}

run_model 返回 logits 后，Sampler 按温度采样出每条序列的下一个 token：

temperatures=[0.6, 0.6] -> token_ids=[151667, 151667] seq_ids=[4, 5]

（这是 prefill 后的第一个 token，两条序列碰巧采样到了同一个 token_id。）

2.7.5 Prefill 和 Decode 的对比小结

PrefillDecode
每序列输入全部 token（或去掉 cached 前缀的部分）只有 last_token（1 个）
input_ids shape(total_tokens,) 一维展平(batch_size,)
序列边界cu_seqlens_q/k 标记不需要（每人恰好 1 token）
KV 来源forward 当场算出的（或 prefix cache）全部从 cache 读（block_tables）
slot_mapping每个新 token 都要写 cache每序列只写 1 个 slot
执行路径Eager（变长不适合固定图）CUDA Graph（shape 固定，减少 launch 开销）

结语

这篇讲了什么

回头看一下，这篇从零开始走了一遍 nano-vllm 的核心链路：

  1. Continuous Batching 的概念：从 static batching 的 Early-Finished / Late-Joining 两个痛点出发，理解 Iteration-Level Scheduling 的核心洞察——在每次 GPU 推理的间隙插入调度
  2. 整体架构：LLMEngine 是指挥官，Scheduler 决定"算谁"，ModelRunner 负责"怎么算"
  3. 端到端 trace：对着真实日志走了一遍两条请求从入队到结束的完整流程——1 步 prefill + 255 步 decode，batch size 从 2 动态降到 1
  4. 引擎初始化：模型加载 → warmup 暴露显存峰值 → 算 KV 块数 → 录 CUDA Graph，理解了那个 num_kvcache_blocks=3052 是怎么来的
  5. Scheduler 细节：两阶段调度（先 prefill 后 decode）、三个停止条件、preempt 机制、postprocess 的结束判断
  6. ModelRunner 细节：prefill 展平 + cu_seqlens 描述边界，decode 每人一个 token + block_tables 读 cache，Eager vs CUDA Graph 的分支

到这里，你应该能对着日志说出每一步是 prefill 还是 decode、engine_step 和 decode_step_index 是什么关系、某条序列为什么在这一步结束。

这篇没讲什么（下一篇的方向）

  • Paged Attention / BlockManager 内部：prefix hash 怎么算、块怎么复用、ref_count 怎么管理——这篇只用了接口表，没拆黑盒
  • Tensor Parallel：多卡时 SharedMemory + Event 怎么同步、模型怎么切分
  • Transformer 各层怎么处理 batch：Attention 里 FlashAttention 的 varlen / kvcache 两种 kernel 怎么用 Context 里的元数据
  • CUDA Graph 录制细节：静态 buffer 怎么分配、图怎么捕获

这些留到"入门-2"。

建议

如果你想自己跑一遍，建议：

  1. 用 Qwen3-0.6B 这种小模型，单卡就能跑，日志也短
  2. 自己在关键位置加打印日志——本文贴的那些 TEACHING/... 格式的日志，就是我为了学习在 nano-vllm 源码里自己加的，不是项目原始的输出。学习一个项目的时候，主动加日志是非常有效的方法
  3. 先跑通默认的 example.py，再试着改 max_tokens、改 prompt 数量，观察日志里 batch size 和 step 数的变化
  4. 对着本文的 trace 和你自己的日志做对比——形态一致就对了，具体数字会因采样随机性和硬件不同而不同
avatar
文章
阅读
粉丝