/AI 芯片的软件和硬件设计(14)——Prefill/

翻译自 Bjarke Hammersholt Roune 的博文

1

Prefill

Prefill 所执行的计算与 Decode 本质上非常相似，区别在于：Prefill 处理的是已经存在的 Token，而 Decode 需要逐个生成新的 Token。

因此，在 Prefill 阶段，AI 的任务是理解已经存在的文本或其他输入数据，并根据这些输入生成相应的 Key 和 Value，从而预先填充 KV Cache（Prefilling the KV Cache），而不是生成新的文本。

这一差异非常重要。

在 Decode 阶段，由于采用自回归生成方式，下一个 Token 依赖前一个 Token 的生成结果，因此只能按照：

Token 1→Token 2→Token 3→……

的顺序逐个处理。

而 Prefill 不存在这种生成依赖，因为所有输入 Token 都已经提前给定。

因此，只要芯片内存容量允许，就可以同时处理输入文本中的所有 Token，而不需要像 Decode 一样逐个 Token 进行计算。

换言之，Decode 的计算形式更接近：

一次处理 1 个 Token→生成 1 个 Query→执行 Attention

而 Prefill 则可以变成：

一次处理大量 Token→同时生成大量 Query→并行执行 Attention

因此，相比 Decode 中只有单个 Query 导致矩阵乘法某一维度为 1 的情况，Prefill 能够自然形成规模更大的矩阵乘法，从而更适合映射到大型脉动阵列。

也就是说，Prefill 和 Decode 虽然执行的 Transformer 计算在本质上非常相似，但二者具有完全不同的并行特征：

Decode：输入 Token 逐步产生，存在严格的自回归依赖，因此 Token 级并行度较低。

Prefill：输入 Token 已经全部存在，不存在 Token 生成依赖，因此可以同时处理大量 Token。

这也是 Prefill 通常比 Decode 更加适合大型脉动阵列的重要原因。

这里可以看到，与 Decode 阶段只有一个 Query 向量不同，Prefill 可以同时处理所有已经存在的 Token，因此能够同时得到与 Token 数量相同的 Query 向量。这与推测解码有些类似，但 Prefill 已经确定后续 Token 是什么，因此不存在任何“推测”。例如，输入可以是一个包含数千个 Token 的完整网页，这些 Token 都可以并行处理。

这会显著改善内存带宽效率。每次从内存读取一个 Key，或者首次生成一个 Key 后，该 Key 都可能同时被数千个 Query 使用。因此，相比 Decode，Key 能够获得非常高的数据复用率，也就不存在 Decode 阶段 Query 数量过少所造成的严重问题。

同样的优势也适用于训练（Training）。训练处理的文本或其他数据同样已经提前生成，因此也具有大量可同时处理的 Query。换言之，前文讨论的低 Query 并行度问题主要是 Decode 特有的问题。

2

Prefill 中的 GQA

Prefill 仍然可以从GQA（Grouped Query Attention）中获益，但其作用与 Decode 有所不同。

在 Decode 中，GQA 的重要作用之一是改善矩阵乘法的 N 维度，因为单 Token Decode 原本只有一个 Query，N 非常小。

但在 Prefill 中，这一问题基本不存在。Prefill 已经拥有大量 Query 向量，因此矩阵乘法的 N 维度通常足够大，可以在该维度上充分利用脉动阵列。

GQA 仍然能够降低 Prefill 的内存带宽需求，因为采用 GQA 之后，需要读取的 Key 数量减少。不过，与 Decode 相比，Prefill 本身就不太容易受到内存带宽限制，因此这一优势的重要性相对较低。

GQA 对 Prefill 更重要的作用体现在矩阵乘法的K 维度。

前文已经提到，要使脉动阵列达到较高甚至 100%的利用率，不仅 N 需要足够大，K 也需要足够大。

在 Decode 部分已经提到，GQA 通常会与增大 Key 和 Query 向量宽度结合使用。从矩阵乘法角度看，增大 Attention 向量宽度实际上就是直接增大 K 维度。

因此：

更宽的 Key/Query 向量→更大的 Attention K 维度→更高的脉动阵列利用率

