【infra之路】Prefill和Decode是如何一起计算、为什么可以batch并行计算

因为它们本质上做的是同一件事——把 token 送进 Transformer 做 forward pass。

从模型的角度看

不管 Prefill 还是 Decode，模型执行的计算是一样的：

输入 token → Embedding → 32层 Transformer → LM Head → logits

区别只在于输入长度和 KV Cache 的处理方式：

Prefill: 输入 [token_1, token_2, ..., token_1000]，长度 1000 K, V 存入新分配的 Block Decode: 输入 [token_1001]，长度 1 K, V 追加到已有的 Block

但模型本身（权重、矩阵乘法、LayerNorm、FFN）完全不关心你输入的是 1 个 token 还是 1000 个 token。它只管做矩阵运算，输入的 seq_len 维度多大都行。

具体怎么混在一起

假设当前 iteration 有 2 个 Decode 请求和 1 个 Prefill 请求：

请求 A (Decode): 输入 1 个 token，KV Cache 已有 500 个 请求 B (Decode): 输入 1 个 token，KV Cache 已有 200 个 请求 C (Prefill): 输入 100 个 token，KV Cache 为空 拼成一个 batch: 总输入 = [A的1个token, B的1个token, C的100个token] 总 seq_len = 1 + 1 + 100 = 102

每层的计算过程：

Embedding: [102, 4096] ← 三个请求的 token 一起查表 Self-Attention: Q, K, V = W_q·x, W_k·x, W_v·x ← 都是 [102, 4096]，一次矩阵乘法搞定 但 Attention 计算时，每个请求只看自己的 K, V: A 的 Q（1 个 token）和 A 的 KV Cache（500+1 个 token）做 Attention B 的 Q（1 个 token）和 B 的 KV Cache（200+1 个 token）做 Attention C 的 Q（100 个 token）和 C 自己的 K, V（100 个 token）做 Attention（causal mask） 这里每个请求的 Attention 是独立的，但 Q, K, V 的投影是一起算的 FFN: [102, 4096] → [102, 11008] → [102, 4096] ← 一起算，互不干扰 LM Head: [102, 32000] ← 一起投影 然后拆回各请求： A 取自己的 logits → 采样 B 取自己的 logits → 采样 C 取最后一个位置的 logits → 采样第一个生成 token

为什么能一起算？

关键在于 Transformer 的两个特性：

1. 线性层是 batch 友好的

W · x这个矩阵乘法，不管 x 是 1 个请求的 token 还是 100 个请求的 token 拼在一起，都是同一套权重 W 在算。GPU 的矩阵乘法单元天然擅长处理大 batch。

2. Attention 通过 Block Table 隔离

虽然 Q, K, V 的投影是混在一起算的，但 Attention 计算时，每个请求通过自己的 Block Table 找到自己的 KV Cache，不会和其他请求的数据混在一起：

A 的 Attention: Q_A · K_cache_A^T ← 通过 A 的 Block Table 找到 A 的 Block B 的 Attention: Q_B · K_cache_B^T ← 通过 B 的 Block Table 找到 B 的 Block C 的 Attention: Q_C · K_C^T ← C 刚 Prefill，K 就在当前输入中

vLLM 的 Attention kernel 支持这种混合计算——它能处理一个 batch 中不同请求有不同 seq_len 和不同 KV Cache 长度的情况。
这个问题问到了 GPU 并行计算的核心。

为什么能用同样的参数？

因为推理时模型参数是只读的，不会修改。

训练时: 前向 → loss → 反向 → 更新参数 参数每步都在变，必须串行 推理时: 前向 → 输出 参数固定不变，就是查表 + 矩阵乘法 多少请求同时读同一组权重都没冲突

就像一本教科书放在图书馆里，10 个人可以同时看它的复印件，书本身不会被改变。推理时 W_q、W_k、W_v 这些权重矩阵就是"教科书"，所有请求只是读取它做乘法。

有先后顺序吗？

没有。是真正的同时计算。

这一点和 CPU 的思维方式很不一样。CPU 是"先算 A，再算 B，再算 C"。GPU 是通过矩阵运算把 A、B、C合并成一次操作：

