1. 1. 1. 从 token 到多头 Attention:先看懂模型在算什么
    1. 1.1. 1.1 先固定场景:已知 t1、t2、t3,准备预测 t4
    2. 1.2. 1.2 从文本到向量:Tokenizer 编号,Embedding 查表
    3. 1.3. 1.3 最小的 Head:每个 q、k、v 只有一个数
    4. 1.4. 1.4 把 Head 扩到 128 维:权重仍然按位置分配
    5. 1.5. 1.5 从一个 Head 到 32 个:各自读取,再沿列拼接
  2. 2. 2. 从一层到一次生成:Prefill 与 Decode 怎样执行
    1. 2.1. 2.1 一层 Transformer:Attention 与 MLP 都在层内
    2. 2.2. 2.2 走完所有层,用最后一行预测下一个 token
    3. 2.3. 2.3 一次请求的三轮:Prefill → Decode → Decode
    4. 2.4. 2.4 Prefill:三行一起计算,放大最后一行 t3
    5. 2.5. 2.5 Decode:只新增 t4,历史以 K/V 的形式参与
  3. 3. 3. KV Cache:用空间换掉重复计算
    1. 3.1. 3.1 旧位置的 a 有什么用,追加 token 后为什么不用重算
    2. 3.2. 3.2 为什么只缓存各层 K/V,就能让 Decode 只算新位置
    3. 3.3. 3.3 Context window:新算一行,可以读取多少历史
    4. 3.4. 3.4 与无缓存对照:计算次数少了,读取范围没少
    5. 3.5. 3.5 怎样验证缓存等价
  4. 4. 4. 把一个新 token 放到 GPU:从矩阵读写推导显存与性能
    1. 4.1. 4.1 把教学形状换成真实配置,再跟踪一个新位置
    2. 4.2. 4.2 先看显存:权重、历史 KV 和当前向量是三类数据
    3. 4.3. 4.3 把北京这一轮放到 GPU:每层读什么、写什么
    4. 4.4. 4.4 沿张量形状,一步一步算出 KV 显存
    5. 4.5. 4.5 新增只有一行,为什么 Decode 仍然可能慢
    6. 4.6. 4.6 为什么权重放得下,运行时仍可能 OOM
    7. 4.7. 4.7 最后再计时:首 token、每轮间隔和总吞吐
  5. 5. 5. 从模型计算到推理服务:vLLM 负责什么
    1. 5.1. 5.1 裸跑能生成,为什么还需要推理引擎与服务
    2. 5.2. 5.2 离线生成:先看清接口的输入输出
    3. 5.3. 5.3 在线服务:请求如何进入引擎,输出如何返回
    4. 5.4. 5.4 为什么下一步读 nano-vLLM
  6. 6. 6. nano-vLLM 全貌:组件架构与请求流程
    1. 6.1. 6.1 单请求、单卡:一次 generate() 为什么要反复执行 step()
    2. 6.2. 6.2 单请求、单卡:运行第一轮之前,模型和 KV 池怎样准备好
    3. 6.3. 6.3 单请求、单卡:ModelRunner 怎样把一轮工作交给 GPU
    4. 6.4. 6.4 单请求、单卡:把 Prefill、两轮 Decode 和结束回收连起来
    5. 6.5. 6.5 多请求、单卡:Scheduler 怎样把独立请求组织成一批
    6. 6.6. 6.6 多请求、单卡:BlockManager 怎样让请求共享 KV 显存
    7. 6.7. 6.7 多请求、多卡:一个批次怎样由多个 rank 一起执行
    8. 6.8. 6.8 回到全貌:组件归属、源码与外部依赖
    9. 6.9. 6.9 优化机制分层:调度、缓存、传输与执行各管什么
  7. 7. 7. 请求生命周期:LLMEngine 与 Sequence
    1. 7.1. 7.1 Engine:一次调用,很多轮推进
    2. 7.2. 7.2 Sequence:保存的是请求记录,不是 token 向量
    3. 7.3. 7.3 三种长度:已经知道、已经算完、本轮待算
    4. 7.4. 7.4 候选、接受、结束,是不同的状态变化
  8. 8. 8. Scheduler:怎样让不同请求共享每一轮计算
    1. 8.1. 8.1 Continuous Batching:调度单位是 step,不是整条请求
    2. 8.2. 8.2 Prefill:按预算选择新位置,必要时切分长输入
    3. 8.3. 8.3 Decode:每条请求一个新位置,空间不足则抢占
    4. 8.4. 8.4 回写:先确认缓存完成,再决定是否增长回答
  9. 9. 9. PagedAttention 与 KV Cache:分页存储怎样参与计算
    1. 9.1. 9.1 为什么要分页:请求长度不断增长,显存容量却固定
    2. 9.2. 9.2 从 token 位置到物理槽位
    3. 9.3. 9.3 分配、增长、释放:块表怎样变化
    4. 9.4. 9.4 “空闲”不等于“显存已释放”,也不等于“内容已擦除”
    5. 9.5. 9.5 PagedAttention:K/V 分散在不同页,Attention 怎样算
    6. 9.6. 9.6 分页究竟省了什么,与 FlashAttention 有何不同
  10. 10. 10. Prefix Cache:不同请求怎样复用已算过的前缀
    1. 10.1. 10.1 为什么必须是相同前缀,而不只是相同文字块
    2. 10.2. 10.2 一个完整例子:少算 512 个位置,但仍读取它们
    3. 10.3. 10.3 共享边界:完整前缀只读,最后逻辑块独占
    4. 10.4. 10.4 缓存什么时候可用,什么时候失效
    5. 10.5. 10.5 命中前缀后,首字与后续生成分别省了什么
    6. 10.6. 10.6 把三个机制串起来:谁来算、存在哪里、哪些不用重算
  11. 11. 11. ModelRunner:把调度结果变成 GPU 张量
    1. 11.1. 11.1 一轮执行:工作单怎样变成候选 token
    2. 11.2. 11.2 Prefill 输入:把新位置拼起来,同时保留请求边界
    3. 11.3. 11.3 写地址与读地址:为什么两套信息都需要
    4. 11.4. 11.4 Decode:每条请求只输入一个新 id,历史仍然可读
    5. 11.5. 11.5 本轮 Context 与常驻 KV 池:寿命不同的两种对象
    6. 11.6. 11.6 Pinned Memory:本轮 CPU 输入怎样送到 GPU
  12. 12. 12. Qwen3 前向、权重加载与采样:走完模型内部
    1. 12.1. 12.1 接住 Runner 的输入:220 个位置,最后只需要两份预测
    2. 12.2. 12.2 Attention:新位置生成 Q/K/V,再按权重读取上下文
    3. 12.3. 12.3 残差与 MLP:每行完成本层计算,继续进入下一层
    4. 12.4. 12.4 从最终表示到候选:LM Head 与 Sampler 各做什么
    5. 12.5. 12.5 回看启动:权重文件怎样对应合并投影
  13. 13. 13. GPU 执行:缓存写入、FlashAttention 与 CUDA Graph
    1. 13.1. 13.1 新 KV 写入:地址已经确定,kernel 负责搬入数值
    2. 13.2. 13.2 三种 Attention 路径:新 query 一样,K/V 来源不同
    3. 13.3. 13.3 FlashAttention:分块累积,少搬中间矩阵
    4. 13.4. 13.4 CUDA Graph:输入每轮变化,执行安排可以复用
    5. 13.5. 13.5 形状分桶与收益边界:少提交也可能多算填充
  14. 14. 14. Tensor Parallel:把一次 forward 分给多张 GPU
    1. 14.1. 14.1 先区分多份模型与一份模型的切分
    2. 14.2. 14.2 Column 与 Row 的名字怎样对应 PyTorch 权重
    3. 14.3. 14.3 worker 为什么不需要拿到完整采样状态
    4. 14.4. 14.4 为什么多卡不一定更快:容量、计算与通信的取舍
  15. 15. 15. 验证方法:把原理变成可复查的证据
    1. 15.1. 15.1 先区分三类证据:数值、机制与性能
    2. 15.2. 15.2 数值等价:从单头手算到真实模型
    3. 15.3. 15.3 机制正确:从状态重建本轮到底发生了什么
    4. 15.4. 15.4 性能收益:同一负载,只改变待研究的机制
    5. 15.5. 15.5 现有证据到哪一步:保留已测与未测的边界
  16. 16. 16. 实现边界与排障:先判断出错层次
    1. 16.1. 16.1 按组件看边界:简化发生在哪里
    2. 16.2. 16.2 按请求经过的组件定位故障
    3. 16.3. 16.3 最小复现:一次只恢复一个复杂因素
    4. 16.4. 16.4 显存排查:先对齐观测口径,再判断是否泄漏
  17. 17. 17. 全篇回顾与阅读导航:从这份实现走向生产引擎
    1. 17.1. 17.1 用一次 Decode,把所有组件接回去
    2. 17.2. 17.2 按问题查回正文
    3. 17.3. 17.3 从真实瓶颈选择进阶方向
    4. 17.4. 17.4 资料与版本:带着问题继续读
  18. 18. 18. 动手学习路线:半天入门与两周深入验证
    1. 18.1. 18.0 快速入门与通用环境准备
      1. 18.1.1. 18.0.1 先选路线:读通主流程,不必先完成两周计划
      2. 18.1.2. 18.0.2 准备一份最小调试入口
      3. 18.1.3. 18.0.3 实验 A:一次 Prefill、两次 Decode
      4. 18.1.4. 18.0.4 实验 B:两条请求,为什么会从同批执行变成只剩一条
      5. 18.1.5. 18.0.5 实验 C:第 257 个已知 token 的 KV 究竟写到哪里
      6. 18.1.6. 18.0.6 入门验收与进阶分流
      7. 18.1.7. 18.0.7 通用环境准备:没有现成环境时从这里开始
    2. 18.2. 18.1 Day 1:从文本、token 到采样概率
    3. 18.3. 18.2 Day 2:手算并验证一次因果 Attention
    4. 18.4. 18.3 Day 3:逐轮分清“已知 token”和“已计算 KV”
    5. 18.5. 18.4 Day 4:用两个路径验证 KV Cache 省了哪些计算
      1. 18.5.1. 18.4.1 选做:沿完整小模型追踪无缓存计算
    6. 18.6. 18.5 Day 5:算显存并区分 TTFT、TPOT 与吞吐
    7. 18.7. 18.6 Day 6:配置 GPU,完成第一次真实生成
    8. 18.8. 18.7 Day 7:从日志还原两条请求的一生
    9. 18.9. 18.8 Day 8:验证分块 Prefill 与抢占顺序
    10. 18.10. 18.9 Day 9:从逻辑位置算到真实 cache slot
    11. 18.11. 18.10 Day 10:验证 512-token 前缀命中与引用计数
    12. 18.12. 18.11 Day 11:把 Runner 元数据接到真实模型形状
    13. 18.13. 18.12 Day 12:完成 eager 与 CUDA Graph 的可比较实验
    14. 18.14. 18.13 Day 13:先推导 Tensor Parallel,再按设备选做多卡
    15. 18.15. 18.14 Day 14:交付一份别人能复查的推理实验报告
    16. 18.16. 18.15 选做:体验 vLLM 的离线接口与本机服务
      1. 18.16.1. 18.15.1 独立环境与离线补全
      2. 18.16.2. 18.15.2 本机在线服务与流式返回
  19. 19. 19. 附录:从数字与向量看懂 RMSNorm、SiLU/SwiGLU、GQA 与 RoPE
    1. 19.1. 19.1 Softmax:怎样把分数变成权重
    2. 19.2. 19.2 Norm 与 LayerNorm:归一化究竟在处理什么
    3. 19.3. 19.3 RMSNorm:不减均值,只按均方根调整尺度
    4. 19.4. 19.4 同一个向量,四种“归一化”结果并不相同
    5. 19.5. 19.5 SiLU 与 SwiGLU:MLP 中的非线性和门控
    6. 19.6. 19.6 GQA:多个 Q Head 共享一组 K/V,但各自计算权重
    7. 19.7. 19.7 RoPE:旋转 Q/K 的分量,让匹配分数感知位置
    8. 19.8. 19.8 回到同一轮请求:核对形状与计算轴
  20. 20. 20. 延伸阅读|深入 nano-vLLM 高效推理:授权全文翻译
    1. 20.1. 20.1 引言
    2. 20.2. 20.2 快速开始
    3. 20.3. 20.3 为什么需要 vLLM 这样的推理引擎?
      1. 20.3.1. 20.3.1 KV Cache
      2. 20.3.2. 20.3.2 内存碎片
      3. 20.3.3. 20.3.3 连续批处理
      4. 20.3.4. 20.3.4 Prefill 与 Decode
    4. 20.4. 20.4 架构
    5. 20.5. 20.5 generate() 循环
      1. 20.5.1. 20.5.1 调度器
        1. 20.5.1.1. 20.5.1.1 Prefill 路径
        2. 20.5.1.2. 20.5.1.2 Decode 路径
        3. 20.5.1.3. 20.5.1.3 抢占策略
        4. 20.5.1.4. 20.5.1.4 Decode 期间的 KV 块簿记
      2. 20.5.2. 20.5.2 KV Cache
        1. 20.5.2.1. 20.5.2.1 分页注意力的实现
        2. 20.5.2.2. 20.5.2.2 Flash Attention
    6. 20.6. 20.6 在多个 GPU 上运行模型
      1. 20.6.1. 20.6.1 ModelRunner 的并行执行
      2. 20.6.2. 20.6.2 Embedding 层
      3. 20.6.3. 20.6.3 MLP 的张量并行:列切分与行切分技巧
        1. 20.6.3.1. 20.6.3.1 按列切分
        2. 20.6.3.2. 20.6.3.2 按行切分
      4. 20.6.4. 20.6.4 把 Attention 层切分到多个 GPU
  21. 21. 21. 延伸阅读|Neutree:沿一次请求读懂 nano-vLLM(中文导读)
    1. 21.1. 21.1 这篇文章适合解决什么问题
    2. 21.2. 21.2 把类名放回各自的职责
    3. 21.3. 21.3 带着一个跨块例子阅读
    4. 21.4. 21.4 阅读时保留版本与术语边界
  22. 22. 22. 延伸阅读|Aleksa Gordić:从推理引擎到高吞吐在线服务(中文导读)
    1. 22.1. 22.1 它补上了 nano-vLLM 之外的哪一层
    2. 22.2. 22.2 用“解决什么瓶颈”理解高级特性
    3. 22.3. 22.3 多 GPU 与多副本不是同一件事
    4. 22.4. 22.4 最后用指标判断设计是否有价值
  23. 23. 23. 附录|vLLM 首发博客:中文摘要与工程解读
    1. 23.1. 23.1 原文要点:从 KV 显存管理改善服务吞吐
    2. 23.2. 23.2 工程解读:省显存,为什么可能让服务更快
    3. 23.3. 23.3 工程解读:共享块为什么还需要引用计数与写时复制
    4. 23.4. 23.4 读性能数字:保留测试条件,不把历史结果当保证
    5. 23.5. 23.5 对照 nano-vLLM:继承的是思路,不是全部实现
