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

推理机制动态图解（进阶）：并行、投机解码、PD、FlashAttention、调度

推理过程动态图解（基础） · 框架对比 · 推理优化方法

基础篇讲了张量变化、连续批处理和 PagedAttention。推理里还有很多“文字很难讲清、画出来就懂”的机制，这一页再补 7 个逐步演示，覆盖并行、加速、部署和调度中最常被问到的部分。

演示要回答的问题
1. 张量并行权重怎样切到多张卡？为什么每层要 all-reduce？
2. 专家并行（MoE）token 怎样被发到专家所在的卡？为什么负载不均是核心问题？
3. 投机解码为什么一次前向能产出多个 token？结果为什么不变？
4. PD 分离Prefill 算完的 KV 怎样交给 Decode？什么时候才能释放？
5. FlashAttention不写出完整注意力矩阵，怎样得到同样的结果？
6. 重叠调度CPU 调度开销怎样被藏起来？代价是什么？
7. 抢占KV 块不够时调度器怎样腾空间？为什么抢占多了很浪费？

1. 张量并行：权重怎样切开

GPU 0X[T, 4096]
GPU 1X[T, 4096]
通信—
GPU 0X× W_up 左半Y₀[T, 7168]
GPU 1X× W_up 右半Y₁[T, 7168]
通信不需要：每张卡独立算一半输出列
GPU 0Y₀激活（逐元素）
GPU 1Y₁激活（逐元素）
通信不需要：激活只看自己这一列
GPU 0Y₀× W_down 上半Z₀[T, 4096] 部分和
GPU 1Y₁× W_down 下半Z₁[T, 4096] 部分和
通信两张卡的结果都只是一部分
GPU 0Z = Z₀ + Z₁[T, 4096]
GPU 1Z = Z₀ + Z₁[T, 4096]
通信all-reduce
GPU 0注意力头 0–15KV head 0–3
GPU 1注意力头 16–31KV head 4–7
通信注意力同样按头切开，o_proj 后再一次 all-reduce
  1. 张量并行（TP=2）：每张卡都有完整的输入 X，权重切开放在两张卡上。以 MLP 为例（为简化只画 up 投影）。
  2. 第一个矩阵按“列”切（Column Parallel）：每张卡算出一半的输出列，彼此不需要通信。
  3. 激活函数逐元素计算，各卡只处理自己的那一半，仍然不需要通信。
  4. 第二个矩阵按“行”切（Row Parallel）：每张卡用自己那一半输入乘对应的一半权重，得到的是完整形状但只是部分和。
  5. 一次 all-reduce 把两张卡的部分和相加，两张卡都得到完整结果。每层 MLP 和注意力各需要一次 all-reduce，所以 TP 通常只在 NVLink 连接的同一台机器内使用。
  6. 注意力按头切：32 个 query 头、8 个 KV 头平均分到两张卡，每张卡的 KV Cache 只有一半。如果 KV 头数少于 TP 数，就要复制 KV 头，每卡容量不再减半。

2. 专家并行：token 怎样找到专家

tokent1t2t3t4
路由router 给每个 token 选 top-2 专家
GPU 0（专家 0、1）—
GPU 1（专家 2、3）—
tokent1 → E0, E2t2 → E1, E2t3 → E2, E3t4 → E2, E0
路由选出的专家和权重
GPU 0（专家 0、1）—
GPU 1（专家 2、3）—
tokendispatch：按专家所在的卡发送 token
路由all-to-all
GPU 0（专家 0、1）E0: t1, t4E1: t2
GPU 1（专家 2、3）E2: t1, t2, t3, t4E3: t3
token各专家独立计算
路由—
GPU 0（专家 0、1）算 3 个 token
GPU 1（专家 2、3）算 5 个 token（E2 是热点）
tokencombine：结果发回 token 原来的卡，按权重加权求和
路由all-to-all
GPU 0（专家 0、1）—
GPU 1（专家 2、3）—
  1. MoE 层里每个 token 只用少数几个专家。专家并行（EP）把不同专家放在不同 GPU 上。这里 4 个专家放在 2 张卡上，每个 token 选 2 个专家。
  2. router（一个小的线性层）为每个 token 打分，选出 top-2 专家和对应权重。
  3. dispatch：用一次 all-to-all 通信，把每个 token 发到它所选专家所在的卡。这一步的通信量随 token 数和选中的专家数增长。
  4. 各卡上的专家独立计算。专家 E2 被 4 个 token 选中，是热点，GPU 1 要算得更久，其他卡只能等它——负载不均是 EP 的核心问题，常见对策有负载均衡损失、冗余热点专家、动态调整专家位置。
  5. combine：再一次 all-to-all 把结果发回 token 原来的卡，按 router 权重加权求和。每个 MoE 层两次 all-to-all，所以大规模 EP 对网络延迟和拓扑非常敏感；一张卡故障会影响整个 EP 组。

