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阶段可以并行"的时候，希望你能在心里微微一笑，因为你知道，并行背后也藏着秩序和等待。