Macduan Notes

大模型推理入门与 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 通常包含 AttentionMLP:Attention 让各位置读取可见的上下文,MLP 在每个位置内部变换特征;Norm 调整数值尺度,残差连接把更新加回原表示。上一个 Block 的输出,就是下一个 Block 的输入。

这套网络会反复用于预测,参数保持不变。首次处理整段输入,叫作 Prefill;随后复用历史结果、每轮只新算一个位置,叫作 Decode。 用三个输入 token,就能串起整个过程:

原文配图 1

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]
这里的词表大小和 id 仅作示意

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.75Q 和 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 是投影参数的形状。

原文配图 2

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 个头的输出拼接。

原文配图 3

实现时可以把 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、残差。

原文配图 4
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             # 第二次残差相加,传给下一层
跟踪一个位置;所有中间主线向量均为 [1,4096]

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)
LM Head 负责打分,Sampler 负责选出 token

选取 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。

原文配图 5

轮次

本轮新输入

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]
一个 Head 的三个查询;softmax 沿每一行的可读位置计算

分数矩阵的行表示查询位置,列表示被读取的位置。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;同样的层结构依次执行,最后进入词表输出。

原文配图 6

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 汇合。

原文配图 7

每个 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
跨层依赖;“后处理”包括多头拼接、WO、残差、Norm 和 MLP