这意味着，虽然 Prefill 已经通过大量 Query 解决了 N 维度问题，但 GQA 仍然可以通过增大 K 维度，帮助 Prefill Attention 获得较高甚至 100%的脉动阵列利用率。

因此，GQA 不仅对 Decode 重要，对 Prefill 同样重要。

/Decode-Prefill 的聚合、解耦与受控聚合/

Prefill 和 Decode 可以采用不同的系统部署方式。作者将其分为三类：默认聚合（Default Aggregation）、解耦（Disaggregation）和受控聚合（Managed Aggregation）。

3

默认聚合

最直接的方法，是在同一颗芯片上使用相同的 Kernel 同时执行 Prefill 和 Decode。

从数学运算本身来看，Prefill 和 Decode 之间并没有太大差异。两者最明显的区别，是每个对话一次处理的 Token 数量存在巨大差异。

例如：

Decode：一次约 1～4 个 Token

Prefill：一次约 256 个 Token

因此，两类任务完全可以混合起来同时执行。

不过，在这种默认方式下，Prefill 与 Decode 的比例没有经过专门控制。结果可能是：

某颗芯片主要执行 Prefill，而另一颗芯片主要执行 Decode。

这种分配是动态且不可控的，因此可能导致不同芯片上的计算资源和内存带宽利用率出现明显差异。

4

Decode-Prefill 解耦

Decode-Prefill Disaggregation是指将 Prefill 和 Decode 分别放到不同芯片上执行。

甚至可以为二者使用不同类型的芯片，不过这并不是必须的。例如，由于 Decode 通常需要更高的内存带宽，因此可以为 Decode 配置具有更高带宽的硬件。

这种方案需要额外完成一个重要步骤：

Prefill 芯片生成 KV Cache→将 KV Cache 传输到 Decode 芯片→Decode 继续使用

其主要优势是，可以分别针对 Prefill 和 Decode 开发专用的软件，甚至采用专门优化的硬件。

解耦还可以解决另一个问题：Prefill 与 Decode 混合执行时，某些时刻芯片可能被大量 Prefill 任务占据。

一次 Prefill 可能需要同时处理数千个 Token，因此单次执行时间可能明显长于只处理一个 Token 的 Decode。

这样就可能出现：

单个 Decode Token 等待数千个 Prefill Token 处理完成→Decode 再次获得执行机会

从而增加 Decode 的 Token 延迟。

作者指出，这一点经常被作为 Decode-Prefill 解耦的优势。不过作者同时认为，如果系统能够在构造 Batch 时智能考虑 Prefill 与 Decode 之间的差异，那么这一问题实际上不一定非常严重。

5

Decode-Prefill 受控聚合

作者进一步提出受控聚合（Managed Aggregation）的概念，即不简单地随机混合 Prefill 和 Decode，而是按照经过控制的比例将两类工作负载组合起来执行。

这种方式在以下情况下尤其有价值：

Decode 受到内存带宽限制（Memory-bound）

而：

Prefill 受到计算能力限制（Compute-bound）

Decode 的大量时间可能消耗在等待内存数据，此时脉动阵列存在空闲；Prefill 则可能主要等待脉动阵列计算，此时部分内存带宽存在空闲。

因此，如果按照合适比例将二者混合：

Memory-bound Decode + Compute-bound Prefill

就有可能形成更加平衡的执行状态，使：

内存带宽充分利用+脉动阵列充分利用

并且二者都尽量避免等待对方。

相比先只执行 Prefill、再只执行 Decode，这种方式既可以降低高优先级 Token 的延迟，也可以提高全部 Token 的总体吞吐率。

“Managed Aggregation”是作者在这里临时提出的名称。从实际实现角度看，它与目前所谓的Chunked Prefill非常相似，甚至可能本质上就是同一种方法。

受控聚合还有一个重要的硬件含义：

如果确定系统会使用这种技术，那么芯片可能不需要配置原本那么高的内存带宽。

因此，它的作用并不仅仅是防止芯片被大量 Prefill 任务占满，还可能直接影响芯片的硬件配置。

6

如何达到 Prefill 与 Decode 之间的平衡

整个 AI 系统中的 Prefill 与 Decode 任务总量不一定恰好满足理想比例。因此，并不能保证所有芯片始终获得完全平衡的工作负载。