3. 投机解码：先猜，再验证

已确认今天天气
草稿模型—
目标模型—
已确认今天天气
草稿模型很好吗
目标模型草稿模型很小，连续猜 3 个 token，速度快
已确认今天天气
草稿模型很好吗
目标模型一次前向验证 4 个位置
已确认今天天气
草稿模型很 ✓好 ✓吗 ✗
目标模型第 3 个位置目标模型想要的是“，”
已确认今天天气很好，
草稿模型“吗”及其 KV 被丢弃（回滚）
目标模型一次前向得到 3 个 token
  1. 投机解码：用一个便宜的方法先猜几个 token，再让大模型一次性检查。目标是减少大模型的前向次数。
  2. 草稿模型（或 EAGLE、MTP 头等）连续猜出 3 个 token。
  3. 目标模型把这 3 个猜测连同前文一起做一次前向，同时得到 4 个位置的预测。因为 Decode 受带宽限制，多算几个位置几乎不增加时间。
  4. 从左到右比较：前两个猜对了，第三个和目标模型的预测不同，从这里开始的猜测全部作废。
  5. 接受“很、好”，再加上目标模型在出错位置给出的“，”，这一次前向产出 3 个 token。被拒绝位置的 KV 要回滚。只要用合适的接受规则，输出分布和不用投机解码完全一样；接受率越低、并发越高，收益越小（算账见推理优化）。

4. PD 分离：KV 怎样交接

路由请求到达
Prefill 实例空闲
KV 传输—
Decode 实例运行其他请求的 Decode
路由选择 Prefill 和 Decode 实例
Prefill 实例计算整个 prompt 的 KV采样第 1 个 token
KV 传输—
Decode 实例预留 KV 块
路由—
Prefill 实例KV 留在显存，等待传输
KV 传输RDMA 写到 Decode 的块
Decode 实例接收中，不能使用
路由—
Prefill 实例确认传输完成后释放 KV
KV 传输完成
Decode 实例KV 就绪
路由流式返回输出
Prefill 实例处理下一个请求的 Prefill
KV 传输—
Decode 实例加入 Decode batch，逐个生成
  1. PD 分离：Prefill 和 Decode 跑在不同的实例（机器池）上。请求先到路由。
  2. 路由同时选好 Prefill 实例和 Decode 实例。Prefill 实例算出整个 prompt 的 KV（这里画成 Prefill 侧同时采样出第一个 token；首 token 由哪一侧产生取决于实现，例如 SGLang 的 PD 模式由 Decode 侧在 KV 就位后算出）；Decode 实例提前预留好接收用的 KV 块。
  3. KV 通过 RDMA（例如 Mooncake Transfer Engine、NIXL）从 Prefill 的显存传到 Decode 预留的块。传输完成之前，Decode 不能使用这些块。8K 上下文的示例模型约 1 GiB。
  4. 传输完成后 Prefill 才能释放自己的 KV；Decode 的块变为就绪。如果传输失败或超时，要有明确的处理：重传、重新 Prefill，或者报错。
  5. Decode 实例把这个请求加入自己的批处理，继续生成。Prefill 不再被长 Decode 拖慢，Decode 也不会被长 Prefill 打断，两个池可以按负载分别扩缩容。

5. FlashAttention：分块与在线 softmax