这说明中间层的 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、残差等细节,但它们在实际计算中仍然存在。

原文配图 8

进阶说明:最后一个 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 的输入
Decode 在每一个 Block 内的依赖;展示一个 Head

新 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

缓存对照图|同一条生成路径、相同的可读前缀;有缓存时只对新位置运行整条模型。

原文配图 9

三轮中,无缓存共处理 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 维残差流。

原文配图 10

读图时先沿顶部走完模型,再向下放大一个 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 内进入下一层。

原文配图 11

这里有两种不同的搬运。 一种是 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,未结束就继续下一轮。

原文配图 12

这张图暂不展开请求调度和内存分块。单请求仍会经过这些代码,只是没有请求间竞争;先把“准备本轮输入、完成计算、接收结果”这条主线记住。等加入多条请求,再解释为什么要把这些职责拆成独立组件;完整函数调用图放在第 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 的缓存读写。

原文配图 13

模型内部的计算沿用第 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 确保缓存可用。

原文配图 14

这份实现采用 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。从单请求单卡到多请求多卡,主循环不变,逐步增加的是资源选择、地址管理与并行执行。

原文配图 15

按数据寿命再核对一次: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:七种机制在同一轮推理中的位置——沿主线看数据与工作如何前进,沿分层看各自在省什么。