但即使无法实现完美平衡，受控聚合仍然有意义。

至少可以避免出现：

一颗芯片几乎全部执行 Prefill，而旁边另一颗芯片几乎全部执行 Decode

这种明显不合理的资源分配。

公司还可以通过定价策略在一定程度上调节 Prefill 与 Decode 的比例，尤其对于 Offline Token，即用户并不要求立即获得结果的任务。

实际上，目前许多公司的输入 Token 价格已经明显低于输出 Token 价格。

此外，如果运营方自己拥有可以无限生成的 Decode 或 Prefill 工作负载，也可以利用这些任务动态填补系统中的空闲资源。

例如，可以自动生成一些用于检查 AI 回答一致性的任务，包括判断一个问题在加入否定条件之后，答案是否应该发生相应变化。

系统可以根据当前 Prefill 和 Decode 的负载比例，动态增加其中一类内部任务，从而尽量使整个系统始终保持较高利用率。

除了 Prefill 之外，还可以将 Decode 与训练任务混合，因为训练通常同样偏向 Compute-bound，尽管其计算受限程度不一定像 Prefill 那么明显。

因此，关键并不是一定要实现完美的 Prefill/Decode 比例，而是：

即使无法完全平衡，受控聚合仍然能够改善整体资源利用率。

7

通过 Attention 与 FF 重叠进一步改善平衡

受控聚合会产生一个最低的 Prefill/Decode 比例要求，以保证系统始终处于 Compute-bound，而不是重新落入 Memory-bound。

还可以通过重叠执行 Attention 与 FF进一步放宽这一比例要求。

FF 通常可能存在部分空闲内存带宽，因此可以在执行 FF 的同时，将这些带宽提供给 Attention 使用。

这种方法在一颗芯片包含两个或更多核心时效果更好。在 GPU 中，这类核心通常称为SM（Streaming Multiprocessor）。

如果只有一个核心，却强行让 Attention 和 FF 同时执行，那么两者可能争用 Cache 空间，最终不一定能够获得收益。

为了实现 Attention 和 FF 重叠，可以采用两种方式。

一种方式是使用流水线（Pipelining）：

Attention 和 FF 分别位于不同流水级，并同时处理不同任务。

另一种方式是直接修改模型结构：

让 Attention 与 FF 并行执行，而不是按照传统方式串行执行。

作者指出，后一种方法还能够将 Token 延迟降低约一半。

8

混合部署方案

还可以采用一种混合系统。

在这种方案中，大部分服务器集群按照接近理想的 Prefill/Decode 比例运行受控聚合。

与此同时，再部署少量：

Prefill 专用芯片

或者：

Decode 专用芯片

具体选择哪一种，取决于系统中哪类任务更容易出现过量。

这些专用芯片负责处理无法通过主集群平衡掉的溢出工作负载。

因此，整个系统可以形成：

主集群：受控 Prefill+Decode

辅助集群：处理 Prefill 或 Decode 的不平衡溢出任务

9

受控聚合对芯片内存带宽设计的影响

如果一颗芯片从设计之初就确定主要采用受控聚合，那么理论上可以为它配置比传统方案更低的内存带宽。

因为 Decode 虽然需要大量内存带宽，但可以与 Compute-bound 的 Prefill 混合执行，使内存带宽和计算资源互补。

减少内存带宽配置可以直接降低芯片及系统的制造成本。

不过，作者认为，这是一种相当大胆，同时也具有较高风险的设计选择。

原因在于，并不是所有客户都一定会使用受控聚合。

如果某些客户仍然按照传统方式分别执行 Decode 和 Prefill，那么由于芯片本身配置的内存带宽较低，其 Decode 性能可能明显下降。

这再次说明了作者此前反复强调的一个问题：

不同软件能力的客户，对 AI 芯片硬件配置的需求并不相同。

对于软件能力较弱、无法充分进行系统级优化的客户，芯片往往需要通过配置更多硬件资源来弥补软件不足，因此产品成本也更高。

而对于能够实施受控聚合、Chunked Prefill 以及其他高级调度优化的客户，则可以使用更加激进、成本更低的硬件设计，同时仍然获得较高的系统利用率。