LLM Infer Infra 学习资料库
02 推理运行时 · 专注 KV Cache 存储

推理过程动态图解：张量变化、连续批处理与 PagedAttention

推理运行时 · vLLM 请求源码 · 块管理器深入 · vLLM 与 SGLang 对比

推理里最难想象的三件事：一次前向中张量的形状怎么变、几十个请求怎样在每一轮动态拼成一个 batch、KV Cache 怎样被分页存放和读取。这一页用三个可以逐步播放的演示把它们画出来。点“下一步”或“自动播放”，也可以直接点说明跳到某一步；黄色框是这一步的重点。

读完应该能回答：为什么 Decode 受带宽限制？为什么连续批处理比静态批处理吞吐高？PagedAttention 到底省了什么、代价是什么？

1. 一次推理中张量怎样变化

模型参数取 Llama-3-8B 一类的公开配置，只为说明形状，不涉及具体实现细节。

阶段Prefill：输入 6 个 token
张量token_ids[6]
KV Cache空
阶段Embedding 查表
张量token_ids[6]→X[6, 4096]
KV Cache空
阶段第 1 层：RMSNorm + QKV 投影
张量X[6, 4096]→QKV[6, 6144]拆成Q[6, 32, 128]K[6, 8, 128]V[6, 8, 128]
KV Cache空
阶段RoPE：给 Q、K 加上位置 0–5
张量Q[6, 32, 128]K[6, 8, 128]形状不变，数值按位置旋转
KV Cache空
阶段写入 KV Cache
张量K[6, 8, 128]V[6, 8, 128]按 slot_mapping 写到对应物理槽位
KV Cachepos 0pos 1pos 2pos 3pos 4pos 5
阶段注意力（因果掩码）
张量scores[32, 6, 6]softmax × V →O[6, 32, 128]→o_proj[6, 4096]
KV Cachepos 0pos 1pos 2pos 3pos 4pos 5
阶段残差 + MLP
张量gate / up[6, 14336] ×2激活相乘 →down[6, 4096]+ 残差
KV Cachepos 0pos 1pos 2pos 3pos 4pos 5
阶段重复 32 层 → 采样第 1 个输出 token
张量最后一个位置[1, 4096]lm_head →logits[1, 128256]采样 → token 7
KV Cache每层都存了 pos 0–5
阶段Decode：只输入 1 个新 token
张量X[1, 4096]Q[1, 32, 128]新 K / V[1, 8, 128]→ scoresscores[32, 1, 7]
KV Cachepos 0pos 1pos 2pos 3pos 4pos 5pos 6
  1. 以 Llama-3-8B 类模型为例（hidden 4096、32 层、32 个 query head、8 个 KV head、head_dim 128、MLP 14336、词表 128256）。Prefill 阶段一次输入 6 个 token。vLLM 会把 batch 里所有请求的 token 拼成一维，所以张量的第一维是“本轮 token 总数”，而不是 [batch, seq]。
  2. Embedding 把每个 token id 查表成 4096 维向量，得到 X：[6, 4096]。
  3. RMSNorm 后做一次合并的 QKV 投影：4096 → 6144（Q 4096 + K 1024 + V 1024）。拆开后 Q 有 32 个头，K、V 只有 8 个头——这就是 GQA，每 4 个 query head 共用 1 组 K、V，KV Cache 因此只有 MHA 的 1/4。
  4. RoPE 按每个 token 的位置旋转 Q 和 K。形状不变，但位置信息写进了数值里——所以同样的 token 放在不同位置，K 就不同，不能直接复用。
  5. 把这一层的 K、V 写进 KV Cache。每个 token 写到哪个物理槽位由 slot_mapping 决定（PagedAttention，见第 3 个演示）。V 不经过 RoPE。
  6. 注意力：每个 query 和它之前的所有 key 做点积。scores 形状是 [头数, 查询数, 键数] = [32, 6, 6]，因果掩码把“看未来”的位置置为负无穷。softmax 后乘以 V，再经 o_proj 投影回 [6, 4096]。实际 kernel（FlashAttention 等）分块计算，不会把完整的 scores 写回显存。
  7. 残差相加后进入 MLP：gate 和 up 两个投影各得到 [6, 14336]，激活后相乘，再 down 投影回 [6, 4096]，加上残差。MLP 是参数量和计算量的大头。
  8. 32 层都做完后，只取最后一个位置的隐藏状态，经 lm_head 得到 128256 维的 logits，采样出第一个输出 token。注意：Prefill 也要算所有位置的 K、V 并存进缓存，但只有最后一个位置需要算 logits。
  9. Decode 阶段每步只输入 1 个新 token：X 是 [1, 4096]，新算出的 K、V 追加到缓存（pos 6），注意力要读全部 7 个位置的 K、V，scores 变成 [32, 1, 7]。计算量很小，但每步都要读一遍全部权重和这个请求的全部 KV——所以 Decode 受显存带宽限制，必须靠批处理提高效率。