原文配图 16

请求调度与 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;单请求与多请求经过同一条路径。

原文配图 17

在本文固定版本中,公开 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:固定批次与连续批处理——看每轮成员如何变化,而不是把格子数当作加速倍数。

原文配图 18

每轮被一起处理的请求共享模型参数和 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,仍能命中的完整前缀块则可能省掉部分重算。

状态图:正常生成沿主线前进,抢占把请求送回等待恢复的路径。

原文配图 19

图中的“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 章。

原文配图 20

图使用实际块大小 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——物理位置可以分散,数学上仍是同一个上下文。

原文配图 21

先寻址,再打分。内核根据块表读取物理块 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 留在显存中。

原文配图 22

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 双路径实现。

原文配图 23

源码把 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、权重加载与采样。

模型配置:Qwen3-0.6B 固定 revision

源码: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 图:对比完整中间矩阵的显存往返与分块累积。

原文配图 24

图中的 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 图:先捕获模型主干,再逐轮更新缓冲区并重放。

原文配图 25

本提交捕获 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:上半按输出特征分工,下半把输入分片产生的部分和相加。

原文配图 26

把 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 张量通信。

原文配图 27

源码: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 截图更有意义。下一章用同样的分层方式定位边界与故障。

源码对照:仓库基准测试的统计口径。