Q 块Q 的一块[Bq, d]
K/V 块K/V 块 1K/V 块 2K/V 块 3
片上状态m = −∞（当前最大值）　l = 0（指数和）　O = 0
Q 块Q 的一块
K/V 块K/V 块 1K/V 块 2K/V 块 3
片上状态算 S₁ = Q·K₁ᵀ，得到 m₁、l₁，O = exp(S₁ − m₁)·V₁
Q 块Q 的一块
K/V 块K/V 块 1K/V 块 2K/V 块 3
片上状态新最大值 m₂ 更大 → 旧的 l 和 O 乘以 exp(m₁ − m₂) 重新缩放，再加上块 2 的贡献
Q 块Q 的一块
K/V 块K/V 块 1K/V 块 2K/V 块 3
片上状态同样方式合并块 3
Q 块输出 = O / l
K/V 块全部处理完
片上状态只把最终结果写回显存；完整的注意力矩阵从未写出
  1. FlashAttention 把 Q 和 K/V 切成小块，放进 GPU 片上的高速存储里计算。每个 Q 块维护三个量：当前见过的最大分数 m、指数和 l、未归一化的输出 O。
  2. 处理第一个 K/V 块：算出这一块的注意力分数，用块内最大值做数值稳定的 exp，累加到 O 和 l。
  3. 处理第二块时，如果出现了更大的分数，之前算的 exp 都要按新的最大值缩小——乘以 exp(旧最大值 − 新最大值)。这就是“在线 softmax”：不用先看完所有分数，也能得到和标准 softmax 完全相同的结果。
  4. 依次处理所有 K/V 块。中间结果一直在片上，不写回显存。
  5. 最后 O / l 就是标准注意力的输出。计算量没有减少，省掉的是把 [序列长度 × 序列长度] 的大矩阵写回显存再读出来的开销。PagedAttention 的 kernel 在此基础上再按块表去收集不连续的 K/V 块。

6. 重叠调度：让 GPU 不再等 CPU

不重叠CPU 调度 NGPU 计算 NCPU 处理结果 NCPU 调度 N+1GPU 计算 N+1
GPU 利用CPU 工作时 GPU 空闲，Decode 每轮越短，空闲占比越大
重叠CPU 调度 NGPU 计算 N同时CPU 调度 N+1
GPU 利用GPU 算第 N 轮时，CPU 已在准备第 N+1 轮
重叠GPU 计算 N+1同时CPU 处理结果 N、调度 N+2
GPU 利用GPU 几乎不停
代价第 N+1 轮调度时还不知道第 N 轮的结果
GPU 利用例如某请求在第 N 轮已生成结束符，第 N+1 轮可能多算一步，再丢弃
  1. 不重叠时，每一轮都是“CPU 调度 → GPU 计算 → CPU 处理结果”串行进行。Decode 一轮可能只有几毫秒，CPU 的几百微秒到几毫秒就会占掉很大比例。
  2. 重叠调度（SGLang 的 overlap scheduler、vLLM 的异步调度）让 CPU 在 GPU 计算第 N 轮的同时，准备第 N+1 轮的输入。
  3. 稳定后，GPU 一直有活干，CPU 的调度和结果处理被隐藏在 GPU 计算时间里。
  4. 代价是调度第 N+1 轮时还拿不到第 N 轮的输出，实现上要处理“结果晚一轮到达”：已结束的请求可能多算一步再丢弃，采样结果要在 GPU 上直接传给下一轮（vLLM 源码中对应 stale / in-flight token 的回退处理）。

7. 抢占：显存不够时怎么办

运行中A8 块B6 块C5 块
空闲块1 块
等待队列空
运行中A需要新块BC需要新块
空闲块不够分
等待队列空
运行中AB
空闲块C 释放 5 块
等待队列C（已生成的部分记下来）
运行中AB 结束
空闲块空出更多块
等待队列C
运行中AC 重新 Prefill（prompt + 已生成部分）
空闲块—
等待队列空
  1. 显存里的 KV 块快用完了：A、B、C 都在 Decode，只剩 1 个空闲块。
  2. 下一轮 A 和 C 都需要新块，空闲块不够。
  3. 调度器抢占一个请求（vLLM V1 中通常是优先级最低或最晚到达的）：释放它的全部 KV 块，把它放回等待队列，已经生成的 token 保留下来。
  4. 等其他请求结束、显存空出来。
  5. 被抢占的请求重新被调度时，要把 prompt 和已生成的部分重新 Prefill 一遍（重算）。抢占越频繁，浪费的计算越多——这是显存容量不足的典型信号，可以用 KV 量化、减小 batch、分层缓存等方式缓解。另一种做法是把 KV 换出到主机内存再换回，代价是 PCIe 带宽。

还有一些机制已经在其他页面用动态演示讲过：前缀缓存的命中与淘汰（块管理器深入）、SGLang 基数树（框架对比）、Transfer Engine 的切分与重试（Transfer Engine）、Mooncake Store 的租约与淘汰（Store 内核），以及互动计算器里的前缀命中折叠和超时租约。

输入主题、术语或问题