Transformer Prefill阶段并行计算：本质、流程与思考
"为什么Transformer在Prefill阶段可以并行？并行到什么程度？哪里还需要同步？今天讲清楚！"
引子
在大语言模型（LLMs）爆发的时代，Transformer架构成为了当之无愧的基石。随着模型体积越来越大，推理加速成为性能瓶颈。
在推理过程中，人们常常提到两个阶段："Prefill"和"Decode"。
尤其是Prefill阶段 ，经常被描述为"可以并行处理所有token"，但这个并行到底是怎么实现的？有没有限制？
很多资料一笔带过，今天我们深入浅出，带你直击本质。
Prefill阶段是什么？
简单来说：
- Prefill阶段 就是：
- 模型已经收到了完整的一段输入文本（prompt）， 需要对这段文本进行一次性处理，
- 生成隐藏状态（hidden states）和缓存（KV Cache），以便后续快速推理。
比如输入一句话：
"The quick brown fox jumps over the lazy dog."模型会一次性处理整句话，把每个token的位置都计算好。
Prefill阶段的最大特点
关键词：上下文已知
- 在Prefill阶段，所有token的输入内容都是已知的。
- 每个token在Self-Attention中可以直接看到它应该能看到的所有上下文（比如前面的token，或者全部token）。
这与Decode阶段（一边生成一边喂入）完全不同。
正是因为上下文完整，Prefill阶段可以进行大规模并行计算！
Prefill阶段内部的计算流程
让我们以一个简单的4-token例子来展开：
假设输入是 [A] [B] [C] [D]。
在Transformer内部，每一层主要做两件事：
- Self-Attention 机制
- 前馈神经网络（Feed Forward）
其中Self-Attention最关键，让我们细细拆开看。
（1）首先，每个token独立计算自己的 Q/K/V
- 每个token的hidden state（比如embedding）输入到三个不同的线性层，分别产生：
- Query向量（Q）
- Key向量（K）
- Value向量（V）
这一阶段是完全并行的：
| Token | 生成 | 并行性 |
|---|---|---|
| A | QA、KA、VA | ✅ |
| B | QB、KB、VB | ✅ |
| C | QC、KC、VC | ✅ |
| D | QD、KD、VD | ✅ |
✔️ 完全并行，不需要等待。
（2）然后，需要同步所有token的Q/K，进行Attention打分
接下来，每个token要做Attention：
-
比如token C，要用自己的QC去和所有K做点积：
Attention Score(C, A) = QC · KA
Attention Score(C, B) = QC · KB
Attention Score(C, C) = QC · KC
Attention Score(C, D) = QC · KD -
这就意味着：
- 每个token的注意力得分不仅依赖自己，还依赖所有其他token的K。
所以：
必须同步等待所有token的Q/K都计算完成，才能继续。
🚨 同步点出现了！
（3）Softmax归一化 Attention Scores
- 每个token的Attention得分进行Softmax归一化，得到权重。
这一步可以在每个token内部并行做，但需要确保Attention Scores已经准备好。
（4）根据Attention权重，聚合Value向量
- 每个token根据权重对所有token的V向量加权求和，生成新的hidden state。
这也是并行的（每个token自己处理自己的加权），但同样要等所有Value准备好。
（5）前馈网络（Feed Forward）
- 每个token的新的hidden state经过前馈神经网络（一般是MLP）。
这一阶段又是完全并行的。每个token独立处理自己。
Prefill阶段并行与同步总结
| 阶段 | 并行性 | 是否需要同步？ | 说明 |
|---|---|---|---|
| 生成Q/K/V | 完全并行 | ❌ 不需要同步 | |
| Attention打分（Q×K^T） | 并行矩阵乘法 | ✅ 需要全部token的Q/K | |
| Softmax归一化 | 局部并行 | ✅ 需等待Attention Scores | |
| 聚合V向量 | 局部并行 | ✅ 需完成Softmax | |
| 前馈网络 | 完全并行 | ❌ |
层与层之间：严格串行
虽然单层内部可以高度并行，但不同层之间 存在严格的串行依赖。
原因很简单：
- 下一层的输入是上一层的输出hidden state。
- 没有上一层结果，下一层就无法开始。
所以：
Prefill阶段是层内并行、层间串行。
Prefill阶段的并行上限
Prefill阶段可以并行到什么程度？
具体取决于：
- Token数量（prompt长度）
- 注意力头数
- 批量大小（batch size）
在硬件层面上，Prefill阶段主要是大规模矩阵乘法（比如GEMM），适合GPU并行加速。
这也是为什么：
- A100、H100等GPU在Prefill阶段能展现出超强吞吐能力。
Prefill阶段 vs Decode阶段
| 特性 | Prefill阶段 | Decode阶段 |
|---|---|---|
| 输入 | 完整prompt已知 | 一个一个生成 |
| 计算模式 | token内并行，层间串行 | token串行，层间串行 |
| 并行度 | 高 | 非常低 |
| 速度 | 快 | 慢 |
Prefill阶段基本是Transformer推理中最快的一段，而Decode阶段则是最慢、最痛苦的一段。
小结
🔵 Prefill阶段并行计算特点总结：
- 单层内部：
- Q/K/V生成可以完全并行。
- Attention打分需要同步所有token的Q/K后统一处理。
- 聚合V向量也需要同步。
- 层与层之间：
- 必须等待上一层完全计算完，才能进入下一层。
- 整体：
- 层内高度并行（有同步点），
- 层间严格串行。
结语
了解Transformer Prefill阶段的并行计算本质，不仅有助于我们理解模型的推理性能，还能帮助我们在未来优化部署、提升吞吐率。
Prefill并不是"完全自由飞奔"，它内部也有同步障碍，只不过整体上是高度并行的。
下一次，当你听到"Prefill阶段可以并行"的时候，希望你能在心里微微一笑，因为你知道，并行背后也藏着秩序和等待。