源码:bench.py,L8–28

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 → 多头输出

第 1 章

Prefill / Decode 与完整模型

本轮输入哪些位置,最后怎样选下一个 id

第 2 章

为什么只缓存各层 K/V

因果性、旧结果不变、新位置的直接依赖

第 3 章

显存、带宽与时间指标

权重、KV、中间激活和计时范围

第 4 章

为什么有引擎,组件怎样协作

入口、状态、调度、缓存、执行与回写

第 5 章第 6 章总图

这轮选谁、各算多少

Sequence 三种长度与 Scheduler 工作单

第 7 章第 8 章

KV 放在哪里、何时可以共享

逻辑位置、物理块、引用与完整前缀

第 9 章第 10 章

工作单怎样变成张量与候选

输入准备 → 模型前向 → LM Head → 采样

第 11 章第 12 章

GPU 怎样减少开销、怎样多卡执行

中间搬运、工作提交、特征分片与部分和

第 13 章第 14 章

怎样证明、怎样定位、怎样运行

证据方法、组件边界、统一操作路线

15 / 16 / 18 章

Softmax、temperature、LayerNorm、RMSNorm、SiLU/SwiGLU、GQA 与 RoPE 的基础数值解释统一在第 19 章附录19.8可查形状与计算轴。不必每遇到一个术语就重新走整篇主线。

