大模型推理入门与 nano-vLLM 源码详解
这篇文章从一个问题出发:当我们向大模型提交一段文本,模型和推理引擎究竟做了什么,才让回答逐个 token 地出现?我们会先看清模型内部的计算,再理解这些计算如何被组织成一个能够处理多个请求的推理系统。文章面向有基本编程经验、希望深入理解大模型推理的读者。从 token、向量和 Attention 开始,逐步解释 Prefill、Decode,以及为什么缓存各层 K/V 就能避免历史位置重算;随后结合 Qwen3 和 nano-vLLM,沿请求的执行路径理解调度、缓存管理、GPU 执行与多卡协作。源码用于把原理落到具体实现,实验用于验证理解,附录和延伸阅读则可按需查阅。
大模型推理,是用训练好的参数对输入进行计算、得到输出的过程。本文聚焦于普通因果 Transformer 的自回归生成:根据已有文本预测下一个 token,选出并追加它,再预测下一个,逐步形成回答。token 是模型处理文本的基本单位,可以是字、词或文本片段。
一次预测沿着同一条路径完成:文本经 Tokenizer 转成 token 编号,经 Embedding 查表变成向量,依次通过 Block 1 → Block 2 → … → Block L,最后经 Norm 和 LM Head 得到词表分数,由 Sampler 选出下一个 token。L 是模型架构确定的 Block 数量。
每个 Transformer Block 通常包含 Attention 和 MLP:Attention 让各位置读取可见的上下文,MLP 在每个位置内部变换特征;Norm 调整数值尺度,残差连接把更新加回原表示。上一个 Block 的输出,就是下一个 Block 的输入。
这套网络会反复用于预测,参数保持不变。首次处理整段输入,叫作 Prefill;随后复用历史结果、每轮只新算一个位置,叫作 Decode。 用三个输入 token,就能串起整个过程:
Prefill:处理已知前缀。 输入 t1、t2、t3,在每个 Block 中更新这三个位置,并保存各自的 K/V。每个位置经过 Attention、MLP 等计算后,才进入下一个 Block;最后取 t3 的最终表示,选出 t4。同一层的多个位置可以并行计算,各层仍依次执行。
Decode:每轮只新算一个位置。 下一轮输入刚选出的 t4,在每层读取 t1~t3 的缓存 K/V,并生成、保存 t4 的 K/V。t4 仍完整经过每个 Block 的 Attention、MLP 等计算,最终选出 t5;随后输入 t5,继续生成。
KV Cache 让历史位置不必重算。 因果约束使旧位置看不到后来追加的 token,因此旧位置的结果不变。新位置读取历史时,用 K 计算分配给各位置的权重,再用这些权重对 V 加权求和。保存各层历史 K/V,就能复用这条依赖,不必重跑旧位置的 Attention 和 MLP;第 3 章会展开原因。
阅读主线是:第 1–3 章理解向量、模型执行与缓存;第 4–6 章接入真实模型、GPU 和推理引擎;第 7–14 章沿请求追踪 nano-vLLM 源码。第 15–18 章用于验证、排障与动手实验,第 19–23 章是算子附录与延伸阅读,可按需查阅。
1. 从 token 到多头 Attention:先看懂模型在算什么
下面先拆开模型内部最关键的一步:一个 token 的向量,怎样通过 Attention 读取它能看到的位置?从只有一个数的 Head 开始,再扩展到向量和多头结构。
原理部分采用教学配置:隐藏维度 4096、普通因果 Transformer、Pre-Norm、标准多头注意力。先逐步改变 Head 的数量与宽度来理解机制,再固定为 32 个 Head、每头 128 维。它不是后文 Qwen3-0.6B 的真实配置;默认省略 batch=1。
1.1 先固定场景:已知 t1、t2、t3,准备预测 t4
本章固定一条序列:t1、t2、t3 是从左到右三个已知位置,各存一个 token ID;模型处理它们,准备预测下一个位置 t4。先放大 t3 在某一层中的 Attention,再在第 2 章走完生成过程。
公式下标表示位置,例如 q3 属于 t3;上标 h 表示第几个 Head。向量内部的分量序号则表示特征,三者分别回答“哪个位置、哪个头、哪个数”。
1.2 从文本到向量:Tokenizer 编号,Embedding 查表
Tokenizer 按配套的词表与规则把文本切成 token,再转换成整数 id。一个 token 可能是词、词的一部分、汉字、标点或字节片段。Embedding Table 是模型参数;它按 id 查出一行浮点数,把离散编号变成计算所需的向量。
文本 → Tokenizer → [id(t1), id(t2), id(t3)]
E:Embedding Table [100000,4096]
一个 id=1234 → E[1234] → 向量 [4096],保留位置轴时写为 [1,4096]
三个 id → 三行向量 X [3,4096]4096 就是这条 token 向量的特征分量数:一条向量里有 4096 个数。 例如 x = [0.2, −0.7, 1.3, …],共 4096 项;写成 [1,4096] 时,1 表示一行、一个 token 位置,4096 表示这一行的 4096 个特征分量。三个位置排在一起是 [3,4096]:输入变长只增加行数,不改变每行的特征宽度。4096 这个宽度在模型架构设计时确定,不是每次输入时临时决定的。
“特征分量”是表示向量中的一个数,不是一个预先命名的属性。模型通过训练学习如何用这些数表示信息;一个概念往往由多个分量共同表达,不能简单认定某一列就是“颜色”、另一列就是“地点”。刚查表得到的向量叫初始 embedding;经过各个 Block 更新后,每个位置在 Block 之间传递的向量仍是 4096 维,但数值已经改变,成为融合上下文的隐藏表示(hidden state)。相同 token id 的初始 embedding 相同,在不同位置或上下文中的隐藏表示则可以不同。
“向量宽度”与“数组的轴数”描述不同事情:[4096] 是一个轴的数组,[1,4096] 是两轴行矩阵,都可表示一条含 4096 个分量的向量。本文常保留位置轴,写成 [T,4096];再带上批次轴,可写为 [B,T,4096]。读形状时,要同时看每个轴代表什么。
对象 | 如何确定 | 普通推理时 |
|---|---|---|
Tokenizer、词表与特殊 token | 预训练前学习切分规则或复用现有配置;通常不参与大模型的梯度训练 | 按固定规则编码、解码;必须与模型的 id 约定匹配 |
隐藏宽度、Head 数、Head 宽度 | 架构设计时确定 | 不会随 prompt 长度改变 |
Embedding、投影和 MLP 权重 | 通常通过反向传播学习;微调时可选择冻结 | 保持固定,由不同位置与请求共享 |
Hidden state、Q/K/V、Attention 权重 | 用模型参数对当前输入计算 | 随位置、层和上下文变化 |
不同模型不一定各自训练一套 tokenizer,同系列与微调模型常复用它。向量不保存在 tokenizer 中;词表大小相同,也不代表两套 tokenizer 可以互换。
1.3 最小的 Head:每个 q、k、v 只有一个数
Head 是一条独立的“投影 → 匹配 → 加权汇总”计算分支。 它内部会产生 q、k、v 向量。为便于手算,先只设一个 Head,并让每条 q、k、v 只有一个数。设 u3 是 t3 在本层经过 Norm 后的表示,形状为 [1,4096]。
q3 = u3 × WQ [1,4096] × [4096,1] → [1,1]
k3 = u3 × WK [1,4096] × [4096,1] → [1,1]
v3 = u3 × WV [1,4096] × [4096,1] → [1,1]三个投影并列发生:q 用来查询,k 用来匹配,v 提供被读取的内容。t1、t2 用同一组本层参数各自产生自己的 k、v。WQ、WK、WV 是训练好的参数;下面的 Attention 权重则是本轮现算的结果。
直接给定一组便于手算的投影结果:q3=1,三个位置的 k 分别是 0、1、2,v 分别是 10、20、30。当前 q3 对每个 k 做乘法,得到三个分数 [0,1,2];除以 √1 不改变分数,softmax 再把它们变成总和为 1 的权重:
可读位置 | k | v | q3×k | Attention 权重 |
|---|---|---|---|---|
t1 | 0 | 10 | 0 | 0.0900 |
t2 | 1 | 20 | 1 | 0.2447 |
t3 | 2 | 30 | 2 | 0.6652 |
单头输出 a3=0.0900×10+0.2447×20+0.6652×30≈25.75。Q 和 K 决定各位置读多少,V 提供被加权的内容,a3 是汇总结果。 因此 v3=30 与 a3≈25.75 是不同的量:前者由输入投影产生,后者由 Attention 读取多个位置后产生。
另外两个查询也各有结果:a1 读取 v1,a2 读取 v1、v2。a1、a2、a3 按位置各占一行,组成 [3,1] 的单头输出矩阵,不是把 a1+a2+a3 相加。加权求和发生在每个查询内部。每头扩到 128 维后,三个位置的结果就是 [3,128]。
1.4 把 Head 扩到 128 维:权重仍然按位置分配
现在把同一个 Head 内 q、k、v 的宽度扩到 128。一个位置的 q、k、v 各是 [1,128],即一行 128 个特征分量;这就是“128 维 Head”的含义。4096 是 Block 之间传递的向量宽度,128 是本例单头内部的向量宽度;Head 数量仍是 1,可读位置仍是 3。
从 4096 个分量变成 128 个分量,需要投影矩阵 WQ、WK、WV,各为 [4096,128]。以 WQ 为例,每列用一组权重组合输入的 4096 个数,得到 q 的一个分量;128 列得到完整的 q。k、v 分别由 WK、WV 生成。这是学到的线性组合,不是从输入中截取 128 个数。
u3 [1,4096] × WQ [4096,128] → q3 [1,128]
u3 [1,4096] × WK [4096,128] → k3 [1,128]
u3 [1,4096] × WV [4096,128] → v3 [1,128]
k1、k2、k3 按行放在一起 → K [3,128]
v1、v2、v3 按行放在一起 → V [3,128]q3 与一个 k 点积:128 对分量分别相乘,再求和,得到一个分数。与三个位置的 k 匹配,就得到三个分数;softmax 将它们转换成三个位置权重。
scores = q3 × Kᵀ [1,128] × [128,3] → [1,3]
weights = softmax(scores / √128) [1,3]
a3 = weights × V [1,3] × [3,128] → [1,128]如果 t2 的权重是 0.2,它会乘上 v2 的整条 128 维向量;三个加权向量逐分量相加,得到 a3。K/V 的列表示特征,但 scores 的列表示可读位置,必须结合轴的含义读形状。
除以 √128 是为了控制分数尺度:在分量近似独立、零均值且单位方差的简化假设下,128 项点积之和的标准差约为 √128。不缩放时 softmax 容易过早集中到极少数位置。更具体的数字示例见第 19.1 节。
向量放大图|一格是一个数;输入表示经不同投影得到 q、k、v。128 是一条向量的宽度,4096×128 是投影参数的形状。
1.5 从一个 Head 到 32 个:各自读取,再沿列拼接
32 个 Head 接收同一份 u3,但各有不同的投影参数,各自算出位置权重和单头输出 a3。它们可以学习不同的信息组合,不能预先规定某个头永久只负责语法或事实。标准 MHA 中,每个头都有对应的 K/V。
Head 1 → a3⁽¹⁾ [1,128]
Head 2 → a3⁽²⁾ [1,128]
…
Head 32 → a3⁽³²⁾ [1,128]
沿列拼接:[a3⁽¹⁾ | a3⁽²⁾ | … | a3⁽³²⁾] → [1,4096]
再做输出投影:concat × WO [4096,4096] → o3 [1,4096]拼接是把同一位置的 32 条 128 维向量首尾接成一条 4096 维向量。 WO 再混合各头结果,得到多头 Attention 的输出 o3。这里沿特征方向拼接;不同位置的输出仍分别占一行。
Head 递进图|从 1 维单头,到 128 维单头,再到 32 个头的输出拼接。
实现时可以把 32 组投影合成一次较大的矩阵乘法,再把总 Q 从 [1,4096] 重排为 [1,32,128]。计算含义仍是 32 个头分别匹配、分别汇总;源码中的权重布局与融合实现见第 12 章。
1 维单头、128 维单头、32 个头是为理解机制逐步扩展的教学配置。从下一章起固定为 32 个头、每头 128 维;第 4 章再明确切换到 Qwen3-0.6B 的真实配置。
延伸阅读:3Blue1Brown — Attention in transformers, step-by-step。对照 Q/K 匹配、归一化与加权 V。
2. 从一层到一次生成:Prefill 与 Decode 怎样执行
第一章拆开了 Attention。现在把它放回完整模型:每一层用 Attention 读取其他位置,用 MLP 加工当前特征;所有层结束后,才选择下一个 token。 以下固定隐藏宽度 4096、32 个 Head、每头 128 维。
2.1 一层 Transformer:Attention 与 MLP 都在层内
本文的普通 decoder-only 模型由多个 Transformer Block 依次堆叠;每个 Block 通常有一个 Self-Attention 子层和一个 MLP 子层。多个 Head 是同一 Attention 内的并列分支,不是多个顺序执行的 Block。每个 Block 的 Attention 都要对 V 加权汇总。
第一层输入 x 来自 Embedding,后续层输入来自上一层输出。常见 Pre-Norm 结构先归一化再做 Attention,把更新加回 x;再归一化并经过 MLP,把更新再次加回。以下“层”在跨层讨论中指一个完整 Block。
模型结构图|Embedding → 多层 Transformer → 最终 Norm → LM Head;每层内部是 Norm、Attention、残差,再接 Norm、MLP、残差。
u = Norm1(x)
o = Attention(u) # 内含 Q/K/V、加权汇总、拼接与 WO
h = x + o # 第一次残差相加
z = Norm2(h)
m = activation(z × W_up) × W_down
x_next = h + m # 第二次残差相加,传给下一层Norm 调整这条向量的数值尺度,不把不同位置混在一起;残差是同形状向量逐分量相加。MLP 在每个位置上使用同一套本层参数,完成扩维、非线性变换、缩回。例如 W_up=[4096,11008]、W_down=[11008,4096],一行的路径就是 [1,4096]→[1,11008]→[1,4096]。这是简化的两层 MLP,后文 Qwen3 使用带门控的 SwiGLU。
三个位置进入 MLP 时,输入为 [3,4096],中间结果为 [3,11008],输出仍为 [3,4096]。同一套本层 MLP 参数处理三行,不需要三套 MLP。 Attention 在位置之间读取信息;MLP 按行加工已经包含上下文的表示。
传给下一层的是 x_next,不是本层的 v 或单头输出 a。每一层有自己的参数和 K/V;进入下一层后,要根据更新后的 hidden state 重新计算该层的 Q/K/V。
2.2 走完所有层,用最后一行预测下一个 token
输入 t1~t3 后,各个 Block 逐层更新这三个位置的向量。最后一个 Block 仍要完成多头拼接、WO、残差和 MLP;然后取 t3 的最终向量,经最终 Norm 和 LM Head 转成词表分数,由 Sampler 选出一个 token ID,作为 t4。训练时各位置学习预测紧接着的 token,因此最后位置的输出用于本轮续写。Attention 的中间结果 a3 仍是向量,不是 t4。
t3 的最终向量,经最终 Norm:[1,4096]
LM Head:× W_vocab [4096,Vocab] → logits [1,Vocab]
Sampler:根据 logits 与生成设置 → 选出 token ID(t4)选取 token 有不同策略:贪心选择分数最大的 ID,随机采样按分布抽取,并可通过 temperature、top-k、top-p 调整候选分布。这是通用生成策略;本文 nano-vLLM 实际提供的采样功能见第 12.4 节。Attention 中的 softmax 沿上下文位置分配读取权重,这里的分布则沿词表分配选择概率。
刚选出的 t4 是一个 ID。继续生成时,查 E[id(t4)] 得到它自己的初始向量,再输入模型;历史信息由各层 KV Cache 接入。遇到 EOS 或达到输出长度上限后停止,Tokenizer 再将输出 ID 序列还原成文本。
2.3 一次请求的三轮:Prefill → Decode → Decode
继续沿用三个已知位置 t1、t2、t3,观察一次 Prefill 和两轮 Decode。这里的 prompt 是完成聊天模板与编码后的全部输入,可能包含系统提示和历史消息;先假设没有可复用的前缀缓存。
Prefill 处理已知前缀,并选出第一个新 token;之后每轮 Decode 输入刚选出的 token,再选出下一个。下面按“本轮输入 → 前向后缓存 → 本轮输出”对齐三个时刻。这里的 Decode 是模型执行阶段,Tokenizer 的 decode 则是编号转文本。
三轮时序图|输入 t1~t3 选出 t4,再输入 t4 选出 t5,再输入 t5 选出 t6。
轮次 | 本轮新输入 | Embedding 后 | 前向后各层 KV | 本轮选出 |
|---|---|---|---|---|
Prefill | t1~t3 | [3,4096] | 覆盖 t1~t3 | t4 |
Decode 1 | t4 | [1,4096] | 覆盖 t1~t4 | t5 |
Decode 2 | t5 | [1,4096] | 覆盖 t1~t5 | t6 |
选出 t4 时还没有 t4 的 KV;下一轮实际处理 t4 时才生成它。 若选出的 token 已触发停止条件,就不必再为它额外运行一轮。
2.4 Prefill:三行一起计算,放大最后一行 t3
三个 ID 查表得到 X=[3,4096]。进入某一层后,先做 Norm;对一个 Head,投影得到 Q、K、V,各为 [3,128],每行对应一个位置。K/V 保存到本层缓存;三个查询分别读取因果约束允许的位置:
scores = QKᵀ [3,128] × [128,3] → [3,3]
P = softmax(scores / √128 + causal_mask) [3,3]
A = P × V [3,3] × [3,128] → [3,128]
最后一行:a3 = P₃₁v1 + P₃₂v2 + P₃₃v3 [1,128]分数矩阵的行表示查询位置,列表示被读取的位置。t1 只读 t1,t2 读 t1~t2,t3 读 t1~t3;未来位置的分数在 softmax 前屏蔽为 −∞。同一层的输入已经全部就绪,因此三个查询可以成批计算,t3 不需要等待 t2 在本层的 Attention 输出。
对 t3,把 32 个单头结果 a3 拼成 [1,4096],经 WO 得到 o3,再完成残差、Norm、MLP 和残差,得到传给下一层的 x3_next。t1、t2 也各自完成这条路径;整层输出仍是 [3,4096]。这些旧位置的更新如何参与下一层计算,第 3.1 节会详细解释。
Prefill 向量路径图|主线跟踪 t3 在一层中的变化,旁路展示 t1、t2 提供的 K/V;同样的层结构依次执行,最后进入词表输出。
Prefill 的并行发生在同一层的多个已知位置之间;层与层依次执行。实际 GPU 会进一步切块、融合算子,长 prompt 也可以分块 Prefill,这些执行方式将在引擎章节展开。
2.5 Decode:只新增 t4,历史以 K/V 的形式参与
Prefill 选出 t4 后,下一轮查它自己的 embedding,得到 [1,4096]。进入某一层,先 Norm,再生成本层的 q4、k4、v4;对一个 Head,它们各为 [1,128]。把新 K/V 纳入该层缓存,本轮就能读取四个位置:
K_old、V_old:各 [3,128]
纳入 k4、v4 后:K、V 各 [4,128]
scores = q4 × Kᵀ [1,128] × [128,4] → [1,4]
weights = softmax(scores / √128) [1,4]
a4 = weights × V [1,4] × [4,128] → [1,128]一个查询分别匹配四个 K,得到四个位置权重;每个权重乘对应的 V,再逐分量求和,得到一条 128 维输出。接着,t4 完成多头拼接、WO、残差、Norm 和 MLP,将更新后的一行传给下一 Block。历史 t1~t3 通过缓存 K/V 参与本层 Attention,不再重新经过这些网络计算。
Decode 向量路径图|新 t4 走完整层,历史 t1~t3 从本层 KV Cache 接入;两条路径在 Attention 汇合。
每个 Block 都按同样方式处理 t4,并读取、更新该 Block 自己的缓存。走完全部层后选出 t5,此时缓存覆盖到 t4。Prefill 与 Decode 使用同一套模型;区别是本轮新算哪些位置,以及哪些历史 K/V 已经可用。
延伸阅读:Jay Alammar — The Illustrated GPT-2。沿输入、向量、模型层与词表输出追踪同一个位置。
3. KV Cache:用空间换掉重复计算
理解 KV Cache 要同时回答两个问题:追加 token 后,哪些结果没有变化?新位置后续还会读取哪些旧结果? 因果性回答前者,Attention 的 Q/K/V 依赖回答后者。两者合起来,才说明为什么只保留各层 K/V,就能避免重复计算整个前缀。
3.1 旧位置的 a 有什么用,追加 token 后为什么不用重算
延续 t1、t2、t3 预测 t4 的场景。以两个 Block 为例:Block 2 的 q3 需要读取本层的 k1/v1、k2/v2,而这些 K/V 是由 Block 1 更新后的 t1、t2 表示投影得到的。因此,Prefill 中 Block 1 不仅要计算 a3,也要完成 t1、t2 的 Attention、残差和 MLP。
Block 1:
t1 的各 Head 输出 a1 → 后处理 → 更新后的 x1 → Block 2 的 k1、v1
t2 的各 Head 输出 a2 → 后处理 → 更新后的 x2 → Block 2 的 k2、v2
t3 的各 Head 输出 a3 → 后处理 → 更新后的 x3 → Block 2 的 q3、k3、v3
Block 2:
q3 查询本层 k1、k2、k3,再加权本层 v1、v2、v3 → 本层 a3
本层 a3 → 后处理 → t3 的最终表示 → 最终 Norm、LM Head → 预测 t4这说明中间层的 a1、a2 有用:它们参与形成下一层的历史表示。同一层的 a3 不依赖同一层的 a1、a2,但下一层的 t3 会间接用到它们加工后的信息。 不同层的 a3 通常不同,因为输入、参数及所读 K/V 都不同。
现在选出 t4,序列变成 t1~t4。若完整重算,Block 1 的 t1 仍只读 t1,t2 仍只读 t1~t2,t3 仍只读 t1~t3;追加的 t4 不在它们的可读范围内。因此旧 a1、a2、a3 和完整层输出都不变。Block 2 收到的旧行不变,旧行的可读前缀也不变,结果仍不变。逐层推下去,每一层旧位置的 hidden state、K/V 和 a,都与该层 Prefill 时在数学上相同。
这个结论依赖于权重、已有前缀、位置处理和因果 mask 保持不变。修改前文或模型后,需要重新判断缓存能否复用;不同内核和求和顺序则可能带来浮点误差。接下来要回答的是:旧结果既然不变,究竟需要留下哪一部分?
下图把两个相邻 Block 展开:上半展示旧位置怎样建立更深层的缓存,下半展示新位置怎样直接读取这些缓存。图中省略 Norm、残差等细节,但它们在实际计算中仍然存在。
进阶说明:最后一个 Block 的裁剪。 若只需要本轮最后位置的 logits,在数学上可以只计算最后查询的 Attention 输出及其后续 MLP;但所有已知位置在该 Block 的 K/V 仍需生成,供当前查询与后续生成使用。常规批量实现也可能保留全部行的计算。这个优化不适用于需要为下一 Block 建立历史表示的中间层。
3.2 为什么只缓存各层 K/V,就能让 Decode 只算新位置
输入 t4、预测 t5 时,第 2.5 节已经给出了新位置的完整执行路径。现在只看其中的历史依赖:在每个 Block 内,新查询 q4 要匹配哪些旧数据,又要汇总哪些旧数据?
本层 t4 的输入 → Norm → q4、k4、v4
缓存的本层 K1~3、V1~3 + 新 k4、v4 → K1~4、V1~4
a4 = softmax(q4 × K1~4ᵀ / √d_head) × V1~4
各 Head 的 a4 → 拼接、WO、残差、Norm、MLP、残差
→ t4 在下一 Block 的输入新 q4 匹配历史只用旧 K,汇总历史只用旧 V。 它不直接读取旧 Q、旧 Attention 输出 a 或旧 MLP 输出。后两者在之前的计算中已经形成更深层的输入,并据此生成了更深层的 K/V。因而只要各层 K/V 已保存,新位置就能取得所需的历史信息,无须重建旧位置的整条计算链。
这里要分清层次:第 l 个 Block 的 K/V 来自它的输入,经 Norm 和投影得到;该输入已经经过前面 l−1 个 Block 的加工。因此,高层 K/V 承接了前面各层 Attention、残差和 MLP 处理过的信息,但生成在本层的加权结果 a 之前。第一层的 K/V 则从初始 embedding 出发,结合模型要求的位置处理得到。
K/V 足够,是因为新位置的跨位置计算只依赖它们;不是因为 K/V 能还原所有历史中间结果。新位置的残差和 MLP 使用自己的当前表示,所以仍须执行。“只保存 K/V”指可复用的历史计算状态;引擎还会保留 token ID、位置和请求进度等管理信息。
执行阶段 | 中间 Block 需要更新哪些位置 | 历史 K/V 从哪里来 | 本轮预测 |
|---|---|---|---|
首次 Prefill,无前缀命中 | t1、t2、t3 | 本轮逐层计算并保存 | t4 |
下一轮,不使用 KV Cache | 重新更新 t1、t2、t3、t4 | 从完整前缀逐层重算 | t5 |
下一轮,使用 KV Cache | 每个 Block 只更新 t4 | 本层缓存的 t1~t3,加本轮新算的 t4 | t5 |
在教学 MHA 配置中,单层的 K 和 V 各可组织为 [batch,32,S,128]:32 是头数,S 是已缓存位置数,128 是每行特征宽度。只看一个请求、一个 Head,就是两张 [S,128] 矩阵。必须每层分别保存,不能让所有 Block 共用第一层或最后一层的一份 KV。
“追加一行”是逻辑描述。高效引擎通常预先分配存储空间,把新 K/V 写进对应槽位,不会每轮复制整份矩阵;第 9 章会把这些逻辑位置映射到物理块。
3.3 Context window:新算一行,可以读取多少历史
新增输入长度与可读上下文长度是两个量:处理 t4 时只新算一行,在全注意力下仍可读取 t1~t4。Context window 是模型或服务支持的上下文容量,请求预算包含系统提示、历史、模板标记以及不断追加的生成内容。达到上限后,服务可能拒绝请求、停止生成或按既定策略截断,不能默认它会自动滑动。
只有明确采用“每次最多直接读取 1024 个位置,包含自己”的滑动窗口时,处理 t1025 才会直接读 t2~t1025:读的是 1024 个位置,不是只读 t2。缓存中的高层表示还可能间接携带更早信息。本章其余示例都按全注意力,不主动滑动。
3.4 与无缓存对照:计算次数少了,读取范围没少
不用缓存也能生成:选出 t4 后,重新输入 t1~t4 来预测 t5;再重新输入 t1~t5 来预测 t6。同一条 token 路径下,两种方式的最后位置 logits 应在数值容差内等价,区别在于是否重复执行旧位置的投影、Attention 和 MLP。
本轮选出 | 有缓存:本轮模型输入行数 | 最后查询可读范围 | |
|---|---|---|---|
t4 | 3 | 3 | t1~t3 |
t5 | 4 | 1 | t1~t4 |
t6 | 5 | 1 | t1~t5 |
缓存对照图|同一条生成路径、相同的可读前缀;有缓存时只对新位置运行整条模型。
三轮中,无缓存共处理 3+4+5=12 个输入位置,有缓存共处理 3+1+1=5 个。一般地,prompt 长 S、选择 G 个新 token,无缓存处理 S+(S+1)+…+(S+G−1),有缓存处理 S+G−1。第一个输出由 Prefill 产生,之后需要 G−1 轮 Decode。
这些是位置执行次数,不是加速倍数。缓存路径仍读取模型权重与不断增长的历史 KV;时间取决于并行度、带宽和算子实现。单请求每轮一个新位置是普通自回归 Decode 的设定,合批和推测解码可以在一轮处理多个位置。
3.5 怎样验证缓存等价
验证时固定同一条 token 序列,对每个前缀分别运行全量重算和增量缓存,比较最后位置的 logits。不要让两条路径各自随机采样后继续比较,因为输入一旦不同,输出自然可能不同。记录模型、精度、位置处理与数值容差;数学等价不要求不同 kernel 的浮点结果逐位完全一致。
验证方法见第 15.2 节。动手时,第 18.2 节从单头 Attention 入手,第 18.3–18.4 节核对生成时序与缓存依赖,第 18.6 节用真实模型比较 logits。原有无缓存小模型及附件已集中到第 18.4 节,阅读原理时可先跳过。
从计算走向性能,要同时考虑三类数据:请求共享的模型权重、随上下文增长的 KV、当前轮次的临时激活。Prefill 的多行可以共同使用权重;单请求 Decode 虽然只新增一行,仍要读权重和历史 KV。第 4 章据此估算显存与瓶颈。
原理核对:Hugging Face — How caching works。
4. 把一个新 token 放到 GPU:从矩阵读写推导显存与性能
模型计算落到 GPU 上,要区分两件事:显存保存哪些数据,计算时必须搬运哪些数据。 本章先换成真实配置,再依次估算 KV 容量、计算与带宽成本,最后定义延迟和吞吐指标。
从本章开始切换到真实模型 Qwen3-0.6B,仍沿用三个输入位置。为便于对应文本,采用裸文本“中国的首都是”,在原实验的配套 Tokenizer 中切为“中国的 / 首 / 都是”;假设随后选择“北京”、句号、EOS。这个输出序列用于讲解,不是真实生成记录。以下按单卡、BF16 KV、全注意力估算容量,模型维度不再使用前文的 4096 维教学配置。
4.1 把教学形状换成真实配置,再跟踪一个新位置
前面的 4096=32×128 只是教学配置。Qwen3-0.6B 的 hidden size 是 1024,Q 有 16 个 heads,每头 128 维,总宽度为 2048;K/V 各有 8 个 heads,总宽度为 1024。每两个 Q heads 共用一个 KV head,这就是本例的 GQA。输出投影把 2048 维的 Attention 结果映射回 1024 维残差流。
读图时先沿顶部走完模型,再向下放大一个 Block 和其中的 Attention。形状以 Decode 新输入的一个位置为例,省略 batch 轴;N 表示本层可读取的位置数,包含历史与当前位置。注意:1024 是隐藏表示的宽度,16×128=2048 是各 Q Head 输出拼接后的宽度,两者不必相等。 配置核对自 Qwen3-0.6B 官方 config.json。
配置 | 前文主例 | 本章 Qwen3-0.6B |
|---|---|---|
层数 / hidden size | L 层 / 4096 | 28 层 / 1024 |
Q heads / KV heads | 32 / 32(MHA) | 16 / 8(GQA) |
每个 head 的维度 | 128 | 128 |
MLP | 简化两层网络,扩维示例 11008 | SwiGLU,gate/up 各 3072 维 |
位置与归一化 | 用常见 Pre-Norm 结构说明 | RMSNorm、Q/K head Norm、RoPE |
词表 | 记为 Vocab | 151936,Embedding 与 LM Head 共享权重 |
Prefill 选出“北京”后,下一轮输入它的一行 [1,1024]。按 [位置, Head, 特征分量] 记形状:每层新 Q 为 [1,16,128],新 K/V 各为 [1,8,128];纳入三个历史位置后,可读 K/V 各为 [4,8,128]。每个 Q Head 读取对应 KV Head 的四个位置,得到 [1,128];16 个结果拼成 [1,2048],再经 WO 回到 [1,1024]。
图中几个新算子各有明确职责:RMSNorm 调整数值尺度,Q/K Head Norm 在每头内部归一化,RoPE 给 Q/K 引入位置信息,SwiGLU 是带门控的 MLP。它的 gate/up 分支各为 [1,3072],经 SiLU 与逐元素相乘,再投影回 [1,1024]。直观解释见第 19 章;这里先关注各张量如何占用和读写显存。28 层后得到 [1,151936] 的 logits,假设选出句号;此时缓存覆盖到“北京”。
4.2 先看显存:权重、历史 KV 和当前向量是三类数据
显存负责保存数据,GPU 计算单元负责执行乘加等操作。服务启动时把模型权重从文件加载到 GPU 显存,通常长期保留。请求到来后,还要为历史 KV 和本轮中间结果提供空间。模型空闲时权重依然占着显存,所以显存占用高并不等于正在忙于计算。
数据 | 具体保存什么 | 怎样变化 |
|---|---|---|
模型权重 | Embedding、各层投影与 MLP 权重、Norm 参数、LM Head 等。 | 请求之间共享;普通推理不更新。各层通常有各自参数。 |
历史 KV Cache | 每层每个已处理位置的 K、V。 | 随请求上下文增长;本例缓存位置数 3→4→5。 |
当前激活与工作区 | 本轮 hidden state、Q、Attention 输出、MLP 分支等。 | 随本轮位置数和算子变化;不用保留的部分可释放或复用。 |
输入、元数据与运行时 | token ids、positions、缓存槽位、CUDA 上下文及框架缓冲等。 | 输入和元数据按轮更新,运行时占用依框架与配置而定。 |
数组形状决定元素数量,dtype 决定每个元素占多少字节。BF16 的每个数占 2 bytes:北京的一行 hidden state [1,1024] 仅 2 KiB;三行 Prefill 为 6 KiB。这只是当前一种激活,不代表整个推理只需要这么多显存。MLP 扩维、其他层权重、缓存和工作区还要另算;部分累加也可能使用更高精度。
4.3 把北京这一轮放到 GPU:每层读什么、写什么
CPU 先准备新 token id、position 和缓存元数据,把它们交给 GPU。GPU 查 Embedding,然后逐层执行。每层使用本层的权重生成当前 Q/K/V;Attention 读取本层历史 K/V 并纳入当前位置;当前 K/V 被写入缓存槽位,供后续轮次使用。当前 Attention 输出继续经过 WO、残差、Norm 和 MLP,得到下一层输入。
GPU 数据流图|左侧是 CPU 与 GPU 的输入输出交接;GPU 内部区分显存中保存的数据与执行当前层的计算单元。每层各自读取权重和 KV,向量在 GPU 内进入下一层。
这里有两种不同的搬运。 一种是 CPU 内存与 GPU 显存之间传输 ids、元数据或结果;另一种是 GPU 内部将显存中的权重、KV 和激活分块送到计算单元附近。数据可能经过 L2 缓存、寄存器和共享内存等层次。计算单元处理当前数据块后,结果留在片上或写回显存。不能把“权重已加载到 GPU”理解为“每轮计算不用再读取权重”。
例如本层 Q 投影可以写成 [1,1024]×[1024,2048]→[1,2048]。GPU 需要读取输入和权重块,执行乘加;之后的 Norm、Attention 和 MLP 继续处理这些数字。层之间有依赖,但向量不需要每层送回 CPU。某些算子会融合,部分中间值可直接留在寄存器中,所以逻辑上的每一步不一定都对应一次独立显存写入。
KV Cache 和 GPU 的 L2 硬件缓存也是两回事:KV Cache 是推理软件保存的模型中间张量,通常主要位于显存;L2 是硬件自动管理的数据缓存。Attention 读 KV 时也可能受益于硬件缓存,但两者的含义、容量与生命周期不同。
nano-vLLM 在采样后把新 token id 取回 CPU,更新请求状态并组织下一轮;模型的层间计算留在 GPU。其他引擎可以把更多调度或采样留在 GPU,具体边界取决于实现。
4.4 沿张量形状,一步一步算出 KV 显存
先算最小单位:一个位置、一个层、一个 KV head。K 有 128 个 BF16 数,V 也有 128 个,因此占 2×128×2=512 bytes。再乘 8 个 KV heads,是每位置每层 4096 bytes,即 4 KiB;再乘 28 层,是每位置跨全模型 114688 bytes,即 112 KiB。
这就是公式 2(K/V)×28(层)×8(KV heads)×128(head_dim)×2(bytes)。这里必须数 8 个 KV heads,不能拿 16 个 Q heads 代入:多个 Q heads 共享 K/V,并不会让缓存多存一份。这个单位成本不包含权重、临时激活、池预留或块尾空槽。
已缓存的范围 | 单卡 BF16 逻辑 KV 容量,无前缀共享 |
|---|---|
本例 Prefill 后:3 个位置 | 3×112 KiB=336 KiB |
处理北京后:4 个位置 | 4×112 KiB=448 KiB |
处理句号后:5 个位置 | 5×112 KiB=560 KiB;EOS 刚选出,不再额外计算。 |
1 条请求:4096 个位置 | 4096×112 KiB=448 MiB |
8 / 16 条请求,各 4096 个位置 | 3.5 GiB / 7 GiB |
一个 256-token 缓存块,跨 28 层 | 256×112 KiB=28 MiB |
一般公式为 KV bytes=2×层数×KV head 数×head_dim×每元素字节数×已缓存位置总数。多条不同长度请求要把位置数相加。共享前缀、滑动窗口、多卡切分或其他注意力结构需要按实际保留的数据调整;公式中的容量和本轮“实际从显存读了多少 bytes”也是不同口径。
逻辑数据量还不等于物理分配量。后文 nano 默认每块容纳 256 个位置,3-token 请求也至少占一个块,对应本配置跨全层 28 MiB 的物理槽位容量,只有一小部分已经写入有效 KV。引擎通常预先申请更大的 KV 池,请求是在池内领取块,而非每增长一个 token 就重新申请一段 GPU 内存。
权重另外按“参数个数×每参数字节数”粗估:标称 0.6B、每参数 2 bytes,约 1.2 GB 或 1.12 GiB;这不是精确模型文件或进程占用。共享权重只算一次,量化还会有 scales 等附加数据。KiB/MiB/GiB 按 1024 进位,GB 按 1000 进位;权重降精度也不会自动把 KV 同比例缩小。
4.5 新增只有一行,为什么 Decode 仍然可能慢
Prefill 与 Decode 使用同一套权重,变化的是本轮输入行数。对某个投影,Prefill 是 X[T,1024]×W[1024,2048],单请求 Decode 是 X[1,1024]×同一个 W。T 较大时,加载到计算单元附近的权重块可以服务更多输入行,形成较大的矩阵运算;只有一行时,仍须使用这份大权重,却只有较少计算可分摊读取成本。
以一般矩阵乘法 [T,d]×[d,m] 为例,按乘与加各算一次,计算量约为 2Tdm。权重本身的元素数量仍是 dm;在其他条件相同时,更多行能够提高每次读取所完成的计算量。真实读取还取决于分块、缓存和融合,不能直接把模型文件大小认定为每轮必然传输量。
场景 | 本轮在算什么 | 可能形成的限制 |
|---|---|---|
较长 Prefill | 很多已知位置共同使用权重,各自计算因果 Attention。 | 较大矩阵计算量;长序列 Attention 与临时空间也可能成为重点。 |
单请求 Decode | 新增一行,仍执行全部层并读取历史 KV。 | 权重/KV 读取、GPU 利用率及许多小算子的提交开销。 |
很长上下文 Decode | 新查询数量仍少,可读 K/V 的位置更多。 | Attention 的历史读取与计算随长度增加。 |
多请求合批 Decode | 多条请求各新增一行,共同使用模型权重。 | 权重利用率可能提高,但各请求的 KV、调度等待和延迟仍需考虑。 |
这解释了为什么“小 batch Decode 常受显存带宽影响,较长 Prefill 常更容易体现计算能力限制”,但它不是所有模型、长度、设备都成立的固定标签。粗估一段 GPU 工作可比较“运算量/有效算力”和“实际搬运量/有效带宽”;两者可能重叠,不宜无条件相加,还要考虑提交、通信、同步等开销。
不同优化针对不同成本:合批提高权重复用,FlashAttention 减少 Attention 中间矩阵搬运,CUDA Graph 减少重复提交。第 6.9 节给出完整的机制分层,第 7–14 章再对照源码展开。
4.6 为什么权重放得下,运行时仍可能 OOM
权重成功加载,只说明这部分数据放得下。运行时还要容纳 KV、当前激活峰值、算子工作区和运行时开销。增加并发或上下文会增大 KV 需求,较大的 Prefill 可能提高激活与工作区峰值。模型声明的上下文容量,不保证某台机器在任意并发下都有足够显存。
可以先按“可用预算−权重−运行时和峰值工作区−必要余量”估算 KV 预算,再按每块成本估算可容纳多少块,最后用实测修正。KV 池的预分配和请求在池内的使用不要重复相加;GPU 上有其他进程时,也不能把标称总显存当成自己的剩余预算。
排查 OOM 时,先确认是哪部分预算被用尽,再检查请求长度、并发与缓存池。不同显存统计口径,以及“请求结束后显存为什么不下降”,集中见第 16.4 节。
4.7 最后再计时:首 token、每轮间隔和总吞吐
从用户视角展开一次请求:发送→网络与排队→输入处理→Prefill→选择首 token 并返回→多轮 Decode→最后输出返回。权重加载通常发生在服务启动时;若测到冷启动,应单独说明。不同位置的时间回答不同问题。
指标 | 怎么计算或观察 | 回答什么问题 |
|---|---|---|
TTFT | 请求发出到客户端收到首个输出 token 的时间。 | 多久开始看到回答;端到端口径包含排队、网络、Prefill 等。 |
ITL | 相邻输出 token 的到达时间差。 | 输出是否平滑;也可能受流式发送合并影响。 |
TPOT | (最后 token 时间−首 token 时间)/(输出 token 数−1),至少 2 个输出。 | 首 token 之后平均每个 token 用多久。 |
输出吞吐 | 观察区间内输出 token 总数/区间时长。 | 系统总体产出,不能直接替代单请求等待时间。 |
用教学时间戳验算:请求在 0 秒发出,4 个输出分别在 0.12、0.16、0.21、0.25 秒到达,则 TTFT=120 ms,ITL=[40,50,40] ms,TPOT=(250−120)/3≈43.33 ms。它们不是某张 GPU 的实测性能。增加 batch 可能提高总吞吐,也可能使某个请求因排队或每轮工作变多而更晚收到输出,两者要分别测量。
离线 generate 一次返回整批文本,仅凭总调用耗时无法得到客户端的 TTFT 或各次 ITL。具体的同步边界、GPU 计时和对照条件见第 15.4 节,操作步骤见第 18 章;本章先用这些指标区分“个人等待多久”和“系统完成多少”。
5. 从模型计算到推理服务:vLLM 负责什么
5.1 裸跑能生成,为什么还需要推理引擎与服务
前四章的生成链本身不依赖 vLLM:模型得到词表分数,采样器选出 token,外层循环将它送回模型,直到停止。加载模型后,Hugging Face Transformers 的 generate() 就能完成这条链。推理引擎要解决的是另一个问题:怎样让持续到来的多个请求高效共享有限的 GPU 资源?
Transformer 是模型架构,Transformers 是提供模型实现与生成接口的 Python 库。后者已经支持 KV Cache、批量生成和多种优化后端。学习原理、调试模型或处理低并发任务时,直接使用它往往就够了。
问题出现在多个请求持续到来时。假设 A、B 正在生成,A 先结束,而 C 此时刚提交一段长 prompt:A 的 KV 占用何时释放?C 能否立即进入,还是必须等 B 结束?C 的 Prefill 要一次算完,还是分段执行,以免 B 长时间收不到下一个 token?这些不是 Attention 公式的问题,而是有限的计算能力和显存该分给谁的问题。
推理引擎负责组织每一轮执行。 它选择本轮处理哪些请求、各计算多少新位置,为 KV 分配和回收空间,再调用计算后端执行模型。把多条请求的新位置合批,可以更充分地利用 GPU;请求进度、缓存和停止状态仍需分别维护。专用引擎把这些机制做成可复用的系统。
推理服务负责让客户端使用生成能力。 它通过 HTTP 等接口接收请求,将任务交给引擎,把结果一次性或流式返回,并处理取消与异常。鉴权、限流、路由和监控可由服务层及外围系统承担。引擎也可以只在本地程序中运行,不必开放网络接口。
三层职责可以概括为:模型负责怎么算,引擎负责本轮算谁及怎样执行,服务负责客户端怎样调用。 小规模在线应用可以先封装 Transformers;并发、延迟或成本要求提高后,再引入专用引擎。
vLLM 同时提供离线 Python 接口和在线服务入口。下面用它区分两种使用方式,再通过 nano-vLLM 打开引擎内部,观察调度、缓存管理与执行之间的关系。
5.2 离线生成:先看清接口的输入输出
使用引擎不等于必须部署服务。离线程序可以直接把一批 prompt 交给 vLLM,由引擎完成调度与生成,再取回结果;整个过程不需要 HTTP 请求。下面先用这种最短路径,看清引擎接口的输入和输出。
离线接口可以概括为:一批 prompts 与采样参数 → 引擎调度生成 → 每条请求的输出对象。普通字符串用于文本补全;聊天消息需通过 chat template 组织成模型输入。完整 Python 示例与独立环境安装步骤见第 18.15.1 节。
5.3 在线服务:请求如何进入引擎,输出如何返回
当调用方变成浏览器或远程程序,路径变为:客户端请求 → 服务层解析与排队 → 引擎生成 → 响应返回客户端。模型加载后常驻,客户端不必各自加载权重。聊天 messages 先转换成模型输入,生成结果再恢复为对应请求的文本响应。流式接口允许逐步接收输出,但一个网络数据块不一定对应一个 token。第 18.15.2 节提供仅本机可访问的体验步骤。
5.4 为什么下一步读 nano-vLLM
调用 vLLM 能看到输入输出,但要回答“A 结束后 C 怎样进来、KV 空间怎样复用”,还需要打开引擎内部。nano-vLLM 用较少的 Python 代码展示了完整执行链:调度器决定谁进入本轮,缓存管理模块决定 KV 放在哪里,ModelRunner 组织模型在 GPU 上执行。后文沿这条链,把本章的问题落实到函数和变量。
nano-vLLM 有自己的实现选择与边界,并非从某一版生产 vLLM 原样裁剪出的源码子集。本文先用它学习请求、缓存和 GPU 执行,再回到 vLLM 理解生产系统如何扩展这些机制。
第 6 章先沿单请求、单卡建立最小闭环,再逐步加入多请求与多卡。第 6.9 节统一说明连续批处理、分页 KV、前缀缓存、Pinned Memory、FlashAttention、CUDA Graph 与 Tensor Parallel 各自解决什么问题;随后在组件章节深入。
阅读时把通用机制与具体实现分开:vLLM 的功能组合取决于版本、模型、硬件和后端,nano-vLLM 则有自己的精简边界。下一章会先固定源码版本,再跟踪一条请求。
vLLM 官方 Quickstart(核对日期:2026-09-06)
6. nano-vLLM 全貌:组件架构与请求流程
前五章解释了模型怎样生成 token,以及为什么需要推理引擎。本章换一个视角:沿请求经过的路径,看 nano-vLLM 怎样组织这些计算。先只放一条请求、一张 GPU,把启动、Prefill、Decode 和结束回收讲完整;再加入多条请求,理解调度与缓存管理;最后增加 GPU,理解多个进程怎样合作执行同一批工作。第 7–14 章在这条主线上继续拆解实现细节;各项优化分别处在哪个层面,可先看本章末尾的第 6.9 节分层总览。
本文固定使用上游提交 bb823b3e06983d71485a8e1f23715ebd87d98ef8(包版本 0.2.0),模型例子采用 Qwen3-0.6B。它提供 Python 推理引擎,不是完整在线服务:公开的 generate() 先加入整批请求,循环计算,再返回结果;这个版本没有完整 HTTP 服务或逐 token 流式 API。下面讲的是这份实现,不把它的调度策略当作所有引擎的通则。
6.1 单请求、单卡:一次 generate() 为什么要反复执行 step()
先假设只有一条请求,prompt 编码后为 t1、t2、t3,模型完整放在一张 GPU 上。t1~t3 表示先后排列的三个 token,而不是三个用户请求。用户想要一段回答,但模型每轮只能给出下一个 token 的分布,因此引擎必须不断执行“准备输入 → 运行模型 → 接受输出 → 准备下一轮”。generate() 管整次生成,step() 管其中一轮。
先只认识两个角色:Engine 在 CPU 上组织生成,模型在 GPU 上完成计算。Engine 用 Tokenizer 把文本变成 token ids,并用一份请求记录保存已经有哪些 token、哪些位置已经算过,以及什么时候停止。这份记录在代码中叫 Sequence;它保存的是编号和进度,不是 token 向量。
第一次,Engine 安排 t1~t3 进入模型;GPU 完成 Prefill,保存各层 K/V,并选出 t4。Engine 接收 t4,如果还没结束,下一轮只安排 t4 进入同一个模型,复用缓存来预测 t5。如此反复,直到满足停止条件。模型负责算下一步,Engine 负责让这个过程持续、正确地推进。
图 6-1 先看最小闭环:一条请求、一个 Engine、同一张 GPU;新 token 回到 Engine,未结束就继续下一轮。
这张图暂不展开请求调度和内存分块。单请求仍会经过这些代码,只是没有请求间竞争;先把“准备本轮输入、完成计算、接收结果”这条主线记住。等加入多条请求,再解释为什么要把这些职责拆成独立组件;完整函数调用图放在第 7.1 节。
请求结束后,Engine 收集输出 ids,generate() 用 tokenizer.decode() 将其转成文本,返回包含 text 和 token_ids 的结果。这里的 tokenizer.decode 是“编号转文字”,与模型逐轮生成的 Decode 阶段不是同一件事。
6.2 单请求、单卡:运行第一轮之前,模型和 KV 池怎样准备好
模型只在启动时加载,不随每个 token 重新加载。调用方传入模型目录和引擎配置,Config 用 AutoConfig 读取目录中的模型配置。配置描述层数、hidden_size、Q/KV 头数、head_dim 等结构;权重文件提供具体参数;Tokenizer 文件描述文本怎样切分及 token 与编号怎样对应。这三者必须配套。
本版本的 Runner 明确创建 Qwen3ForCausalLM,再由 loader 将 safetensors 文件中的权重装入各层;Engine 从同一目录加载 Tokenizer。三者的关系是:代码决定支持哪种架构,配置决定张量尺寸,checkpoint 提供参数数值。AutoConfig 负责读取配置,不负责自动适配任意架构。
以 Qwen3-0.6B 为例,配置指定 hidden_size=1024、28 个 Block。输入 ids 是整数数组,查 Embedding table 后,每个 token 才得到一个 1024 维向量,再依次经过这 28 个 Block。这里用的是真实模型配置,不是前文 4096 维的教学例子;Q/K/V 的具体尺寸已在第 4 章展开,本章不重复列举。
准备设备、加载权重,以及分配 KV 存储空间,由下一节要展开的 ModelRunner 负责。它先预热,估算模型运行时还需要多少显存,再利用可用容量建立常驻 KV 池,给每层 Attention 分配相应切片。启动时只是备好存储空间,还没有用户请求的历史 K/V;这些数值必须在请求实际前向时生成并写入。
实际缓存容量确定后,Engine 才建立管理请求和缓存分配的组件,因为它们必须知道可用空间有多少。配置允许时,启动阶段还会准备 CUDA Graph,减少后续部分 Decode 计算的启动开销;这是第 13 章的执行优化,不改变本章的生成流程。
启动完成后有两类寿命不同的数据:模型权重和 KV 池随引擎常驻;请求占用的块、进度和本轮输入随生成过程变化。一个请求结束只是释放它的块引用,不会卸载模型或销毁整个 KV 池。下一条请求继续使用同一套计算设备和存储池。
6.3 单请求、单卡:ModelRunner 怎样把一轮工作交给 GPU
ModelRunner 接收本轮选中的 Sequence,返回采样候选 ID。 中间需要把请求记录整理成 GPU 张量:准备 token ID、位置和缓存地址,构造执行 Context,调用模型与 Sampler,再将候选交回 Engine。它连接 CPU 上的调度结果与 GPU 上的计算。
第一次 Prefill,Runner 取 t1~t3 的 ids,得到形状 [3] 的 input_ids,以及 positions=[0,1,2];查 Embedding 后是 [3,1024]。下一轮 Decode 输入刚选出的 t4,input_ids 变为 [1],positions=[3],Embedding 后是 [1,1024]。输入只剩一行,表示本轮只新算一个位置,不表示上下文只剩一个 token。
只有新 token 的 id 还不够。输入 t4 时,Attention 必须知道它位于第 4 个位置,哪些缓存属于这条请求,以及新 K/V 应写在哪里。Runner 把这些本轮执行信息放入 Context,供模型内部读取。Context 描述怎样读写,历史 K/V 数值则保存在常驻缓存池中;具体的长度、边界与地址字段留到第 11 章逐项展开。
图 6-2 Runner 执行链:工作单 → 输入与 Context → 模型 → 采样 → 候选 ids;下方展开每层 Attention 的缓存读写。
模型内部的计算沿用第 4.1 节:新位置经过全部 28 个 Block,每层读取自己的历史 K/V,写入新 K/V,并完成 Attention 与 MLP。此处关注 Runner 的交接边界:本轮输入和 Context 随调用准备,模型权重与 KV 池则跨轮保留。
全部 Block 完成后,LM Head 得到 logits,Sampler 选出候选 ID。Runner 返回候选并清理本轮 Context;Engine 再把结果交给 Scheduler 更新请求进度。Runner 完成计算,Scheduler 接受结果并决定请求是否继续。
6.4 单请求、单卡:把 Prefill、两轮 Decode 和结束回收连起来
现在走完同一条请求。输入为 t1、t2、t3,缓存充足、没有前缀复用,Prefill 一轮完成;最多输出 3 个 token,并假设中途没有 EOS。以下 t4~t6 是模型依次选出的 token 的示意名称;用 cached 表示已经完成前向、形成各层 K/V 的位置数。
轮次 | Runner 新输入 | 模型完成后,各层有效 K/V | 接受输出及后续处理 |
|---|---|---|---|
1:Prefill | t1~t3,位置 0~2 | t1~t3,共 3 个位置 | 接受 t4;总长 4、cached=3,继续生成 |
2:Decode | t4,位置 3 | t1~t4,共 4 个位置 | 接受 t5;总长 5、cached=4,继续生成 |
3:Decode | t5,位置 4 | t1~t5,共 5 个位置 | 接受 t6;达到 3 个输出,随后结束并释放引用 |
第三轮接受 t6 后,已经得到 3 个输出,生成结束。引擎释放这条请求的缓存引用,但保留输出 ids,用于返回 t4、t5、t6 对应的文本。t6 不会再作为下一轮输入,因此没有产生自己的 K/V。上表最后一行描述的是释放前的计算结果;释放后,请求不再持有有效缓存映射,cached 也会清零,具体状态变化见第 7 章。
所以,引擎至少要分清“序列已有多长”与“已经算到哪里”。正常 Decode 中,刚接受的新 token 要等下一轮才形成 K/V,因此活跃请求的总长通常比 cached 多 1。第 7 章再把这两种进度和“本轮准备算多少”对应到 Sequence 字段。
6.5 多请求、单卡:Scheduler 怎样把独立请求组织成一批
现在加入请求 B,仍然只用一张 GPU。如果先把 A 整段生成完再处理 B,B 就一直等待,而 A 的每轮 Decode 只有一个新位置,通常也不能充分利用 GPU。能否把 A、B 当前要算的位置放在同一轮?这就是批处理要解决的问题。两条请求共享模型权重和设备,但各有进度、采样参数与上下文;合并计算不等于合并历史,A 的 query 不能读取 B 的 K/V。
先看一个不受预算或容量限制的小例子:A 输入 3 个 token,要求输出 3 个;B 输入 2 个 token,要求输出 2 个;都不提前遇到 EOS。第一轮 Prefill 共新算 5 个位置,Runner 将 ids 打包成 [5],Embedding 后为 [5,1024],并用请求边界把两段 Attention 分开,各采样一个输出。第二轮 Decode 各取一个新 token,输入 [2]、Embedding 后 [2,1024];A 和 B 分别读取自己的 4 个、3 个上下文位置。第二轮结束,B 已得到 2 个输出并退出;第三轮只留下 A 的一个新位置。
合批并行处理的是不同请求当前这一轮的工作;每条请求的后续输出仍有先后依赖。更多输入行可以提高权重读取的复用,但也会增加本轮工作量,因此需要同时考虑缓存容量、吞吐与请求延迟。
有的请求刚开始,有的在 Decode,有的已经结束;一轮能容纳的计算量和缓存空间又有限,Engine 就不能始终沿用同一个请求列表。Scheduler 因此要在每个 step 边界重新决定“算谁、算多少”。它用 waiting 管理待处理或待恢复的请求,用 running 管理已接纳、继续生成的请求;再把本轮选择交回 Engine,由 Engine 调用 Runner 执行。执行后也由 Engine 把候选交给 Scheduler,推进进度并移出完成请求。
图 6-3 从单请求扩展到多请求:Scheduler 先尝试 Prefill,否则选择 Decode,并通过 BlockManager 确保缓存可用。
这份实现采用 Prefill 优先:预算和缓存容量允许时先安排 Prefill,否则安排 Decode;同一轮不混合两种阶段。长输入可拆成多轮 Prefill,中间轮推进缓存,完整前缀算完才接受首个输出。下面用一个固定例子观察这种选择,具体队列与预算字段在第 8 章展开。
例如后续章节统一使用 A=700、B=32、Prefill 预算=512:第一轮只算 A 的前 512 个位置,不接受中间候选;第二轮算 A 剩余 188 个位置和 B 的 32 个位置,共 220 个,各接受首输出;第三轮 Decode 才是两条请求各一个新位置。切分的是一条请求本轮新计算的范围,不是把它变成多条互不相关的请求。第 8 章会继续追踪这一例子的队列和计数。
6.6 多请求、单卡:BlockManager 怎样让请求共享 KV 显存
多个请求长度不同、不断增长,也会在不同时间结束。如果每条请求都预留最大上下文长度的连续缓存,大量空间会暂时用不上,还容易因为连续空闲空间不足而难以接纳新请求。nano-vLLM 将 KV 池分成固定大小的物理块,用每条 Sequence 的 block_table 记录“逻辑上的第几个块,放在池里的哪一个物理块”。请求的逻辑顺序连续,物理块不必连续。
默认一块可放 256 个 token 位置。若 A 的块表为 [17,5,23],就表示它的前三个逻辑块分别放在池里的物理块 17、5、23;B 可以使用另一个块 11。比如 A 的零基位置 300 位于第二个逻辑块,应读取物理块 5、块内偏移 44。同一映射用于各模型层,但每层在自己的 K/V 切片中保存不同数值。第 9 章再展开具体寻址公式。
BlockManager 管归属和地址,Attention 读写 K/V 数值。 BlockManager 在 CPU 上维护空闲块、引用计数和前缀索引,把块表交给请求;Runner 据此生成本轮的写入槽位与读取映射,Attention 才能在 GPU 缓存池中访问数据。Scheduler 因而可以在执行前确认:本轮请求有地方保存新的计算结果。
Prefill 先建立请求所需的块映射;后续 Decode 使用尾块的空槽,跨过块边界时才申请新块。分块 Prefill 的计算预算与整段已知序列的存储需求是两个量,具体分配规则见第 9 章。
后来请求若具有相同上下文前缀,可以复用符合条件的完整 KV 块,省掉重复 Prefill 与存储。前缀索引负责查找,引用计数负责记录谁还在使用;不同请求后续生成的内容仍相互隔离。匹配与失效规则见第 10 章。
某条请求结束时,管理器减少它占用块的引用;只有不再被任何请求引用的块,才可以重新分配给其他用途。释放的是请求的占用,不是销毁 GPU 上的整个缓存池。因此请求可以不断进入和退出,而模型与 KV 存储池持续服务后续工作。
缓存不足时,Scheduler 可暂停请求并释放它的块占用;请求 ID 序列仍保留,恢复时重新 Prefill,并尝试复用可命中的前缀。这份实现通过重新计算恢复,而非把整份 KV 暂存到 CPU。抢占与恢复的顺序见第 8 章。
6.7 多请求、多卡:一个批次怎样由多个 rank 一起执行
批处理提高了单卡利用率,但没有增加单卡的显存和算力。换成更大的模型、需要更多缓存或承担更重的计算时,就要考虑如何利用多张 GPU。先区分两条路线:每张卡放一份完整模型、分流不同请求;或者把同一个模型的参数和计算分到多张卡,一起处理同一批请求。本版本主要实现后一种,即单机张量并行 Tensor Parallelism(TP),不是多副本请求路由。
TP=1 时只有 rank 0;TP=P 时,Engine 启动 rank 1~P−1 的 worker 进程,并在主进程建立 rank 0 的 Runner,每个 rank 对应一张 GPU。各 Runner 按模型层的并行实现持有相应权重与 KV 分片。主进程仍只有一套 Scheduler/BlockManager 决定本轮批次,不是每张卡自己挑一批请求。增加卡数也不要求必须有多个请求:单条请求同样可以走 TP,只是这里沿多请求场景继续讲。
一轮工作先经过控制通路:rank 0 通过 SharedMemory 传递命令与请求执行信息,用 Event 通知 worker。worker 被唤醒后准备本轮输入与 Context,执行对应的模型分片。权重与历史 K/V 已在各卡常驻,不随每轮命令重新传输。
模型前向经过 GPU 通信通路:各 rank 计算同一层中自己负责的分片,并在需要汇合结果时调用 torch.distributed/NCCL。SharedMemory/Event 传达本轮算什么,NCCL 交换模型计算中的张量。 本例是同一层内部的张量切分;权重如何切、何处通信,见第 14 章。
例如 Qwen3-0.6B 使用 TP=2 时,每张卡负责部分 Attention 头,并保存相应的 KV 分片;两张卡都参与请求 A 和 B 的计算,不是一张卡只服务 A、另一张卡只服务 B。每个 head 的特征宽度并未因此改变。具体权重切分、缓存形状和层内通信放到第 14 章;这里先看清“同一份工作,多张卡协作完成”。
模型得到词表分数后,只由 rank 0 执行采样、把候选 ids 返回 Engine,再由同一个 Scheduler 回写状态、安排下一轮。多卡改变的是 Runner 内部如何完成一次模型计算,不改变 Engine 的 schedule → run → postprocess 循环。它可以分担参数、缓存与运算,但会增加通信和同步成本,所以多一张卡并不保证单请求时延按比例缩短。
6.8 回到全貌:组件归属、源码与外部依赖
走完三个场景,再看完整架构就不必先背类名了:Engine 持有 Tokenizer、Scheduler 和 rank 0 Runner;Scheduler 持有 waiting/running 队列及 BlockManager;Sequence 把某条请求的状态带过一轮轮调用;Runner 持有模型、Sampler 和常驻 KV 池。多卡时再增加 worker Runner,而不是增加另一套请求调度逻辑。
图 6-4 完整架构:上半部分汇总主进程、组件与 GPU 数据;下半部分展开可选的单机 TP。从单请求单卡到多请求多卡,主循环不变,逐步增加的是资源选择、地址管理与并行执行。
按数据寿命再核对一次:token ids、状态、块表和前缀索引属于 CPU 控制信息;input_ids、positions 和 Context 中的张量服务本轮执行;模型权重与各层 KV 池常驻 GPU。权重对请求共享,KV 按请求映射使用,符合条件的完整前缀可以共享。调度器知道算谁,块管理器知道放哪,Runner 知道怎样调用,模型才真正算出向量和词表分数。
沿请求要找什么 | 源码入口(nanovllm/ 下) | 后续深入章节 |
|---|---|---|
入口、生成循环与请求状态 | llm.py、engine/llm_engine.py、engine/sequence.py | 第 7 章:ids、cached、scheduled 与生命周期 |
本轮选谁、何时回写或抢占 | engine/scheduler.py | 第 8 章:队列、预算、阶段与回写 |
KV 块分配、寻址与共享 | engine/block_manager.py | 第 9–10 章:映射、引用、前缀 hash |
工作单变成张量与执行元数据 | engine/model_runner.py、utils/context.py | 第 11 章:输入形状、边界、读写地址 |
配置、参数与实际模型运算 | config.py、utils/loader.py、models/qwen3.py、layers/ | 第 12–14 章:模型内部、执行优化、张量并行 |
项目源码短,是因为它主要写“这些部件怎样接起来”,大量底层能力来自外部库。Transformers 在这里提供配置读取与 Tokenizer,并不代替 nano 自己的 Qwen3 前向;safetensors 读取 checkpoint,loader 再把参数装入模型;PyTorch 提供张量、模型层、GPU 执行与分布式通信接口;FlashAttention 提供实际 Attention 内核;Triton 实现定制的 KV 写入内核。CPU 上的 multiprocessing、SharedMemory/Event 管进程与通知,xxhash/NumPy 帮助构造前缀块索引。读懂这些依赖的边界,比把每个依赖源码一并展开更有助于理解 nano-vLLM。
组件划分及完整架构、Engine、Scheduler、Runner 图参考了 CalvinXKY 的 Structure of Nano-vLLM,并按本文固定的上游版本重画;开头的单请求简图用于建立阅读顺序。具体调用关系和实现行为以上游源码为准。
6.9 优化机制分层:调度、缓存、传输与执行各管什么
组件回答“谁负责”,优化机制回答“减少哪一种开销”。把已经认识的组件放回一轮推理,可以分清三个层面:调度与 KV 管理决定本轮算谁、哪些位置需要新算、历史存在哪里;传输与 GPU 执行负责把工作高效完成;多卡并行决定这份计算怎样由多张 GPU 分担。这是理解职责的分层,不是七个必须依次调用的算法。
图 6-5:七种机制在同一轮推理中的位置——沿主线看数据与工作如何前进,沿分层看各自在省什么。
请求调度与 KV 管理。Continuous Batching 让 Scheduler 在每轮边界调整请求集合,及时利用完成请求留下的容量;Prefix Cache 识别已经算过的公共前缀,减少重复 Prefill;分页管理让 BlockManager 用块表组织可以分散存放的 K/V。PagedAttention 还跨到 GPU 内核:不仅要在 CPU 分配块,Attention 也必须按块表读取,所以它不只是一个内存分配器。连续批处理见第 8 章,分页存储与访问见第 9 章,前缀复用见第 10 章。
主机传输与单卡执行。Runner 用 Pinned Memory 准备锁页的 CPU 输入缓冲区,支持把新 id、位置和地址元数据高效传到 GPU;FlashAttention 在 GPU 内分块融合 Attention,减少中间分数和权重矩阵的显存往返;CUDA Graph 在适用的 Decode 路径复用工作提交安排,减少 CPU 反复提交小操作的开销。三者分别针对主机到设备传输、GPU 内部读写和执行提交,不是同一种“缓存”。Pinned Memory 见第 11.6 节,FlashAttention 与 CUDA Graph 见第 13 章。
多卡计算与通信。Tensor Parallel 把同一层的权重和计算分给多个 rank,各卡执行本地分片,并在必要处通过 NCCL 合并结果。它是贯穿模型执行的分工方式,不是模型算完后再追加的一步,也不是给每个请求单独分配一张卡;每卡仍会用到自己的输入传输、Attention 内核与 KV 分片。切分和通信成本见第 14 章。后续阅读每种机制时,都可以回到这张图,先判断它改变的是工作选择、数据位置,还是执行方式。
读到这里,应当能独立复述一条请求从文本到输出、从入队到释放缓存的全过程,并解释加入更多请求或更多 GPU 后哪里发生变化。第 7–11 章不再承担第一次介绍组件的任务,而是沿 A=700、B=32、Prefill 预算=512 的例子追踪具体状态、地址与形状;验证方法集中在第 15 章,安装与实验步骤集中在第 18 章。不运行实验,也能先读通这一章。
7. 请求生命周期:LLMEngine 与 Sequence
第 6 章已经把整条路径连起来。本章只拆开入口与请求状态:Engine 负责不断推进工作,Sequence 负责记住每条请求走到了哪里。它们本身不做 Attention,却决定哪些 token 是输入、哪些是已接受输出、哪些位置已经有了可复用的 K/V。
7.1 Engine:一次调用,很多轮推进
LLM 是 LLMEngine 的薄封装。收到文本时,Engine 先用 Tokenizer 得到 token ids,再结合 SamplingParams 创建 Sequence;若输入本来就是 ids,则直接建立请求。每条 Sequence 被赋予标识并加入 Scheduler 的 waiting 队列。此时只登记了请求,还没有为它执行模型前向。
Engine 的循环以 step 为单位,而不是以“完整生成一条请求”为单位。每个 step 先让 Scheduler 选择本轮工作,再让 ModelRunner 执行,最后让 Scheduler 回写状态。Engine 收集本轮结束的请求;只要还有 waiting 或 running 请求,就继续下一轮。因此,A 的长输入和 B 的短输入可以在不同轮共享模型,而不必等 A 全部生成完才开始 B。
图 7-1 完整 Engine 循环:把第 6.1 节的单请求简图展开为具体调用。Engine 依次调用 schedule、run、postprocess;单请求与多请求经过同一条路径。
在本文固定版本中,公开 generate 先加入给定的一批 prompts,完成整个循环后才返回结果,按 seq_id 顺序组织为包含 text 和 token_ids 的字典列表。内部可以逐轮更新,并不等于公开接口已经提供逐 token 流式返回。
7.2 Sequence:保存的是请求记录,不是 token 向量
Sequence 是 CPU 上的状态对象。token_ids 是整数编号,block_table 是整数地址映射;两者都不是 embedding,也不包含历史隐藏向量或实际 K/V 数值。真正的 KV 在 GPU 池中,Sequence 只记录自己引用了哪些块、有效内容有多长。
字段 | 它回答的问题 | A 刚入队时 |
|---|---|---|
seq_id / token_ids | 这是哪条请求?目前有哪些 token? | A 的标识;700 个 prompt ids |
num_prompt_tokens | 原始输入有多长?后续保持不变 | 700 |
num_tokens | 输入加已接受输出,总共有多长? | 700 |
num_cached_tokens | 前多少个位置已有可用的各层 KV? | 0 |
num_scheduled_tokens | 本轮安排新计算多少个位置? | 尚未调度,为 0 |
last_token | 当前序列的最后一个 id 是什么? | prompt 的第 700 个 token |
block_table | 逻辑块对应 GPU 池的哪些物理块? | 尚未分配,为空 |
status / is_prefill | 请求在哪种状态?阶段相关处理如何进行? | WAITING / True |
temperature / max_tokens / ignore_eos | 怎样采样,生成到何时停止? | 来自这条请求的生成参数 |
7.3 三种长度:已经知道、已经算完、本轮待算
沿用 A=700、B=32、Prefill 预算=512 的主例,只看 A。假设没有前缀命中、没有抢占,也未达到停止条件。下面表中的“总长”和“已缓存”均是本轮回写后的状态;“本轮新算”是刚执行的工作量,不是回写后已经清零的 scheduled 字段。
时点 | 本轮新算 | 总长 | 已缓存 | 差额表示什么 |
|---|---|---|---|---|
刚入队 | — | 700 | 0 | 700 个已知位置还没计算 |
第 1 轮后 | 512 | 700 | 512 | 188 个 prompt 位置还没计算 |
第 2 轮后 | 188 | 701 | 700 | 首个输出已接受,但还没有自己的 KV |
第 3 轮后 | 1 | 702 | 701 | 第二个输出已接受,留待下一轮计算 |
所以,“已有 token”不等于“已有这个 token 的 KV”。采样得到一个新 id,只是决定下一步输入什么;要等它作为输入穿过各层,才会形成自己的各层 K/V。稳定 Decode 的活跃请求通常因此满足“总长比已缓存长 1”。Prefill 未完成、被抢占或已经结束时,不能机械套用这个关系。
7.4 候选、接受、结束,是不同的状态变化
第一轮只计算 A 的前 512 个位置。Runner 可以从这一段末尾的表示采出候选,但 A 后面还有 188 个已知 prompt token,不能把候选插到它们中间。Scheduler 因此只推进缓存进度,不追加该候选。第二轮补齐整个 prompt 后,候选才成为正式的首个输出;之后每次 Decode 通常接收一个新输出。
接受输出后,Scheduler 检查 EOS 和输出数量上限。结束时,它把请求标为 FINISHED 并释放块引用,Engine 则收集 completion_token_ids。释放缓存会清空该请求的 block_table 和 cached 计数,但不等于丢掉生成结果。
也不要把“Decode 不必保存旧输入向量”理解成“引擎连旧 token ids 都不要了”。旧隐藏向量无需跨轮保留,是因为新位置直接依赖各层历史 K/V;CPU 上的 ids 仍用于结果还原、前缀识别,以及抢占后重新建立缓存。这两种保存解决的是不同问题。
到这里,Sequence 提供了完整的进度记录,但它不会自己决定何时执行。下一章的 Scheduler 将用这些状态回答:本轮选谁,以及每条请求算多少个新位置。
源码对照:入口循环、初始化与 Sequence 字段定义。
源码:nanovllm/engine/llm_engine.py,L49–55
源码:nanovllm/engine/sequence.py,L18–31
源码:nanovllm/engine/llm_engine.py,L17–35
8. Scheduler:怎样让不同请求共享每一轮计算
有了 Sequence,引擎知道每条请求还欠哪些计算;但 GPU 一轮容量有限,KV 空间也有限。Scheduler 的工作,就是在这些约束下生成本轮工作单:选出的请求、Prefill 或 Decode 阶段,以及各请求的新计算位置数。它不算向量,而是保证交给 Runner 的工作有顺序、有预算、有缓存位置。
8.1 Continuous Batching:调度单位是 step,不是整条请求
Continuous Batching(连续批处理)的核心,是每轮重新组批,而不是等整批回答全部生成完。一条请求的生成会经过很多次 step;两次 step 之间,已完成的请求可以退出,等待请求可以在容量允许时加入。调度单位因此从“整条请求”变成了“这一轮需要计算的位置”。
先区分三件事。逐条执行是 A 全部结束才开始 B;固定批次是 A、B 一起开始,但这批结束前不接纳 C;普通动态组批通常是在执行前短暂收集到达的请求。连续批处理进一步把组批边界放到生成过程内部:同一条长请求还没有结束,和它一起计算的其他请求已经可以换一批。
例如 A、B 已在 Decode,A 还需 1 轮,B 还需 3 轮;第一轮结束时 A 完成,C 到达。固定批次要让 C 等 B 结束。连续批处理可以在后续轮次接纳 C,但 C 必须先做自己的 Prefill,不能直接拿着原始 prompt 加入 Decode。以 nano 本版本为例,第二轮先 Prefill C,让 B 暂等一轮;第三轮才让 B、C 各输入一个新 token,一起 Decode。
图 8-1:固定批次与连续批处理——看每轮成员如何变化,而不是把格子数当作加速倍数。
每轮被一起处理的请求共享模型参数和 GPU 执行机会,不共享彼此的上下文。B 的 query 仍只读取 B 的 K/V,C 只读取 C 的 K/V;请求边界、各自长度与块表保证 Attention 不会串到另一条请求。单条请求内部也仍然是自回归的,下一轮必须等待上一轮选出新 token。
nano 的 Scheduler 持有 waiting 和 running 两个队列,以及负责缓存容量的 BlockManager。waiting 包括尚未完成 Prefill 的请求与被抢占后需要恢复的请求;running 里的请求通常等待 Decode,但最后一段 Prefill 被选中时也会提前进入这里。状态表示调度阶段,不能单凭 RUNNING 判断 GPU 已经执行完成。
本版本先尝试 waiting 中的 Prefill。只要本轮选出了任意 Prefill,就返回这一批,不混入 Decode;没有选出 Prefill 工作时才尝试 Decode。Prefill 优先只是让已有 Decode 请求等到后续 step,不会自动把它们移回 waiting,也不打断正在运行的 GPU 内核。真正的抢占发生在 Decode 准备新位置却缺少缓存空间时,详见 8.3。这个优先级是实现策略,不是连续批处理的定义;持续加入新输入可能拉长已有请求的输出间隔。
收益来自更及时地利用容量,以及把多条请求的新位置合并计算;特别是 Decode 中,批量矩阵运算有机会提高模型权重读取的复用和 GPU 利用率。它不是删掉每条请求的必要计算,也不保证所有延迟都下降。让新请求早做 Prefill 可以缩短其首 token 等待时间(TTFT),却可能拉长已有请求的相邻输出间隔(ITL);较大的批次也可能让单轮更慢。预算、切分长 Prefill 与优先级,都是在吞吐和延迟之间取舍。
图里的“C 中途到达”描述调度能力:需要上层在 step 之间调用 add_request。仓库的 generate() 则先将传入的 prompts 全部入队,再循环 step,最后统一返回;它不是现成的在线 HTTP 或流式服务。max_num_seqs 限制的是本轮选入的请求数,也不能直接解释成整个引擎所有存活请求的总上限。
完整调度决策图已放在第 6.5 节(图 6-3)。下面沿 A=700、B=32 的例子,进一步追踪预算、队列迁移、抢占和回写,而不是重新介绍一次组件。
8.2 Prefill:按预算选择新位置,必要时切分长输入
Scheduler 先确认缓存能否分配,并减去可复用的前缀位置,再用剩余 token 预算确定本轮新算多少。切块改变的是单轮计算量,而不是把原 prompt 改成多条请求。每段仍属于原来的 Sequence,后段继续读取前段已经形成的 K/V。
设 A 的 prompt 长度为 700,B 为 32,按 A、B 顺序进入;Prefill 预算为 512,max_num_seqs 足够大,显存充足且没有前缀命中。两条请求都不触发停止条件,执行过程如下:
轮次 | 工作单 | 回写后的缓存进度 | 接受输出后的总长 |
|---|---|---|---|
1:Prefill | A 的位置 0~511,共 512 个 | A=512,B=0 | A=700,B=32;不接受 A 的中间候选 |
2:Prefill | A 的 512~699,加 B 的 0~31,共 220 个 | A=700,B=32 | A=701,B=33;各接受首个输出 |
3:Decode | A 的位置 700,B 的位置 32,各 1 个 | A=701,B=33 | A=702,B=34;各接受第二个输出 |
这里有两条实现边界。第一,本批只有第一条请求可以被切块;若已经选了其他请求,下一条的剩余输入放不进剩余预算,就结束选取。第二,max_num_batched_tokens 约束这段 Prefill 选择循环,max_num_seqs 约束选入的请求数;不能把 Prefill 的预算循环直接套到此版本的 Decode 路径。
8.3 Decode:每条请求一个新位置,空间不足则抢占
Decode 选中一条 running 请求后,需要为它刚接受的最后一个 token 确保缓存槽位。尾块还有空间时直接使用;跨入新块时,则需要空闲物理块。Scheduler 通过 BlockManager 检查和扩展映射,再把这条请求本轮新计算数设为 1。
若没有足够的空闲块,本实现先抢占 running 队尾的其他请求,必要时抢占当前请求。抢占会释放该请求的块引用,清空映射与缓存进度,重设为 Prefill,并放回 waiting 前端。它不是把 KV 搬到 CPU 暂存;恢复时要利用保留的 token ids 重新 Prefill,仍能命中的完整前缀块则可能省掉部分重算。
状态图:正常生成沿主线前进,抢占把请求送回等待恢复的路径。
图中的“Prefill 完成后进入 RUNNING”表达逻辑阶段;源码在选定最后一段 Prefill 时提前迁移队列,实际有效缓存长度随后由 postprocess 推进。被抢占的请求不仅保留原 prompt,也保留已经接受的输出,恢复时不能丢掉这部分上下文。
8.4 回写:先确认缓存完成,再决定是否增长回答
Runner 返回后,GPU 已完成本轮新位置的计算。Scheduler 先用“旧 cached + 本轮 scheduled”确定刚完成的完整块,登记它们的前缀 hash;再增加 cached,并清零 scheduled。hash 登记放在计数更新之前,是因为这份实现使用这两个旧字段来界定本轮完成的范围,而不是因为 hash 会参与 Attention。
若 Prefill 尚未覆盖整个已知序列,回写到此为止,不接受候选;否则追加候选,并检查 EOS 或 max_tokens。结束请求释放引用并移出 running,未结束请求等待后续轮次。这个顺序把“缓存已经算到哪里”和“回答已经增长多少”分开维护,也解释了第一轮 A 没有新输出、第二轮才有首输出。
Scheduler 最终交出的仍只是 CPU 上的工作单和映射。第二轮 A 的 188 个新位置、B 的 32 个新位置具体写到哪里,由下一章的 BlockManager 与分页寻址解释;它们怎样组成 GPU 输入,则留到第 11 章。
源码对照:Prefill 选择、Decode 与抢占、postprocess 回写。
源码:nanovllm/engine/scheduler.py,L30–55
源码:nanovllm/engine/scheduler.py,L58–79
源码:nanovllm/engine/scheduler.py,L81–92
源码:engine/llm_engine.py,add_request、step 与 generate 的调用边界
9. PagedAttention 与 KV Cache:分页存储怎样参与计算
第二轮要新算 A 的 188 个位置、B 的 32 个位置。A 的前 512 个位置已经有 K/V,但“cached=512”只说明进度,不能告诉 GPU 到哪里读取。BlockManager 负责缓存空间和地址映射;分页让请求使用逻辑上连续、物理上可以分散的存储。缓存解决重复计算,分页解决这些缓存怎样被分配和共享。
9.1 为什么要分页:请求长度不断增长,显存容量却固定
若每条请求都按最大长度预留连续缓存,短请求会浪费大量空间;若只分配当前长度,请求增长时又可能找不到足够的相邻空间。分页把固定数量的 token 位置组成逻辑块,再将逻辑块映射到常驻 KV 池中的物理块。请求只需要知道自己的块表,不要求相邻逻辑块在显存里相邻。
这里的存储 block 不是 Transformer Block,也不是一个 head。一个物理块编号在各层对应同一组 token 槽位;每个槽位仍分别存放各层、各 KV head 的 K 与 V。分页改变存储组织,不改变模型的位置编号、因果约束或 Attention 公式。
BlockManager 在 CPU 侧维护物理块元数据、空闲队列、使用集合和前缀索引。Scheduler 向它申请空间;Runner 根据 Sequence 的 block_table 计算读写元数据;模型与 GPU kernel 产生并写入真正的 K/V。分配块、生成地址、写入向量,是三个不同动作。
9.2 从 token 位置到物理槽位
先用 block size=4 的纸面例子:10 个 token 占 3 个逻辑块,覆盖位置 0~3、4~7、8~9;block_table=[17,5,23]。下标是逻辑块号,数组值是物理块号。本文 nano 实际默认每块 256 个 token,配置要求是 256 的倍数;4 只用于帮助看清地址换算。
slot = block_table[position // block_size] × block_size + position % block_size
零起始位置 | 逻辑块号 | 块内偏移 | 物理块号 | slot |
|---|---|---|---|---|
0 | 0 | 0 | 17 | 17×4=68 |
6 | 1 | 2 | 5 | 5×4+2=22 |
8 | 2 | 0 | 23 | 23×4=92 |
图 9-1:上半部分看逻辑块如何映射到物理块,下半部分的前缀共享留到第 10 章。
图使用实际块大小 256,并用独立的 620-token 请求展示 [17,5,23] 三个块。位置 300 的逻辑块号为 1、块内偏移为 44,因此 slot=5×256+44=1324。slot 只是 token 槽位编号,不是字节地址;层号、K/V、KV head 和特征分量还对应缓存张量的其他轴。
回到 A=700、B=32 的主例,假设 A 分到 [17,5,23],B 分到 [11]。第二轮 A 的位置 512~699 写入 slot 5888~6075,B 的位置 0~31 写入 slot 2816~2847。第三轮 Decode 的新位置 700 和 32,则分别写到 6076 和 2848。编号只是手算示意,真实分配器可以选其他空闲块。
9.3 分配、增长、释放:块表怎样变化
首次分配。BlockManager 检查当前序列需要几个逻辑块、哪些完整前缀可复用,以及空闲容量是否足够。可复用块建立引用,其余块从空闲队列取出,最终形成 Sequence 的 block_table。初次分配按请求当前总长度需要的块数进行;A 虽然第一轮只计算 512 个位置,仍会为当前 700-token 序列建立 3 块映射。
Decode 增长。已有尾块能容纳新位置时,不增加块数。序列长度从 256 变成 257 后,刚接受的新 token 位于位置 256,才跨入新逻辑块。因此源码在 len(seq) % block_size == 1 时申请新块。申请发生在这轮前向之前,此时新 token 已存在,但它的 K/V 还没写入。
序列总长度 | 需要的逻辑块数 | Decode 时尾部意味着什么 |
|---|---|---|
255 / 256 | 1 / 1 | 最后位置仍在第一块内 |
257 | 2 | 位置 256 需要新块 |
512 | 2 | 最后位置 511 仍在第二块内 |
513 | 3 | 位置 512 需要第三块 |
结束或抢占。BlockManager 对请求引用的各块减少 ref_count,只有降为 0 才放回空闲队列;随后清空该请求的块表与 cached 计数。只要还有别的请求引用某个共享块,就不能把它当作可覆盖空间。这也是前缀共享需要引用计数的原因。
9.4 “空闲”不等于“显存已释放”,也不等于“内容已擦除”
请求释放的是块的使用权,不是整个 GPU KV 池。池通常常驻,物理块可以很快交给下一条请求,因此请求结束后显存占用不一定明显下降。ref_count=0 表示当前没有活跃持有者,旧 K/V 和前缀元数据却可能仍然存在;只要未被覆盖,就可能在第 10 章的前缀查找中再次命中。
分页也不是零浪费:最后一块可能只有少量有效位置。更大的块减少块表和管理开销,但尾部空余和前缀复用粒度也更大。读缓存时必须同时使用块表与有效长度,不能把已分配容量都当作有效上下文;“有这个槽位”与“这个槽位已写好 K/V”仍是两回事。
9.5 PagedAttention:K/V 分散在不同页,Attention 怎样算
只把显存切成块,还没有完成 PagedAttention。CPU 侧负责分配物理块、维护块表;GPU 侧的 Attention 内核还必须能按块表直接读取分页 K/V。否则,每轮先把分散的历史搬回一个连续大矩阵,仍会产生额外复制。分页注意力把“逻辑连续、物理分散”的缓存组织接入实际计算,而不是改变模型公式。
沿用 620 个位置、block size=256、block_table=[17,5,23] 的例子。现在只看某一层、某一个 head,假设 Decode 正在处理从 0 编号的位置 619:已有历史是 0~618,当前新 query 是 q₆₁₉,形状为 [1,128]。本轮先把新 k₆₁₉、v₆₁₉ 写到物理块 23 的偏移 107,即 slot=23×256+107=5995;此后本次 Attention 的有效上下文才是 620 个位置,包含当前自己。
图 9-2:分页 K/V 进入 Attention——物理位置可以分散,数学上仍是同一个上下文。
先寻址,再打分。内核根据块表读取物理块 17 的 256 个位置、块 5 的 256 个位置、块 23 的前 108 个位置;尾部其余 148 个槽位不参与。对每个有效位置 j,用 q₆₁₉ 与 kⱼ 点积并缩放,得到 sⱼ=q₆₁₉·kⱼ/√128。逻辑上相当于 [1,128] × [128,620] → [1,620],但不要求先构造一份新的连续 K 矩阵。
再统一归一化,最后读取 V 加权求和。令 m 是这 620 个分数的最大值,则 wⱼ=exp(sⱼ−m) / Σᵣ exp(sᵣ−m),分母覆盖全部 620 个有效位置。每个 wⱼ 是一个标量,用它乘对应的 128 维 vⱼ,再把结果逐分量相加:a₆₁₉=Σⱼ wⱼvⱼ,输出仍为 [1,128]。K 决定“读多少”,V 提供“读出的内容”;这个 a 继续进入多头拼接、输出投影和后续网络,不是直接生成的 token。
不能每页各做一次 softmax,再把各页结果简单相加。举例说,若两页在同一数值尺度下的指数和分别是 2 和 6,它们应分到总权重的 1/4 和 3/4;独立归一化却会让每页权重都加起来等于 1。内核可以分块计算并合并最大值、指数和、加权分子等统计量,得到跨页统一归一化的结果,无须把完整分数和权重都写回显存。图中分阶段展示的是数学关系,不限定内核必须拆成多个操作。
同样的寻址也用于有历史缓存的 Prefill,只是这时有多个新 query;每个 query 仍只能看到不晚于自己的位置。换到下一层,要读取下一层自己的 K/V,而不是复用上一层的向量。
9.6 分页究竟省了什么,与 FlashAttention 有何不同
PagedAttention 主要减少按最大长度预留和连续空间分配造成的浪费,使有限 KV 容量能容纳更多请求,并为按块共享前缀提供方便。它不是稀疏 Attention:相同上下文下,该读取的历史 K/V 仍需参与计算,q·k 和权重乘 V 不会因为“分页”自动消失。尾块空余、块表维护与寻址开销也仍存在;只有一个短请求时,分页未必让单步更快。
可以借操作系统分页理解“逻辑页 → 物理页”,但这里通常是引擎维护的 KV 块表和 GPU 内核寻址,不是直接依赖操作系统缺页中断;这份 nano 实现也不会因此自动把 KV 换出到 CPU。
PagedAttention 关注缓存如何布局和寻址,FlashAttention 关注 Attention 如何分块融合计算、减少显存中间读写;两者可以配合使用。本版本 nano 先由 store_kvcache 写入新 K/V;有缓存的 Prefill 使用带 block_table 的 flash_attn_varlen_func,Decode 使用 flash_attn_with_kvcache。不要因为它支持分页,就以为它直接照搬了最初 vLLM 的那套 CUDA 内核。
源码:layers/attention.py,缓存写入与分页 Attention 调用
源码对照:块引用回收与 Decode 跨块增长。完整分配与前缀索引见下一章。
源码:nanovllm/engine/block_manager.py,L94–108
10. Prefix Cache:不同请求怎样复用已算过的前缀
同一请求在下一轮读取自己的旧 KV,是跨轮缓存;另一条请求复用已经计算过的公共前缀,则是 Prefix Cache。nano 没有为它增加一个独立执行组件,而是在 BlockManager 中用前缀索引寻找已有物理块,再通过块表与引用计数共享。底层仍是上一章的那一份 GPU KV 池。
10.1 为什么必须是相同前缀,而不只是相同文字块
对固定模型、权重与计算语义,同样的 token 前缀在同样的位置会产生可复用的历史 K/V。因果模型的旧位置不会受后来追加内容影响,因此公共系统提示词、公共文档前缀等不必为每条新请求重新计算。
但是,中间一段文字相同还不够:它的高层表示受更早上下文影响;位置不同也可能改变含位置信息的 K。需要匹配的是“从开头直到这里的前缀”,不是在文本中任意查找一段相同句子。共享范围也限定在当前引擎的模型与缓存语境内,不能跨不同权重直接套用。
本实现为完整块建立滚动 hash:首块由 token ids 计算,后续块把前一个前缀 hash 与当前块 token ids 一起计算。hash_to_block_id 找到候选物理块后,还检查该块保存的 token ids。索引保存的是“去哪里找”的信息,不是 K/V 本身。
10.2 一个完整例子:少算 512 个位置,但仍读取它们
这里暂离 A=700、B=32 的调度主例,改用 P、Q 两条各长 620 的请求:前 512 个 token 相同,最后 108 个不同。先让 P 完成相关前缀计算和 hash 登记,再在同一个引擎中接纳 Q;假设公共块没有被覆盖。
步骤 | BlockManager 与 Sequence | 模型需要做什么 |
|---|---|---|
查找 Q 的前缀 | 命中两个完整的 256-token 块 | 复用这些块在所有层的 K/V |
建立 Q 的块表 | 前两项指向共享块,尾块独占;cached=512 | 前 512 个位置不重跑前向 |
执行 Q 的剩余 Prefill | 若预算足够,scheduled=108 | 新算位置 512~619,并写入 Q 的尾块 |
Attention 读取 | 按块表找到公共前缀与自己的尾部 | 108 个新 query 读取长达 620 的上下文,仍受因果约束 |
第 9 章分页图的下半部分展示的就是这种关系:不同请求的块表可以指向同一组完整前缀块,但各自尾部不同。收益是少算前缀位置的投影、Attention、MLP 等前向,不是把前缀从上下文删掉。Q 的 query 长度是 108,key 上下文总长是 620;第一个新 query 实际只能看位置 0~512,最后一个才能看 0~619。
10.3 共享边界:完整前缀只读,最后逻辑块独占
这个版本初始查找只遍历 num_blocks−1 个块,始终跳过请求的最后逻辑块,即使它恰好写满。620-token 请求有 3 块,可检查前 2 块;512-token 请求有 2 块,只检查第 1 块。不要把“只有完整块才可共享”误写成“所有完整块都会被这份实现查找”。
所以,即使两次输入完全相同,也不能推导出“第二次不用跑模型,直接返回答案”。本版本对完全相同的 512-token 输入最多复用前 256 个位置,仍计算最后一块。KV Cache 保存的是各层 K/V,不是词表 logits 或最终回答;下一 token 的选择还需要最后位置的输出与采样。其他引擎可以采用不同的全命中处理方式,但那是额外的实现选择,不能只凭 KV 命中就假定能直接复用答案。
共享前缀后续不再修改,新位置写入独占尾部;因此这条路径不需要给共享尾块做 copy-on-write。这里的安全性来自明确的共享与写入边界,而不是简单地认为“复制永远没有必要”。其他引擎若允许共享可写尾块,就需要另外的隔离机制。
匹配到块也仍需通过容量检查。正在被其他请求使用的命中块可增加引用,不额外占用一个空闲块;已经空闲的命中块则要从空闲集合重新占用。尾部和未命中部分仍需分配,所以“命中了一段前缀”不保证整条请求立刻就能被接纳。
10.4 缓存什么时候可用,什么时候失效
完整块的索引在对应计算完成后的 postprocess 中登记。若 P、Q 一起进入最初同一批 Prefill,调度查找发生在 GPU 执行和 hash 登记之前,不能假设 Q 会命中 P 尚未算出的内容。这也是示例明确使用“先完成前缀、后接纳另一请求”的原因。
P 结束后,块引用数可以降为 0,但旧 hash 与 K/V 不会立即全部擦除。Q 若及时命中,就能重新激活这些空闲块;若它们已经被分配给新内容,旧索引会在相应重用流程中清理,旧前缀便不能再使用。它是可被覆盖的进程内缓存,不是持久化保存;重启引擎后不能继续依赖它。
10.5 命中前缀后,首字与后续生成分别省了什么
计算上,主要省的是重复 Prefill。Q 命中前 512 个位置后,不再为这些位置重算各层投影、Attention 和 MLP;本次新计算位置由 620 减到 108。但 512/620≈82.6% 是复用位置的比例,不是耗时下降比例:108 个新 query 仍要读取公共前缀,调度等待、缓存查找和其他开销也没有按这个比例消失。是否明显缩短首 token 时间,要看公共前缀长度、实际命中率与当时负载。
空间上,可以少存重复的公共 KV;上下文长度却不会缩短。在上述 P、Q 都还持有 620 个位置的时刻,各自复制需要 6 个物理块,共享前两块后只需 4 块:2 个公共块,加 P、Q 各 1 个尾块。但 Q 后续 Decode 仍然面对自己的完整历史,不能把“共享了 512 个位置”理解成“只需对剩下 108 个位置做 Attention”。因此 Prefix Cache 不保证每个后续 token 都更快。
实际使用时,要区分冷缓存与热缓存:第一次请求负责算出并登记前缀,之后只有在内容匹配且块尚未被覆盖时才能复用。公共系统提示词或固定文档放在开头、保持 tokenizer 与 chat template 一致,有利于形成稳定的 token 前缀;在开头插入不同的时间戳或请求信息,则可能让后面很长的相同文字也无法复用。判断收益,应同时看复用了多少前缀 token、减少了多少 Prefill 工作和首 token 延迟,而不只看“请求有没有命中过”。
10.6 把三个机制串起来:谁来算、存在哪里、哪些不用重算
机制 | 核心问题 | 主要收益与边界 |
|---|---|---|
Continuous Batching | 这一轮让哪些请求一起计算? | 及时接纳和移除请求、提高资源利用率;不取消每条请求的自回归依赖。 |
PagedAttention | K/V 存在哪里,内核怎样找到它? | 降低预留和碎片浪费、支持分页读取;不减少相同上下文所需的 Attention 数学工作。 |
Prefix Cache | 哪些前缀已经算过,可以直接复用? | 省掉重复 Prefill,并可共享物理 KV;不把公共前缀从 Attention 上下文删掉。 |
在 nano 中,新请求进入 waiting 后,Scheduler 通过 BlockManager 查找可复用前缀、检查容量,并确定本轮的新计算位置;Runner 把工作单和块表交给 GPU,内核按分页地址写入新 K/V、读取有效上下文。采样与回写结束后,完成请求退出、释放块引用,未完成请求留待下一轮。三个机制分别管理计算时机、存储地址与复用范围,结合起来形成推理引擎;它们不是三个同义词,也不是缺一个就必然无法实现另一个。
前缀命中后,Scheduler 少安排一些新位置,ModelRunner 却仍要保留完整的可读上下文。下一章就把这种“输入短、上下文长”的工作单变成实际张量。
源码对照:前缀查找、空闲块重用与共享分配。
源码:nanovllm/engine/block_manager.py,L58–73
源码:nanovllm/engine/block_manager.py,L43–56
源码:nanovllm/engine/block_manager.py,L75–92
11. ModelRunner:把调度结果变成 GPU 张量
Scheduler 选出了请求,BlockManager 确定了块映射,但这些仍是 CPU 对象。ModelRunner 把“算哪些位置、读写哪些块”翻译成 GPU 可以执行的输入和元数据。它持有模型、Sampler 与 KV 池,负责调用和组织执行;模型内部如何完成 Q/K/V、Attention 和 MLP,留到第 12 章。
11.1 一轮执行:工作单怎样变成候选 token
Engine 将本轮 Sequence 列表与 is_prefill 标记交给 Runner。Runner 选择 prepare_prefill 或 prepare_decode,生成 input_ids、positions 和本轮 Context;rank 0 还准备采样温度。随后 run_model 完成模型前向和词表输出,rank 0 的 Sampler 根据 logits 选出候选 id,再由 Engine 转交 Scheduler.postprocess。是否接受候选、是否结束请求,属于 Scheduler 的职责,不由 Runner 判断。
Runner 本身是 Python 编排器,不是一个全部运行在 GPU 上的黑箱。它先在 CPU 收集 ids、位置和映射,构造锁页张量(Pinned Memory,详见 11.6)并提交到 GPU;GPU 执行模型与采样。返回时取得候选列表、清理本轮 Context,但权重与 KV 池继续保留。
Runner 的完整执行图及 Attention 缓存读写展开见第 6.3 节(图 6-2)。下面不再重复全貌,重点把主例中每一项输入、边界和读写地址落实到具体值与形状。
11.2 Prefill 输入:把新位置拼起来,同时保留请求边界
回到主例第二轮,A.cached=512、scheduled=188;B.cached=0、scheduled=32。Runner 从 A 取位置 512~699 的 ids,再接 B 的 32 个 ids,得到 input_ids 形状 [220]。这是包含 220 个整数的单轴数组;进入 Qwen3-0.6B 后查 Embedding table,才变成 [220,1024] 的向量矩阵,一行对应一个本轮新位置。
把 ids 拼接起来可以避免按最长请求补出大量无效输入,但不能因此把两条请求当成同一个上下文。Runner 还必须给出每条请求的位置、query 边界与 key 长度。以下形状均以本章单卡主例为准:
张量 / 字段 | 第二轮的值或形状 | 用途 |
|---|---|---|
input_ids | [220] | A 的 188 个新 id,接 B 的 32 个 id |
positions | [512,…,699, 0,…,31] | 每个位置在自身请求里的编号,供位置编码使用 |
cu_seqlens_q | [0,188,220] | 新 query 的累计边界,两段分别长 188、32 |
cu_seqlens_k | [0,700,732] | 两个独立 key 上下文的累计长度:700、32 |
max_seqlen_q / k | 188 / 700 | 本批最长的新 query 段 / key 上下文 |
slot_mapping | [220] | 按新输入顺序列出每组 K/V 的写入槽位 |
block_tables | [2,3] | A 为 [17,5,23];B 为 [11,−1,−1],后两项为填充 |
cu_seqlens 的 cu 表示累计值,相邻两项的差才是对应请求的长度。732 是 700+32,不是任何 query 都能看到 732 个 token。请求边界先隔开 A 与 B,因果规则再限制请求内部可见的位置:A 的新 query 512 只看 A 的 0~512;B 的 query 0 只看 B 的位置 0。
没有旧缓存时,本轮 query 与 key 长度一致;有前一个 chunk 或公共前缀时,两者不同,Runner 会提供 block_tables 以读取缓存池。例如上一章 Q 的公共前缀命中后,新 query 只有 108 个,key 上下文仍长 620,对应累计长度 [0,108] 与 [0,620]。
11.3 写地址与读地址:为什么两套信息都需要
slot_mapping 回答“本轮新 K/V 逐个写在哪里”。沿用 A=[17,5,23]、B=[11],第二轮前 188 项是 5888~6075,后 32 项是 2816~2847。它与 input_ids 的顺序一一对应,各层使用这些槽位写入自己的新 K/V。
block_tables 配合上下文长度回答“这条请求的 Attention 到哪些物理块读取历史和当前 K/V”。它描述的是完整可读路径,而不是只有新写入的位置。写入只有 220 组新 K/V,读取却包含 A 以前的 512 个位置;因此逐 token 的写地址不能代替按请求的读映射。
块表为了组成 batch 会补到相同宽度,填充列不是有效上下文。GPU 需要结合真实长度使用它,既不能跨到其他请求,也不能把尾块尚未写入的容量当作有效内容。位置、边界、映射、长度必须一起正确,只有形状正确还不够。
11.4 Decode:每条请求只输入一个新 id,历史仍然可读
第二轮回写后,A、B 各接受了一个首输出,总长分别为 701 和 33。第三轮只取各自 last_token:input_ids 的形状为 [2],positions=[700,32],context_lens=[701,33];查 Embedding 后是 [2,1024]。这表示两条请求各有一行新输入,不是上下文只剩两个 token。
新位置的 slot_mapping 为 [6076,2848],块表仍分别指向 A、B 的完整缓存。每层 Attention 先形成并写入本轮 K/V,再读取包含当前位置在内的上下文,所以 context_lens 包含这两个新位置。A 的旧 700 个位置不再重跑投影和 MLP,新 query 却仍需读取它们的 K/V。
Prefill 与 Decode 因而不是两套模型,而是两种输入组织方式。前者每条请求可能新算很多位置,后者通常各新算一个位置;新位置都必须穿过完整模型。Runner 可以为合适的 Decode 批次复用 CUDA Graph 以降低执行开销,具体机制放在第 13 章。
11.5 本轮 Context 与常驻 KV 池:寿命不同的两种对象
Context 是本轮执行信息的容器,供 Attention、LM Head 等读取阶段、边界和缓存地址,避免把所有元数据逐层传递。容器在 CPU 侧,里面可以持有 GPU 张量引用。run 结束时 reset_context 只是清理这一轮的元数据,不是删除 KV,也不是清空聊天历史;这个版本的全局 Context 与顺序执行相配合,扩展并发执行时需要重新考虑隔离。
常驻 KV 池在初始化时分配。Runner 加载权重、完成预热并估算可用容量,再建立 [2,L,num_blocks,block_size,Hkv_per_rank,head_dim] 的大张量,给每层 Attention 绑定对应切片。单卡 Qwen3-0.6B、块大小 256 时为 [2,28,num_blocks,256,8,128]。
轴 | 本例长度 | 含义 |
|---|---|---|
K / V | 2 | 分别保存 key 与 value |
模型层 | 28 | 每层有自己的历史 K/V |
物理块 | num_blocks | 整个引擎缓存池的容量,不是某条请求长度 |
块内位置 | 256 | 每块可放多少个 token 槽位 |
KV head | 8 | 每个位置的 KV 头数 |
每头特征 | 128 | 一个 k 或 v 向量包含的特征分量数 |
这里 [2,28,…] 描述数组的 6 个轴,最后的 128 才是每个 head 向量的特征维度。池内哪些块归某条请求使用,由 block_table 决定;哪些位置已有效,由缓存进度与本轮长度决定。gpu_memory_utilization 参与池容量估算,不是每条请求分配显存的比例,也不意味着数值越大就一定越快。
11.6 Pinned Memory:本轮 CPU 输入怎样送到 GPU
前面已经准备好 input_ids、positions 和缓存地址,但它们最初来自 CPU 上的 Python 列表,GPU 不能直接把这些列表当作模型输入。Runner 还要把数据整理成张量,并从主机内存传到显存。这段 CPU → GPU 的传输称为 H2D(Host to Device);Pinned Memory 优化的是这段传输,不是 Attention 计算,也不是历史 KV 的复用。
“Pinned”指锁页的 CPU 内存,不是显存。普通主机内存由操作系统分页管理;被锁定的页在使用期间保持驻留,不被换出,从而为设备传输提供稳定的主机内存页。普通可分页内存上传到 GPU 时,运行时通常需要先暂存到锁页缓冲区,再由 DMA(直接内存访问)传输到设备。直接准备 pinned 缓冲区可以避免这段临时暂存复制,并为高效异步传输提供条件;DMA 搬运数据时,不需要 CPU 逐个数值执行复制,但数据仍须经过实际的传输链路。
图 11-1:CPU 准备输入、锁页缓冲区承接传输、GPU 消费输入;模型权重和历史 KV 留在显存中。
nano 每轮上传的是新 id 和执行元数据,不是整份历史 K/V。以第三轮两请求 Decode 为例,input_ids 只有 A、B 各一个新 id;positions=[700,32]、context_lens=[701,33]、slot_mapping=[6076,2848],以及两条请求的 block_tables 一起送到 GPU。源代码在 prepare_prefill、prepare_decode、prepare_block_tables 和 prepare_sample 中,用 pin_memory=True 构造主机张量,再通过 cuda(non_blocking=True) 提交传输。新 id 到达 GPU 后才查 Embedding 并生成各层 Q/K/V;模型权重与历史 KV 已常驻 GPU,不会随这些小张量每轮重新上传。
non_blocking 描述 CPU 的等待方式,不表示 GPU 已经拿到数据。对这里的锁页主机张量,提交异步拷贝后,CPU 可以继续安排后续工作,不必在每次提交后都等待该次传输完成;但依赖这些输入的 GPU 计算必须等拷贝完成。CUDA stream 可以理解为一条有序的 GPU 工作队列。本版本没有专门建立跨轮传输流水线,同一 stream 上的拷贝和后续计算按顺序执行,保证模型不会读到尚未上传好的输入。走 CUDA Graph 时,还会把本轮 GPU 输入更新到固定地址的 GPU 缓冲区,再 replay;固定的是地址与执行安排,不是输入数值。
异步提交,不等于传输与计算自动重叠。若要让一份数据的 H2D 与另一份数据的计算同时进行,还需要硬件支持、合适的 stream/事件依赖、互不冲突的缓冲区,以及确实独立的工作。当前输入的计算不能抢在自己的传输前面;自回归的下一轮输入也要等本轮采样结果。这里不能只看到 pin_memory=True 和 non_blocking=True,就宣称 nano 已实现双缓冲或跨轮传算重叠。返回路径同样需要同步:rank 0 将采样结果转成 CPU 可读的 Python 列表时,必须取得实际完成的结果,不能把 run() 理解成全程不等待。
锁页也有成本:分配或注册缓冲区需要时间,大量锁页会减少系统可灵活管理的主机内存;很小的输入还可能主要受 Python 打包、张量分配和提交开销影响。因此收益应看整段“准备 → 传输 → 执行”的耗时,不能只比较一次拷贝,也不能认为锁页越多越快。使用可复用缓冲区时,传输未结束前不得修改其源数据;跨 stream 使用目标张量也要建立正确依赖。
最后区分两种“页”:Pinned Memory 锁定的是操作系统管理的主机内存页;Paged KV 把若干 token 的 K/V 组织成引擎管理的显存块。两者的单位、位置和目的都不同。传输完成后,GPU 内部如何用 FlashAttention 减少中间读写、用 CUDA Graph 减少重复提交,继续看第 13 章。
源码:model_runner.py,prepare_* 中的锁页张量与 H2D 提交;Graph 输入更新与采样结果返回。
延伸阅读:PyTorch 官方指南:non_blocking 与 pin_memory 的正确用法,包含传输与同步语义、性能测量及缓冲区安全边界。
至此,控制部分与计算部分已经接上:Engine 驱动循环,Sequence 记录进度,Scheduler 选择工作,BlockManager 管地址和复用,Runner 准备输入并调用 GPU。下一章进入这次调用内部,沿一行 token 向量看模型怎样产出词表分数。
源码对照:run、Prefill / Decode 输入准备与 KV pool 分配。
源码:nanovllm/engine/model_runner.py,L214–220
源码:nanovllm/engine/model_runner.py,L138–148
源码:nanovllm/engine/model_runner.py,L172–188
源码:nanovllm/engine/model_runner.py,L106–121
12. Qwen3 前向、权重加载与采样:走完模型内部
第 11 章已经准备好本轮输入。现在进入 ModelRunner 调用的模型内部,沿着同一批数据走到候选 token:Embedding → 28 个 Decoder Block → 最终 Norm → LM Head → Sampler。权重在启动时就已加载,本轮不会重新加载。本章先走完执行主线,最后再解释文件中的权重如何对应这些计算。
12.1 接住 Runner 的输入:220 个位置,最后只需要两份预测
沿用 A=700、B=32 的第二轮:A 新算 188 个位置,B 新算 32 个位置,共 T=220,本轮请求数 N=2。input_ids 先查 Embedding,得到 [220,1024]。这里每行是一个本轮新位置,1024 是 Qwen3-0.6B 的隐藏特征宽度,不是前置教学例子中的 4096。
模型内部始终保留请求边界与位置语义。220 行可以一起执行线性层,但 Attention 不会让 A 读取 B 的上下文。下一轮 Decode 时只需把 T 换成 2:每条请求各输入上轮刚接受的一个 token,仍走同一套模型。
形状表:先看各步骤怎样改变特征宽度,再逐段解释计算。
单卡阶段 | 张量形状 | 为什么 |
|---|---|---|
Embedding 输出 | [T,1024] | 每个 token 得到 hidden_size 维向量 |
qkv_proj 输出 | [T,4096] | 拼接 Q 的 2048 与 K/V 各 1024 |
Q / K / V reshape | [T,16,128] / [T,8,128] / [T,8,128] | 显式区分 Q heads 与 KV heads |
Attention → flatten | [T,16,128] → [T,2048] | 拼回所有 Q heads 的输出 |
o_proj | [T,2048] → [T,1024] | 回到残差流宽度 |
gate_up_proj → SwiGLU → down_proj | [T,6144] → [T,3072] → [T,1024] | 两个 3072 维分支做门控,再投影回去 |
最终 Norm → LM Head | [T,1024] → [N,151936] | Prefill 先选每条请求最后一个 query 位置 |
本模型有 16 个 Q heads、8 个 KV heads,每头 128 维。因此 Q 的总宽度是 2048,而 K、V 各为 1024;合并 QKV 投影输出 4096 维。hidden_size 不必等于 Q heads×head_dim。Attention 后的 o_proj 把 2048 维映射回残差流的 1024 维。这也说明,同样出现“4096”,可能表示完全不同的张量。
12.2 Attention:新位置生成 Q/K/V,再按权重读取上下文
每层先对本轮隐藏向量做输入 RMSNorm,再通过 qkv_proj 一次算出 Q/K/V 三段。合并只是执行组织:逻辑上仍是三组不同的投影权重。将 [220,4096] 拆开并重排,得到 Q [220,16,128]、K/V 各 [220,8,128];没有重新生成 A 前 512 个位置的 Q/K/V。
Qwen3 对每个 Q、K 向量沿 128 个特征分量做各自的 RMSNorm,然后按 positions 应用 RoPE;V 不经过这一步旋转。写入缓存的 K 因而已经带有相应的位置信息。RoPE 的成对旋转细节见附录 19.7,这里记住 positions 必须是请求中的真实位置,而不是拼接后的行号。
对 A 的某一个 Q head,本轮 Q 是 [188,128],对应的 K/V 上下文各是 [700,128]。QKᵀ 得到 [188,700] 的匹配分数;除以 √128、施加因果 mask,再沿可读位置做 softmax,得到注意力权重。权重乘 V,将每个 query 可见的位置加权求和,输出仍为 [188,128]。
GQA 中每 2 个 Q heads 共享一组 K/V,但各自的 Q 不同,权重和输出也分别计算。16 个头的输出沿特征方向拼成 [220,2048],再经 o_proj 回到 [220,1024]。缓存读写和分块计算怎样落到 GPU,由第 13 章展开;这里的数学关系始终是“Q/K 决定权重,权重加权 V”,不是用权重再生成 V。
12.3 残差与 MLP:每行完成本层计算,继续进入下一层
设本层输入为 x,将多头 Attention 与 o_proj 合称 Attn,将 gate/up、SwiGLU 与 down_proj 合称 MLP。一层的数学关系是:
h = x + Attn(Norm1(x))y = h + MLP(Norm2(h))
h 表示 Attention 分支与残差相加后的整行表示,y 是本 Block 的输出;沿用第 2.1 节的命名。两条分支都回到 1024 维,才能在同一位置、同一特征上与旁路逐元素相加。单头加权输出仍记为 a。
MLP 将归一化后的 [T,1024] 投影成 gate、up 两个 [T,3072] 分支;合并执行时先得到 [T,6144],拆开后计算 silu(gate)×up,再由 down_proj 回到 [T,1024]。它对每个位置独立执行,跨位置的信息已经由 Attention 混入。Norm 与 SwiGLU 的数值例子见附录 19.2–19.5。
残差图:上半是数学关系,下半是 nano 的 h/r 双路径实现。
源码把 hidden_states 与 residual 分开携带:Attention 后,在 post_attention_layernorm 中先相加再归一化;MLP 后返回两部分,等下一层 input_layernorm 再相加。最后一层的相加由最终 Norm 完成。这是将残差相加与 Norm 融合执行,不是省掉残差,也不是缓存历史 token 的中间激活。
上述结构重复 28 层,每层权重不同,同层各位置共享权重。第二轮的 220 个新位置都继续经过各层;Decode 时只有两个新位置走这些分支,历史只通过每层缓存的 K/V 被读取。
12.4 从最终表示到候选:LM Head 与 Sampler 各做什么
走完 28 层与最终 Norm 后,第二轮仍得到 [220,1024]。但我们只预测每条请求的下一个 token,因此 LM Head 先选择每段最后一个 query。由 cu_seqlens_q=[0,188,220] 减去各段末尾的 1,得到拼接张量的行号 [187,219],取出 [2,1024];它们对应 A 的原位置 699 和 B 的原位置 31。
LM Head 将这两行投影为 [2,151936] 的词表分数 logits。它还不是两个 token,也不是两个 embedding:每行有 151936 个候选词表项的分数。Sampler 再按各请求温度调整分布,从每行选择一个 id,返回两个候选 token。
本提交把 logits 转为 float、除以 temperature、做 softmax,然后将各概率除以独立的指数随机数,对随机分数取 argmax。虽然最后用了 argmax,它不是对原 logits 取最大值的 greedy。公开参数要求 temperature>1e-10,没有独立的 greedy、top-k 或 top-p 接口;概率和温度的直觉见附录 19.1。
候选返回 Scheduler 后,才决定接受和停止。完整 Prefill 或 Decode 接受输出;中间 Prefill chunk 的候选被丢弃。至此,模型内部的计算重新接回第 8 章的状态回写,而不是模型一采样就意味着回答已经增长。
12.5 回看启动:权重文件怎样对应合并投影
上面的计算依赖启动时已装好的权重。Hugging Face 文件分别命名 q_proj、k_proj、v_proj,nano 通过 packed_modules_mapping 把它们放入 qkv_proj 的不同区间;gate_proj 与 up_proj 同样合入 gate_up_proj。loader 遍历 safetensors,找到本地参数,再由该参数的 weight_loader 放入正确位置。
按 PyTorch 的 [输出特征,输入特征] 存储约定,单卡 qkv_proj 权重为 [4096,1024]:前 2048 行来自 Q,随后 1024 行来自 K,最后 1024 行来自 V。数学写法 XW 的 W 则是它的转置。多卡时,各源投影先按本 rank 的 heads 分片,再放进本地对应区间,不是把整个合并矩阵随意对半切开。
本模型还将 Embedding 与 LM Head 绑定到同一份权重存储。共享的是参数,不是操作:Embedding 按 id 取行,LM Head 用最终隐藏向量计算整份词表分数。它们发生在请求链的两端,作用不同。
配置对照:这些值决定启动时的参数与缓存形状,完整模型结构见第 4.1 节。
配置项 | 本文模型值 | 影响 |
|---|---|---|
hidden_size / layers | 1024 / 28 | 残差流宽度与模型深度 |
Q heads / KV heads | 16 / 8 | 每 2 个 Q heads 对应一组 K/V heads |
head_dim | 128 | 每个注意力头的向量宽度 |
intermediate_size | 3072 | MLP 的中间宽度 |
vocab_size | 151936 | Embedding 与 logits 的词表维度 |
dtype / tied embedding | BF16 / True | 每元素 2 字节;Embedding 与 LM Head 共享权重存储 |
读到这里,已经能从输入行数和模型配置推导本轮主要张量。下一章不再改变这些数学含义,而是解释 GPU 怎样写缓存、减少中间搬运,并降低重复提交的开销。
源码对照:模型配置、Qwen3 前向、Norm、权重加载与采样。
源码:nanovllm/layers/rotary_embedding.py,L6–14
源码:nanovllm/models/qwen3.py,L77–88
源码:nanovllm/models/qwen3.py,L146–183
源码:nanovllm/layers/layernorm.py,L28–50
源码:nanovllm/models/qwen3.py,L186–203
源码:nanovllm/utils/loader.py,L12–28
源码:nanovllm/layers/linear.py,L114–128
源码:nanovllm/layers/sampler.py,L5–12
13. GPU 执行:缓存写入、FlashAttention 与 CUDA Graph
上一章解释了模型要算哪些数,本章解释这些数怎样在 GPU 上流动。缓存写入解决“新 K/V 放哪里”,FlashAttention 减少中间矩阵的显存往返,CUDA Graph 减少重复提交 GPU 工作的开销。三者可以配合,但不能互相替代,也不改变 Prefill/Decode 的数学目标。
13.1 新 KV 写入:地址已经确定,kernel 负责搬入数值
第二轮某一层生成的 K、V 各为 [220,8,128]。写入 kernel 将每个位置的 8 个 heads 视为 1024 个连续数值,从 Context 读取与输入顺序一致的 slot_mapping。一个 Triton program 负责一个新位置,分别把它的 1024 个 K 和 1024 个 V 写入该层缓存切片。
例如 A 的位置 512 对应 slot=5888,写入起点是在该层 K 或 V 切片内的第 5888×1024 个元素。5888 是 token 槽位号,不是层号或 head 号;选哪一层、K 还是 V,已由 Attention 持有的切片决定。slot=−1 则是无效填充,必须跳过写入。
这一步把三个组件的职责连起来:BlockManager 分配物理块,Runner 算出本轮槽位张量,GPU kernel 写入模型刚产生的数值。映射正确但还没执行 kernel,槽位就还没有本轮 K/V;kernel 正常执行但映射错误,则可能污染别的请求。
13.2 三种 Attention 路径:新 query 一样,K/V 来源不同
本轮场景 | Attention 使用的数据 | nano 的调用路径 |
|---|---|---|
没有历史缓存的 Prefill | Q/K/V 都直接使用本轮张量;同时写 KV 池,供以后复用 | flash_attn_varlen_func |
已有旧 chunk 或公共前缀的 Prefill | Q 是新位置;K/V 从分页池读取历史与刚写入部分 | flash_attn_varlen_func,携带 block_table |
Decode | 每请求一个新 Q;按上下文长度读取包含当前位置的缓存 | flash_attn_with_kvcache |
因此,“使用 KV 缓存”不等于每种 Prefill 都从缓存池读回本轮 K/V。普通首段 Prefill 可以直接使用刚算出的张量;A 的第二段 Prefill 必须把前 512 个位置和本轮 188 个位置接起来,因此改用分页池。Decode 也先写本轮 K/V,再按完整长度读取。请求边界和 causal mask 在这些路径里都必须保持。
13.3 FlashAttention:分块累积,少搬中间矩阵
只看 A 的一个 Q head:Q=[188,128],K/V 各为 [700,128]。逻辑上仍要得到 [188,700] 的分数、缩放并做因果 softmax,再加权 V,得到 [188,128]。朴素实现可能把完整分数矩阵和概率矩阵写到显存,后续算子再读回来;这些中间往返不是最终输出,却消耗带宽与空间。
FlashAttention 每次取一小块 Q,依次读取 K/V 的小块,在片上计算临时分数并累积结果。每行维护当前最大值、指数和与未归一化的加权输出;读到新块时按新的最大值重新缩放旧累积量。所有可见块处理完成后,再得到最终归一化输出。因此它不是对每个小块单独 softmax 后直接相加,而是在分块执行中保持同一次全局 softmax 的含义。
FlashAttention 图:对比完整中间矩阵的显存往返与分块累积。
图中的 bq、bk 是一次处理多少个 query/key 位置,不是 head 的特征宽度 128。A 的 query 512 只看位置 0~512,query 699 才能看 0~699;B 则完全独立。减少中间矩阵物化,不等于不计算匹配分数、不读取历史 KV,也不把一般全注意力变成线性运算量。
Paged KV 与 FlashAttention 处在不同层面:前者让历史数据可以放在不连续物理块中,后者优化 Attention 的计算和搬运。具体后端利用块表找到数据后,仍需完成对应的加权计算。
13.4 CUDA Graph:输入每轮变化,执行安排可以复用
即使搬运已经高效,短 Decode 仍可能有很多小操作,需要 CPU 反复调度并提交。CUDA Graph 先捕获一段 GPU 工作及其依赖关系,之后将本轮输入写到固定地址的缓冲区,再 replay,以降低重复提交开销。
主例第三轮输入 A/B 的两个新 token,positions=[700,32],context_lens=[701,33];下一轮位置和长度继续增长,id、写入 slot 也改变。Graph 复用的是执行安排与静态缓冲区地址,不是上一轮的值。矩阵乘法、Attention、MLP 都仍要计算本轮的新位置。
CUDA Graph 图:先捕获模型主干,再逐轮更新缓冲区并重放。
本提交捕获 self.model(input_ids, positions),即 Embedding、所有 Decoder Blocks 与最终 Norm;compute_logits、Sampler、Scheduler 与状态回写在这张图之外。KV Cache 保存历史计算的数值,Graph 复用 GPU 工作的提交安排,两者减少的是不同开销。部分算子的 torch.compile 也是另一层机制,不能与 Graph 视为同义词。
13.5 形状分桶与收益边界:少提交也可能多算填充
Graph 使用固定形状,所以 nano 为若干 batch 大小预先捕获图:1、2、4、8,以及从 16 起按 16 递增的桶,上限受 max_num_seqs 与 512 限制。实际有 9 条请求时选择 16 的图,7 个额外位置是填充;slot_mapping=−1 禁止写真实 KV,context_lens=0 表示无有效上下文,最终只取前 9 行输出。
这个版本在 Prefill、enforce_eager=True 或 Decode batch>512 时走普通路径,Graph 只覆盖合适的 Decode 批次。分桶也未覆盖任意配置:例如 max_num_seqs=20 时可能没有能覆盖 17~20 的桶。这是 nano 的实现限制,不是 CUDA Graph 的通用要求。
Graph 的收益取决于原本的提交开销占比、填充浪费和实际 GPU 工作量;还要付出初始化捕获与静态缓冲区的成本。FlashAttention 与 Graph 都不保证在每种负载下同样有效。如何固定变量、区分初始化与稳定性能,见第 15 章;具体对照操作统一在第 18.12 节。
源码对照:缓存写入、Attention 路径、Graph 运行与捕获。
源码:nanovllm/layers/attention.py,L21–30
源码:nanovllm/layers/attention.py,L59–75
源码:nanovllm/engine/model_runner.py,L195–212
源码:nanovllm/engine/model_runner.py,L223–256
14. Tensor Parallel:把一次 forward 分给多张 GPU
14.1 先区分多份模型与一份模型的切分
前面用一张 GPU 执行第二轮的 220 行输入。现在只改变设备组织:让两张 GPU 共同执行同一层。这不是把 A 发给卡 0、B 发给卡 1;在 Tensor Parallel 中,两张卡都参与这两条请求,只是各自负责部分特征与权重。先看矩阵怎样分,再看进程怎样协同。
Data Parallel 通常让不同设备上的模型副本处理不同请求;Tensor Parallel 将同一个层的权重和计算切给多个设备协同完成。nano 的多卡路径是单机 TP。较大模型需要多卡共同容纳时,首先解决的是容量;是否变快还取决于本地计算和通信的比例。
14.2 Column 与 Row 的名字怎样对应 PyTorch 权重
按数学写法 Y=XW,Column Parallel 切输出特征,让各卡计算不同输出列;Row Parallel 切输入特征,让各卡计算完整输出的部分和,再相加。PyTorch Linear 实际保存的权重是 [out_features,in_features],所以源码中 Column 的 tp_dim=0、Row 的 tp_dim=1。名称来自数学布局,存储转置后轴号会变化。
先看 Column:把数学权重 W 横向分成 [W0,W1],两张卡读取相同 X,各自得到 XW0 与 XW1;它们是不同输出特征,可以继续在本卡计算,不必立即拼起来。QKV 层会对 Q、K、V 各自分片,所以 TP=2 时每卡持有 8 个 Q heads、4 个 KV heads,不是把合并 QKV 张量随意从中间切成前后两段。
再看真实的 o_proj:完整 Attention 输出是 [T,2048],但每卡只拥有其中 [T,1024] 的本地 heads。把它们记为 X0、X1,数学权重对应地纵向切成 W0、W1,每份 [1024,1024]。卡 0 算 Z0=X0W0,卡 1 算 Z1=X1W1;两者虽都是 [T,1024],却各自只包含一半输入特征的贡献。完整结果必须是 Y=Z0+Z1,all-reduce 做的是逐元素求和并让两卡都拿到 Y,不是把两个结果拼成 2048 维。
图 6:上半按输出特征分工,下半把输入分片产生的部分和相加。
把 T=220 代入,图下半每卡的部分和都是 [220,1024],求和后仍是 [220,1024];第三轮 Decode 只需把 T 换成 2。MLP 的 gate/up → down 遵循相同的“先分输出、再合部分和”结构:每卡 gate/up 各 1536 维,down 后 all-reduce 回到完整残差流。矩阵尺寸描述外形,是否包含完整信息还取决于计算用了哪些分片。
Qwen3-0.6B,TP=2 | 每卡本地形状或工作 | 何时通信 |
|---|---|---|
Embedding | 保存 75968×1024 的词表 shard;非本 shard token 的输出置 0 | all-reduce 得到完整 embedding |
QKV 投影 | 本地 Q 8×128;K/V 各 4×128,合计输出 2048 维 | 此处无需立刻合并,各卡做本地 heads |
Attention 与 KV | 本地 KV heads=4,cache 按头切分 | 本地 Attention 不需要合并历史 KV |
o_proj | 本地输入 1024 维,计算一个 [T,1024] 部分和 | all-reduce 合并两个部分和 |
gate/up → down | gate/up 各 1536 维;down 生成 [T,1024] 部分和 | down 后 all-reduce |
LM Head | 每卡产生 [N,75968] 局部 logits | gather 到 rank 0,拼成完整词表后采样 |
RowParallelLinear 的 bias 若存在只在 rank 0 加一次,避免 all-reduce 后重复累加。QKV、MLP 与词表的分片还要求相关维度能够整除 TP size;本实现不提供任意 KV heads 复制等兼容路径。
源码:nanovllm/layers/linear.py,L152–156
源码:nanovllm/layers/embed_head.py,L56–66
14.3 worker 为什么不需要拿到完整采样状态
rank 0 位于主进程,负责 tokenizer、Scheduler、采样与结果整理;其他 rank 是 worker。SharedMemory 和 Event 传递方法名、Sequence 等 CPU 元数据;各 rank 真正的张量合并通过 NCCL 完成。不能把两类通信混为一条“传模型权重的消息队列”。
图 7:TP 的 CPU 控制通信与 GPU 张量通信。
源码:nanovllm/engine/model_runner.py,L41–89
主进程把 run 方法和本轮 Sequence 信息写入共享内存,唤醒 worker;所有 rank 按一致顺序进入模型计算与 collective。Prefill 序列化包含完整 token_ids,Decode 只需 last_token 加长度、block_table 等元数据。非零 rank 不负责 tokenizer 或最终采样,完整词表 logits 聚合到 rank 0。
源码:nanovllm/engine/sequence.py,L72–83
CPU 控制消息不搬运每层 KV 数值:KV 按本地 heads 常驻各卡,各 rank 用一致的请求位置和块映射访问自己的分片。张量合并则要求所有 rank 按相同顺序进入 collective;某个 worker 提前失败时,其他 rank 可能表现为等待。控制与数值通信的分工,是理解多卡执行和定位问题的关键;固定资源名与容量限制集中见第 16 章。
14.4 为什么多卡不一定更快:容量、计算与通信的取舍
TP 首先让一份模型的权重与 KV 按分片分布到多张卡,从而扩大可容纳的模型和上下文规模;它不自动保证时延下降。每卡本地乘法少了,但 o_proj、down_proj 等处的通信仍要发生。一个位置必须等需要的部分和合并完,才能进入依赖完整残差流的后续计算。
当模型或批量足够大,本地计算的减少可能覆盖通信成本;当模型很小、Decode 每轮只有少量位置时,collective 启动、数据传输和进程协同反而可能占主要时间。卡间互联也会改变这个平衡。因此增加卡数,需要分别回答“容量是否够”和“同一负载是否更快”。
Qwen3-0.6B 很适合看清分片形状,却未必能展示 TP 的速度优势。判断结果时,同时看每卡权重与 KV、每层通信点、端到端耗时;不能把 rank 0 的显存当作多卡总量,也不能把随机输出文本相同当成数值等价。验证方法见第 15 章,具体操作统一放在第 18.13 节。
至此,推理主线完整闭合:前面决定本轮新算哪些位置,这里决定这些位置由一张卡还是多张卡执行。调度、缓存和并行改变执行组织,新位置仍然必须经过完整的模型计算。
15. 验证方法:把原理变成可复查的证据
理解了引擎怎样工作,还需要区分“解释合理”与“证据充分”。本章不再提供另一套操作路线,而是回答:要验证哪个判断,保持哪些条件相同,观察什么,以及结果不能证明什么。具体安装、命令和日志采集统一在第 18 章,附件也集中在那里。
15.1 先区分三类证据:数值、机制与性能
数值对照问“优化前后是否计算了同一个函数”;机制对照问“调度和缓存状态是否按预期推进”;性能对照问“同一负载下是否减少时间或资源”。脚本返回 PASS,只能覆盖它实际检查的断言,不能自动跨越这三类问题。
对照表:每项判断都先固定不应变化的条件,再确定观察量。
对照 | 必须保持相同 | 应检查什么 |
|---|---|---|
有 / 无 KV Cache | 每一步输入对应的完整 token 路径、权重与位置语义 | logits 误差、形状、输入位置数量 |
整段 / chunked Prefill | 同一个长 prompt 与后续 token 路径 | 完整 prompt 末位置 logits、缓存进度、输出接受时点 |
冷 / 热 prefix cache | 待测请求内容相同;用其他预热请求建立相同前缀 | 命中数与 Prefill 减少量;相同位置 logits |
eager / Graph 或 TP=1 / TP=2 | 权重、输入、dtype、生成设置与工作负载 | logits 误差、完成数量、稳定吞吐、显存和额外开销 |
尤其要固定同一条 token 路径。若两次独立随机采样很早就选出不同 token,后续 logits 不同可能只是输入已经不同;输出文本偶然相同,也不足以证明实现等价。完整上下文、权重、位置语义、dtype 与后端都应记录,数值差异需要结合具体计算解释。
15.2 数值等价:从单头手算到真实模型
实验 01 验证最小因果 Attention:改变最后位置,较早位置的输出应不变;追加新位置后,“全部重算的最后一行”应与“新 Q 读取历史 K/V”一致。另用 V 与 Q/K 不同的例子,确认权重由 Q/K 决定,而加权内容来自 V。手算、代码和旧运行输出统一保留在18.2。单层单头相等,不代表真实模型的多层缓存、RoPE 与 GQA 已全部验证。
实验 02 则固定真实模型的输入和后续 token,对比全量前向与 past_key_values 前向。prompt-len=64、共 8 次前向时,无缓存输入宽度为 64~71,有缓存为 64 加七个 1,合计分别处理 540 和 71 个输入位置。这是工作量的一个可核对指标,不是 FLOPs 或速度比;新位置的 Attention 仍然读取增长的上下文。
应同时比较最后位置 logits 的最大绝对误差、相对 L2 误差与 top-1 一致步数。浮点运算顺序、dtype 和 Attention 后端会带来误差,不宜给所有模型套同一个万能阈值。差异异常时,先确认缓存是否正确推进、每次对照是否重新初始化、positions 是否一致,再缩小到同一后端或更高精度判断。CPU 原理对照见18.4,真实模型路径见18.6。
15.3 机制正确:从状态重建本轮到底发生了什么
实验 03 围绕同一条组件链收集证据:Scheduler 选了谁,Runner 新算多少位置,KV 写在哪里,候选是否被接受,以及引用何时释放。下面三组输入用于回答不同问题,不应混用参数:
机制案例表:分别观察请求生命周期、分块 Prefill 和跨请求前缀复用。
模式 | 预期观察 | 重点解释 |
|---|---|---|
lifecycle | 64 与 96 token prompt 首轮 Prefill;随后每条请求每轮输入一个 token,较短输出的请求先结束 | num_tokens、num_cached_tokens 与结束后清零 |
chunk | 700 token 请求先计算 512;第二轮计算剩余 188 与另一请求的 32 | 中间 chunk 不追加输出;第二轮 query 总数为 220 |
prefix | 第二条 620 token 请求命中 512 个缓存位置,只 Prefill 剩余 108 | 同进程、完整块已登记、尾块独占 |
sampled 是候选,after_postprocess 才是接受与回写后的状态;total−prompt 才是已经增长的输出数。用请求 id 串联同一条记录,不要求物理块号永远等于纸面示例,因为分配顺序可能不同。应核对逻辑位置到物理槽位的关系,以及释放后引用是否正确,而不是只核对某几个整数。
这些日志能证明状态与地址按预期变化,不能单独证明 chunk/prefix 与基线 logits 等价。严格数值验证仍需在固定 token 路径上采集对应 logits。日志还会将 GPU 元数据取回 CPU、干扰执行时间,因此它不是性能测试。三个操作入口分别是18.7 生命周期、18.8 分块、18.10 前缀复用。
15.4 性能收益:同一负载,只改变待研究的机制
实验 04 分别对照 eager/Graph、TP=1/2。先固定请求数、输入输出长度、生成设置、设备与软件,再只改变一项机制。基线每轮 16 条请求,每条输入 128、输出 32,应有 512 个输出 token;固定输出长度避免 EOS 让某一组少做工作。每轮使用不同前缀,避免原本要测生成吞吐,却变成反复命中前缀的测试。
同一组对照使用相同 seed 序列,分别预热,并将初始化、编译、捕获与稳定生成计时分开。至少保留多轮结果及中位数,不只挑最快的一次。eager/Graph 操作见18.12,TP 对照见18.13。正式计时不启用机制追踪日志。
测 GPU 操作还要处理异步执行。 CPU 提交 CUDA 工作后,可能先继续运行;同一 stream 中按依赖执行,并不需要每个算子后都回 CPU 等待。只用 CPU 时钟包住一次异步调用,可能主要测到提交时间。要测 GPU 完成一段工作的耗时,应使用合适的同步边界或 CUDA Events;CPU 真正读取 GPU 结果时,也要等结果就绪。
本实验和仓库 bench.py 的离线吞吐口径是总输出 token 数除以整批墙钟时间,包含 Prefill 与 Decode;不能直接称为客户端 TTFT 或 TPOT。脚本的显存峰值仅采集 rank 0,多卡总量需要逐卡统计。比较 eager 与 Graph 也不能同时声称单独测出了 FlashAttention 的收益。
结论应落到“这个模型、长度分布、批量与 GPU 上”的实测变化,并同时报告初始化成本与显存。Graph 没有变快、小模型 TP 反而变慢,都是可以解释的结果;不能删掉不符合预期的轮次,也不能把单次吞吐推广到所有在线负载。
15.5 现有证据到哪一步:保留已测与未测的边界
实验包 02 实现缓存与无缓存的真实模型对照,03 侧重机制状态,04 侧重吞吐。要证明 chunk、prefix、Graph 或 TP 的严格数值一致性,还需按前面的控制条件增加 logits 对照。下方保留原始验证记录;正文修订不会自动把“未运行”升级成“已通过”。
历史验证记录:
验证记录(2026-09-06):实验 01 已在本机执行并通过因果 mask 与缓存等价性断言;四个脚本均通过 Python 语法检查。当前写作环境没有 PyTorch/NVIDIA GPU 运行环境,02–04 尚未做 GPU 实跑。本文不给出虚构的 GPU 跑分;下文注明“预期”的日志与状态值来自固定源码和输入的推演。
一份可复查的结论至少说明版本与设备、输入与控制变量、原始观察量、误差或计时口径,以及尚未覆盖的条件。记录里明确“静态推演”“CPU 已测”“GPU 未测”,比仅给出一张 PASS 截图更有意义。下一章用同样的分层方式定位边界与故障。
源码对照:仓库基准测试的统计口径。
16. 实现边界与排障:先判断出错层次
nano-vLLM 展示了推理引擎的核心组件,但不是完整生产服务。读到限制时,先问它属于哪个组件、影响哪一步,再区分通用机制与本版本的简化选择。出错时也沿这张组件地图定位,不把所有异常都归为“模型有问题”。
16.1 按组件看边界:简化发生在哪里
组件 | 固定提交的实现选择 | 影响与使用边界 |
|---|---|---|
入口与模型支持 | 主要支持 Qwen3 causal LM;没有完整 HTTP、流式接口和模型注册兼容层 | generate 是批量结果接口;生产前端能力不能直接假定存在 |
输入与采样参数 | 公开 temperature、max_tokens、ignore_eos;temperature>1e-10;逐请求校验有限 | 无公开 greedy/top-k/top-p;调用前检查非空输入、max_tokens≥1、总长度不越界,以及参数列表与请求数一致 |
Scheduler | Prefill 优先;同一轮不混合 Decode;只允许本批第一条 Prefill 请求切块 | 这是调度策略,不是 Transformer 的要求;不自动具备生产级公平性与服务时延保障 |
BlockManager | 块大小为 256 的倍数;初始前缀查找跳过最后逻辑块;抢占释放引用后重新 Prefill | 教学块大小 4 不能当配置;命中范围和恢复成本要按此实现推导,不能假定有 CPU KV swap |
ModelRunner / Graph | 全局 Context 配合顺序执行;Graph 桶不覆盖任意 max_num_seqs | 不能无隔离地并发复用本轮上下文;部分非对齐上限可能找不到覆盖图 |
多卡控制与通信 | 单机 TP;固定 localhost:2333;共享内存名 nanovllm、容量 1 MiB;TP=1 也初始化 NCCL | 多实例可能冲突;控制 payload 有容量边界,worker 故障恢复有限,不是多机调度系统 |
资源生命周期 | 依赖 atexit 清理,exit 非幂等 | 重复手动清理可能报错;扩展服务时需要明确资源归属与幂等退出 |
同样不应假定这个提交已提供 LoRA、量化、推测解码或多机调度。实际运行时还必须满足模型维度与 TP 分片的整除条件。以上边界限定于本文固定提交,不能推广成所有版本的 nano 或生产 vLLM 都如此。
16.2 按请求经过的组件定位故障
先确定失败发生在启动还是逐轮执行。启动阶段依次涉及依赖、配置与权重、预热、KV 池和 Graph;执行阶段再沿“输入 → 调度 → 缓存映射 → Runner → 模型/采样 → 回写”排查。能找到最早出现偏差的边界,通常比从最后一个报错反推整套系统更有效。
现象 | 先核对的数据或条件 | 对应组件 / 入口 |
|---|---|---|
CUDA 或扩展导入失败 | GPU 可见性、驱动、CUDA 版 torch、扩展二进制组合 | 环境与依赖;此时尚未到 Scheduler |
模型加载失败 | config、tokenizer、safetensors、架构与依赖 API 是否匹配 | Config / loader |
初始化 OOM | 失败在加载、warmup、KV pool 还是 Graph 捕获;其他进程占用 | Runner 初始化路径 |
本轮选择或输出数量异常 | waiting/running、cached/scheduled、参数列表长度、候选接受与停止条件 | Engine / Scheduler / postprocess |
前缀没有复用 | 是否同进程、完整块是否已算完、完整 token 前缀是否相同、块是否被覆盖 | BlockManager 的查找与 hash 登记 |
位置、shape 或 logits 异常 | 新输入范围、positions、请求边界、slot_mapping、有效上下文长度 | Runner 准备 / Attention / 模型层 |
Graph 失败,eager 正常 | 静态缓冲区、桶覆盖、填充槽位与长度元数据 | capture_cudagraph / run_model |
worker 等待或多卡卡住 | 端口、worker 最早异常、是否按同样顺序进入 collective | 控制消息 / ModelRunner / NCCL |
多卡或 Graph 更慢 | 负载与计时口径、启动成本、通信占比、填充计算和设备互联 | 先回到第 15 章的性能对照,不先判定计算错误 |
16.3 最小复现:一次只恢复一个复杂因素
先用一两条请求、短输出、单卡 eager 缩小范围,在固定 token 路径上确认输入、状态、地址与输出;再依次恢复长上下文、更多请求、Graph 或 TP。对照时保留原始失败证据,不同时修改模型、软件版本、批量和执行模式。
记录至少包含固定源码与模型版本、设备依赖、输入长度、失败阶段和关键状态;涉及性能则额外标明预热与统计口径,涉及数值则保留 logits 误差。环境重装、实验命令和资源清理操作见第 18 章。确认进程与资源归属后再处理,不能为了排障结束别人的任务。
16.4 显存排查:先对齐观测口径,再判断是否泄漏
观察工具时区分三个口径:memory_allocated 表示 PyTorch 张量当前占用;memory_reserved 是分配器保留的空间,通常包含 allocated;nvidia-smi 的进程或设备占用还可能包含 CUDA 上下文、其他库与其他进程。allocated 和 reserved 不能直接相加。
一个预先创建的 KV 池即使有很多未分给请求的槽位,也可能已经计入 allocated。请求结束后只把块还给池,nvidia-smi 不下降是常见现象;需结合活跃请求、可用块和池大小判断,不能仅凭显存不下降认定泄漏。empty_cache 也不能释放仍被引用的权重或 KV 池。
17. 全篇回顾与阅读导航:从这份实现走向生产引擎
到这里,可以把模型数学与引擎工程合成同一幅图:模型决定如何从上下文算出下一个 token;引擎决定本轮算哪些位置、缓存放在哪里,以及如何组织设备执行。先用一轮请求收束全篇,再按问题查索引或继续深入。
17.1 用一次 Decode,把所有组件接回去
主例第二轮结束后,A、B 各接受首个输出,Sequence 总长为 701/33,已缓存为 700/32。下一轮 Scheduler 为两条请求各安排一个新位置,BlockManager 确保尾部槽位可用;Runner 取两个 last_token,准备 positions=[700,32]、context_lens=[701,33] 与相应块表和写地址。
GPU 先查 Embedding 得到两行向量。每一层为这两行生成 Q/K/V,写入新 K/V,用新 Q 读取该层历史与当前 K/V,完成加权、输出投影、残差与 MLP。旧位置不再重跑完整前向,新位置却一层也不能跳过。最终 Norm、LM Head 与 Sampler 给出下一批候选;Scheduler 接受输出、推进状态,并在结束时释放引用。
如果使用 FlashAttention,变化的是中间结果如何分块累积与搬运;如果使用 CUDA Graph,变化的是怎样提交这些 GPU 操作;如果使用 TP,变化的是哪些特征与部分和由哪张卡计算。它们都围绕同一条请求链工作。理解推理引擎的核心,就是同时跟住新计算的位置、可复用的历史、数据地址和执行依赖。
17.2 按问题查回正文
现在想弄清什么 | 关键关系 | 入口 |
|---|---|---|
token、向量、head 与维度 | 离散 id → 特征向量 → Q/K/V → 多头输出 | |
Prefill / Decode 与完整模型 | 本轮输入哪些位置,最后怎样选下一个 id | |
为什么只缓存各层 K/V | 因果性、旧结果不变、新位置的直接依赖 | |
显存、带宽与时间指标 | 权重、KV、中间激活和计时范围 | |
为什么有引擎,组件怎样协作 | 入口、状态、调度、缓存、执行与回写 | |
这轮选谁、各算多少 | Sequence 三种长度与 Scheduler 工作单 | |
KV 放在哪里、何时可以共享 | 逻辑位置、物理块、引用与完整前缀 | |
工作单怎样变成张量与候选 | 输入准备 → 模型前向 → LM Head → 采样 | |
GPU 怎样减少开销、怎样多卡执行 | 中间搬运、工作提交、特征分片与部分和 | |
怎样证明、怎样定位、怎样运行 | 证据方法、组件边界、统一操作路线 |
Softmax、temperature、LayerNorm、RMSNorm、SiLU/SwiGLU、GQA 与 RoPE 的基础数值解释统一在第 19 章附录,19.8可查形状与计算轴。不必每遇到一个术语就重新走整篇主线。
17.3 从真实瓶颈选择进阶方向
遇到的问题 | 继续研究什么 | 先确认的代价或条件 |
|---|---|---|
权重或长上下文 KV 放不下 | 权重量化、KV 精度、并行与缓存管理 | 先区分哪类数据占显存,再检查数值影响与额外计算 |
单请求 Decode 间隔仍高 | 执行优化、推测解码及验证机制 | 判断带宽、提交还是计算受限;推测接受率与验证成本影响收益 |
长短请求互相拖慢 | 调度、公平性、分块策略与 Prefill/Decode 部署组织 | 明确到达模式、首 token 与输出间隔目标,不只看离线吞吐 |
多卡通信成为主要开销 | 并行策略、拓扑与通信组织 | 定位具体 collective,并比较计算减少量和通信增加量 |
从教学引擎走向在线服务 | 服务前端、准入与取消、隔离、监控和故障恢复 | 功能完备、可靠性和性能是不同目标,不能只验证能生成文本 |
回到生产 vLLM 时,仍用“请求 → 调度 → 缓存 → Runner → 模型 → 采样”定位组件,再阅读当前版本为多模型、多设备和服务可靠性增加的工程层。先定义负载与目标指标,再选技术;不要因名称更新,就忽略它最终改变了哪段计算或哪类数据。
17.4 资料与版本:带着问题继续读
资料入口 | 带着什么问题阅读 |
|---|---|
从向量到整层与词表输出;注意 GPT-2 结构不能直接套到 Qwen3 | |
投影、匹配和加权 V 的直觉 | |
因果性、逐层缓存与更新方式 | |
计算、带宽与容量估算 | |
动态 KV 分配和跨请求共享 | |
对照本文实际分析的组件与实现选择 | |
核对真实维度、层数、GQA 与参数文件 | |
区分当前生产接口、设计与 nano 的简化版本 | |
原文参考文章;结合固定源码核对版本差异 |
本文 nano 源码固定为 bb823b3e06983d71485a8e1f23715ebd87d98ef8,模型 revision 为 c1899de289a04d12100db370d81485cdf75e47ca。入门资料沿用 2026-09-09 的核对记录,vLLM stable 介绍沿用 2026-09-06 的记录;这些可变链接以后可能更新,实际运行应记录安装版本。手算与源码推演不是 GPU 实测,验证状态以第 15 章记录为准。
阅读外部公式时也要核对约定:有些教程把 token 向量放在列中,本文按行书写,使用 QKᵀ 并沿可读位置归一化。只要能把位置、特征、权重和数据流一一对应起来,就能将这套理解迁移到其他模型与引擎。
18. 动手学习路线:半天入门与两周深入验证
本章提供两条共用同一套环境的动手路线。已经理解 Attention、Prefill/Decode 和 KV Cache 时,先做 18.0.1–18.0.6 的三个小实验,用半天到一天读通请求、调度和分页缓存的主流程;环境下载与编译另计。需要进一步验证原理、性能和多卡执行时,再选择后面的 Day 1–14,每天约 60–90 分钟。两周计划不是读懂这份精简项目的前置门槛。第 5–6 章用于建立组件全貌,第 15 章用于核对实验的证据范围。
已有环境时,先核对固定源码、模型 revision、env.sh 目录和实际导入路径,再进入快速实验。没有环境时,先按 18.0.7 准备目录、附件与工具链,再跳到 18.6 完成 GPU 环境和短生成,不必按日期从 Day 1 顺序做起。vLLM 接口体验是独立选做,不能把它的依赖装进 nano 环境。下方原实验附件及其历史验证记录保持不变;新加的 quick_walk.py 在 18.0.2 给出完整代码,不在旧附件中。
配套附件:nanovllm-e2e-guide-v2.zip。包含原两周深入路线的离线 README.md、原有 4 个实验脚本、分词/原理练习、源码导航、日志验收器、环境检查与报告生成工具;Day 7 的 trace 脚本包含可选形状 hook。快速路线复用其中的环境文件,新入口 quick_walk.py 按 18.0.2 单独创建。附件里的旧章节导读与线上内容不一致时,以线上导航为准;无需再下载第 15 章的旧实验包。
验证边界:文中的“预期值”是验收目标。CPU 原理练习和真实 Scheduler/BlockManager 的 CPU 测试已在写作机器验证;GPU 安装、模型 forward、CUDA Graph 与 TP 命令已对照固定源码检查,但写作机器没有 NVIDIA GPU,未把这些结果标成实测。附件 VALIDATION.md 记录具体覆盖范围。
18.0 快速入门与通用环境准备
18.0.1 先选路线:读通主流程,不必先完成两周计划
如果已经理解 token、Attention、Prefill/Decode 和 KV Cache,可以直接从源码开始。第一轮目标是用半天到一天跟通一次请求,而不是立即掌握 FlashAttention 内核、CUDA Graph 和多卡通信。这里的时间用于已有环境下的阅读与调试,模型下载、依赖编译和设备准备另计。
项目自带的 example.py 用于运行生成示例,bench.py 用于测批量吞吐;本章固定版本没有配套教学 Notebook 或专门的 tests/ 目录。下面的 A/B/C 是本文设计的学习实验,不是项目官方测试。代码量小使主链路容易读通,但一行 FlashAttention 或分布式调用背后的外部实现并不计入这份源码的行数。
顺序与时间参考 | 这一段只解决什么问题 | 完成标准 |
|---|---|---|
准备:20–30 分钟 | 确认固定版本、单卡环境,读 example.py 与 LLMEngine.generate/step | 知道请求在哪里入队,下一轮在哪里开始 |
A:45–60 分钟 | 3 个输入位置,生成 3 个 token,跟一次 Prefill 和两次 Decode | 能区分 token 已生成、KV 已计算、缓存已释放 |
B:45–60 分钟 | 两条请求分别生成 2/5 个 token | 能解释批次如何从两条变成一条 |
C:45–60 分钟 | 256 个输入位置,观察下一轮跨 KV 块边界 | 能算出新 token 的物理槽位 |
回顾:30 分钟 | 把输入、调度、缓存、模型计算和状态回写串起来 | 对照日志与断点说明,而不是只记住类名 |
主线从 example.py 进入 LLMEngine.generate;每轮 step 依次调用 Scheduler.schedule、ModelRunner.run、Scheduler.postprocess。Sequence 保存请求状态,BlockManager 被调度器用于管理 KV 块;Runner 再进入 Qwen3 和 Attention。第 6 章已有完整组件图,读到不熟悉的细节再查第 7—14 章,不要求先通读全部正文。
18.0.2 准备一份最小调试入口
先确认环境。已有本章环境时,直接复用 env.sh、.venv-nano、NANO_DIR 和 MODEL_DIR。尚未准备时,先执行本节后面的“通用环境准备”,再跳到 18.6 完成 GPU 环境、权重下载及 smoke.py;不必先做 Day 1–5,也不必先跑当天的缓存性能对照。没有兼容 NVIDIA GPU 时,只能做源码推演或已有 CPU 练习,不能把本调试入口当作 CPU 推理程序。
源码固定为 bb823b3e06983d71485a8e1f23715ebd87d98ef8;模型固定为 Qwen/Qwen3-0.6B revision c1899de289a04d12100db370d81485cdf75e47ca。安装应使用本章的 editable 模式,确保断点所在源码就是实际导入的源码。不要为了跟某篇旧文章而修改 rope_scaling,也不要把 vLLM 安装进这个 nano 环境。
创建入口。在 NANO_DIR 仓库根目录新建 quick_walk.py,粘贴下面完整代码。它只组织输入并调用原来的 generate,不修改引擎源码;旧版附件不含这个新增文件。为固定输入长度,实验直接传入合法 token ID,跳过文本分词与 chat template。它们不用于评估回答质量;temperature 使具体生成 ID 可能变化,但 ignore_eos=True 固定了输出轮数。
"""Minimal reading harness for nano-vLLM commit bb823b3 (requires a CUDA environment)."""
import argparse
import os
def cases(name):
if name == "A":
return [[100, 200, 300]], [3]
if name == "B":
return [[100, 200, 300], [400, 500, 600, 700]], [2, 5]
return [list(range(100, 356))], [3]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default=os.environ.get("MODEL_DIR"))
parser.add_argument("--case", choices=["A", "B", "C"], required=True)
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()
if not args.model:
parser.error("Use --model or set MODEL_DIR to the local model directory")
from nanovllm import LLM, SamplingParams
import nanovllm.engine.llm_engine as engine_module
llm = LLM(
args.model,
enforce_eager=True,
tensor_parallel_size=1,
max_model_len=1024,
max_num_seqs=2,
max_num_batched_tokens=512,
kvcache_block_size=256,
gpu_memory_utilization=0.5,
)
prompts, limits = cases(args.case)
sampling = [
SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=n)
for n in limits
]
print("Loaded source:", engine_module.__file__)
print("Prompt lengths:", [len(p) for p in prompts], "Output limits:", limits)
# Break after initialization, so model warmup does not pollute the walkthrough.
if args.debug:
breakpoint()
outputs = llm.generate(prompts, sampling, use_tqdm=False)
for label, output, limit in zip(["R1", "R2"], outputs, limits):
print(label, "output_count=", len(output["token_ids"]),
"token_ids=", output["token_ids"])
assert len(output["token_ids"]) == limit
if __name__ == "__main__":
main()三个案例都用单卡、256-token 块、512-token Prefill 预算和短上下文,分别启动独立进程;不要并发运行。enforce_eager=True 关闭 CUDA Graph 回放,但仍使用 KV Cache、FlashAttention 等实现,不等于关闭全部优化。程序在初始化完成后才进入断点,避免把 warmup 的假请求误认成实验请求;不要再次手动调用 llm.exit(),本版本会在进程退出时清理。
source "$HOME/llm-inference-lab/env.sh"
source "$LAB_ROOT/.venv-nano/bin/activate"
cd "$NANO_DIR"
CUDA_VISIBLE_DEVICES=0 python quick_walk.py --case A --model "$MODEL_DIR"
CUDA_VISIBLE_DEVICES=0 python quick_walk.py --case A --model "$MODEL_DIR" --debug第一次不加 --debug,确认 output_count=3;第二次会在初始化后进入 Python 自带的 pdb。交互调试时不要把输入输出接到 tee;断点与张量打印会影响计时,不用这次运行测性能。也可以用 IDE,在相同文件和语句处设置断点,并以 NANO_DIR 为工作目录运行相同参数。
18.0.3 实验 A:一次 Prefill、两次 Decode
步骤 1|先写下预测。输入 ID 列表有 3 个元素,依次记为 t1、t2、t3。预期 Prefill 选出 t4,Decode 1 输入 t4 并选出 t5,Decode 2 输入 t5 并选出 t6。输出是 3 个 token,但只有后两轮属于 Decode;t6 被选出后达到上限,不会再输入模型计算自身 KV。
步骤 2|设置两个断点。在初始 pdb 提示符输入以下命令。行号仅对应固定提交;52 行位于 schedule 已返回、模型尚未运行的位置,54 行位于 postprocess 已返回的位置。不要编辑引擎文件后继续照抄旧行号。
b nanovllm/engine/llm_engine.py:52
b nanovllm/engine/llm_engine.py:54
c在 52 行查看下面两项。为避免 warmup 导致编号偏移,认准同一个 seq_id,不假定它从 0 开始。各字段依次为请求 ID、总 token 数、已缓存位置数、本轮待算位置数、物理块表。
p is_prefill
p [(s.seq_id, s.num_tokens, s.num_cached_tokens, s.num_scheduled_tokens, list(s.block_table)) for s in seqs]第一轮应为 True、total=3、cached=0、scheduled=3。用 n 执行模型调用,停到 53 行时查看 token_ids:候选已经产生,但尚未执行 postprocess,所以 Sequence 的 total 仍为 3。再执行 n,到 54 行查看回写后的状态。这是区分“算出候选”和“把候选追加进请求”的最直接方法。
p [(s.seq_id, s.num_tokens, s.num_cached_tokens, s.num_completion_tokens, s.status.name, list(s.block_table)) for s in seqs]轮次 | 模型本轮处理的位置 | postprocess 后应看到什么 |
|---|---|---|
1:Prefill | t1、t2、t3;positions=[0,1,2] | total=4,cached=3,输出数=1,RUNNING |
2:Decode | t4;positions=[3] | total=5,cached=4,输出数=2,RUNNING |
3:Decode | t5;positions=[4] | total=6,输出数=3,FINISHED;缓存已释放,cached=0、blocks=[] |
最后一轮模型实际已算到 t5 的 KV,但随后请求完成,deallocate 清掉该请求的缓存计数和块表。因此 54 行看到 cached=0,并不是模型没计算 KV,也不是把 GPU 显存逐元素清零。未结束的普通 Decode 请求通常满足 total=cached+1;完成后的清理状态不适用这个关系。
步骤 3|再看 Runner 的实际输入。重跑 A,在初始 pdb 只设置下面这个断点。216 行位于 prepare_prefill/prepare_decode 返回之后,能够同时看到输入张量和已准备好的 Attention 上下文。随后按 c 在三轮之间继续。
b nanovllm/engine/model_runner.py:216
c
p input_ids.shape, positions.tolist()
from nanovllm.utils.context import get_context
p get_context().slot_mapping.tolist()
p None if get_context().context_lens is None else get_context().context_lens.tolist()A 的 input_ids 形状依次为 [3]、[1]、[1];Decode 的 context_lens 依次为 [4]、[5],包括当前即将写入的 KV 位置。进入模型后,Qwen3-0.6B 的 Embedding 输出分别是 [3,1024]、[1,1024]、[1,1024];各层 Q 为 [T,16,128],K/V 各为 [T,8,128]。要查看 Attention,重点是新 Q 如何读该层历史和当前 K/V,再做加权求和,不是重新处理全部历史 Embedding。
验收问题:为什么生成 3 个 token 只有 2 轮 Decode?为什么刚采样出的 token 比 KV Cache 多一个位置?为什么最后一轮 cached 又变成 0?能用上述断点说明这三件事,就进入 B。
18.0.4 实验 B:两条请求,为什么会从同批执行变成只剩一条
步骤 1|运行案例 B。输入长度分别为 3/4,输出限制分别为 2/5。启动一个新进程,沿用实验 A 的 52/54 行断点;预热之外的第一轮应该只有这两条请求,且没有前缀缓存命中。
source "$HOME/llm-inference-lab/env.sh"
source "$LAB_ROOT/.venv-nano/bin/activate"
cd "$NANO_DIR"
CUDA_VISIBLE_DEVICES=0 python quick_walk.py --case B --model "$MODEL_DIR" --debug步骤 2|看本轮批次与下一轮队列。52 行的 seqs 是本轮工作集合;54 行的 seqs 仍包含刚完成的请求,不等于下一轮仍活跃的请求。要判断谁还活跃,另看 scheduler.running 与 waiting:
p [s.seq_id for s in self.scheduler.running]
p [s.seq_id for s in self.scheduler.waiting]
p len(self.scheduler.block_manager.free_block_ids)轮次 | 本轮执行 | 回写后 |
|---|---|---|
1:Prefill | R1 的 3 个位置 + R2 的 4 个位置,共 7 个新位置 | 各生成 1 个 token;total/cached 分别为 4/3、5/4 |
2:Decode | R1、R2 各算 1 个新位置 | R1 输出达到 2,退出 running 并释放其块;R2 输出为 2,继续 |
3、4:Decode | 只有 R2,每轮算 1 个位置 | R2 输出数依次为 3、4 |
5:Decode | 只有 R2,算 1 个位置 | R2 输出达到 5;waiting 与 running 都为空 |
这个案例里两条请求各占一个物理 KV 块且不共享前缀。第二轮结束时空闲块数应增加 1,最后再增加 1;不要假定具体块 ID。完成请求的 deallocate 释放的是引用,不是承诺立即擦除 KV 数值或所有前缀索引。
步骤 3|对比参数。临时把 quick_walk.py 中的 max_num_seqs 改成 1,再用新进程运行 B,先预测每轮选择谁。注意此版本优先选择可调度的 Prefill,不要想当然地认为 R1 一定全部生成完才轮到 R2。看完恢复为 2。这个变体是为了理解调度,不用它得出性能结论。
验收问题:为什么不是 2+5=7 次独立模型运行,而是共 5 轮?为什么 R1 已完成后,当前 seqs 中还能看到它?本实验展示逐轮移除完成请求,但 generate 在这一版本仍是同步返回整个批次结果,不代表已经提供可动态接收网络请求的在线服务。
18.0.5 实验 C:第 257 个已知 token 的 KV 究竟写到哪里
步骤 1|用正好一块的 prompt。案例 C 输入 list(range(100,356)),恰好 256 个合法 token ID,输出上限为 3。块大小仍为 256;不要为了缩短示例把真实 GPU 块大小改成 4,本版本要求它是 256 的倍数。使用新进程避免其他请求或前缀缓存影响观察。
source "$HOME/llm-inference-lab/env.sh"
source "$LAB_ROOT/.venv-nano/bin/activate"
cd "$NANO_DIR"
CUDA_VISIBLE_DEVICES=0 python quick_walk.py --case C --model "$MODEL_DIR" --debug步骤 2|先在 A 的两个断点看整体。Prefill 处理 positions=0…255,写完一块 KV,并选出第 257 个 token;此时 total=257、cached=256、block_table 仍只有一个物理块。Sequence.num_blocks 按已知 token 长度计算会是 2,但 len(block_table) 暂时是 1:最后一个 token 还没有进入模型,所以尚不需要为其 KV 分配第二块。
步骤 3|再跟一次 may_append。重跑 C,在初始 pdb 设置下列两个断点。107 行进入 may_append 的条件判断,216 行在 Runner 已构造输入和槽位后停下。
b nanovllm/engine/block_manager.py:107
b nanovllm/engine/model_runner.py:216
c第一次先在 Runner 中看到 Prefill;继续后,下一轮会进入 may_append。在这里查看 len(seq)、seq.num_cached_tokens、seq.num_blocks 和 seq.block_table,应该是 257、256、2 和一个物理块。因为 257 % 256 == 1,执行分配语句后,块表变成两个物理块。然后继续到 Runner 216 行,查看 positions、context_lens、slot_mapping。
p len(seq), seq.num_cached_tokens, seq.num_blocks, list(seq.block_table)from nanovllm.utils.context import get_context
p positions.tolist()
p get_context().context_lens.tolist()
p get_context().slot_mapping.tolist()时刻 | 位置与物理块 | 应得到的映射/状态 |
|---|---|---|
Prefill | positions=0…255;块表 [p0] | 槽位范围为 p0×256 … p0×256+255;回写后 total=257、cached=256 |
Decode 1 | position=256;先追加物理块 p1 | context_lens=[257];slot_mapping=[p1×256];回写后 total=258、cached=257 |
Decode 2 | position=257;复用已有尾块 p1 | context_lens=[258];slot_mapping=[p1×256+1];选出第三个输出后结束并释放两块 |
统一公式为 slot = block_table[position // 256] × 256 + position % 256。p0、p1 是运行时观察到的物理块编号,不要求相邻,更不能直接把逻辑位置 256 当作物理地址 256。真实 Attention 内,store_kvcache 再把本层新 K/V 写入该槽位;同一个块表映射在各层对应各层自己的缓存切片。
验收问题:为什么 Prefill 结束时已知 token 有 257 个,却只分配一块 KV?为什么第二轮必须分配新块,第三轮不用?为什么最后选出的第 259 个 token 没有自己的 KV?能用块表、position 和 slot 三项数值回答,就已读通分页缓存的主路径。
18.0.6 入门验收与进阶分流
满足这些条件后即可结束第一轮,不必继续堆实验。下一步按问题选择:缓存前后的计算量对照看 18.4/18.6;更完整的生命周期日志看 18.7;分块 Prefill 与抢占看 18.8;共享前缀看 18.10;CUDA Graph 和 TP 看 18.12/18.13。原有 Day 1–14 是进一步验证的路线,不是阅读项目的前置作业。
验证边界:此处预期状态依据固定源码推导;新增 quick_walk.py 已做语法和命令行检查,断点行号已对照固定提交。写作机器没有 NVIDIA GPU,本次未实际执行这些 GPU 实验,不能把表中预期值当作实测结果。运行时遇到 Traceback、块数或轮次不符,应先保存现场,核对版本、导入路径、参数和资源,而不是修改日志去匹配表格。
18.0.7 通用环境准备:没有现成环境时从这里开始
下面保留原有统一目录、版本、设备和附件约定。快速路线与两周路线共用这一套环境,不需要各自安装一次。
步骤 1|准备机器并下载附件。完整路线使用 Linux x86_64、Ampere/Ada/Hopper NVIDIA GPU;推荐独占 A100/A800/H100/H800,显存至少 24 GB、主机内存至少 32 GB、磁盘至少 30 GB。驱动建议 550.54.15 或更新。Day 1–5 及标明“CPU”的推演可在普通 Linux/macOS 完成;没有 GPU 时,GPU 验收保持“未运行”。本指南不包含 GPU 资源申请、驱动或 Docker 的管理员安装。
在飞书点击附件下载。若浏览器在本机、实验在远端,先把文件传过去;下面的 user@gpu-host 要换成自己的登录地址。后续主线路径统一为执行环境中的 $HOME/llm-inference-lab。
# 仅远端实验需要:在下载附件的电脑运行
scp "$HOME/Downloads/nanovllm-e2e-guide-v2.zip" user@gpu-host:~/
ssh user@gpu-host
# 在实验机运行;本机实验时将 ZIP 改为实际下载路径
mkdir -p "$HOME/llm-inference-lab"
unzip "$HOME/nanovllm-e2e-guide-v2.zip" -d "$HOME/llm-inference-lab"
cd "$HOME/llm-inference-lab"
ls README.md env.sh e2e labs步骤 2|确定 GPU 工具链。主线使用 CUDA 12.4.1 devel 环境。已具备 Linux + CUDA 12.4 nvcc + C++ 编译器时直接跳到步骤 3;否则,在已有 Docker 和 NVIDIA Container Toolkit 的 GPU 宿主机执行下面命令。后续步骤都在这个容器中运行。容器不随退出删除,宿主机的附件目录会保存所有输出。
# 在 GPU 宿主机,先确认驱动和 GPU 可见
nvidia-smi
docker run --gpus all --ipc=host --name nanovllm-e2e -it -v "$HOME/llm-inference-lab:/root/llm-inference-lab" nvidia/cuda:12.4.1-devel-ubuntu22.04 bash
# 以下在容器内执行(容器默认 root)
apt-get update
apt-get install -y ca-certificates curl git unzip build-essential
nvidia-smi
nvcc --version
# 以后退出后重进:在宿主机执行 docker start -ai nanovllm-e2e
# 容器仍在运行时另开终端:docker exec -it nanovllm-e2e bash步骤 3|安装独立 Python 3.11.11 与 CPU 环境。下面固定 uv 版本,仅安装到实验目录并创建虚拟环境,不替换系统 Python。使用容器时在容器内执行;纯 CPU 路线在自己的 Linux/macOS 终端执行。不要把在 macOS 创建的虚拟环境直接搬进 Linux。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
mkdir -p tools
case "$(uname -s)-$(uname -m)" in
Linux-x86_64) UV_BUILD=uv-x86_64-unknown-linux-gnu ;;
Darwin-arm64) UV_BUILD=uv-aarch64-apple-darwin ;;
Darwin-x86_64) UV_BUILD=uv-x86_64-apple-darwin ;;
*) echo "本指南未覆盖此平台"; return 1 2>/dev/null || exit 1 ;;
esac
curl -fL --retry 3 "https://github.com/astral-sh/uv/releases/download/0.6.9/$UV_BUILD.tar.gz" -o tools/uv.tar.gz
tar -xzf tools/uv.tar.gz -C tools
"tools/$UV_BUILD/uv" python install 3.11.11
"tools/$UV_BUILD/uv" venv --seed --python 3.11.11 .venv-cpu
source .venv-cpu/bin/activate
python -m pip install -r requirements-cpu.txt
python --version
python -m pip check
python e2e/final_report.py --root "$LAB_ROOT" --init-notes看到 Python 3.11.11、pip check 没有依赖冲突,并生成 notes/day01.md–day14.md 后再继续。requirements-cpu.txt 固定 NumPy 1.26.4、Transformers 4.57.1、tokenizers 0.22.1、huggingface-hub 0.36.0、safetensors 0.6.2、xxhash 3.5.0 与 Jinja2 3.1.6。后续 GPU 环境另建,避免和 vLLM 的依赖互相覆盖。
步骤 4|下载固定源码和分词配置。本章始终使用 nano-vLLM 提交 bb823b3e06983d71485a8e1f23715ebd87d98ef8(0.2.0)与 Qwen/Qwen3-0.6B revision c1899de289a04d12100db370d81485cdf75e47ca。先只下载分词/模型配置;第 6 天再下载权重。下载失败重跑同一命令,不要换到未固定的 main 分支。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
mkdir -p "$NANO_DIR"
curl -fL --retry 3 "https://codeload.github.com/GeeeekExplorer/nano-vllm/tar.gz/$NANO_COMMIT" -o nano-vllm-source.tar.gz
tar -xzf nano-vllm-source.tar.gz -C "$NANO_DIR" --strip-components=1
python e2e/source_walk.py verify --repo "$NANO_DIR" | tee runs/source-version.txt
python e2e/download_model.py --model "$MODEL_DIR" | tee runs/tokenizer-download.log
python -m pip freeze > runs/cpu-environment.txt执行与记录约定:每个日任务都从 source env.sh、进入目录、激活环境开始,可以关闭终端后重新接上。env.sh 启用 pipefail;带 tee 的命令要正常结束且没有 Traceback 才算成功。runs/ 保存实际输出,notes/dayNN.md 保存你的预测和解释;原样保留“未运行”项目。不要用带 trace 的程序测性能。nano 固定使用 localhost:2333 和共享内存名 nanovllm,同一主机不要并发启动多个 nano 实验;TP=2 本身属于一个实验。
附件脚本有 --help。principles.py 是小型原理演示;cpu_scheduler.py 直接加载固定源码里的 Sequence/Scheduler/BlockManager,并提供合成的采样 token,不运行模型或 GPU;analyze.py 读取真正的运行日志,输出 JSON 验收与 TSV 状态表。PASS 仅覆盖脚本声明的断言,不能代替对性能和数值误差的解释。
历史实验包(仅用于追溯)。下方 nanovllm-learning-labs.zip 为旧版,保留原附件便于复查早期记录。首次运行只使用本章的 nanovllm-e2e-guide-v2.zip;新版已包含四个实验与验收工具,不需要同时解压两个包。
18.1 Day 1:从文本、token 到采样概率
今天读第 1.1–1.2 节和第 2.2–2.3 节:从 token id、向量到 logits,再走三轮生成。先回答“一个 token 一定是一个汉字吗”。练习使用三 token 中文示例;完整无缓存对照见第 3.4–3.5 节。
步骤 1|运行真实分词与教学采样。脚本对一段固定中英文文本做 encode/decode,对人为给定的 [2,1,0] 做 softmax,并打印关闭 thinking 的 Qwen3 chat template。此处尚未加载模型权重,三个分数不是 Qwen3 的预测。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/principles.py tokenizer --model "$MODEL_DIR" --out runs/day01-tokenizer.json
cat runs/day01-tokenizer.json步骤 2|对照输出写记录。打开 notes/day01.md,抄下真实 ids、token_pieces 和 decoded;看到单个中文 token piece 像乱码时,对照完整 decode,token 的内部字节表示不一定是独立可读字符。记录 chat_template_text 如何包含角色边界与生成提示。
预期与验收:decoded 与原文逐字相同;probabilities 约为 [0.66524,0.24473,0.09003],greedy_index=0;随机采样频率接近概率但不要求完全相等。输出 check=PASS。补写两句:权重推理时通常固定,上下文/缓存随生成增长;greedy 每次取最大值,采样按分布抽取。卡住时:缺 tokenizer.json 就重跑 18.0 的分词下载;不要用不存在的模型路径继续。
18.2 Day 2:手算并验证一次因果 Attention
今天读第 1.3–1.5 节与第 2.4 节:从单头到多头,再理解因果 mask。下面实验采用 Q=K=V=X;本节末尾另保留不同 V 的手算对照,第 15.2 节解释这些对照能证明什么。先预测修改未来位置是否影响过去位置,再运行验证。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python labs/01_attention_numpy.py | tee runs/day02-original.log
python e2e/principles.py attention --out runs/day02-attention.json
sed -n '1,150p' e2e/principles.py步骤 2|手算最后一行。X=[[1,0],[0,1],[1,1]]。写出 Q/K/V=[3,2]、score=[3,3]、output=[3,2];最后一行 score=[1/√2,1/√2,2/√2],softmax 后权重约 [0.2483,0.2483,0.5035],加权输出约 [0.7517,0.7517]。把过程写入 notes/day02.md。
步骤 3|验证 mask 的作用。同一脚本将最后一个位置改为 [9,-5],分别保留和去掉 causal mask,自动对比前两个位置。验收:masked_earlier_delta=0,unmasked_earlier_delta>0(本输入约 8.15455),各行权重和为 1,check=PASS。能解释为什么不允许较早位置读取未来信息。卡住时:ModuleNotFoundError: numpy 说明没有激活 .venv-cpu 或安装未完成。
补充对照与历史输出。以下保留原实验 01 的运行记录,以及 V 与 Q/K 不同的独立手算代码。前一份输出来自 Q=K=V 的旧实验;后一段代码是另一个局部例子,两者不要当作同一次运行的输入输出。
实验 01 使用三位置、二维的独立例子 Q=K=V=X,验证因果 mask 和缓存等价性:修改最后位置,检查较早输出不变;追加一个向量,比较全量重算与新 Q 读取旧 K/V。它不使用第一章 4096 维教学模型的实际参数。
attention weights:
[[1. 0. 0. ]
[0.3302 0.6698 0. ]
[0.2483 0.2483 0.5035]]
last output: [0.7517 0.7517]
future change leaves earlier positions unchanged: True
cached last output: [0.6304 0.8374]
max error: 0.0
query positions: full=4, cached=1; KV positions=4先预测去掉 causal mask 后哪个断言会失败,再自行修改观察。验收标准是能解释相等的原因,而不仅是看到 max error=0。这个教学例子只验证单层单头计算思想,不能替代真实模型各层缓存、RoPE、GQA 与数值精度的验证。
补充手算:让 V 与 Q/K 不同。 仍用三个位置、每头二维,直接给定 Q=K=[[1,0],[0,1],[1,1]],V=[[1,0],[0,2],[1,2]]。q3=[1,1] 的三个缩放分数为 [0.7071,0.7071,1.4142],softmax 约为 [0.248255,0.248255,0.503490]。
按这些权重加权 V,输出约为 [0.751745,1.503490]。这里 Q/K 决定“如何匹配”,V 决定“混合什么数”;V 刻意取成不同数值,便于区分职责。下面只验证这段局部 Attention,不是一次真实语言模型生成。
import numpy as np
q = np.array([[1., 0.], [0., 1.], [1., 1.]])
k = q.copy()
v = np.array([[1., 0.], [0., 2.], [1., 2.]])
scores = q @ k.T / np.sqrt(2.)
future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
scores = np.where(future, -np.inf, scores)
e = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = e / e.sum(axis=-1, keepdims=True)
print(weights[-1]) # [0.24825508 0.24825508 0.50348984]
print((weights @ v)[-1]) # [0.75174492 1.50348984]18.3 Day 3:逐轮分清“已知 token”和“已计算 KV”
今天回看第 2.3 节的三轮时序,再读第 3.1–3.4 节。用第 4 章开头的中文示例,预测输入 t1~t3、输出 t4~t6 需要几轮 forward,并分别记录每轮已知 token 数与已计算 KV 数。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/principles.py timeline --out runs/day03-timeline.json
cat runs/day03-timeline.json步骤 2|逐项核对三轮。第 1 轮 Prefill 输入 t1,t2,t3,KV 覆盖 t1–t3,预测 t4;第 2 轮 Decode 输入 t4,KV 到 t4,预测 t5;第 3 轮输入 t5,KV 到 t5,预测 t6。把 rounds 整理成“forward 输入 / KV 覆盖范围 / 刚输出 token”三列表。
验收:每轮 known_token_count=kv_count+1;第一轮输入宽度 3,后两轮各 1;check=PASS。能回答“输出 t4 时为什么还没有 t4 的 KV”。这是符号推演,不是模型实测。卡住时:若把新输出 token 立刻算进 KV,回到生成循环,区分 forward 与采样这两个动作。
18.4 Day 4:用两个路径验证 KV Cache 省了哪些计算
今天读第 3.1 节的因果性与第 3.4–3.5 节的缓存对照。先在 CPU 上验证等价关系与位置计数,第 6 天再用真实模型比较 logits。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/principles.py cache --out runs/day04-cache.json
python labs/01_attention_numpy.py | tee runs/day04-numpy.log步骤 2|核对输入位置数。S=64、生成 G=8 时,不缓存路径的每轮输入是 64–71,求和 540;缓存路径是 [64,1,1,1,1,1,1,1],求和 71。脚本用同一固定输入比较“完整重算最后位置”与“只算新 query 并读取历史 K/V”。
验收:max_error=0 或浮点舍入量级,full_last 与 cached_last 一致,位置总数为 540/71,check=PASS。在 notes/day04.md 写清:历史 K/V 仍被读取;历史 Q 通常不需缓存;540÷71 不是 GPU 提速倍数。卡住时:不要比较两次独立随机生成的文本来判断缓存正确性,必须固定输入 token 路径。
18.4.1 选做:沿完整小模型追踪无缓存计算
先对照两段逻辑伪代码:无缓存每轮输入完整前缀,有缓存只在首轮输入 prompt,之后输入刚选出的 token。两条路径都应使用同一组模型参数、位置语义和停止条件。
ids = tokenizer.encode(prompt)
new_ids = []
for _ in range(max_new_tokens):
logits = model_forward(ids) # 每轮重算全部已知位置
next_id = choose(logits[-1])
ids.append(next_id)
new_ids.append(next_id)
if next_id == eos_id:
break
answer = tokenizer.decode(new_ids)prompt_ids = tokenizer.encode(prompt)
cache = None
pending_ids = prompt_ids
new_ids = []
for _ in range(max_new_tokens):
logits, cache = model_forward(pending_ids, cache)
next_id = choose(logits[-1])
new_ids.append(next_id)
if next_id == eos_id:
break
pending_ids = [next_id] # 下一轮只计算刚选出的新位置
answer = tokenizer.decode(new_ids)下面保留原有 00_uncached_walkthrough.py,作为可选的完整网络对照:4 维表示、2 个单头 Block、RMSNorm、加性位置向量、ReLU MLP 和固定随机权重。模型未经训练,6 个 ID 不对应中文,也没有 EOS;固定三轮输入长度为 3→4→5,共处理 12 个位置。它便于逐层追踪向量,不用于评价语言能力。
完成第 18.0 节 CPU 环境后,从下方保留的 2026-09-09 实验附件取出脚本,解压至实验目录再运行。它不需要 GPU,也不下载语言模型。阅读脚本时按 Embedding→Norm→Q/K/V→Attention→残差→MLP→LM Head 的顺序追踪同一位置。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source .venv-cpu/bin/activate
ZIP_PATH="$HOME/Downloads/llm-inference-pedagogy-20260909.zip"
mkdir -p pedagogy
unzip -o "$ZIP_PATH" -d pedagogy
python pedagogy/00_uncached_walkthrough.py | tee runs/uncached-walkthrough.txt验收先看 full input ids 长度是否为 3、4、5,hidden state 行数相应变化、列数始终为 4,Attention 分数形状分别为 [3,3]、[4,4]、[5,5]。末尾应出现 total input positions = 12 和 CHECK: PASS。causal prefix error 应接近 0,表示追加未来位置不改变旧前缀表示(允许浮点误差)。No module named numpy 时先检查 CPU 虚拟环境是否激活。
入门实验附件(2026-09-09):下方 llm-inference-pedagogy-20260909.zip 包含 00_uncached_walkthrough.py、运行步骤和 CPU 验证记录。包内历史配图不代表当前版本,配图以正文在线画板为准。第 18 章的两周实验包继续用于完整实验。
18.5 Day 5:算显存并区分 TTFT、TPOT 与吞吐
今天的问题与阅读:读第 4 章。今天所有数字都是估算或教学时间戳,尚无 GPU 性能结论。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/principles.py memory --out runs/day05-memory.json
cat "$MODEL_DIR/config.json"步骤 2|从配置代入公式。KV bytes/token=2(K/V)×28 层×8 KV heads×128 head_dim×2 bytes(BF16)=114688 bytes=112 KiB。单请求 4096 个已缓存位置=448 MiB,8 请求为 3.5 GiB;一个 256-token block 跨全层的 KV 为 28 MiB。权重、激活、工作区、Graph、块尾浪费需另计。
步骤 3|算一条教学时间线。请求发出为 0 秒,收到四个 token 的时间是 [0.12,0.16,0.21,0.25] 秒:TTFT=120ms,ITL=[40,50,40]ms,TPOT=(250−120)/(4−1)≈43.33ms。写出“离线输出 token 总数 / generate 总耗时”的吞吐口径为何不同。
验收:day05-memory.json 中 112、448、28、3.5 均吻合,check=PASS;notes/day05.md 同时写出显存估算的假设和指标起止点。卡住时:先检查是否误用了 16 个 Q heads;KV 容量要用 8 个 KV heads,KiB/MiB 按 1024 换算。
18.6 Day 6:配置 GPU,完成第一次真实生成
今天的问题与阅读:读第 5–6 章:先区分 vLLM 的服务能力与 nano 的教学引擎;本章必做主线使用 nano。第 5 章的 vLLM 在线体验留作完成主线后的扩展。
步骤 1|检查工具链并新建 GPU 环境。以下是明确指定的参考软件组合,不是已在写作机器通过的 GPU 镜像:Python 3.11.11、PyTorch 2.6.0+cu124、Triton 3.2.0、FlashAttention 2.7.4.post1、Transformers 4.57.1。FlashAttention 可能需要编译,耗时显著长于普通 pip 包。先看到 GPU 和 nvcc 12.4 再安装。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
nvidia-smi
nvcc --version
g++ --version
python -m venv .venv-nano
source .venv-nano/bin/activate
python -m pip install --upgrade pip
python -m pip install setuptools==80.9.0 wheel==0.45.1 packaging==24.2 ninja==1.11.1.3
python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
python -m pip install -r requirements-cpu.txt einops==0.8.1 psutil==7.0.0
MAX_JOBS=4 python -m pip install flash-attn==2.7.4.post1 --no-build-isolation
python -m pip install --no-deps --no-build-isolation -e "$NANO_DIR"
python -m pip check
python -m pip freeze > runs/gpu-environment.txt
CUDA_VISIBLE_DEVICES=0 python e2e/doctor.py | tee runs/day06-doctor.json
python e2e/download_model.py --model "$MODEL_DIR" --weights | tee runs/weights-download.log步骤 2|先验证短生成。smoke.py 使用 chat template、单卡 eager、max_model_len=1024、token budget=512、gpu_memory_utilization=0.5,并设置 ignore_eos=True 强制生成 2 token。不要以这 2 token 是否组成完整回答作为验收。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
CUDA_VISIBLE_DEVICES=0 python e2e/smoke.py --model "$MODEL_DIR" 2>&1 | tee runs/day06-smoke.log
CUDA_VISIBLE_DEVICES=0 python labs/02_cached_forward.py --model "$MODEL_DIR" --device cuda --prompt-len 64 --steps 8 --repeats 3 2>&1 | tee runs/day06-cached-forward.log
python e2e/analyze.py cache --log runs/day06-cached-forward.log --out runs/day06-cache-check.json预期与验收:doctor 输出 check=PASS;smoke 返回 1 个结果、2 个 token id、正数 kv_blocks;真实模型 cache 对照仍为 540/71,数值误差为有限值。记录 max_abs_logit_error、relative_l2_logit_error、same_top1_steps 及两条 decode 时间;BF16 不强求逐位相等,分析器的 PASS 不等于已证明所有误差都可接受。若 top-1 不一致或误差异常,先保留原始结果,单独排查,不写“等价已通过”。
卡住时按顺序处理:CUDA 不可用先检查容器 GPU 挂载和驱动;没有 nvcc 时回到 18.0 的 devel 环境;flash-attn 出现 undefined symbol 时确认 torch=2.6.0+cu124 后,用 MAX_JOBS=2 python -m pip install --force-reinstall --no-deps --no-cache-dir --no-build-isolation flash-attn==2.7.4.post1 重装;编译被 Killed 时先降低 MAX_JOBS 并检查主机内存。KV blocks 断言或 OOM 时先看 nvidia-smi,停止自己占用的其他实验,再在独占设备重跑;不要直接杀别人的进程。
补充 API 示例与兼容性说明(从第 6 章集中到此处)。以下补全示例便于查看 generate 的返回结构;正式验收仍以上面的 smoke.py 为准。
import os
from nanovllm import LLM, SamplingParams
def main():
model_dir = os.environ["MODEL_DIR"]
llm = LLM(model_dir, enforce_eager=True, tensor_parallel_size=1,
max_model_len=1024, max_num_seqs=8,
max_num_batched_tokens=512, gpu_memory_utilization=0.5)
outputs = llm.generate(["The sky was"],
SamplingParams(temperature=0.6, max_tokens=2),
use_tqdm=False)
print(outputs[0]["token_ids"])
print(outputs[0]["text"])
if __name__ == "__main__":
main()nano 的 generate 返回字典列表,读取 outputs[0]["text"];与 vLLM 的 RequestOutput 对象不同。示例的 main guard 为 spawn 多进程保留入口,本提交通过 atexit 清理资源,不需重复调用 exit()。
正式 smoke 按第 18.6 节使用包内脚本:它应用 chat template,并用 ignore_eos=True 固定输出两个 token。此处的简短补全示例允许提前遇到 EOS,max_tokens=2 是上限。生成可读文本只证明主链能跑,正确性与性能仍要分别验证。
单张兼容的 NVIDIA GPU 足以学习正文主线。A100/H800 适合这些实验;GPU 的架构、驱动、CUDA 版 PyTorch 与 FlashAttention 扩展需要匹配。真实 TP 通信需要同机多卡。没有 GPU 时,可以先完成 NumPy 实验、手算和源码推演,GPU 实验再补;本提交依赖 CUDA/FlashAttention,不能直接当作 CPU 推理引擎运行。
排查环境按依赖顺序推进:GPU/驱动与 CUDA 工具链 → PyTorch、FlashAttention、Triton → 模型配置和路径 → 单卡 eager 短生成。依赖下限不保证任意组合兼容;undefined symbol、no kernel image 应先查扩展与二进制版本。工具能导入,也不代表 GPU kernel 已运行成功。
ModelRunner 读取 hf_config.dtype。如果依赖能导入,但配置对象没有这一属性,应核对 Transformers 与源码的版本配合,不要把配置 API 不匹配当成模型权重损坏。
18.7 Day 7:从日志还原两条请求的一生
今天的问题与阅读:读第 7 章,重点 add_request、step、generate、Sequence.append_token。今天启用形状 hook,日志同时留给第 11 天使用。
步骤 1|预测后运行。两条请求输入长 64/96,输出限制为 2/3,ignore_eos=True。先在 notes/day07.md 写下每轮哪些请求还活跃,再运行:
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
python e2e/source_walk.py engine --repo "$NANO_DIR" > runs/day07-source.txt
CUDA_VISIBLE_DEVICES=0 python labs/03_trace_nanovllm.py --model "$MODEL_DIR" --case lifecycle --shapes 2>&1 | tee runs/day07-lifecycle.log
python e2e/analyze.py lifecycle --log runs/day07-lifecycle.log --out runs/day07-check.json
cat runs/day07-check.tsv步骤 2|按同一个 id 读四种事件。before_schedule 看队列;scheduled 看本轮 cached/scheduled/blocks;gpu_inputs 看送进 GPU 的输入;after_postprocess 看 token 接受、状态变化和资源回收。seq_id 因 warmup 而变化,不要假定从 0 起;物理块 id 也不固定。
预期与验收:三轮依次 Prefill、Decode、Decode;首次 postprocess 后 total=65/97,cached=64/96;短请求先完成,最终输出长 2/3。结束时 cached=0、blocks=[] 是 deallocate 的结果;未完成请求通常 total=cached+1。day07-check.json 为 PASS,TSV 能逐行解释。卡住时:分析器报告缺事件先检查原始 log 的 Traceback,不要手工补 JSON。
18.8 Day 8:验证分块 Prefill 与抢占顺序
今天的问题与阅读:读第 8 章 schedule、preempt、postprocess。先手写 700/32 token 输入在 budget=512 下的前三轮选择。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
python e2e/source_walk.py scheduler --repo "$NANO_DIR" > runs/day08-source.txt
CUDA_VISIBLE_DEVICES=0 python labs/03_trace_nanovllm.py --model "$MODEL_DIR" --case chunk 2>&1 | tee runs/day08-chunk.log
python e2e/analyze.py chunk --log runs/day08-chunk.log --out runs/day08-check.json
cat runs/day08-check.tsv预期观察:第 1 轮只处理 A 的 512 个位置,total 仍为 700、cached 变为 512;中间 chunk 虽有采样结果,postprocess 不接受为输出。第 2 轮处理 A 的 188+B 的 32=220,cu_q=[0,188,220]、cu_k=[0,700,732];第 3 轮进入 Decode。分析器输出 PASS。
步骤 2|用 CPU 精确触发抢占。教学 block size=4、共 4 块,A/B/C 各 4 token;首轮占 3 块并各采样 1 token。Decode 时 A 获得最后一块,B 需要扩块,于是从 running 尾部抢占 C。这个受控例子直接执行真实 Scheduler,不依赖把 GPU 显存耗尽。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/cpu_scheduler.py preempt --repo "$NANO_DIR" --out runs/day08-preempt.json验收:selected 是 A/B,waiting 是 C,C.cached=0、blocks=[],check=PASS。notes/day08.md 解释 C 保留了 token 序列,重新入场要用 Prefill 重建已被释放的 KV。CPU 教学块大小 4 不能直接传给真实 nano Config,真实路径要求 256 的倍数。
18.9 Day 9:从逻辑位置算到真实 cache slot
今天的问题与阅读:读第 9 章与 prepare_prefill/prepare_decode 的 slot 公式。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/source_walk.py blocks --repo "$NANO_DIR" > runs/day09-source.txt
python e2e/principles.py slots --out runs/day09-slots.json
python e2e/analyze.py slots --log runs/day07-lifecycle.log --out runs/day09-real-slots.json步骤 2|先算教学例,再对照真实日志。B=4、block_table=[17,5,23],位置 0/6/8 对应 slot 68/22/92。统一公式为 block_table[position//B]×B+position%B。真实实验 B=256,分析器使用每个 scheduled 事件的 block_table 重算各轮前 8 个 slot,与 gpu_inputs 逐个相等才 PASS。
步骤 3|解释边界。长度 255/256/257 对应 1/1/2 个逻辑块;采样出第 257 个 token 时该 token 的 KV 尚未计算,下一轮 forward 前才需 may_append 分配新块。验收:两个 JSON 均 PASS,并在 notes/day09.md 写出逻辑块、物理块、块内偏移的区别。无 GPU 日志时只完成教学计算,真实 slot 项保持未运行。
18.10 Day 10:验证 512-token 前缀命中与引用计数
今天的问题与阅读:读第 10 章,重点 can_allocate、hash_blocks、allocate/deallocate。
步骤 1|执行先 A 后 B 的受控前缀实验。A/B 长度都为 620,前 512 token 相同、后 108 不同。A 完全生成结束后再提交 B;这与两个请求同时首次入队不同。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
CUDA_VISIBLE_DEVICES=0 python labs/03_trace_nanovllm.py --model "$MODEL_DIR" --case prefix 2>&1 | tee runs/day10-prefix.log
python e2e/analyze.py prefix --log runs/day10-prefix.log --out runs/day10-check.json
cat runs/day10-check.tsv预期观察:prefix_second_request 后 B 首次 scheduled 为 cached=512、scheduled=108;GPU 的 cu_q=[0,108]、cu_k=[0,620],positions 从 512 开始,block_tables 非空。由此解释“本轮只算 108 个 query,却读到 620 个位置的 K/V”。
步骤 2|执行 CPU 引用计数实验。同样调用固定源码 BlockManager,受控登记完整块的 hash,再执行释放、重新激活、并发共享、逐个释放;这里没有计算真正 KV 张量。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/cpu_scheduler.py refs --repo "$NANO_DIR" --out runs/day10-refs.json验收:两个共享块的 ref_count 依次 [0,0]→[1,1]→[2,2]→[1,1]→[0,0];620 长度命中 2 块,恰好 512 长度最多检查并命中 1 块(can_allocate 使用 range(num_blocks−1))。两项均 PASS。notes/day10.md 解释 ref_count=0 不代表 GPU 显存池已归还,也不代表 hash 立即失效;旧块被重新分配覆盖时会清理旧映射。
18.11 Day 11:把 Runner 元数据接到真实模型形状
今天的问题与阅读:读第 11–12 章,先用正文的第二轮 220 个新位置手算一遍,再复用第 7/10 天的 GPU 日志核对同一套规则,不必为此重复初始化模型。注意 Day 7 的 lifecycle 输入是 64/96,总长 160;正文的 chunk 主例是 700/32,在第二轮新算 188+32=220。这是两组刻意不同的实验参数,下面的 160 不能改成 220。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/source_walk.py runner --repo "$NANO_DIR" > runs/day11-runner-source.txt
python e2e/source_walk.py model --repo "$NANO_DIR" > runs/day11-model-source.txt
python e2e/principles.py shapes --tp 1 --out runs/day11-static.json
python e2e/analyze.py runner --log runs/day07-lifecycle.log --out runs/day11-runner.json
python e2e/analyze.py shapes --log runs/day07-lifecycle.log --out runs/day11-shapes.json
cat runs/day10-check.json步骤 2|对照三组元数据。lifecycle Prefill 的 input_shape=[160],cu_q/cu_k 都是 [0,64,160],没有历史 block_tables;首次 Decode 输入 [2],positions=[64,96]、context_lens=[65,97];prefix 的 query=108、key=620,需从页表读取历史缓存。
步骤 3|核对 layer 0 的 hook。首次 Prefill:qkv_proj 输出 [160,4096];Q=[160,16,128],K/V=[160,8,128];o_proj 回到 [160,1024];gate_up=[160,6144];down_proj=[160,1024];LM Head 的 logits=[2,151936]。T=160 是输入位置数,N=2 是本轮序列数;Prefill 在 LM Head 中挑每条序列的最后位置。首次 Decode 的 qkv 输出 [2,4096]。
把实验数据与正文连接起来:lifecycle 的末行索引为 [63,159],正文 chunk 第二轮则为 [187,219],它们都来自 cu_seqlens_q[1:]−1,最终都是两行 logits。再对照第 12.3 节的残差图,在 notes/day11.md 写出 a=x+Attn(Norm1(x))、y=a+MLP(Norm2(a)),并指出代码中哪次 Norm 接收 h/r、在哪个时点完成相加。
步骤 4|定位权重如何装进去。在 day11-model-source.txt 搜 packed_modules_mapping、load_model,再读 layers/linear.py 的 QKVParallelLinear.weight_loader。TP=1 时 Q/K/V 分别写入 qkv 权重的行区间 [0,2048)、[2048,3072)、[3072,4096)。不是把 hidden_size=1024 简单乘三得到 3072。验收:三个 JSON 均 PASS,在 notes/day11.md 记录一条真实形状与对应源码函数。如果提示缺 model_shape,用第 7 天的 --shapes 命令重跑日志。
18.12 Day 12:完成 eager 与 CUDA Graph 的可比较实验
今天的问题与阅读:读第 13 章 Attention.forward、store_kvcache_kernel、run_model、capture_cudagraph。今天启用不带 trace/hook 的 benchmark。
运行前分别对照第 13.3、13.4 节的两张图回答两句:FlashAttention 为什么不必把完整分数矩阵保存到显存;CUDA Graph 为什么每轮仍要更换输入并重新计算。这一天的计时实验比较 eager 与 Graph,不是单独测量 FlashAttention 的收益,不能从同一份结果同时宣称两种机制各自提高了多少。
步骤 1|两个独立进程顺序运行。固定 16 请求、输入 128 token、输出 32 token,ignore_eos=True,每轮应有 512 输出 token。正式测量前预热 2 次,测量 3 次;两组只改 enforce_eager。运行期间保持同一 GPU、没有其他负载。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
python e2e/source_walk.py kernels --repo "$NANO_DIR" > runs/day12-source.txt
CUDA_VISIBLE_DEVICES=0 python labs/04_benchmark_nanovllm.py --model "$MODEL_DIR" --eager --seqs 16 --max-num-seqs 32 --input-len 128 --output-len 32 --repeats 3 2>&1 | tee runs/day12-eager.log
CUDA_VISIBLE_DEVICES=0 python labs/04_benchmark_nanovllm.py --model "$MODEL_DIR" --seqs 16 --max-num-seqs 32 --input-len 128 --output-len 32 --repeats 3 2>&1 | tee runs/day12-graph.log
python e2e/analyze.py compare --variable eager --log runs/day12-eager.log --other runs/day12-graph.log --out runs/day12-comparison.json步骤 2|读结果与源代码。分别记 init_seconds、三轮吞吐、median_tokens_per_second 和 KV block 数;比较文件给出 Graph/eager 的吞吐比。定位 graph_bs=[1,2,4,8,16,32]:batch=9 选择 16 桶,填充的 slot_mapping=-1 使 KV 写入 kernel 跳过;Graph 捕获主要是 Decode 的 model forward,LM Head 与采样在捕获外。
验收:每轮 output_tokens=512、配置仅 eager 不同、分析器 PASS。写一条带设备、软件、负载、重复次数和实际数字的结论;Graph 没变快也可以完成验收,不能删掉不符合预期的轮次。该指标是包含 Prefill/Decode 的离线输出吞吐,不是 TTFT 或 TPOT。卡住时:先确认第 6 天 eager smoke 成功;Graph 初始化失败保留日志,本日标未完成;小于 16 的 max_num_seqs 不属于本实验覆盖范围,不随意改桶参数。
18.13 Day 13:先推导 Tensor Parallel,再按设备选做多卡
今天的问题与阅读:读第 14 章;TP 形状推导为必做,实际 TP=2 需要同机两张兼容 GPU。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/source_walk.py tp --repo "$NANO_DIR" > runs/day13-source.txt
python e2e/principles.py shapes --tp 2 --out runs/day13-static.json步骤 2|核对每张卡的形状。TP=2 时每 rank 有 8 个 Q heads、4 个 KV heads;qkv 输出宽度 2048;o_proj 输入宽度 1024,其输出虽为 [T,1024],仍只是部分和,需要 all-reduce;gate_up 宽度 3072,down_proj 输入宽度 1536;词表切为每卡 75968,logits gather 到 rank 0。KV 每 token 每 rank 56 KiB;两卡总和仍为 112 KiB。
步骤 3|有两卡再做性能对照。以下两条命令也必须顺序运行,参数保持一致;两组仅改变 TP。终端 B 可用 nvidia-smi -l 1 观察每卡占用,用 Ctrl+C 结束观察。benchmark 的 peak 字段仅来自 rank 0,不能据此声称多卡总显存。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-nano/bin/activate"
CUDA_VISIBLE_DEVICES=0 python labs/04_benchmark_nanovllm.py --model "$MODEL_DIR" --eager --tp 1 2>&1 | tee runs/day13-tp1.log
CUDA_VISIBLE_DEVICES=0,1 python labs/04_benchmark_nanovllm.py --model "$MODEL_DIR" --eager --tp 2 2>&1 | tee runs/day13-tp2.log
python e2e/analyze.py compare --variable tp --log runs/day13-tp1.log --other runs/day13-tp2.log --out runs/day13-comparison.json验收:day13-static.json 为 PASS;在 notes/day13.md 标“静态推演”或“含双卡实测”。若实测,比较文件应 PASS,并解释小模型 TP=2 可能受通信开销影响而变慢。SharedMemory 传递方法名和序列控制元数据,模型张量通信走 NCCL。卡住时:NCCL 等待时先检查可见 GPU 数、两卡是否可用和前一实验是否退出,不要并发重试制造更多进程。没有双卡则跳过步骤 3,无需伪造多卡结果。
18.14 Day 14:交付一份别人能复查的推理实验报告
今天的问题与阅读:回看第 6 章的架构与请求流程、第 7 章的请求状态,以及第 15–17 章实验/排障索引。今天验收的是解释与证据是否接得起来。
步骤 1|完成学习记录。用任意文本编辑器打开 notes/day01.md–day14.md,将 TODO 换成自己的内容;未运行项目明确写“未运行及原因”,不要填预期数值代替观察。Day 14 写完整调用链 prompt → Sequence → Scheduler → BlockManager → Runner → Qwen3 → Sampler → postprocess,并给每一步标出一个源码函数或实际日志字段。
步骤 2|写三条有证据的结论。至少覆盖缓存位置数、调度/前缀缓存、性能对照;每条注明实测/CPU 源码执行/静态估算、文件路径、具体数值与适用条件。另写一个只改变单一变量的后续实验设计,例如在线 TTFT;本章的离线 benchmark 不能直接给出在线 TTFT。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source ".venv-cpu/bin/activate"
python e2e/final_report.py --root "$LAB_ROOT"
cat runs/final-report.md
tar -czf learning-evidence.tar.gz README.md VALIDATION.md env.sh requirements-cpu.txt e2e labs runs notes最终验收:runs/final-report.md 列出 Day 1–13 必要输出是否存在、JSON 检查状态和文件摘要,Day 1–14 笔记是否仍有 TODO。完整单卡路线应无 MISSING/INVALID JSON,笔记无 TODO,evidence_and_notes_ready_for_review=true;此值只表示证据和笔记已就绪,仍需人工审核内容。无 GPU 路线允许报告保留 MISSING,结论明确停留在 CPU/静态范围。双卡实测为选学,不影响单卡路线完成。
把 learning-evidence.tar.gz 留存或交给同伴复查。随机抽一行调度 TSV,应能解释为什么选择这些序列、写到哪些 slot、模型输入形状是什么、为什么接收或忽略本轮 sampled token;再抽一条性能结论,应能找到原始日志、环境版本与测量口径。能做到这两点,才算把原理、源码和运行结果连起来。
18.15 选做:体验 vLLM 的离线接口与本机服务
这一节承接第 5 章,只用于观察同一生成能力怎样被不同接口调用,不是 nano-vLLM 源码学习的前置条件。先完成第 18.0 节环境与第 18.6 节模型准备;完整步骤集中在这里,避免打断前面的原理主线。
18.15.1 独立环境与离线补全
以下复用 env.sh 中的 MODEL_DIR,指向第 18.6 节准备的固定模型目录。vLLM 使用独立环境,不覆盖 nano-vLLM 实验环境。
安装方式参考 vLLM 官方 Quickstart;实际运行前按所用版本的安装页核对 GPU、驱动与依赖兼容性。下面在 .venv-vllm 中安装和运行。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
export MODEL_DIR
# 先安装 uv,或使用第 18.0 节下载到 tools/ 的 uv 可执行文件
python -m pip install uv
uv venv .venv-vllm --python 3.12 --seed
source .venv-vllm/bin/activate
uv pip install vllm --torch-backend=autoimport os
from vllm import LLM, SamplingParams
model_dir = os.environ["MODEL_DIR"]
llm = LLM(model=model_dir, max_model_len=2048,
gpu_memory_utilization=0.5, generation_config="vllm")
params = SamplingParams(temperature=0.6, max_tokens=32)
for result in llm.generate(["The sky was", "The river was"], params):
print(result.prompt, result.outputs[0].text)在已激活的 vLLM 环境中运行上面的 Python 示例。它调用文本补全接口;若要按聊天消息问答,应使用聊天接口或先应用相同的 chat template。完成离线体验后先退出该进程,释放其 GPU 资源,再启动下面的本机服务。
18.15.2 本机在线服务与流式返回
服务仅监听 127.0.0.1;在同一台 GPU 机器的另一终端发送请求。它是学习用入口,不是公网生产部署。
source "$HOME/llm-inference-lab/env.sh"
cd "$LAB_ROOT"
source .venv-vllm/bin/activate
vllm serve "$MODEL_DIR" \
--served-model-name qwen3-demo \
--host 127.0.0.1 --port 8000 \
--max-model-len 2048 --gpu-memory-utilization 0.5 \
--generation-config vllmcurl -N http://127.0.0.1:8000/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{"model":"qwen3-demo","messages":[{"role":"user","content":"用一句话解释什么是推理。"}],"temperature":0.6,"max_tokens":64,"stream":true,"chat_template_kwargs":{"enable_thinking":false}}'这里的 stream=true 让客户端逐步接收结果,适合观察首输出和输出间隔;网络 chunk 不一定严格对应一个模型 token。Qwen3 的 thinking 设置会影响输入模板与输出形式,示例显式关闭 thinking,便于先观察简短回答。离线与在线的模型前向原理相同,在线路径额外涉及请求排队、序列化和传输。
19. 附录:从数字与向量看懂 RMSNorm、SiLU/SwiGLU、GQA 与 RoPE
正文主要沿着 token 的执行过程展开,这里把反复出现的小算子单独拆开。遇到不熟悉的名字时,可以只读对应一节,再回到正文。所有小向量都是为了手算构造的示例,不是模型实测值;算子改变的是数值,是否改变形状要单独检查。
先把这些名字放回第 4.1 节的模型结构图:RMSNorm 调整输入尺度;Attention 内部用 GQA 安排 Q 与 K/V 的共享关系,用 RoPE 引入位置信息,再用 Softmax 把匹配分数变成权重;MLP 内部用 SiLU 做非线性,用 SwiGLU 把两路特征门控组合。它们不是五个依次串起来的新 Block,而是 Block 内不同位置的计算或结构选择。
以下真实形状以 Qwen3-0.6B 为准:隐藏宽度 1024,Q Head 数 16,KV Head 数 8,每头 128 维,MLP 中间宽度 3072。前文的 4096 是另一套教学配置,不能混用。T 表示本轮新计算的位置数,N 表示某条请求当前可读的位置数(含历史与当前);D 表示特征维度,采样温度记为 τ。除特别说明外省略 batch 轴,以单卡路径说明。
19.1 Softmax:怎样把分数变成权重
假设有三个候选,分数 z=[1,2,3]。分数只能告诉我们相对偏好,还不是概率:它们可能为负,也不保证和为 1。Softmax 对每个分数取指数,再除以所有指数之和:
p_i = exp(z_i) / Σ_j exp(z_j)
z = [1, 2, 3]
softmax(z) ≈ [0.09003, 0.24473, 0.66524]exp(z) 就是 e 的 z 次方,e 约为 2.718。分数越大,指数越大;再除以共同的总和,就得到非负且和为 1 的权重。输入有三个数,输出仍有三个数,不会把向量变成一个 token。最后还要采样或取最大值,才选出候选 id。
实际计算通常先减去最大分数:把 [1,2,3] 改成 [-2,-1,0],指数约为 [0.13534,0.36788,1],再相除,结果不变。这避免了大正数取指数时溢出;它不是近似地改变模型偏好。causal mask 则把不允许读取的位置在 Softmax 前设为负无穷,使其权重为 0;不能仅设为 0,因为 exp(0)=1。
Attention 和采样用的是同一个数学算子,但候选轴不同。Attention 对一条 query 的各个 key 位置分配权重,再用权重加权汇总 V;LM Head 之后则对词表里的 token 候选分配概率。前者是在选择“读哪些上下文位置”,后者是在选择“输出哪个 token”。QK 的点积先产生分数,Softmax 不产生 V。
Temperature 改的是分数差距。采样时通常计算 softmax(logits/τ),τ 必须为正。对同一组 [1,2,3],τ=0.5 得到约 [0.01588,0.11731,0.86681],τ=2 得到约 [0.18632,0.30720,0.50648]:较低温度更集中,较高温度更平缓,但它不保证答案更正确。Attention 中除以 √head_dim 是控制点积分数尺度的模型运算,不是用户设置的采样 temperature。
对应正文:第 1.3–1.4 节的 Attention 权重、第 2.2 节的词表输出、第 12.4 节的采样实现。
19.2 Norm 与 LayerNorm:归一化究竟在处理什么
先区分同一个词的两种用法。数学里的 norm 常指“范数”,例如向量的 L2 范数 ||x||₂=√Σx_i²,是衡量向量大小的一个数。网络结构图里的 Norm 通常指某种 normalization layer,即归一化层;到底是 LayerNorm 还是 RMSNorm,要看具体实现,不能都理解成“把向量长度变成 1”。
以一个 token 的三个特征 x=[1,2,3] 为例,LayerNorm 先计算这一行的均值,再看各分量偏离均值的程度:
μ = mean(x) = (1+2+3)/3 = 2
σ² = mean((x−μ)²) = (1+0+1)/3 = 2/3
x_hat = (x−μ) / √(σ²+ε)
y_i = γ_i × x_hat_i + β_i先忽略很小的 ε,并取 γ=[1,1,1]、β=[0,0,0],输出约为 [-1.22474,0,1.22474]。先减均值让这行围绕 0,再除以标准差调整尺度。对于非恒定向量,在这些简化条件下,归一化结果的均值为 0、方差为 1;加入 ε 和学习到的 γ/β 后,最终输出不必严格满足这两个数值。
γ(gamma)是逐特征缩放参数,β(beta)是逐特征偏移参数,标准 LayerNorm 可学习它们。它们在模型训练中更新、普通推理中固定;μ 和 σ² 则每次由当前输入重新计算。ε(epsilon)是很小的稳定项,避免零或过小的分母引发数值问题,不是额外的特征维度。
在本文的 Transformer 隐藏状态上,LayerNorm 通常沿每个 token 自己的 D 个特征计算,而不是把不同 token 混在一起。输入 [220,1024] 时,各行分别得到均值和方差,统计量可保留成 [220,1],再广播回该行的 1024 个数;输出仍是 [220,1024]。名字里的 Layer 不表示“把这一层所有 token 合在一起求均值”。
19.3 RMSNorm:不减均值,只按均方根调整尺度
RMSNorm 是 Root Mean Square Normalization,均方根归一化。它的任务是先衡量当前向量的整体数值大小,再用这个尺度调整各个分量,使后续计算面对较稳定的输入尺度。它不负责读取历史 token,不产生 Attention 权重,也不压缩向量:一行 [1,1024] 输入后仍是一行 [1,1024]。
RMS(均方根)的计算顺序是:每个数平方 → 取平均 → 开平方。RMSNorm 让同一行的各个分量除以共同的尺度,再乘各自可学习的缩放参数 γ。与 LayerNorm 的关键区别是,它不先减均值:
mean_square = mean(x²)
RMS(x) = √mean_square
y_i = γ_i × x_i / √(mean_square+ε)
x = [1, 2, 3]
mean_square = (1+4+9)/3 = 14/3
RMS(x) ≈ 2.16025
忽略 ε 且 γ=1:y ≈ [0.46291, 0.92582, 1.38873]三个数都除以同一个尺度,所以这个例子仍保留 1:2:3 的比例,均值也不会被移到 0。忽略 ε 且 γ 全为 1 时,非零向量的输出均方根为 1,但 L2 范数是 √D,不是 1;使用学习到的逐特征 γ 后,也不必保持原比例或单位均方根。RMSNorm 的输出可以为负,也可以大于 1,它不是概率。
为什么有用?若把 [1,2,3] 整体放大成 [10,20,30],它的均方根也会放大十倍;忽略 ε 时,归一化后两者相同。这使后续子层不必直接承受输入整体幅度的变化,也省去了 LayerNorm 的均值居中步骤。γ 在训练中学到、推理时固定;分母则每次按当前 token 的当前向量重新计算,不是训练时存好一个全局平均值。RMSNorm 是训练时选定的模型结构,不能在已训练模型里随意替换 LayerNorm,并假定输出不变。
本文固定源码的核心计算如下。特别注意变量虽然叫 var,实际计算的是 mean(x²),不是减过均值的方差;rsqrt(a) 表示 1/√a。源码先转成 float 做统计和归一化,再转回原 dtype 并乘权重:
x = x.float()
var = x.pow(2).mean(dim=-1, keepdim=True)
x.mul_(torch.rsqrt(var + self.eps))
x = x.to(orig_dtype).mul_(self.weight)在 Qwen3 中,残差流进入 Attention 前、进入 MLP 前各有一次 RMSNorm,走完所有 Block 后还有最终 RMSNorm。源码名字 input_layernorm、post_attention_layernorm 里的 layernorm 不代表其类型是 LayerNorm,实际构造的都是 RMSNorm。它们沿每个 token 的 1024 个特征计算;Attention 内另外的 Q/K Norm 则在拆头之后,分别沿每个 Head 的 128 个特征计算,不把 16 个 Q Head 或不同 token 混成一组。本文 V 不经过对应的 Q/K Norm。
源码:RMSNorm 与融合的 add_rms_forward;源码:Qwen3 中各 Norm 的类型与位置。
19.4 同一个向量,四种“归一化”结果并不相同
把 x=[1,2,3] 放进不同算子,差异就很直观。下表统一忽略 ε,LayerNorm/RMSNorm 的 γ 全为 1,LayerNorm 的 β 全为 0;数值保留五位小数。
操作 | 输出 | 满足什么条件 |
|---|---|---|
Softmax | [0.09003,0.24473,0.66524] | 非负,三个数之和为 1 |
L2 归一化:x/||x||₂ | [0.26726,0.53452,0.80178] | 平方和为 1,即 L2 范数为 1 |
LayerNorm | [-1.22474,0,1.22474] | 均值为 0,方差为 1 |
RMSNorm | [0.46291,0.92582,1.38873] | 平方的平均值为 1,即均方根为 1 |
这四种操作都可以输入 [1,3]、输出 [1,3],但解决的问题不同。不能只看到 shape 没变,就认为它们可以互换;也不能把“归一化”一概理解成“所有数变到 0–1 之间”。
19.5 SiLU 与 SwiGLU:MLP 中的非线性和门控
先分清层次:SiLU 是作用于单个数的激活函数,SwiGLU 是用它构成的门控结构。 SiLU 的全称是 Sigmoid Linear Unit,也等价于固定 β=1 的 Swish。它对输入数 x 计算 x×sigmoid(x),在向量上则逐个分量独立应用。
线性投影负责把特征重新组合,激活函数负责引入非线性。如果几次线性投影之间没有非线性,它们仍可合并成一次线性变换。SiLU 让 MLP 能表达更复杂的关系;它不会把一行特征变成概率,也不参与不同 token 之间的加权读取:
sigmoid(x) = 1 / (1+exp(−x))
SiLU(x) = x × sigmoid(x)
SiLU([-1,0,1]) ≈ [-0.26894,0,0.73106]sigmoid(x) 在 0 到 1 之间,但乘回 x 后,SiLU 的输出可以为负,也可以大于 1。与把负数全部清零的 ReLU 不同,SiLU 是平滑变化的;它保持输入形状。[1,3072] 经过 SiLU,仍是 [1,3072]。
SwiGLU 可以理解为 Swish 门控线性单元(GLU,Gated Linear Unit)的一个变体。在本文模型里,MLP 取 Attention 与残差相加、再经 RMSNorm 后的当前 token 表示 x,用两组不同权重投影出 gate 与 up。gate 经 SiLU 后,逐特征调制 up,最后再投影回隐藏宽度:
x [1,1024]
gate = x @ W_gate W_gate [1024,3072] → [1,3072]
up = x @ W_up W_up [1024,3072] → [1,3072]
z = SiLU(gate) ⊙ up → [1,3072]
out = z @ W_down W_down [3072,1024] → [1,1024]gate 与 up 都由同一个 x 计算得到,但投影权重不同。门控让 up 中每个分量的贡献随输入内容变化;这里的“门”是连续数值调制,不是只取 0/1 的开关,更不是一组和为 1 的 Attention 权重。上面写的是行向量右乘权重的数学形状,PyTorch Linear 的权重存储通常采用转置布局。
用三个数手算:若 gate=[-1,0,1]、up=[2,3,4],SiLU(gate)≈[-0.26894,0,0.73106],逐元素相乘后 z≈[-0.53788,0,2.92423]。注意此时还不是 MLP 的最终输出,仍要经过 W_down。源码把 gate/up 两次投影合并为 gate_up_proj,先得到 [T,6144],再拆成两个 [T,3072];相乘仍为 [T,3072],最后映射回 [T,1024]。6144 变成 3072 来自两路拆分与组合,不是 SiLU 把特征数减半。
SiLU/Softmax 这些函数本身没有这里需要训练的权重;gate/up/down 的投影权重、Norm 的 γ 等才是模型参数。同一层 MLP 对不同 token 使用同一组投影权重,非线性不会让它自动读取其他 token 的特征。
源码:SiluAndMul 的拆分与逐元素乘法。对应正文第 12.3 节。
19.6 GQA:多个 Q Head 共享一组 K/V,但各自计算权重
GQA 是 Grouped-Query Attention,分组查询注意力。它改变的是Q Head 与 KV Head 的配对关系,不是另一套 Attention 公式。在普通 MHA 中,每个 Q Head 都有对应的一组 K/V;GQA 则把 Q Head 分组,同组的几个 Q Head 读取同一组 K/V。这里“一组 K/V”是某个 KV Head 在所有可读位置上的 K 矩阵和 V 矩阵,不是把多个 token 合成一个向量。
以 Qwen3-0.6B 为例:16 个 Q Head、8 个 KV Head,每头都是 128 维。Q Head 1、2 共享 KV Head 1,Q Head 3、4 共享 KV Head 2,依此类推。对 Decode 新输入的一个位置,Q 的形状是 [1,16,128];本轮新产生的 K、V 各为 [1,8,128]。若连同当前共有 N 个可读位置,该层可读 K、V 各为 [N,8,128]。
只放大第一组,并省略位置下标:q⁽¹⁾、q⁽²⁾ 各为 [1,128],共享的 K⁽¹⁾、V⁽¹⁾ 各为 [N,128]。共享的是 K/V,两个 Q Head 仍各自计算权重与输出。 它们的 Q 来自不同投影,所以即使读取相同 K/V,也能得到不同的加权结果:
w⁽¹⁾ = softmax(q⁽¹⁾ @ K⁽¹⁾ᵀ / √128 + mask) # [1,N]
a⁽¹⁾ = w⁽¹⁾ @ V⁽¹⁾ # [1,128]
w⁽²⁾ = softmax(q⁽²⁾ @ K⁽¹⁾ᵀ / √128 + mask) # [1,N]
a⁽²⁾ = w⁽²⁾ @ V⁽¹⁾ # [1,128]
16 个 a 沿特征轴拼接:[1,2048]
再乘 W_O [2048,1024]:[1,1024]这里省略位置下标,上标表示 Head:两个不同的 Q Head 共享第一个 KV Head,但分别计算权重和输出;q、k 已完成模型要求的 Norm 与 RoPE。沿用行向量右乘权重的数学写法,本模型 W_Q 为 [1024,2048],W_K、W_V 各为 [1024,1024]。投影决定各头总宽度,Q 的总宽度可与隐藏宽度 1024 不同。
为什么这样设计?Decode 每轮都要访问历史 K/V,缓存容量与 KV Head 数成正比。在层数、位置数、每头维度和数据类型相同的条件下,8 个 KV Head 的 K/V 数据量是 16 个 KV Head 的一半,也有助于减少历史数据读取。它不是把 16 个 Q Head 的计算全部减半;输出仍有 16 个 Head,质量和速度也需要具体验证。GQA 是训练时采用的结构,不能把一个已训练 MHA 模型的 KV Head 随意删掉,就假定结果不变。
MHA、GQA、MQA 可以放在同一条线上理解:MHA 每个 Q Head 对应一个 KV Head;GQA 每组 Q Head 共享一个 KV Head;MQA 则让所有 Q Head 共享唯一一个 KV Head。本文模型采用中间的 GQA。配置与实现可对照 Qwen3-0.6B 配置和固定版本的 Qwen3Attention。
19.7 RoPE:旋转 Q/K 的分量,让匹配分数感知位置
RoPE 是 Rotary Position Embedding,旋转位置编码。只有内容投影得到的 q、k,还没有这两个向量之间明确的距离信息;因果 mask 只规定能否看见某个位置,不直接告诉模型“相隔几个位置”。RoPE 根据位置编号改变 Q/K 的方向,使后续点积分数能够同时利用内容和相对位置。它不把位置编号作为第 129 个特征,也不是给 token embedding 直接加一行位置向量。
先把 128 维降到两个数:取 q 或 k 中的一对分量 [a,b],把它看成平面上的一个箭头。位置为 p 时,按角度 φ=p×ω 旋转,ω 是这一对分量使用的角频率。旋转会混合这两个分量,但不改变这一对的长度,也不改变整个向量的特征数:
a′ = a cos(φ) − b sin(φ)
b′ = a sin(φ) + b cos(φ)
教学示例:[1,0] 旋转 90° → [0,1]
长度仍为 1,但与另一个向量的点积会受双方方向影响。真实的 128 维 Head 分成 64 对,每对使用不同频率,相当于以不同的“转速”表达位置。本文固定源码采用前后半段配对:零起始下标为 (0,64)、(1,65)、…、(63,127),而不是相邻的 (0,1)、(2,3)。不同实现可以采用不同的分量布局,必须与权重约定一致。上面的 90° 只是帮助理解旋转,不是说真实模型每个 token 都固定转 90°。
为什么这样就能表达相对位置? 对同一对分量,位置 p 的 q 旋转 p×ω,位置 j 的 k 旋转 j×ω;点积中的旋转关系最终取决于角度差 (j−p)×ω。因此同样的两段内容,间隔不同也可能得到不同匹配分数。它不是“越远权重必然越小”的固定规则,最终权重仍由内容、位置、缩放、mask 和 Softmax 一起决定。
在本文 Qwen3 路径里,顺序是:隐藏表示投影出 Q/K/V → Q、K 各自做 Head RMSNorm → 按位置对 Q、K 做 RoPE → 用旋转后的 Q/K 算分数与权重 → 加权汇总 V。V 不做这一步 RoPE。 Q 仍是 [1,16,128],新 K 仍是 [1,8,128];变化的是分量数值,不是 Head 数量或向量宽度。
这也解释了它与 KV Cache 的关系:本文缓存的是已做 Head Norm 和 RoPE 的 K,以及对应的 V。Decode 只按新位置编号旋转新 Q/K;历史 K 已按各自原位置旋转过,读取时不需要随着新 token 到来重新旋转。缓存物理槽位与逻辑位置编号不是一回事,不能因为数据换了存储位置,就重新给它编号。这里假设使用本文固定的位置编码配置。
标准 RoPE 的正弦、余弦来自位置与频率规则,不是为每个位置训练一条 embedding。本文配置 rope_theta=1000000 控制频率分布,rope_scaling=null;这个数不是上下文窗口长度,也不意味着随意增大它就能可靠处理更长上下文。实现见固定版本的 RotaryEmbedding,与第 12.2 节的执行顺序对应。
19.8 回到同一轮请求:核对形状与计算轴
代码中的 dim=-1 表示最后一维,不固定等于“特征轴”或“token 轴”。要先认出张量各维的含义。沿用正文第二轮 T=220、两条请求的例子:
位置与输入 | 在哪一组数上计算 | 输出与含义 |
|---|---|---|
残差流 RMSNorm:[220,1024] | 每个 token 自己的 1024 个特征 | [220,1024],调整每行特征尺度 |
Q 的 RMSNorm:[220,16,128] | 每个 token、每个 Q head 的 128 个特征 | [220,16,128],不把不同头或 token 混在一起 |
K 的 RMSNorm:[220,8,128] | 每个 token、每个 KV head 的 128 个特征 | [220,8,128];本文 V 不经过对应的 Q/K Norm |
RoPE:Q [220,16,128]、K [220,8,128] | 每个位置、每个 Head 内的 64 对分量,按该位置编号旋转 | Q/K 形状不变;V 不做这一步旋转 |
GQA:16 个 Q Head,8 个 KV Head | 同一请求内,每两个 Q Head 读取同一 KV Head;各自沿可读位置计算权重 | 仍得到 16 个 Head 输出,各 128 维;拼接宽度为 2048 |
A 的单个 head 分数:[188,700] | 每条 query 对应的 key 位置,应用 causal mask | [188,700] 的权重;FlashAttention 不必完整物化它 |
词表 logits:[2,151936] | 每条请求的 151936 个候选 token | [2,151936] 的概率,再采样成两个 token id |
SwiGLU 的两路输入:各 [220,3072] | 相同位置、相同特征逐元素相乘 | [220,3072],随后 down_proj 回到 1024 维 |
最后把残差也分清:x+branch(x) 是两个同形张量逐元素相加;它不是 Softmax,也不是 Norm。融合的 Add+RMSNorm 只是把“先相加、再归一化”放在一起执行,数学关系仍要按第 12.3 节的两条路径理解。
20. 延伸阅读|深入 nano-vLLM 高效推理:授权全文翻译
本章译自 Moncef Abboud 的 Deep Dive into Efficient LLM Inference with nano-vLLM(2026-04-05)。原文由作者以 CC BY 4.0 授权发布。本章为中文翻译并增加独立标注的译者说明,整理日期为 2026-09-15;不表示原作者对本译文或本文作出背书。
阅读定位:第 1—19 章是本文自己的原理与固定版本源码解读;第 20 章保留这篇外部文章的论述顺序、示例与代码,便于换一个视角理解同一套系统。第 21、22 章则是另外两篇文章的中文导读,不是全文翻译。不同文章中的版本、术语和实现差别会单独说明。
代码说明:以下代码全部摘自原文,保留英文标识符和注释;其中包含作者主动省略的片段以及少量笔误,并非本次验证过的完整运行教程。中文译文沿用作者口吻,新增解释均标为“译者注”。
20.1 引言
vLLM 是一个开源的大语言模型推理引擎,也是一款令人印象深刻的软件。
本文将深入 nano-vLLM:它是一个保留核心思想的轻量级重新实现。它能用分页注意力运行 Qwen 模型,支持多 GPU 张量并行,还针对高效处理多个请求做了优化,而代码库本身十分容易入手。
原文嵌入视频:YouTube 视频(保留入口,不转录视频内容)。
20.2 快速开始
# install nano-vllm
git clone https://github.com/GeeeekExplorer/nano-vllm.git
cd /root/nano-vllm/
# fix some quirky bug with rope_scaling default param (dict instead of None)
sed -i 's/rope_scaling=rope_scaling/rope_scaling=None/' nanovllm/models/qwen3.py
pip install .
# install HF CLI
curl -LsSf https://hf.co/cli/install.sh | bash
# add it to path e.g. :
export PATH="/root/.local/bin:$PATH"
# download model weights
export model=Qwen3-0.6B
hf download --force-download "Qwen/$model" --local-dir "/root/huggingface/$model/"译者注:以上命令按原文保留,未在本次翻译中执行或验证。特别是把 rope_scaling 改为 None 的命令,是作者当时环境的处理办法,不应直接套用到本文固定版本或其他模型配置。代码中的安装、下载与路径也请按实际环境核对。
在 Python 交互环境中运行下面的示例:
from nanovllm import LLM, SamplingParams
llm = LLM("/root/huggingface/Qwen3-0.6B/", enforce_eager=True, tensor_parallel_size=1)
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
prompts = ["The sky was"]
outputs = llm.generate(prompts, sampling_params)
outputs[0]["text"]20.3 为什么需要 vLLM 这样的推理引擎?
从高层看,一个大语言模型“不过是”带有 forward() 方法的 Python 类。
我们定义模型类,加载权重;这些权重就是张量,也就是一些数。把权重装进模型,调用 forward(),就能为每个输入 token ID 得到 logits,即对下一个 token 的候选分数。把这些分数转成概率,采样出一个 token,以自回归的方式追加到序列中,再继续重复,直到满足停止条件:遇到 EOS,或者达到最大 token 数、上下文长度限制。
PyTorch 已经让这件事变得清晰而直接。那为什么还需要推理引擎?
嗯……因为酷家伙们都在用,我们也想酷一点!
认真说,推理引擎解决了若干重要问题。下面来看其中几个核心问题。
20.3.1 KV Cache
由于模型结构及自回归机制——把输出再次作为输入——大语言模型会反复做同样的计算。处理 token 2 时,需要 token 0 和 token 1 在所有注意力层中的 K、V 向量;处理 token 3 时,需要 token 0、1、2 的 K、V,依此类推。
最朴素的做法会反复运行网络的大部分计算,包括 FFN 层中代价很高的矩阵乘法。但真正需要读取历史 token 信息的地方只有注意力层,具体说就是历史 K、V 向量。
所以,与其重新计算,不如把它们缓存下来。
保存每个历史 token 的 K/V 张量之后,就可以避免重算,只让新生成的 token 通过网络。对于已有 KV 缓存的 token,其余层的计算,包括昂贵的 FFN,都可以跳过。这样能大幅减少计算量。
译者注:这里跳过的是历史位置的重复计算;新位置仍须经过每一层的 Attention、MLP、Norm 等完整计算。各层缓存互不替代。
20.3.2 内存碎片
缓存又带来了一个新问题:内存管理。
我们事先并不知道生成的序列会有多长。那么,应该给 KV Cache 分配多少内存?
- 如果按最大可能序列长度预先分配,就会产生内部碎片,也就是已分配区域内部的空间浪费。
- 如果根据需要逐步增加分配,则可能留下大小不一的空隙,形成外部碎片。
这正是 vLLM 创作者提出的分页注意力(Paged Attention)要解决的问题。它用固定大小的块,也就是“页”,管理 KV Cache,既减少碎片,也允许序列灵活增长。内存浪费从超过 50% 降到不足 5%,于是能够同时服务更多请求,提高吞吐量。
译者注:这里的“超过 50% 到不足 5%”是原文给出的概括,不是本文环境的实测结论,也不是对任意负载的保证。固定大小分页主要消除块级分配中的外部碎片;每条序列的尾块仍可能有内部浪费。
20.3.3 连续批处理
如果加载一个 14B 或 70B 模型,却只对一个请求执行推理,算术强度就会很低。算术强度是浮点运算量与内存访问量之比,也就是每访问一个字节能执行多少次运算;这是提高 GPU 效率的关键。
如果让多个请求使用同一组权重一起推理,权重仍只需要加载一次,却能完成更多计算,算术强度便提高了。因此有一点很明确:批处理是有益的。
不过,朴素的批处理还不够。在静态批处理中,我们把 N 个请求凑成一批,一起执行,并等待最长的请求完成。这并不高效:如果一个请求在第 1,000 个 token 就结束,另一个要生成到第 50,000 个 token,前者就会空等 49,000 步,造成大量浪费。
此外,请求会在不同时间到达,一个紧急请求可能只能等待当前批次完成。
推理引擎提供了更灵活的机制。批处理发生在每一步执行的粒度上:持续加入新请求,移除已完成的请求,甚至可以抢占低优先级请求。这样能更好地利用计算和内存资源,提高整体吞吐量。
20.3.4 Prefill 与 Decode
大语言模型推理有两个不同阶段:
- Prefill:首次提交 prompt 时,让全部 prompt token 通过模型,计算并保存 KV Cache。这个阶段受计算能力限制。
- Decode:Prefill 完成并保存 KV Cache 后,每次生成一个 token。这个阶段对计算能力的依赖相对较弱,对延迟更加敏感。
两个阶段的特征不同,适合的优化也不同,例如采用专门的注意力内核。推理引擎了解这些差别,并分别处理。
译者注:Prefill 通常更偏计算受限,Decode 通常更偏权重/KV 的内存带宽受限,但具体瓶颈仍取决于模型、序列长度、批量、硬件和内核,不能把它们视为无条件结论。
除了上述问题,推理引擎还提供多 GPU 支持、分块 Prefill、推测解码等功能。每个话题本身都很有意思。
20.4 架构
从前面的示例可以看到,LLM 是入口。创建实例时传入模型路径和选项,其中尤其重要的是 tensor_parallel_size,它决定是否把权重切分到多个 GPU 上;模型装不进单张 GPU 时,这很有用。后面会再讲它。
该路径指向 safetensors 格式的权重,nano-vLLM 会把它们加载进 models/qwen3.py 中的 Qwen3ForCausalLM。从概念上看,模型如下:
class Qwen3ForCausalLM:
def forward(input_ids, pos):
x = Embedding(input_ids) # token -> vector
for _ in layers:
x = x + Attention(x, pos) # self-attention (needs KV Cache)
x = x + MLP(x) # nonlinear feature transform
x = FinalNorm(x) # stabilize outputs
return LMHead(x) # logits over vocab (then softmax + sampling)译者注:这是理解计算主线的伪代码,不是完整类实现;省略了 Block 内的 Norm 等细节。Attention 输出还要经过输出投影、残差和 MLP,最终表示再经 LM Head 得到 logits,不能把某个头的加权结果直接当作下一个 token。
本例使用的 Qwen3-0.6B 权重可以在 Hugging Face 上找到。
原文链接:https://huggingface.co/Qwen/Qwen3-0.6B?show_file_info=model.safetensors
调用 Qwen3ForCausalLM(input_ids, pos),然后采样,就是“大语言模型推理”。vLLM 一类的引擎则增加了让这件事在大规模场景下高效运行的配套能力。
译者配图|一次请求经过的组件。下图复用本文第 6 章的原创可编辑画板,不是原文插图;它按本文固定版本补充组件之间的关系,原文调度细节以本章译者注为界。
20.5 generate() 循环
入口是 generate()。下面是一段经过删减的代码:
class LLMEngine: # really the engine
def generate(
self,
prompts: list[str] | list[list[int]],
sampling_params: SamplingParams | list[SamplingParams], ...
) -> list[str]:
if not isinstance(sampling_params, list):
sampling_params = [sampling_params] * len(prompts)
for prompt, sp in zip(prompts, sampling_params):
self.add_request(prompt, sp)
outputs = {}
while not self.is_finished():
t = perf_counter()
output, num_tokens = self.step()
for seq_id, token_ids in output:
outputs[seq_id] = token_ids
if use_tqdm:
pbar.update(1)
outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())]
outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
return outputs
def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
if isinstance(prompt, str):
prompt = self.tokenizer.encode(prompt)
seq = Sequence(prompt, sampling_params)
self.scheduler.add(seq)这里具体发生了什么?
Prompt 连同 sampling_params 一起进入引擎。temperature 用于调节“创造性”,max_tokens 限制输出长度。每个 prompt 成为一个请求:先分词,再包装成 Sequence。
LLMEngine 持有 Scheduler。Scheduler 用两个双端队列 deque 跟踪序列状态,并决定下一步运行哪些序列。
class Scheduler:
def __init__(self, config: Config):
self.max_num_seqs = config.max_num_seqs
self.max_num_batched_tokens = config.max_num_batched_tokens
self.eos = config.eos
self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)
self.waiting: deque[Sequence] = deque()
self.running: deque[Sequence] = deque()
def add(self, seq: Sequence):
self.waiting.append(seq)
def is_finished(self):
return not self.waiting and not self.runningadd_request() 把请求放进 waiting。只要任意一个队列中还有序列,主循环就继续逐步执行。
一个 prompt,也就是一个序列,最初进入 waiting,等待 Prefill。被接纳执行后,它进入 running,逐个 token 地 Decode。这个“生命周期”发生在 LLMEngine.step() 中:
def step(self):
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这就是引擎的核心循环:
- schedule() 选择下一批序列,并告知这一步执行 Prefill 还是 Decode。
- model_runner 运行模型。
- postprocess() 更新序列,例如追加 token,或者在遇到 EOS 时停止。
较新版本的 vLLM 可以在同一步中混合 Prefill 与 Decode。nano-vLLM 为了简单,将两者分开处理。
译者注:本译文忠实保留原文所分析的调度器。这里“Prefill/Decode 分开、Prefill 优先”的描述不应替换本文第 6—11 章对固定提交 bb823b3e06983d71485a8e1f23715ebd87d98ef8 的分析;该版本涉及 chunked prefill,判断行为应以对应版本源码为准。
文章写完了,今天就到这里吧。
还在看?好,那我们继续。
20.5.1 调度器
def schedule(self) -> tuple[list[Sequence], bool]:
# prefill
scheduled_seqs = []
num_seqs = 0
num_batched_tokens = 0
while self.waiting and num_seqs < self.max_num_seqs:
seq = self.waiting[0]
if num_batched_tokens + len(seq) > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq):
break
num_seqs += 1
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
# decode
while self.running and num_seqs < self.max_num_seqs:
seq = self.running.popleft()
while not self.block_manager.can_append(seq):
if self.running:
self.preempt(self.running.pop())
else:
self.preempt(seq)
break
else:
num_seqs += 1
self.block_manager.may_append(seq)
scheduled_seqs.append(seq)
assert scheduled_seqs
self.running.extendleft(reversed(scheduled_seqs))
return scheduled_seqs, False讨论 schedule() 之前,先说明它要保障什么,会更容易理解。
之所以需要调度,是因为 KV Cache 是稀缺资源。使用分页注意力时,不会给每个请求分配一整块连续的 KV 缓冲区,而是给它一张 block table,也就是固定大小 KV 块的列表;默认 kvcache_block_size=256。这样既能减少碎片,也能支持前缀缓存:如果两个序列的完整块所对应的前缀相同,就能借助哈希值和引用计数共享 KV 块。
nano-vLLM 也在调度器层面区分 Prefill 与 Decode:
- waiting:尚未被接纳执行,或者被抢占的序列。
- running:已经被接纳,可以逐个 token 执行 Decode 的序列。
Scheduler.schedule() 返回 (scheduled_seqs, is_prefill)。
20.5.1.1 Prefill 路径
Prefill 时,持续从 waiting 左侧,也就是最早进入队列的一端取出序列,把尽可能多的 prompt 组成一个批次。
遇到以下限制时停止:
- 超过 max_num_seqs。
- 超过 max_num_batched_tokens。
- 或者 block_manager.can_allocate(seq) 判断不能为这个序列分配 KV 块。这表示 KV Cache 被 running 中的序列占用,需要先对那些序列执行 Decode。
接纳一个 waiting 中的序列时:
- block_manager.allocate(seq) 为它分配 seq.block_table;如果哈希匹配,也可以复用已缓存的块。
- 执行 num_batched_tokens += len(seq) - seq.num_cached_tokens。
后一行很重要:如果命中了前缀缓存,一部分 prompt token 已经有 KV,所以这里只累计当前这一步真正需要 Prefill 的新 token,也就是尚未缓存的 token。
只要成功安排了至少一个等待中的序列,就立即返回 (scheduled_seqs, True)。因此 Prefill 具有优先权。
20.5.1.2 Decode 路径
如果没有安排任何 Prefill,就从 running 中选择序列执行 Decode。
这里的 Decode 是指:每个序列只为一个 token 位置执行一次前向计算。把 seq.last_token 输入模型,采样出下一个 token。真正的 seq.append_token(token_id) 会在后面的 postprocess() 中发生,而不是在 schedule() 中。
因此,在 Decode 阶段,调度器的职责是:
1)最多选出 max_num_seqs 个序列执行本轮 Decode;2)确保每个入选序列都为即将计算的 token 准备好有效的 KV 存储槽位;3)如果没有足够空闲 KV 块,就抢占其他序列。
can_append(seq) 的实现是:
return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)表达式 len(seq) % block_size == 1 会被当成 1 或 0。
- 如果 len(seq) % block_size != 1,这个序列本轮 Decode 不需要分配一个全新的 KV 块,因此即使空闲块数量是 0 也可以执行。
- 如果 len(seq) % block_size == 1,序列当前的最后一个 token 位于一个新块的第 0 个位置。这通常发生在 postprocess() 追加 token 后刚好跨过块边界时。此时必须分配一个新块,模型才能写入这个 token 的 KV。
这个判断与 ModelRunner.prepare_decode() 紧密相关:后者用 seq.block_table[-1] 和 seq.last_block_num_tokens 构造 slot_mapping。如果不在 Decode 前分配新块,这次 KV 写入就没有可以映射的位置。
20.5.1.3 抢占策略
Decode 循环通过 seq = self.running.popleft() 先取出最老的序列。如果 can_append(seq) 为假,就尝试通过抢占释放块:
- 如果 running 中还有其他序列,优先抢占 self.running.pop(),即最新的序列。
- 如果只剩 seq 自己,就抢占它,并放弃在本轮执行它。
因此,这个策略本质上是在内存压力下保护较老的序列,牺牲较新的序列。
抢占时,将状态改回 WAITING,调用 block_manager.deallocate(seq) 释放 KV 块并清除 block table,然后把序列放回 waiting 的前端,以便尽快重试。
20.5.1.4 Decode 期间的 KV 块簿记
一旦满足 can_append(seq),调度器就在安排该序列执行 Decode 前调用 block_manager.may_append(seq)。
may_append 基本上就是“完成这个序列当前所需的 KV block table 维护工作”:
如果 len(seq) % block_size == 1:
- 分配一个新 KV 块。
- 把它的 ID 追加进 seq.block_table。
如果 len(seq) % block_size == 0:
- 最后一个块已经满了,因此计算它的哈希,并记录到 hash_to_block_id 中。
- 这样,完整块之后就有资格参与前缀缓存复用。
其他情况:
- 暂时不需要分配新块,也不需要计算哈希。
最后,收集完 scheduled_seqs 后,调度器按原顺序把这些序列放回 running 左侧,让它们在后续轮次中继续保持活跃。
20.5.2 KV Cache
注意力层是大语言模型的关键部分,它让模型能够根据序列中此前的所有 token 预测下一个 token。
每个 token 被映射成:
- 一个或多个 Query 向量:多头注意力 MHA 中是一个,分组查询注意力 GQA 中是多个。
- 一个 Key 向量。
- 一个 Value 向量。
译者注:这组三项应按“每组 KV 头”理解,而非每个 token 在整层只有一个 K/V。MHA 每个 Q 头有对应的 K/V 头;GQA 则让一组多个 Q 头共享一个 K/V 头。本文 Qwen3-0.6B 是每个 token 有 16 个 Q 向量、8 个 K 向量、8 个 V 向量,每个向量 128 维。
要计算 token t 的输出,需要知道 [0, t−1] 范围内所有 token 的 K、V 向量。可以在每次推理时重新计算,也可以保存所有历史 token 的 K、V 向量;这个存储区域就是 KV Cache。
译者注:[0, t−1] 指需要复用的历史位置。普通因果自注意力计算位置 t 时,还会包含当前位置新算出的 K_t、V_t,因此完整可读范围通常为 [0, t]。
对于一个 13B 参数模型,可以有以下配置:
- 40 层。
- 40 个注意力头。
- 每个头 128 维。
- FP16 精度。
每个 token 所需的 KV Cache 为:2(K 和 V)× 40(层)× 40(头)× 128(维度)× 2(FP16 字节数)= 每 token 800 KB。
译者注:上式精确结果为 819,200 字节,即 800 KiB,约 819.2 kB;原文用“800 KB”表达。40 个头在此按 40 个 KV 头计数;若使用 GQA,应按 KV 头数而不是 Q 头数估算 KV 容量。
更大的模型可能每个 token 就需要数 MB。对于包含数千个 token 的序列,仅 KV Cache 就要占用数 GB 显存。
现在假设要给一个请求分配内存,并且知道最大序列长度是 2k。我们可以直接分配 2k × 800 KB,约 1.6 GB,然后就不用再管。但如果请求只生成几百个 token,超过一半的空间就浪费了。这就是内部碎片。
为了避免这个问题,假设改成按需要分配较小的内存区域,并在需要时扩容。每个请求占用一块长度可变的区域,完成后释放。运行一段时间后,空闲内存虽然变多了,却不连续,因此不能真正满足所需分配。这就是外部碎片。
操作系统在给不同进程分配 RAM 时,早已解决了这种变长内存分配问题:方法就是虚拟内存和分页。
进程看到的是连续的内存,称为虚拟内存。在幕后,操作系统把实际内存分成固定大小的块,也就是页,例如 4 KB 或 8 KB。进程虚拟内存中的每一页都映射到一个物理块,并在使用完后释放。物理页不必按顺序排列,也不必连续。因此,虚拟块 1 可以映射到物理块 233,虚拟块 2 可以映射到物理块 4。
分页几乎消除了内部碎片:最多浪费不足一页的空间。由于任意虚拟页都能用任意空闲物理页满足,也就没有外部碎片。
同样的技术可以用于 KV Cache。把 KV Cache 的连续内存划分成块,每块存放 block_size 个 token 的 K、V 向量,比如 16 或 256 个。额外的复杂性在于:必须为每个序列维护虚拟块到物理块的映射。
把 KV Cache 划分成块,还带来了一项很有意思的能力:前缀缓存。
实际场景中,很多请求拥有相同的开头,例如相同的系统提示词,往往还包括相同的工具定义。对于这些开头的 token,KV Cache 是一样的。所以,不必让每个请求都重新计算共享前缀的 KV;计算一次、保存下来,再复用即可。
具体来说,算完某段前缀的 KV Cache 后,保留对应 KV 块的引用,并使用前缀 token 的哈希作为键,把它们记录进查找表。新请求到来时,对其前缀求哈希并查表;如果匹配,就复用已算好的块,从这个位置继续注意力计算,跳过共享前缀的全部重复工作。
一个细微但重要的地方是:不能孤立地只对一个块求哈希。为了保证正确,整个前缀都必须匹配,因此某个块的哈希还需要反映它之前的 token,因为这个块表示的是某个特定序列中的特定位置区间。因此,键实际上是覆盖整个前缀的链式哈希,而不是独立的块内容哈希,即:
# recursive relationship
Hash_Block(0) = Hash(tokens[0] + -1)
Hash_Block(i+1) = Hash(tokens[i+1] + Hash_Block(i))这种方法效果相当好,因为真实助手流量中的共享前缀非常常见。
20.5.2.1 分页注意力的实现
初始化 ModelRunner 时,会调用 allocate_kv_cache。ModelRunner 就是负责加载模型和运行推理的类:
def allocate_kv_cache(self):
config = self.config
hf_config = config.hf_config
free, total = torch.cuda.mem_get_info()
used = total - free
peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
num_kv_heads = hf_config.num_key_value_heads // self.world_size
head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)
block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.torch_dtype.itemsize
config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
assert config.num_kvcache_blocks > 0
self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.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]
module.v_cache = self.kv_cache[1, layer_id]
layer_id += 1这里做了不少事情,我们拆开看。首先,计算一个 KV 块,也就是页,需要多少字节。每个块容纳 block_size 个 token;每个 token 都需要保存 num_hidden_layers 层的 KV,每层有 num_kv_heads 个头,每个头的维度是 head_dim。因此,一个块的大小为:
2 × num_hidden_layers × block_size × num_kv_heads × head_dim × dtype_size
最前面的 2 表示同时保存形状相同的 K 和 V。hf_config.torch_dtype.itemsize 给出每个元素的字节数,例如 FP16 是 2 字节,INT8 是 1 字节。
通过 gpu_memory_utilization,可以控制为 KV Cache 预留的 GPU 内存预算。默认值是 0.9,用来留出余量,而不是耗尽全部可用内存。
译者注:gpu_memory_utilization 是引擎可使用的总显存预算比例,并非该比例全部专供 KV Cache;代码还会扣除权重、已用内存和预估运行峰值。
然后,用估计可用的内存除以 block_bytes,得到能放下多少块。可用空间粗略为 total × util − used。used 是 GPU 当前已经使用的内存,包括模型权重、CUDA 内核和库的缓冲区等。
调用 allocate_kv_cache 之前,会先调用 warmup_model(),按最大批量与序列长度运行推理,迫使模型分配满负荷运行所需的内存,再把剩余预算用于 KV Cache。
译者注:warmup 的实际输入形状由具体实现决定,不能把它理解为把所有独立上限同时取满。预热峰值是容量估计依据,不保证覆盖所有未来请求形状。
除了当前已用 GPU 内存,还要减去 peak − current,也就是代码中的 − peak + current。这里的 current 和 peak 来自 PyTorch 自身的统计,而不是整个 CUDA GPU 的统计:current 是 PyTorch 当前使用量,peak 是历史峰值。减去这个差值,相当于假设 PyTorch 还可能再次达到高于当前值的历史峰值,因此提前给这部分增长留出空间。
知道 num_kvcache_blocks 之后,就能分配缓存:
self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim)
对 K、V 分别而言,每个隐藏层都分配 num_kvcache_blocks 个块,每块容纳 block_size × num_kv_heads × head_dim 个数值。
最后的循环很有意思。对于带有 k_cache 和 v_cache 的模块——这里只有 Attention 模块——把 self.kv_cache[0, layer_id] 赋给 k_cache,把 self.kv_cache[1, layer_id] 赋给 v_cache。这一点很重要:一个“块”不是横跨整个模型的一整段连续内存。块 i 分散在多个层中,每层对应一个切片。在同一层内,相邻块之间相隔 block_size × num_kv_heads × head_dim 个元素;同一层的块 i 与下一个块是连续的。
译者注:原文末尾的“Block i+i”按明显笔误理解为相邻的 Block i+1。块 ID 在各层对应同编号的切片,但这些跨层切片不构成单一连续大块。
如果使用多个 GPU,每张 GPU 只负责一部分 KV 头:num_kv_heads = num_key_value_heads // world_size。
这种内存布局,以及同一个 KV 块跨层时并不连续这一事实,是最让我意外的地方之一。我也不知道自己为什么原本以为它必须连续。
Attention 模块如下:
class Attention(nn.Module):
def __init__(self, num_heads, head_dim, scale, num_kv_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = head_dim
self.scale = scale
self.num_kv_heads = num_kv_heads
self.k_cache = self.v_cache = torch.tensor([])
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
context = get_context()
k_cache, v_cache = self.k_cache, self.v_cache
if k_cache.numel() and v_cache.numel():
store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
# call flash_attn注意,k_cache 和 v_cache 一开始只是占位张量,之后才会在 allocate_kv_cache 中设置。这个 Attention 模块是在计算 q、k、v 的模块之后调用的。因此每次模型运行时,新 k、v 都已经准备好,再通过 store_kvcache 写入缓存。
先退一步看。我们需要把新算出的 k、v 存进分页缓存,那么怎么知道每个新 k、v 应该写到哪里?前面说过,需要把逻辑连续的虚拟块映射到物理块。
每个序列都有一个 block_table 属性,本质上是列表。如果 block_table = [22, 4, 43],就表示虚拟块 0 映射到物理块 22,虚拟块 1 映射到物理块 4,虚拟块 2 映射到物理块 43。假设每块容纳 256 个 token,这也是 nano-vLLM 的默认值。
这个序列中编号为 257 的 token,其 KV 应该存在哪里?它属于虚拟块 1,因此对应物理块 4;又因为它是块中的第二个 token,所以位置是:物理块 4,块内索引 1。
译者注:这里 token 编号从 0 开始。索引 257 = 1 × 256 + 1,所以是第二个逻辑块的第二个槽位;若说自然语言中的“第 257 个 token”,它的零基索引应是 256。
因此,只要知道 block_table 中的虚拟到物理映射,就能定位某个 token 的 KV 在物理内存中的位置。这个映射预先计算并放入 slot_mapping,后者是 store_kvcache 的最后一个参数,用来告诉它新 K、V 应该写在哪里。
理解虚拟到物理映射与块大小的作用后,下面构造 slot_mapping 的代码就容易理解了:
def prepare_prefill(self, seqs: list[Sequence]):
# ...
slot_mapping = []
for seq in seqs:
for i in range(seq.num_cached_blocks, seq.num_blocks):
start = seq.block_table[i] * self.block_size
if i != seq.num_blocks - 1:
end = start + self.block_size
else:
end = start + seq.last_block_num_tokens
slot_mapping.extend(list(range(start, end)))
set_context(..., slot_mapping=slot_mapping, ...)seq.block_table[i] 取出存放虚拟块 i 的物理块编号。乘以 block_size,得到该块第一个 token 的物理槽位索引,再用 range 顺序填入其他槽位。注意那个判断是否到达最后一块的 if:最后一块可能未满,因此用 seq.last_block_num_tokens 确定结束位置。
Decode 更简单,因为每个序列只需要计算一个槽位:
def prepare_decode(self, seqs: list[Sequence]):
# ...
slot_mapping = []
for seq in seqs:
slot_mapping.append(
seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1
)
# ...
set_context(..., slot_mapping=slot_mapping, ...)prepare_prefill 和 prepare_decode 都在运行模型之前调用。它们把用于存储 KV 的 slot_mapping 放进全局 context,Attention 模块再通过 get_context() 取出。
现在已经知道每个 token 的 KV 应该存放的物理槽位,需要把该 token 的 Key、Value 向量写进去。为了高效完成这一步,我们不使用普通的 PyTorch 张量操作,而是更进一步,用 Triton 做较底层的 GPU 操作。Triton 位于高层 PyTorch 与低层 C/C++ CUDA 之间。
def store_kvcache_kernel(
key_ptr,
key_stride,
value_ptr,
value_stride,
k_cache_ptr,
v_cache_ptr,
slot_mapping_ptr,
D: tl.constexpr,
):
idx = tl.program_id(0)
slot = tl.load(slot_mapping_ptr + idx)
if slot == -1:
return
key_offsets = idx * key_stride + tl.arange(0, D)
value_offsets = idx * value_stride + tl.arange(0, D)
key = tl.load(key_ptr + key_offsets)
value = tl.load(value_ptr + value_offsets)
cache_offsets = slot * D + tl.arange(0, D)
tl.store(k_cache_ptr + cache_offsets, key)
tl.store(v_cache_ptr + cache_offsets, value)
def store_kvcache(
key: torch.Tensor,
value: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
slot_mapping: torch.Tensor,
):
N, num_heads, head_dim = key.shape
D = num_heads * head_dim
assert key.stride(-1) == 1 and value.stride(-1) == 1
assert key.stride(1) == head_dim and value.stride(1) == head_dim
assert k_cache.stride(1) == D and v_cache.stride(1) == D
assert slot_mapping.numel() == N
# Launch N programs: one program per token
store_kvcache_kernel[(N,)](
key,
key.stride(0),
value,
value.stride(0),
k_cache,
v_cache,
slot_mapping,
D=D,
)Python 函数 store_kvcache 会在 GPU 上启动 N 个 store_kvcache_kernel 程序实例,N 是 token 数。每个程序处理一个 token:读取刚算出的 Key、Value 向量,然后写入 slot_mapping 指定的物理缓存槽位。
- idx = tl.program_id(0) 是这个 token 在 Key、Value 输入张量中的索引。
- slot = tl.load(slot_mapping_ptr + idx) 读取 token idx 对应的物理槽位。若 slot == −1,就跳过写入。
- idx * key_stride + tl.arange(0, D) 算出 token idx 的 Key 中 D 个元素的偏移;Value 也按同样的方式处理。
- cache_offsets = slot * D + tl.arange(0, D) 算出该物理槽位在缓存中的目标偏移。
- tl.store(k_cache_ptr + cache_offsets, key) 和 tl.store(v_cache_ptr + cache_offsets, value) 把向量写入缓存。
简单说,就是把 Key、Value 从一处 GPU 内存复制到另一处,重复 N 次,每个 token 一次;slot mapping 决定每个 token 的 KV 实际存放在哪里。
20.5.2.2 Flash Attention
KV 已经存进缓存,接下来要真正计算注意力。别忘了,缓存采用分页布局,因此注意力函数必须支持这种布局。
注意力模块中的这一区分非常明确:
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
# ...
if context.is_prefill:
if context.block_tables is not None: # prefix cache
k, v = k_cache, v_cache
o = flash_attn_varlen_func(q, k, v,
max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
softmax_scale=self.scale, causal=True, block_table=context.block_tables)
else: # decode
o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
cache_seqlens=context.context_lens, block_table=context.block_tables,
softmax_scale=self.scale, causal=True)
return oflash_attn_varlen_func 和 flash_attn_with_kvcache 是著名的 Flash Attention 库中的两个函数。简而言之,Flash Attention 用非常巧妙而高效的方式计算注意力:不把完整注意力矩阵实际存储出来,因此注意力计算所需内存大致随序列长度线性增长,而不是像朴素实现那样呈平方增长。
原文链接:https://github.com/dao-ailab/flash-attention
译者注:这里减少的是中间注意力矩阵的存储及读写,不是把普通全注意力的理论计算量从 O(L²) 改成 O(L)。
但即便使用 Flash Attention,Prefill 的 varlen_func 与 Decode 的 with_kvcache 也有不同的行为和前提。Prefill 中有大量 Query 位置需要计算注意力,因此 flash_attn_varlen_func 通常计算量较大;训练也会使用这类路径,不过训练不在推理引擎的讨论范围内。flash_attn_with_kvcache 用于 Decode:通常只有极少的新 Query,往往就是一个 token,此时很看重延迟,瓶颈经常是如何高效读取缓存中的 K/V。
Flash Attention 是独立的代码库,拥有高度优化的 CUDA 内核,也支持分页。注意 context 中传入的 block_tables:只要提供它,Flash Attention 函数就知道正在处理分页 KV,并把 k_cache、v_cache 参数解释为完整的分页缓存内存;block_table 告诉内核如何为每个序列解释这些内存。注意,varlen_func 的参数虽然叫 k、v,但在这种情况下传入的其实是 k_cache、v_cache。
从概念上看,如果序列的 block table 是 [4, 18, 3, ...],就表示逻辑块 0、1、2、……分别映射到物理缓存块 4、18、3、……。
用于构造 block_tables 的 prepare_blocks 会在 prepare_prefill 和 prepare_decode 中调用,代码如下:
def prepare_block_tables(self, seqs: list[Sequence]):
max_len = max(len(seq.block_table) for seq in seqs)
block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs]
block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
return block_table译者注:原文片段定义的是 prepare_block_tables,前文写成 prepare_blocks;片段末尾 return block_table 与局部变量 block_tables 也不一致。此处保留原代码,阅读时应识别这些简写或笔误,不能把所有摘录直接当作可运行程序。
把各序列的 seq.block_table 放在一起,再用 −1 填充较短的表,让每条序列的块表长度一致。−1 代表分页内核接口约定的“未使用”项。
Q、K 都需要最大长度,是因为 CUDA 内核需要 max_seqlen_q、max_seqlen_k 这样的上界,用于特化和边界处理。
顾名思义,varlen_func 支持不同长度的序列,因此需要 cu_seqlens_q,即各条 Query 序列长度的累计值。每个需要计算注意力的输入 token 都会贡献 Query 向量。例如三条序列长度是 4、8、3,那么 cu_seqlens_q 就是 [0, 4, 12, 15],长度为 B+1。cu_seqlens_k 同理,记录 K/V 的累计长度,也就是每条序列可以读取多少个缓存位置。简单情况下,K/V 位置数与序列或上下文长度一致;涉及前缀缓存或截断时,把它理解为“实际参与注意力计算的上下文长度”更稳妥。
因此,我们实际上是在为 Flash Attention 准备元数据,让它高效地完成分页注意力。它的内部实现,也很适合未来再写一篇文章来讨论 :D。
20.6 在多个 GPU 上运行模型
nano-vLLM 支持把模型切分到多个 GPU。如果模型权重装不进单张 GPU,这就至关重要。前面的示例使用 Qwen 0.6B。以 FP16,也就是每个权重 2 字节计算,每 1B 个权重需要 2 GB。所以一个 70B 参数模型需要 140 GB 显存;这还没有计入 KV Cache,只是权重本身。
量化可以帮忙:不使用 FP16 的 2 字节表示,而把权重映射到精度较低、占用较小的表示,例如 INT8,每个元素 1 字节,或者更激进的 INT4,后者的内存占用是完整 FP16 的四分之一。
总之,利用多个 GPU 很重要。nano-vLLM 支持多 GPU,但仅限单节点。完整的 vLLM 可以跨多个节点运行,每个节点再包含多张 GPU,不过这会增加复杂性,带来更多分布式系统挑战。
那么,nano-vLLM 是怎么实现单节点多 GPU 的呢?
首先,要把模型权重切分到多个 GPU,这叫张量并行。还有其他并行方式,例如数据并行:模型能装进单张 GPU,在多个 GPU 上分别运行它的副本,处理不同数据。它有助于加速训练,nano-gpt 就采用了这种方式。
另一种是层并行:不是切分每一层内部,而是把完整的层分配出去。还有专家并行,把 MoE 中的专家放到各自的 GPU 上。
这里使用的是张量并行 TP,也就是在每一层内部,把权重切分到多个 GPU。
PyTorch 中的 world_size 表示 GPU 或进程的数量,两者一一对应。每个进程,也就是 GPU,都有一个 rank。使用 8 张 GPU 时,world_size 为 8,rank 从 0 到 7。进程 0 通常是主进程,负责簿记和协调。
20.6.1 ModelRunner 的并行执行
通过 tensor_parallel_size 配置进程与 GPU 数量。模型权重切分到 N 张 GPU,由 N 个进程管理。编号 0 是主进程,只有它持有 LLMEngine 实例,并通过 generate 接收请求;其他进程只有 ModelRunner 实例。执行推理时,也由主进程发起,再触发其他进程各自在所负责的模型分片上运行。
有些层需要协调,各个模型分片会等待彼此、传递结果并做归约。例如,分片后的 Embedding 层中,每张 GPU 持有 1/N 的词嵌入向量。我们执行 dist.all_reduce 把结果汇合起来:严格说这是求和,但每张 GPU 都把不属于自己范围的 token 对应向量置零,所以求和能得到相当于汇集后的结果。RowParallelLinear 等层也是类似的道理。
不过,拿到全部 logits 并转成下一个 token 的概率之后,最终采样只发生在主进程。
主进程使用 multiprocessing 的两个基础组件向其他进程派发工作:Event 和 SharedMemory。稍后会详细讨论,先看 LLMEngine 的构造函数:
# wrapper around Python's built-in `multiprocessing`
import torch.multiprocessing as mp
class LLMEngine:
def __init__(self, model, **kwargs):
config_fields = {field.name for field in fields(Config)}
config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
config = Config(model, **config_kwargs)
self.ps = []
self.events = []
ctx = mp.get_context("spawn")
for i in range(1, config.tensor_parallel_size):
event = ctx.Event()
process = ctx.Process(target=ModelRunner, args=(config, i, event))
process.start()
self.ps.append(process)
self.events.append(event)
self.model_runner = ModelRunner(config, 0, self.events)对于 tensor_parallel_size 中的每张非主 GPU,都创建一个 Event 和一个运行 ModelRunner 的 Process。每个 Runner 接收自己的 Event 实例以及作为 rank 的 i;之后会根据 rank 加载自己负责的权重分片。
循环从 range(1, config.tensor_parallel_size) 开始,看起来少了一个,那就是主进程。它紧接着在后面创建,并接收所有 Event。
进程调用 Event.wait() 后会阻塞,直到另一个进程对这个 Event 调用 set()。等待方被唤醒后执行工作,再用 clear() 重置 Event,于是下次 wait() 又会阻塞。
因此,每当有工作需要执行,例如在模型分片上运行推理,主进程就对其他进程各自的 Event 调用 set()。它们此前一直在等待,完成工作后,又回到等待新工作的状态。
还差一块拼图:每个进程怎么知道具体要做什么?这就要用 SharedMemory。它是操作系统提供的共享内存机制,多个进程都能读写同一块内存。因此,主进程可以把工作描述写入共享内存,让其他进程读取。
class ModelRunner:
def __init__(self, config: Config, rank: int, event: Event | list[Event]):
self.config = config
hf_config = config.hf_config
self.block_size = config.kvcache_block_size
self.enforce_eager = config.enforce_eager
self.world_size = config.tensor_parallel_size
self.rank = rank
self.event = event
dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)
torch.cuda.set_device(rank)
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(hf_config.torch_dtype)
torch.set_default_device("cuda")
self.model = Qwen3ForCausalLM(hf_config)
load_model(self.model, config.model)
self.sampler = Sampler()
self.warmup_model()
self.allocate_kv_cache()
if not self.enforce_eager:
self.capture_cudagraph()
torch.set_default_device("cpu")
torch.set_default_dtype(default_dtype)
if self.world_size > 1:
if rank == 0:
self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
dist.barrier()
else:
dist.barrier()
self.shm = SharedMemory(name="nanovllm")
self.loop()def exit(self):
if self.world_size > 1:
self.shm.close()
dist.barrier()
if self.rank == 0:
self.shm.unlink()
if not self.enforce_eager:
del self.graphs, self.graph_pool
torch.cuda.synchronize()
dist.destroy_process_group()def loop(self):
while True:
method_name, args = self.read_shm()
self.call(method_name, *args)
if method_name == "exit":
breakdef read_shm(self):
assert self.world_size > 1 and self.rank > 0
self.event.wait()
n = int.from_bytes(self.shm.buf[0:4], "little")
method_name, *args = pickle.loads(self.shm.buf[4:n+4])
self.event.clear()
return method_name, argsdef write_shm(self, method_name, *args):
assert self.world_size > 1 and self.rank == 0
data = pickle.dumps([method_name, *args])
n = len(data)
self.shm.buf[0:4] = n.to_bytes(4, "little")
self.shm.buf[4:n+4] = data
for event in self.event:
event.set()def call(self, method_name, *args):
if self.world_size > 1 and self.rank == 0:
self.write_shm(method_name, *args)
method = getattr(self, method_name, None)
return method(*args)说不清为什么,我觉得这段代码尤其简洁优美。每个 ModelRunner 实例都创建自己的模型实例,初始化自己的 KV Cache。在初始化末尾,主进程创建 SharedMemory,其他进程则打开它并进入 self.loop。顾名思义,这是一个持续运行的循环,直到收到 exit 调用才结束。
循环里调用 read_shm,它只应在非主进程执行。该函数会阻塞在 self.event.wait()。搜索代码库会发现,event.set() 只出现在主进程执行的 write_shm 中。
执行方式是:先在主进程上调用 call——是的,名字就叫 call!在 call 内部,如果 world_size > 1,说明还有其他进程,就调用 self.write_shm,把方法名和参数写进共享内存,再通过 set() 唤醒其他进程。就是这样。
还有一个有意思的细节:Sequence 类定义了 __getstate__ 和 __setstate__。因为写入 run 的参数时会用 Pickle 序列化,而这些参数包含 Sequence 实例,所以需要这两个方法。
因此,借助 Event 和 SharedMemory,把方法名与参数写入公共区域,让其他进程持续循环等待工作,就能在 N 个进程及 GPU 上调用同一个方法。
这些分布式方法必须仔细编写,才能正确协调,而 PyTorch 的 dist 帮助完成了这一点。同时,这些方法也必须能在只有一个 GPU 或进程时工作,所以代码会检查 world_size。这种设计确实很漂亮。
20.6.2 Embedding 层
Embedding 层本质上是一张查找表,把每个 token ID 映射到一个词嵌入向量。在 TP 中,对应代码是:
self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)希望每张 GPU 负责词表的一部分。如果有 8 张 GPU,每张就加载 vocab_size / 8 个词嵌入。前向计算时,确保每张 GPU 只处理属于自己区间的 token:GPU 0 负责 [0, vocab_size/8),GPU 7 负责 [7*vocab_size/8, vocab_size)。
class VocabParallelEmbedding(nn.Module):
def __init__(self, num_embeddings: int, embedding_dim: int,):
super().__init__()
self.tp_rank = dist.get_rank()
self.tp_size = dist.get_world_size()
assert num_embeddings % self.tp_size == 0
self.num_embeddings = num_embeddings
self.num_embeddings_per_partition = self.num_embeddings // self.tp_size
self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank
self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition
self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim))
self.weight.weight_loader = self.weight_loader
def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
param_data = param.data
shard_size = param_data.size(0)
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
param_data.copy_(loaded_weight)
def forward(self, x: torch.Tensor):
if self.tp_size > 1:
mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
x = mask * (x - self.vocab_start_idx)
y = F.embedding(x, self.weight)
if self.tp_size > 1:
y = mask.unsqueeze(1) * y
dist.all_reduce(y)
return y每张 GPU 负责 num_embeddings_per_partition 个词嵌入,通过 weight_loader 只加载这一部分权重。
前向计算时,从 Embedding 表中查出向量,同时构造 mask:对于不属于当前 GPU 区间的 token,mask 为 false,也就是 0。y = mask.unsqueeze(1) * y 保证这些越界输入对应的 y 都为 0。
最后调用 dist.all_reduce(y),在所有 GPU 间求和,其默认操作是 SUM。对于某个 token,除负责其词表区间的 GPU 外,其他 GPU 的向量都是 0,因此求和后的结果包含全部输入 token 的词嵌入,即使 Embedding 表本身分散在不同 GPU 上。
注意,这个模块在 tp_size = 1 时也能正常工作。此时只有一张 GPU,整张 Embedding 表都在它上面,num_embeddings_per_partition == num_embeddings,全部权重都加载到这张 GPU;前向计算则跳过 mask 和 all_reduce。
这是并行设计的一项关键要求:数量为 1 的基本情况必须始终成立。
20.6.3 MLP 的张量并行:列切分与行切分技巧
这可能是最令人印象深刻的技巧之一。
一个 MLP 层可以写成 Y = XW,W 是权重矩阵。由于矩阵有两个维度,自然就有两种切分方式:把列分配到多个 GPU,或者把行分配到多个 GPU。
译者注:Y=XW 描述的是线性投影,不是完整 MLP。实际 Qwen3 MLP 还包含 gate/up 两路投影、SiLU 与逐元素门控,以及 down 投影,详见本文第 12 章和第 19 章。
分别想一想这两种情况。
20.6.3.1 按列切分
GPU0:X × W0 = Y0;GPU1:X × W1 = Y1;依此类推。
每张 GPU 都需要完整的输入 X。输出中的每个元素,是 X 的一行与 W 的一列做点积得到的。因此,切分 W 的列以后,每张 GPU 负责输出中的一部分列。
这样得到 Y0、Y1、……,每份都包含输出列的 1/N。
但这些输出还要成为下一层的输入,下一层通常又是矩阵乘法。因此,在继续之前,需要通过 all_gather 把完整 Y 拼回来。
译者注:只有下一算子要求完整输入时才需要先 all_gather。下文列并行接行并行的技巧,正是利用下一层可直接消费分片,省掉这次通信。
20.6.3.2 按行切分
现在换一种方式,每张 GPU 持有 W 中 1/N 的行。
情况就不同了。每个输出值都依赖 W 的完整一列;只持有一部分行,就只能算出该点积的一部分和。
举个例子:
# X = [ 1 2 3 4 ]
# W = [ 1 2
# 3 4
# 5 6
# 7 8 ]
X * W = [ 50 60 ]如果按行切分:
W_top:
[ 1 2 ]
[ 3 4 ]
W_bot:
[ 5 6 ]
[ 7 8 ]
x_left = [ 1 2 ]
x_right = [ 3 4 ]上半部分:
x_left * W_top =
[ (1*1 + 2*3) (1*2 + 2*4) ]
= [ 7 10 ]下半部分:
x_right * W_bot =
[ (3*5 + 4*7) (3*6 + 4*8) ]
= [ 43 50 ]把它们相加:
[ 7 10 ] + [ 43 50 ] = [ 50 60 ]因此按行切分时,每张 GPU 算出的局部 Y 都与最终输出形状相同,但每个元素只包含完整点积中约 1/N 的贡献。最后需要 all_reduce 求和,才能得到正确结果。
对比这两种方法。
列并行:
- 每张 GPU 生成输出列的 1/N,其中每个值都已经完整算出。
- 需要 all_gather 才能拼出完整 Y。
行并行:
- 每张 GPU 都生成完整形状的输出,但每个值只是局部贡献。
- 需要 all_reduce 求和才能合成最终结果。
真正巧妙的地方来了:
- 行并行只需要输入中对应的切片,例如 x_left、x_right。
- 列并行恰好生成这种分片输出,也就是输出的一部分列。
所以,列并行的输出形式,正好就是行并行所需的输入形式。
如果把 MLP 中的层按下面的方式排列:
[ column-parallel ] -> [ row-parallel ]那么:
- 第一层生成分片输出,不必聚合。
- 第二层直接消费这些分片。
- 只需要在最后通过 all_reduce 通信一次。
这样,两层都完成了分布式计算,却只付出一次通信的代价。
这就是其中的技巧。
译者配图|列并行与行并行。下图复用本文第 14 章的原创可编辑画板,不是原文插图。重点看“输出分片”和“部分和”的区别,以及末尾的 AllReduce。图中用 Attention 的 o_proj 展示行并行;MLP 的列接行采用相同的分片衔接原则。
来看实际代码。
class LinearBase(nn.Module):
def __init__(self, input_size: int, output_size: int, tp_dim: int)
self.tp_dim = tp_dim
self.tp_rank = dist.get_rank()
self.tp_size = dist.get_world_size()
self.weight = nn.Parameter(torch.empty(output_size, input_size))
self.weight.weight_loader = self.weight_loaderColumnParallelLinear 是采用张量并行的线性层。关键的列切分发生在 __init__ 中。注意 super().__init__(input_size, divide(output_size, tp_size), bias, 0):各 GPU 的行数,也就是输入维度保持相同,而列数,也就是输出维度,缩小为 divide(output_size, tp_size)。这里 tp_size 就是 world size,即 GPU 数量。
还要注意传给父类构造函数的最后一个参数 tp_dim,它表示张量并行的切分轴。列并行时 tp_dim 为 0。由于权重矩阵形状是 (output_size, input_size),沿第 0 维切分,就是把行分到不同 GPU。原文接着将其描述为“从每一行中取 1/N 的列”。
译者注:原文在此混用了数学写法与 PyTorch 存储布局,最后一句与前一句矛盾。数学 Y=XW 中 W 为 [input, output],“列并行”切输出列;PyTorch F.linear 存 weight=[output, input],计算 X·weightᵀ,所以同一件事对应切存储矩阵的第 0 轴,即行。每个分片保留全部输入列。
这一点也体现在 weight_loader 中。加载权重时,沿 tp_dim 执行 loaded_weight.narrow(self.tp_dim, start_idx, shard_size)。也就是说,每张 GPU 沿第 0 维加载一段连续的行分片。GPU 0 从索引 0 开始,编号为 N 的 GPU 则从 shard_size × N 开始。
最后注意 forward:这里只执行矩阵乘法,没有调用 dist,因此没有 gather 或 reduce。这是因为每张 GPU 只生成输出的一部分特征。通常假定后面会跟着 RowParallelLinear,或者其他能正确合并结果的操作。
说到这里,就继续看行并行。
class RowParallelLinear(LinearBase):
def __init__(self, input_size: int, output_size: int, bias: bool = False):
tp_size = dist.get_world_size()
super().__init__(divide(input_size, tp_size), output_size, bias, 1)
def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
param_data = param.data
shard_size = param_data.size(self.tp_dim)
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
param_data.copy_(loaded_weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
if self.tp_size > 1:
dist.all_reduce(y)
return y它与 ColumnParallelLinear 很像,但有三个关键区别:
- 这里用 input_size 除以 tp_size;列并行则是用 output_size 除以 tp_size。
- 传给父类构造函数的最后一个参数是 tp_dim = 1。因为权重矩阵形状是 (output_size, input_size),沿第 1 维切分就是把列分到不同 GPU。原文接着将其描述为“每张 GPU 持有每一列中 1/N 的行”。
译者注:行并行同理:数学 W=[input, output] 切输入行;PyTorch weight=[output, input] 则切第 1 轴,即输入列。原文最后一句又混入了另一种布局,应以前面的形状与 tp_dim=1 为准。
- forward 执行 F.linear 后,调用 dist.all_reduce(y),把所有 GPU 的局部输出相加。每张 GPU 只看到输入的一部分,因此只能计算完整输出的一部分贡献。求和后,每张 GPU 都得到完整结果。
必须做这次归约,因为下一阶段需要完整的输出张量,而不是局部贡献。
20.6.4 把 Attention 层切分到多个 GPU
模型的核心是注意力层:
Given input X:
Q = X @ W_Q
K = X @ W_K
V = X @ W_V
Attention scores (causal / with masking):
S = Q @ K^T
Scaled scores:
S_scaled = S / sqrt(d_k)
Apply softmax:
A = softmax(S_scaled)
Output:
O = A @ V因此,从技术上说,这里关心三组权重:W_Q、W_K 和 W_V。
上面的数学过程还不完全等于实际实现。每个权重矩阵实际上会分成 NH 部分,每部分对应一个头。各个头独立完成计算,直到最后算出 O = A @ V 后,才把每个头的输出拼接起来得到结果。
前三个矩阵乘法,也就是把 X 分别投影到 W_Q、W_K、W_V,是否拆分计算并不改变结果。但实际注意力计算、Softmax 和每个头的输出计算,都发生在头内部。第 i 个头的 Query 只关注对应头的 Key 和 Value。
译者注:这里“对应头”在 MHA 中是一一对应,在 GQA 中是多个 Q 头对应同一个 KV 头。不是把全部 Q/K/V 都按同样的头数拆分;也不是跨头先做 Softmax 再相加。
前三次矩阵乘法及其权重,可以视为普通线性层。由于投影结果之后会拆成多个头,而各头之间的计算相互独立,所以很适合列并行:每个 ColumnParallel 模块输出 1/N 的特征列,这些列对应独立的头。
接下来稍微复杂一些。Hugging Face 权重文件中,注意力权重分开存放在不同层里:
原文链接:https://huggingface.co/Qwen/Qwen3-0.6B?show_file_info=model.safetensors
model.layers.0.self_attn.k_proj.weight [1 024, 1 024]
model.layers.0.self_attn.o_proj.weight [1 024, 2 048]
model.layers.0.self_attn.q_proj.weight [2 048, 1 024]
model.layers.0.self_attn.v_proj.weight [1 024, 1 024]注意,这个 Qwen 模型中,每个 K 头和 V 头对应两个 Q 头,所以 q_proj.weight 的形状是 [2 048, 1 024]。
但在 vLLM 中,为了提高效率,QKV 会合并到同一层:
class QKVParallelLinear(ColumnParallelLinear):
def __init__(
self,
hidden_size: int,
head_size: int,
total_num_heads: int,
total_num_kv_heads: int | None = None,
bias: bool = False,
):
tp_size = dist.get_world_size()
total_num_kv_heads = total_num_kv_heads or total_num_heads
self.head_size = head_size
self.num_heads = divide(total_num_heads, tp_size)
self.num_kv_heads = divide(total_num_kv_heads, tp_size)
output_size = (total_num_heads + 2 * total_num_kv_heads) * self.head_size
super().__init__(hidden_size, output_size, bias)
def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str):
param_data = param.data
assert loaded_shard_id in ["q", "k", "v"]
if loaded_shard_id == "q":
shard_size = self.num_heads * self.head_size
shard_offset = 0
elif loaded_shard_id == "k":
shard_size = self.num_kv_heads * self.head_size
shard_offset = self.num_heads * self.head_size
else:
shard_size = self.num_kv_heads * self.head_size
shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size
param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
param_data.copy_(loaded_weight)每个 QKV 分片,也就是每张 GPU 的分片,包含 num_heads 个 Q 头,以及各 num_kv_heads 个 K 头和 V 头。从 weight_loader 可以看出打包顺序:
at offset "0": Q weights
at offset "num_heads * head_size" : K weights
at offset "num_heads * head_size + num_kv_heads * head_size": V weights因此,每张 GPU 的 QKV 分片分别保存 Q、K、V 层的 1/N。weight_loader 计算偏移,再把权重放进对应区域。
注意力部分还只讲了开头:以上只是投影,后面还要执行 Softmax、计算输出并拼接。相应代码如下:
class Qwen3Attention(nn.Module):
def __init__(self, hidden_size:int, num_heads:int, num_kv_heads:int, max_position:int=4096*32,...) -> None:
# ...
self.qkv_proj = QKVParallelLinear(hidden_size, self.head_dim, num_heads, num_kv_heads, bias=qkv_bias)
self.o_proj = RowParallelLinear(num_heads*self.head_dim, hidden_size, bias=False)
self.rotary_emb = get_rope(self.head_dim, rotary_dim=self.head_dim, max_position=max_position,
base=rope_theta, rope_scaling=rope_scaling)
self.attn = Attention(self.num_heads, self.head_dim, self.scaling, self.num_kv_heads)
def forward(self, positions:torch.Tensor, hidden_states:torch.Tensor) -> torch.Tensor:
q,k,v = self.qkv_proj(hidden_states).split([self.q_size,self.kv_size,self.kv_size], dim=-1)
q,k,v = q.view(-1,self.num_heads,self.head_dim), k.view(-1,self.num_kv_heads,self.head_dim), v.view(-1,self.num_kv_heads,self.head_dim)
if not self.qkv_bias: q,k = self.q_norm(q), self.k_norm(k)
q,k = self.rotary_emb(positions, q, k)
return self.o_proj(self.attn(q,k,v).flatten(1,-1))这段代码经过删减压缩,但核心逻辑都在。用列并行的 QKVParallelLinear,在同一层里得到 q、k、v 投影。
接着拆分并 reshape,确保后续计算按各个头分别进行。通过 RoPE 纳入位置信息,然后就能真正计算注意力:缩放、Softmax,以及对 Value 加权。Attention,也就是 self.attn,是调用 KV Cache 和 Flash Attention 的地方,分页注意力真正的优化也在这里发生。
最后一步是 self.o_proj,它是普通线性层,不过使用 RowParallelLinear,因为这种分布式线性层会在末尾执行归约,提供完整结果,这与 ColumnParallel 不同。
别忘了,每张 GPU 都有自己的进程,运行一个 ModelRunner 实例,KV Cache 也在 ModelRunner 中管理。这很合理:缓存管理的是显存,每张 GPU 都有自己的显存。而且,KV Cache 是在 self.attn(q, k, v) 内写入的;如前所述,每张 GPU 用自己的那些头处理自己负责的部分。
完整的 Qwen 模型就是这些层依次组合而成的;前面的讨论省略了少数层:
class Qwen3DecoderLayer(nn.Module):
def __init__(self, config: Qwen3Config) -> None:
super().__init__()
self.self_attn = Qwen3Attention(
config.hidden_size, config.num_attention_heads, config.num_key_value_heads,
config.max_position_embeddings, config.rms_norm_eps,
getattr(config, 'attention_bias', True),
getattr(config, 'head_dim', None),
getattr(config, "rope_theta", 1_000_000),
getattr(config, "rope_scaling", None),
)
self.mlp = Qwen3MLP(config.hidden_size, config.intermediate_size, config.hidden_act)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(self, positions, hidden_states, residual=None):
hidden_states, residual = (
(self.input_layernorm(hidden_states), hidden_states)
if residual is None else self.input_layernorm(hidden_states, residual)
)
hidden_states = self.self_attn(positions, hidden_states)
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
return self.mlp(hidden_states), residual
class Qwen3Model(nn.Module):
def __init__(
self,
config: Qwen3Config,
) -> None:
super().__init__()
self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
self.layers = nn.ModuleList([Qwen3DecoderLayer(config) for _ in range(config.num_hidden_layers)])
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
) -> torch.Tensor:
hidden_states = self.embed_tokens(input_ids)
residual = None
for layer in self.layers:
hidden_states, residual = layer(positions, hidden_states, residual)
hidden_states, _ = self.norm(hidden_states, residual)
return hidden_states把所有部分组合起来,就能看到:Attention、MLP、RMSNorm 和残差连接共同组成一个 Decoder Block,再堆叠多个这样的 Block,构成模型的核心。
译者说明(全文结束):以上保留原文全部实质段落与 39 段代码;代码及代码注释保持原样,网页导航、推荐文章、分享按钮等站点界面不属于译文。原文嵌入视频仅保留链接。标题做了中文翻译与连续编号,明显拼写问题作了说明;所有额外的技术澄清均以“译者注”标出。
21. 延伸阅读|Neutree:沿一次请求读懂 nano-vLLM(中文导读)
来源:Neutree,Understanding LLM Inference Engines: Inside Nano-vLLM (Part 1),2026-02-01;副标题为 Architecture, Scheduling, and the Path from Prompt to Token。本章是结合本文源码分析写成的原创导读,不是全文翻译,也未复制原文配图。原文中的代码与图示请通过链接阅读。
21.1 这篇文章适合解决什么问题
理解 Attention 和 KV Cache 后,仍可能不知道一个推理引擎为什么要有那么多类。这篇文章最值得参考的是它的组织方式:跟随一个请求,从入口走到调度、显存管理、GPU 执行,再回到输出。读它时,不必先记住所有类名,而应不断追问:这个组件持有什么状态,接收什么输入,把什么交给下一步?
结合本文,可以把整条链分成两种工作。模型计算决定“下一个 token 的分数是多少”;引擎管理决定“这一轮让谁计算、在哪里读写缓存、算完后谁继续谁结束”。后者不改变模型的基本数学,却决定 GPU 能否持续处理足够多的有效工作。
21.2 把类名放回各自的职责
组件 | 应该用它回答的问题 | 与本文的衔接 |
|---|---|---|
LLMEngine / Tokenizer | 文本怎样变成 token ID 和请求对象?什么时候进入生成循环? | 入口组织请求;不是在这里执行全部模型数学 |
Sequence | 输入、已生成 token、状态和块表,属于哪一个请求? | 请求级状态,不是 GPU 上的完整中间激活 |
Scheduler | 本轮选哪些请求、处理多少新位置、预算不足怎么办? | 读本文固定版本时,注意 chunked prefill 的具体策略 |
BlockManager | 请求的逻辑 KV 块对应哪些物理块,何时复用或释放? | 管理块 ID、引用计数、哈希等元数据,不执行 Attention |
ModelRunner | 如何准备输入、位置、槽位映射和执行上下文,并发起模型运行? | 连接 CPU 调度结果与 GPU 计算 |
模型、Attention 与采样器 | 如何读写 KV、得到 logits,再选出新 token? | 新 token 回到 Sequence,下一轮继续调度 |
这张职责表可以与第 6 章的请求流程图、第 20.4 节的译者配图对照阅读。最关键的边界是:BlockManager 知道“这个请求用了哪些块”,ModelRunner 和 Attention 才会处理“GPU 缓存中的 K/V 数值”。块表不是 KV 数据本身。
21.3 带着一个跨块例子阅读
建议在阅读时只跟踪一条即将跨过块边界的请求。上一轮采样得到一个 token,Sequence 已经把它追加进去,但它的 KV 还没有通过下一轮模型计算写入。于是下一轮先检查是否需要新块,再构造 slot_mapping,随后才执行各层投影和 KV 写入。这样就能把“追加 token ID”“分配槽位”“写入 KV”三个不同时刻分开。
再加入第二条请求,观察两条请求如何共用一次模型执行。它们共享模型权重和执行批次,但不应读取彼此的上下文;序列边界、长度信息和 block tables 会让注意力内核区分请求。连续批处理的意义也就不再只是“把列表变长”,而是在每一轮重新组织有效工作。
21.4 阅读时保留版本与术语边界
文章适合建立架构直觉,但“production-grade”等定位不宜直接套到本文所分析的精简版本。是否具备在线流式接口、服务治理、特定调度策略或任意模型适配,需要另外核对实现。本文固定版本的 Runner 直接构造 Qwen3ForCausalLM,不能因为它读取了 AutoConfig,就推断为能够自动运行任意模型架构。
还有两个术语值得校准:logits 是尚未归一化的候选分数,不是概率分布;Event.wait() 是阻塞等待通知,不能简单理解为持续忙轮询。分清这些边界后,这篇文章很适合作为第 6—11 章的第二遍阅读材料。
22. 延伸阅读|Aleksa Gordić:从推理引擎到高吞吐在线服务(中文导读)
来源:Aleksa Gordić,Inside vLLM: Anatomy of a High-Throughput LLM Inference System,2025-08-29。作者明确以 vLLM 提交 42172ad(2025-08-09) 的 V1 引擎为分析对象。本章是原创中文导读与对照分析,不是全文翻译;没有复制其全部代码或配图。原文的交互示例、系统图和参考文献请在原站阅读。
22.1 它补上了 nano-vLLM 之外的哪一层
nano-vLLM 帮助我们看清引擎的最小骨架;这篇文章则把镜头拉远,展示它怎样发展成可接收并发网络请求的完整系统。阅读价值不在于多记一些高级特性,而在于分清三条边界:会运行模型,不等于会高效安排多个请求;会批量推理,不等于提供在线服务;能跑出高吞吐,也不等于满足用户的延迟要求。
可以把原文当作第 5 章“为什么需要推理框架与服务”的扩展阅读。底部是模型前向计算,中间是调度器、执行器与 KV 管理,上面才是 API、异步输入输出和副本路由。一个 HTTP 请求往下走时逐步变成模型可执行的输入,结果往上返回时则重新变成带有请求身份、结束原因和用量信息的响应。
22.2 用“解决什么瓶颈”理解高级特性
技术 | 主要改变什么 | 不能因此推断什么 |
|---|---|---|
Chunked Prefill | 把长 prompt 的新计算分摊到多轮,让调度器安排其他工作 | 不是只处理 prompt 的最后一个 token,也不是省掉其余 token 的计算 |
Prefix Cache | 复用已经算过的相同前缀 KV,减少重复 Prefill | 不会让后续 Query 不再读取前缀 K/V;匹配条件也不只是当前块文本相同 |
约束解码 | 根据格式或语法状态限制允许采样的候选 token | 格式合法不等于事实正确,也不是重新训练模型 |
推测解码 | 先提出多个候选,再由目标模型批量验证,争取一次确认多个 token | 不是未经验证就采用小模型答案;收益取决于候选质量、验证成本和负载 |
P/D 分离 | 让 Prefill 与 Decode 使用不同执行资源,并转移所需 KV | 不是没有传输成本,也不是每一轮 Decode 都要重新搬运完整历史 |
这些技术不在同一层面上:分块和前缀缓存改变要安排多少工作;约束解码改变可选输出;推测解码改变候选与验证流程;P/D 分离改变工作部署在哪里。按瓶颈分类,比把它们当作一串默认全部开启的开关更容易理解。
尤其要区分 P/D 分离与本文前面的 KV Cache:缓存首先是同一个模型执行过程中避免历史重算的机制;跨实例转移 KV 则是额外的系统能力,还需要传输协议、缓存布局兼容、资源管理和调度配合。原文的共享存储示例用于解释机制,不等于所有生产部署都采用同一种传输方式。
22.3 多 GPU 与多副本不是同一件事
张量并行把同一个模型的一层切到多张 GPU 上共同计算,通信发生在模型内部;流水线并行按层或阶段分工;数据并行副本则接收不同请求,主要依靠路由分摊流量。模型装不下与请求太多,通常是不同的问题,不能只用一个“多卡”概念概括。
读原文的在线服务部分时,建议追踪请求身份与队列,而不是只盯住 GPU:API 接入后,请求如何进入合适的引擎;引擎的输出如何返回对应连接;取消、结束与异常又如何清理资源?这些是裸模型 forward() 不负责、精简推理引擎也未必完整提供的能力。原文通过异步客户端、引擎进程和执行器说明了这种分层;类名与进程组织方式则有明确的版本范围。
22.4 最后用指标判断设计是否有价值
推理优化必须先回答“为谁优化”。交互用户在意多久出现第一个字,以及后续是否持续流畅;离线批量任务往往更看重总完成量。TTFT 衡量首 token 等待时间,ITL 衡量相邻输出 token 的间隔,端到端延迟衡量整次请求完成时间。吞吐量统计单位时间处理量,而 goodput 只统计满足服务目标的有效完成量。
因此,评估一个调度策略不能只报 tokens/s。还要说明输入输出长度分布、并发或到达速率、缓存命中情况,以及延迟分位数。批量变大可能更充分地利用 GPU,却也可能加重排队或单步延迟;具体拐点取决于硬件、模型与工作负载,不能套用一条固定经验。
推荐阅读顺序是:先用本文第 1—19 章建立计算与 nano-vLLM 的具体模型,再读第 20 章译文核对分页缓存和 TP,接着用第 21 章串起请求链路,最后到这篇 vLLM 文章看生产系统扩展。遇到具体 API、默认值、支持范围或性能数字时,回到对应提交与实验条件,而不把旧版文章当成当前产品说明书。
23. 附录|vLLM 首发博客:中文摘要与工程解读
来源:vLLM 团队,vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention,2023-06-20。本附录提供原文要点的中文摘要,以及结合本文的原创工程解读,不是全文译文;原文图示、完整基准结果和发布背景请在原站阅读。下文将原文报告的结果与我们的教学算例分开,nano-vLLM 的实现仍以本文固定提交 bb823b3e06983d71485a8e1f23715ebd87d98ef8 为准。
23.1 原文要点:从 KV 显存管理改善服务吞吐
这篇发布文章的核心观点是:服务很多条不断增长的序列时,KV Cache 不只是“保存计算结果”,还是限制并发的重要资源。为每条请求预留大块连续空间,容易产生未使用的预留容量与碎片。PagedAttention 将逻辑上的连续上下文映射到可分散存放的物理块,并让 Attention 内核按块表读取,从而更有效地使用 KV 显存。
原文还介绍了多输出生成中的缓存共享:相同 prompt 可以共用物理块,用引用计数管理持有者,用写时复制保护需要分叉的内容。作者通过当时的吞吐测试和服务部署说明其价值。它改善的是整个推理服务的资源使用,不是把 Transformer 换成另一套预测公式。
23.2 工程解读:省显存,为什么可能让服务更快
先把“容量”和“带宽”分开。容量决定有多少条请求的历史 K/V 能同时留在 GPU;带宽影响每轮把权重、K/V 等数据读到计算单元要多久。分页首先改善前者,不能单凭块分配更灵活,就推断每次读取都更快。
下面是本文自设的容量算例,不是博客实验。假设模型与数据类型固定,除去权重和其他开销后,KV 池能容纳 8192 个 token 位置;每个位置的容量包含该模型所有层需要的 K/V。各请求当前都有 400 个有效位置,最大允许长度为 1600。暂不考虑前缀共享、元数据开销和后续增长,只比较当前这一时刻的容量:
分配方式 | 每条请求占用 | 池中最多容纳多少条 |
|---|---|---|
按最大长度预留连续空间 | 1600 个位置,其中 1200 个尚未使用 | ⌊8192 / 1600⌋ = 5 条 |
按 256 个位置一块分配 | 2 块,即 512 个位置;尾部空余 112 个 | ⌊8192 / 512⌋ = 16 条 |
如果有足够多的待处理请求,调度器就有机会把更多请求的新位置放进同一轮。Decode 从很少的行变成较多的行,往往能让一次权重读取服务更多位置,并提高计算单元的利用率。因此,更少的容量浪费 → 更多可驻留请求 → 更充分的批处理 → 更高的总吞吐,才是这里的因果链。调度器必须把新增容量用起来;只有一个请求,或瓶颈已经转到计算、KV 读取、通信时,这条链的收益就会减弱。
16 条对 5 条只是该算例的容量差,不能直接当作吞吐倍数。请求增长还会继续申请块,调度器必须处理容量不足;batch 更大也可能增加单步耗时或排队等待。要判断体验是否改善,应同时看首 token 延迟、后续 token 间隔与吞吐,而不是只看能装下多少条请求。
这个例子也说明分页不是零浪费:每条请求的尾块空余占已分配容量的 112/512,约 21.9%。它与原文报告的低浪费比例并不矛盾,因为比例取决于块大小和请求长度分布。具体地址换算与分页 Attention 计算见第 9 章。
23.3 工程解读:共享块为什么还需要引用计数与写时复制
共享不只是把两张块表填成相同编号。假设为了手算,块大小取 4,两个输出分支目前共享一个已有 3 个位置的块。如果它们接下来把不同的新 K/V 都写到第 4 个槽位,就会互相覆盖。引用计数只能告诉我们“有几个持有者”,不能自动防止这种写冲突。
写时复制(Copy-on-Write)的含义是:读取时共享,需要独立修改时再分离。在上述教学场景中,一个分支先取得独占块,复制需要保留的已有 K/V,更新自己的块表,然后写入新位置;其他分支继续引用原块。若共享的是已经填满、以后不再修改的前缀块,分叉后的新位置直接使用各自的新块,就不需要为了分叉复制这些完整前缀。
由此可以区分三件事:块表决定请求读哪份数据,引用计数决定何时可以回收,写入策略决定共享数据会不会被破坏。原文展示的多输出共享,是由同一 prompt 派生多个输出;跨请求 Prefix Cache 则还要识别后来请求是否具有相同的、已经算好的前缀。两者可以利用相同的分页存储基础,但不是同一个功能。
这份 nano-vLLM 采用更简单的边界:前缀查找不包含请求的最后逻辑块,可复用的是更早的、已计算完成的完整前缀;最后逻辑块由请求独占。因此不要因为看到了 ref_count,就写成“nano 已实现与原始 vLLM 相同的通用写时复制”。其完整前缀匹配、回收和末块独占规则见第 10.3 节。
23.4 读性能数字:保留测试条件,不把历史结果当保证
博客的吞吐对比采用 LLaMA-7B / NVIDIA A10G,以及 LLaMA-13B / NVIDIA A100 40GB;请求的输入输出长度采样自 ShareGPT,并分别考察每请求一个输出和三个并行输出。它描述的是 2023 年报告中的系统、模型与负载,不是本文 Qwen3-0.6B 的实测,也不是当前各框架版本的横向结论。
原文报告 | 应该怎样理解 |
|---|---|
相对 HF 最高 24 倍、相对 TGI 最高 3.5 倍吞吐 | 来自文中不同配置与输出数量的服务吞吐对比;不是单次 q·k 快了这些倍数,也不代表每条请求延迟同比下降。 |
既有系统因碎片与过度预留浪费 60%–80% 内存 | 这是作者针对当时所考察系统的观察,不能套到任何框架、任何工作负载,更不能当成所有 GPU 总显存的固定比例。 |
PagedAttention 实践中浪费低于 4% | 对应原文的 KV 存储利用情境,不是数学上保证每种块大小和请求长度都低于 4%;本附录容量算例就有更高的尾块空余。 |
这些数据证明了该系统设计在相应条件下的价值,但仅凭引擎间的总体吞吐比较,不能把全部提升精确归因于某一个开关。今天复现时,还应固定模型、精度、硬件、版本、输入输出长度、到达负载和缓存命中条件;记录延迟与吞吐,再做控制变量对照。本文的验证方法见第 15.4 节。
23.5 对照 nano-vLLM:继承的是思路,不是全部实现
读完原文后,可以回到源码逐项核对。下表描述本文固定的 nano-vLLM 提交,不用于概括所有同名 fork 或后续版本。
要核对的机制 | 本文 nano-vLLM 的做法 | 阅读落点 |
|---|---|---|
分页分配与增长 | CPU 管理块表;首次按已知序列长度建立所需映射,Decode 跨块时再追加。不是预留整个最大上下文,也不是只给当前 Prefill chunk 分配空间。 | BlockManager;第 9.2–9.3 节 |
GPU 怎样读分页 K/V | 用 Triton 内核写入新 K/V,再调用支持块表的 FlashAttention 接口;不是直接照搬首发 vLLM 的 Attention 内核。 | layers/attention.py;第 9.5–9.6 节、第 13 章 |
共享与写入安全 | 前缀 hash 查找、完整块复用和引用计数;最后逻辑块独占。没有在共享末块上实现通用的复制后写入路径。 | BlockManager;第 10 章 |
怎样利用更多容量 | Scheduler 在 step 边界调整 batch;本版本 Prefill 优先,同一轮不混合 Prefill 与 Decode。 | Scheduler;第 8 章 |
引擎与在线服务 | generate() 提交一批请求并循环返回结果;这份项目不是包含 HTTP 接入、逐 token 流式返回等能力的完整在线服务。 | LLMEngine;第 5–7 章 |
源码核对:block_manager.py、layers/attention.py、scheduler.py、llm_engine.py。
回到“分页本来就是经典技术,为什么重要”这个问题:经典思想的价值,要看它是否解决了真实瓶颈,以及能否贯穿分配、寻址、共享、调度和实际 GPU 执行。只换一个分配器,内核却仍要求完整连续 K/V,或调度器无法利用腾出的容量,都不能自动得到原文的收益。反过来,也不应把 vLLM 后续全部成功只归因于分页;这篇首发文章适合解释其起点,生产系统的扩展继续看第 22 章,七种优化的职责分层可回看第 6.9 节。