对比Prefill（6 个 token）Decode（1 个 token）
输入张量[6, 4096][1, 4096]
注意力 scores[32, 6, 6]，因果掩码[32, 1, 7]，读全部历史 KV
每个权重被复用6 次（每个 token 一次）1 次
主要瓶颈算力显存带宽（读权重 + 读 KV）
KV Cache一次写入 6 个位置每步追加 1 个位置

所以提高 Decode 效率的核心是让更多请求一起算，把“读一遍权重”的成本分摊掉——这就引出了第 2 个演示。原理推导见 GPU 与 kernel 基础。

2. 连续批处理：每一轮怎样组 batch

演示的是 vLLM V1 一类“统一 token 预算”的调度思路：每轮先满足正在运行的请求，剩余预算再给新的 Prefill，长 Prefill 按预算切块。SGLang 的思路类似，但它还会参考前缀缓存决定谁先进（见对比页）。

本轮 batchAprefill 8 / 10
token 预算 8已用 8
等待队列Bprompt 3
完成—
本轮 batchAprefill 2（剩余）Bprefill 3
token 预算 8已用 5
等待队列空
完成A、B 各采样出第 1 个 token
本轮 batchAdecode 1Bdecode 1Cprefill 4
token 预算 8已用 6
等待队列空（C 刚到达就被加入）
完成—
本轮 batchAdecode 1Bdecode 1Cdecode 1
token 预算 8已用 3
等待队列空
完成B 生成结束，释放 KV 块
本轮 batchAdecode 1Cdecode 1Dprefill 6 / 12
token 预算 8已用 8
等待队列空
完成B 已离开
本轮 batchAdecode 1Cdecode 1Dprefill 6（剩余）
token 预算 8已用 8
等待队列Eprompt 5
完成E 要等下一轮有预算才能进来
  1. 每一轮（一次前向）都有一个 token 预算，这里设为 8（vLLM 对应 max_num_batched_tokens）。A 的 prompt 有 10 个 token，一轮放不下，就切出 8 个先算——这就是 chunked prefill。B 暂时等待。
  2. 下一轮先把 A 剩下的 2 个算完，B 的 3 个也放进来。两个请求的 token 拼成一个 5 个 token 的一维 batch，不需要补齐到同样长度。Prefill 算完时，各自采样出第一个输出 token。
  3. A、B 进入 Decode，每轮各贡献 1 个 token。新请求 C 到达，不用等别人结束，这一轮就被加入。调度器先保证正在运行的请求（Decode 优先），剩余预算再给新的 Prefill。
  4. B 生成到结束条件，立刻离开 batch 并释放它的 KV 块。下一轮就可以把空出来的显存给新请求——这是“连续批处理”和传统静态批处理的根本区别：按轮（iteration）而不是按请求组调度。
  5. 新请求 D 的 prompt 有 12 个 token，这一轮只剩 6 个预算，就先算 6 个。Decode 的请求每轮都在推进，不会因为 D 的长 Prefill 被卡住，ITL 更平稳。
  6. D 继续算完剩下的 6 个，预算再次用满，新到的 E 要等下一轮。预算设得太大，Decode 的每轮时间变长（ITL 变差）；设得太小，Prefill 要拆很多轮（TTFT 变差）。这就是调参时要权衡的地方。