17.3 从真实瓶颈选择进阶方向

遇到的问题

继续研究什么

先确认的代价或条件

权重或长上下文 KV 放不下

权重量化、KV 精度、并行与缓存管理

先区分哪类数据占显存,再检查数值影响与额外计算

单请求 Decode 间隔仍高

执行优化、推测解码及验证机制

判断带宽、提交还是计算受限;推测接受率与验证成本影响收益

长短请求互相拖慢

调度、公平性、分块策略与 Prefill/Decode 部署组织

明确到达模式、首 token 与输出间隔目标,不只看离线吞吐

多卡通信成为主要开销

并行策略、拓扑与通信组织

定位具体 collective,并比较计算减少量和通信增加量

从教学引擎走向在线服务

服务前端、准入与取消、隔离、监控和故障恢复

功能完备、可靠性和性能是不同目标,不能只验证能生成文本

回到生产 vLLM 时,仍用“请求 → 调度 → 缓存 → Runner → 模型 → 采样”定位组件,再阅读当前版本为多模型、多设备和服务可靠性增加的工程层。先定义负载与目标指标,再选技术;不要因名称更新,就忽略它最终改变了哪段计算或哪类数据。

17.4 资料与版本:带着问题继续读

资料入口

带着什么问题阅读

Jay Alammar — The Illustrated GPT-2

从向量到整层与词表输出;注意 GPT-2 结构不能直接套到 Qwen3

3Blue1Brown — Attention in transformers

投影、匹配和加权 V 的直觉

Hugging Face — How caching works

因果性、逐层缓存与更新方式

Kipply — Transformer Inference Arithmetic

计算、带宽与容量估算

vLLM — PagedAttention 介绍

动态 KV 分配和跨请求共享

nano-vLLM 固定源码基线

对照本文实际分析的组件与实现选择

Qwen3-0.6B 固定模型与配置

核对真实维度、层数、GQA 与参数文件

vLLM Quickstart / 官方文档

区分当前生产接口、设计与 nano 的简化版本

Deep Dive into Efficient LLM Inference with nano-vLLM

原文参考文章;结合固定源码核对版本差异

本文 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 章的旧实验包。

nanovllm-e2e-guide-v2.zip

验证边界:文中的“预期值”是验收目标。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()
quick_walk.py:新增的独立学习入口,不改引擎源码

三个案例都用单卡、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
在实验机运行;先检查输出中的 Loaded source 路径

第一次不加 --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
pdb 命令:c 继续,n 执行当前行,s 进入函数,p 查看表达式,q 退出