不是这样（串行）: step 1: W · x_A → y_A step 2: W · x_B → y_B step 3: W · x_C → y_C 而是这样（并行）: 一步完成: W · [x_A, x_B, x_C] → [y_A, y_B, y_C]

具体到矩阵乘法的维度：

W 的形状: [4096, 4096] 单个输入 x_A: [1, 4096] batch 3 个请求: X = [x_A; x_B; x_C] 形状: [3, 4096] Y = X · W^T 形状: [3, 4096] 这是一次 cuBLAS sgemm 调用 GPU 内部用几千个 CUDA 线程同时算 三个请求的结果在一次 kernel 执行中同时产出

GPU 是怎么做到的？

用一个具体的 4×4 矩阵乘法举例：

W (权重，固定): [w00 w01 w02 w03] [w10 w11 w12 w13] [w20 w21 w22 w23] [w30 w31 w32 w33] X (3 个请求的输入拼在一起): [a0 a1 a2 a3] ← 请求 A [b0 b1 b2 b3] ← 请求 B [c0 c1 c2 c3] ← 请求 C Y = X · W^T: [a0*w00+a1*w01+a2*w02+a3*w03, ...] ← A 的结果 [b0*w00+b1*w01+b2*w02+b3*w03, ...] ← B 的结果 [c0*w00+c1*w01+c2*w02+c3*w03, ...] ← C 的结果

GPU 会把这个计算分配给大量线程：

线程(0,0): 算 Y[0][0] = a0*w00 + a1*w01 + a2*w02 + a3*w03 线程(1,0): 算 Y[1][0] = b0*w00 + b1*w01 + b2*w02 + b3*w03 线程(2,0): 算 Y[2][0] = c0*w00 + c1*w01 + c2*w02 + c3*w03 ... 这些线程在 GPU 上是真正同时运行的 不存在"先算 A 再算 B"的顺序

一个类比

想象一个巨大的乘法表：

CPU 思维（串行）: 老师给 3 个学生各出一道题 学生 A 做完 → 学生 B 做 → 学生 C 做 总时间 = 3 × 单题时间 GPU 思维（并行）: 一面墙上挂着乘法表（权重） 3 个学生同时站在墙前 每人看自己需要的那部分，同时算 总时间 = 1 × 单题时间

batch size 越大，GPU 越高兴——因为矩阵乘法的维度越大，GPU 的几千个计算单元越能被充分利用。这也是为什么 Continuous Batching 能提升吞吐量：更大的 batch 意味着更大的矩阵，GPU 的计算资源被更充分地填满。

那 Attention 部分呢？不同请求的 KV Cache 长度不同怎么办？

Attention 计算确实每个请求的 K、V 长度不一样，这里不能直接拼成一个整齐的矩阵乘法。实际有两种处理方式：

Padding 方式（简单但浪费）：把所有请求 pad 到相同长度，空位填 0，然后一起算。浪费计算。

Padded + FlashAttention 方式（现代方案）：vLLM 用 FlashAttention 的 varlen 模式——传入每个请求的实际长度，kernel 内部根据长度分别处理，但依然在同一个 GPU kernel 调用中完成。不同请求的 Attention 在 GPU 的不同 thread block 上并行执行。

所以整体流程是：线性层部分（QKV 投影、FFN、LM Head）完全合并成一个大矩阵乘法；Attention 部分在同一个 kernel 中按请求分别处理但并行执行。没有先后顺序。

如果不混在一起会怎样？

方案1: Prefill 和 Decode 分开跑 Iteration 1: 只做 Prefill（请求 C） A 和 B 等着，GPU 算力没被充分利用 Iteration 2: 只做 Decode（A, B, C） 正常 问题: A 和 B 的 Decode 被 C 的 Prefill 阻塞了一次 用户感知到 A 和 B 的生成突然停顿了一下

混在一起的好处就是 GPU 在一次 forward pass 中同时推进了所有请求，没有等待时间。

所以本质上，Prefill 和 Decode 的区别只是"输入长度不同"和"KV Cache 操作不同"，但底层模型计算完全一致，GPU 可以在同一次 forward pass 中一起处理。