方式怎么组 batch问题
静态批处理凑齐一组请求一起跑，等最长的那个结束才换下一组；序列要补齐到同样长度短请求早早结束后，GPU 空等；补齐浪费算力
连续批处理每轮重新选择参与的请求；结束的立刻离开，新的立刻加入；所有 token 拼成一维，不补齐需要分页的 KV 管理和能处理不等长序列的注意力 kernel
+ chunked prefill长 Prefill 按预算切块，和 Decode 混在一轮里Prefill 被拆成多轮，单个请求的 TTFT 可能变长；预算需要调

3. PagedAttention：块表与物理块

逻辑 token请求 A：10 个 token（pos 0–9），每块 4 个 token
块表 A逻辑块 0 → 7逻辑块 1 → 2逻辑块 2 → 5
GPU 物理块物理块 2物理块 5物理块 7物理块 9
说明三个物理块不连续也没关系
逻辑 tokenpos 9 → 逻辑块 9 // 4 = 2，块内偏移 9 % 4 = 1
块表 A逻辑块 0 → 7逻辑块 1 → 2逻辑块 2 → 5
GPU 物理块物理块 2物理块 5物理块 7物理块 9
说明slot = 5 × 4 + 1 = 21
逻辑 tokenDecode 生成 pos 10、11：写入物理块 5 的偏移 2、3
块表 A逻辑块 0 → 7逻辑块 1 → 2逻辑块 2 → 5（满）
GPU 物理块物理块 2物理块 5物理块 7物理块 9
说明块写满之前不用分配新块
逻辑 token生成 pos 12：需要逻辑块 3
块表 A逻辑块 0 → 7逻辑块 1 → 2逻辑块 2 → 5逻辑块 3 → 9
GPU 物理块物理块 2物理块 5物理块 7物理块 9
说明按需从空闲队列取一个块，最多浪费不到一个块
逻辑 token注意力计算：pos 12 的 query 要读 pos 0–12 的 K、V
块表 A74→24→54→94
GPU 物理块物理块 7物理块 2物理块 5物理块 9
说明kernel 按块表顺序逐块读取，再拼成完整的注意力
逻辑 token请求 B 与 A 的前 8 个 token 相同
块表 B逻辑块 0 → 7逻辑块 1 → 2逻辑块 2 → 4
GPU 物理块物理块 7物理块 2ref = 2，两个请求共享，不拷贝
说明B 只为不同的部分分配新块 4
  1. PagedAttention 借鉴了操作系统的虚拟内存：每个请求的 KV 按固定大小切成“逻辑块”（这里每块 4 个 token；vLLM 默认 16），通过块表映射到 GPU 上任意位置的“物理块”。不用为每个请求预留一整段连续的最大长度显存。
  2. 找一个 token 的 K、V 在哪：逻辑块号 = 位置 ÷ 块大小，块内偏移 = 位置 % 块大小，再查块表得到物理块。vLLM 在每轮调度后算好 slot_mapping，模型写 KV 时直接用。
  3. Decode 每生成一个 token 就写进当前最后一个块的下一个空位，块写满之前不需要分配。
  4. 写满后需要新块时，才从空闲队列取一个。所以每个请求最多浪费不到一个块的空间，而传统做法要按最大长度预留，浪费可能达到一大半。
  5. 注意力 kernel 拿到块表，按顺序把各物理块里的 K、V 读进片上存储计算。因为块不连续，kernel 要做“按块收集”，这也是 PagedAttention 需要专门 kernel 的原因。
  6. 前缀共享：B 和 A 的前 8 个 token 相同，B 的块表直接指向 A 的物理块 7、2，引用计数变成 2，不需要拷贝；只有不同的部分才分配新块。块怎么被命中、释放和淘汰，见块管理器深入。

代价也要知道：注意力 kernel 要按块表收集不连续的数据，比读连续内存复杂；块太小会让块表变长、kernel 效率下降，块太大又会增加尾部浪费并降低前缀命中精度。vLLM 默认块大小是 16 个 token。

4. 继续深入

输入主题、术语或问题