在 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]
断点 52:此轮要算什么

第一轮应为 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]
断点 54:此轮完成后请求变成什么

轮次

模型本轮处理的位置

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()
pdb:每轮输入及 KV 地址元数据

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)
在 may_append 中查看;用 n 逐步执行到块追加完成
from nanovllm.utils.context import get_context
p positions.tolist()
p get_context().context_lens.tolist()
p get_context().slot_mapping.tolist()
在 ModelRunner.run 的 216 行查看

时刻

位置与物理块

应得到的映射/状态

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;新版已包含四个实验与验收工具,不需要同时解压两个包。

nanovllm-learning-labs.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]
局部 Attention 手算;3 个位置、每头 2 维

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)
有缓存逻辑伪代码;model_forward 返回本轮各位置的 logits 及更新后的各层 KV

下面保留原有 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
原有无缓存小模型实验;目录约定沿用第 18.0 节

验收先看 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 章的两周实验包继续用于完整实验。

llm-inference-pedagogy-20260909.zip

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()
API 阅读示例:在已准备好的 nano 环境运行;完整 smoke 以第 18.6 节为准

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 不匹配当成模型权重损坏。

源码:pyproject.toml,L6–20

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=auto
选做:复用实验目录,但使用独立的 vLLM 环境
import 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)
选做:离线补全;MODEL_DIR 来自第 18 章 env.sh

在已激活的 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 vllm
选做:在 GPU 机器启动仅本机可访问的学习服务
curl -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]
一组输入分数,对应一组权重;i、j 都是候选下标

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
这里按 D 个特征计算总体方差,分母为 D,不是 D−1

先忽略很小的 ε,并取 γ=[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]
本文 Qwen3 使用的 RMSNorm 具有逐特征权重 γ,没有 β 偏移项

三个数都除以同一个尺度,所以这个例子仍保留 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)
摘自 nano-vLLM 的 RMSNorm.rms_forward

在 Qwen3 中,残差流进入 Attention 前、进入 MLP 前各有一次 RMSNorm,走完所有 Block 后还有最终 RMSNorm。源码名字 input_layernormpost_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 独立作用于每个数,不像 Softmax 那样在候选之间分配总和

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]
Qwen3-0.6B 的单 token MLP;@ 是矩阵乘法,⊙ 是逐元素乘法

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 的特征。

原文配图 28

源码: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]。

原文配图 29

只放大第一组,并省略位置下标: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]
同一位置、同一 KV 组内两条读取路径;上标表示 Head

这里省略位置下标,上标表示 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,但与另一个向量的点积会受双方方向影响。
只看一对分量;这里的 a、b 不是 Attention 输出或 Value
原文配图 30

真实的 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 章的原创可编辑画板,不是原文插图;它按本文固定版本补充组件之间的关系,原文调度细节以本章译者注为界。

原文配图 31

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.running
原文代码/示例

add_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 之间。

@triton.jit
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 o
原文代码/示例

flash_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":
                break
原文代码/示例
def 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, args
原文代码/示例
def 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 的列接行采用相同的分片衔接原则。

原文配图 32

来看实际代码。

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_loader
原文代码/示例

ColumnParallelLinear 是采用张量并行的线性层。关键的列切分发生在 __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.pylayers/attention.pyscheduler.pyllm_engine.py

回到“分页本来就是经典技术,为什么重要”这个问题:经典思想的价值,要看它是否解决了真实瓶颈,以及能否贯穿分配、寻址、共享、调度和实际 GPU 执行。只换一个分配器,内核却仍要求完整连续 K/V,或调度器无法利用腾出的容量,都不能自动得到原文的收益。反过来,也不应把 vLLM 后续全部成功只归因于分页;这篇首发文章适合解释其起点,生产系统的扩展继续看第 22 章,七种优化的职责分层可回看第 6.9 节