/AI 芯片的软件和硬件设计(13)——Decode/

翻译自 Bjarke Hammersholt Roune 的博文

1

Decode

Decode 是模型需要真正“输出内容”的阶段，例如生成文本、图像，或者任何采用所谓自回归（Autoregressive）方式生成输出的任务。

在推理的两种主要阶段中，Decode 是硬件支持起来更加棘手的一类。Decode 的计算特征天然并不十分适合脉动阵列。不过，通过足够多的 AI 算法研究和软件优化，仍然可以使其较高效地运行在脉动阵列上。因此，对于 AI 芯片设计者而言，必须设法解决这一问题。

前文已经提到，Transformer 主要包含两类计算模块：

• 注意力（Attention）；
• 前馈网络（FF）。

在 Decode 阶段，这两类计算都会遇到问题。不过相比之下，Decode Attention 的问题更加突出。

1.1

Decode 的核心问题：一次只能生成一个 Token

Decode 的主要困难在于，如果采用最直接的实现方式，模型每次只能生成一个输出 Token。

这类似于人写作时通常需要逐个词向后写，而不能在尚未确定前一个词的情况下直接确定后面的内容。

这种生成方式称为自回归（Autoregression）。

其基本含义是：模型在决定前一个 Token 之前，无法开始生成下一个 Token。只有前一个 Token 确定之后，模型才能将这个 Token 作为新的输入，继续计算下一个 Token。

也就是说：

生成 Token t→将 Token t 作为输入→生成 Token t+1→将 Token t+1 作为输入→……

“Autoregression”这一名称也来源于此：

• Auto表示“自身”；
• Regression在这里表示模型需要基于此前的信息继续进行预测。

换言之，模型必须先看到自己刚刚生成的输出，才能继续生成后续内容。

这一严格的前后依赖关系限制了不同 Token 之间的并行计算，也是 Decode 难以高效映射到大规模脉动阵列的重要原因。

为了理解为什么这一点会产生问题，接下来可以观察自回归 Decode 过程中 Attention 的具体计算方式。

此前生成的每一个 Token 都会产生一个Key。Key 是一个向量，用于描述对应 Token 所包含的信息。由于这些 Key 已经在此前的计算过程中生成，因此可以直接从所谓的**KV Cache（Key-Value Cache，键值缓存）**中读取。这里暂时不讨论其中 Value 的具体作用。

当模型需要处理一个新的 Token 时，例如预测下一个词可能是“jumped”，模型会生成一个Query向量，并将这个 Query 分别与此前所有 Token 对应的 Key 进行比较，即执行点积运算。关于 Transformer Attention 的完整计算过程，可以参考相关 Transformer 基础资料。

这里真正的问题在于：这一过程本质上对应一次矩阵乘法，但由于当前只处理一个 Token，因此只有一个 Query。

也就是说，其中一个矩阵维度不仅很小，而且实际上只有1。

例如，可以将这一计算简化理解为：

1×K 的 Query 向量×K×M 的 Key 矩阵→1×M 的 Attention 结果

对于大型脉动阵列而言，这种矩阵形状非常不利。由于其中一个维度只有 1，绝大多数计算单元无法得到有效利用。因此，如果采用最直接的方式执行 Decode，脉动阵列利用率会非常低，甚至可能低于 1%。

1.2

长上下文进一步带来巨大的内存带宽压力

另一个问题是，KV Cache 中可能存在大量 Key。

假设用户已经与 AI 持续交流了两年。如果使用支持超长上下文窗口的模型，并且模型需要保留完整的历史对话，那么 KV Cache 中可能存在多达100 万个 Key。

为了生成当前的一个 Token，模型可能需要：

读取 100 万个 Key→分别与当前唯一的 Query 进行比较

即使绝大多数历史信息实际上与当前 Token 并没有关系，也仍然可能需要将这些 Key 从内存中读取出来。

这会产生巨大的内存带宽需求。

更严重的是，从内存读取如此大量的数据之后，每个 Key 只需要与一个 Query 进行少量计算。因此，整个过程表现出非常低的计算/访存比：

大量读取 KV Cache→只执行少量计算

结果就是，芯片中的大量计算资源实际上处于等待状态，主要时间都消耗在等待内存将大量 Key 送入计算单元。

1.3

Transformer 多层结构会进一步放大这一问题

Transformer 并不只有一个 Attention 层，而是通常包含多个层。

例如，一个模型可能包含 32 层。

这意味着，为了确定下一个 Token 是“jumped”，模型可能需要在32 个 Transformer 层中反复执行类似的 KV Cache 读取和 Attention 计算。

当最终确定“jumped”之后，整个过程并没有结束。

为了生成再下一个 Token，例如“over”，模型又需要从头开始执行下一轮 Decode：

生成“jumped”→重新执行整个模型→生成“over”→重新执行整个模型→生成下一个 Token……

因此，自回归 Decode 具有非常明显的串行特征：

一次生成一个 Token，并且每生成一个 Token，都需要重新执行整个 Transformer 模型。

1.4

Batch 为什么不能很好地解决 Decode Attention 问题？

这里可能会产生一个自然的问题：Transformer 不是可以使用Batch吗？

也就是说，一颗芯片可以同时处理多个不同用户的对话。这样看起来似乎能够将多个 Token 组合起来执行，从而提高脉动阵列利用率。

Batch 确实非常重要，尤其能够提高Decode 阶段 FF 层的计算效率。

但是，作者指出，Batch 对 Attention 的帮助并不大，因此无法从根本上解决这里讨论的问题。

Batch 可以让多个不同对话共享部分权重，从而减少权重重复读取。例如：

Conversation A、B、C、D→使用相同模型权重

因此，在 FF 等以权重矩阵乘为主的计算中，Batch 可以提高权重数据的复用率，并扩大矩阵乘法规模。

但是，Decode Attention 当前面临的主要问题并不是权重读取，而是Key 的读取。

不同对话具有完全不同的历史 Token，因此也具有不同的 KV Cache：

Conversation A→KV Cache A

Conversation B→KV Cache B

Conversation C→KV Cache C

这些对话之间的 Key 无法相互复用。

因此，即使通过 Batch 同时处理多个对话，每个对话仍然需要分别读取自己的大量 Key。Batch 虽然增加了同时处理的 Query 数量，但也相应增加了需要读取的 KV Cache 数据量，并没有像权重那样形成有效的数据共享。

因此，Decode Attention 仍然面临两个核心问题：

第一，单个 Token 只有一个 Query，导致矩阵维度极小，脉动阵列利用率非常低；

第二，Query 需要与大量历史 Key 进行比较，导致 KV Cache 读取量巨大，使计算容易受到内存带宽限制。

这也是为什么相比 Decode FF，Decode Attention 通常是 Decode 阶段更加棘手的硬件加速问题。

下面继续按前文一致的第三人称视角翻译，并保留原文关于 GQA、推测解码、脉动阵列利用率和 Decode 带宽需求之间的逻辑关系。

2

多头注意力与 GQA

Transformer 通常采用多头注意力（Multi-Head Attention，MHA）。这意味着前述 Attention 计算并不是每层只执行一次，而是每一层最多可能执行 128 次，每个 Attention Head 分别执行一次。

不同 Attention Head 拥有各自独立的 Key 和 Query 向量，用于描述 Token 的不同特征。

仅仅增加 Attention Head 本身并不能解决 Decode 阶段效率低的问题，因为每个 Head 仍然只有一个 Query。但如果与分组查询注意力（Grouped Query Attention，GQA）结合，就可以显著改善这一情况。

GQA 的核心思想是：让多个 Attention Head 共享同一组 Key。

假设模型共有 128 个 Attention Head，并将其划分为 G 组，则每组包含：

128/G 个 Attention Head

同一组内的 Attention Head 共享相同的 Key。

原本每个 Attention Head 只有一个 Query，并且每个 Head 都拥有自己独立的 Key。而经过 GQA 分组之后，同一组 Key 可以同时对应：

128/G 个 Query

因此，对于同一组 Key 而言，矩阵乘法中原本只有 1 的维度现在增加到了：

128/G

这能够显著提高脉动阵列的利用率。

与此同时，原本需要读取 128 套 Key，现在只需要读取 G 套，因此 Key 相关的内存带宽需求降低为原来的：

G/128

也就是带宽需求改善：

128/G 倍

除此之外，由于每组 Key 现在能够被重复使用 128/G 次，因此可以适当增大 Key 和 Query 向量的长度，从而提高其中包含的信息量。

从矩阵乘法角度看，这相当于增大了 K 维度。而前文已经指出，增大 K 同样有助于提高 Attention 阶段的脉动阵列利用率。

只要 Key 向量长度的增加倍率小于 128/G，那么即使 Key 本身变大，总体内存带宽仍然能够获得净收益。

因此，作者认为，GQA 是 Decode 阶段非常重要的一项优化。

3

推测解码进一步增加并行度

还可以进一步改善这一问题。

实际上，Decode 并不一定严格限制为每次只处理一个 Token，因为可以采用推测解码（Speculative Decoding）。

其核心思想是：首先利用一个计算成本较低的方法，对接下来可能生成的几个 Token 进行快速、有依据的猜测。

例如，可以一次猜测后续 3 个 Token。

假设当前上下文中，模型预测接下来可能出现：

“jumped over the”

由于已经提前获得了这 3 个候选 Token，主模型就可以在验证“jumped”“over”“the”是否正确的同时，继续计算“the”之后的 Token。

这些工作可以在一次处理中完成。

因此，原本一次只能处理：

1 个输出 Token

现在可以一次处理：

4 个输出 Token

这样可以进一步改善 KV Cache 内存带宽效率和脉动阵列利用率，理论上额外获得约4 倍收益。

当然，推测结果并不一定正确。

例如，快速模型可能预测：

“jumped over pineapple”

其中“pineapple”实际上是错误的。

此时，不仅处理“pineapple”所消耗的计算被浪费，基于“pineapple”继续预测第 4 个 Token 的计算同样无效，因为这些计算建立在错误的前序 Token 之上。

因此，推测解码本质上存在一定风险。

不过，只要快速模型预测正确的概率足够高，整体上仍然可以同时改善内存带宽效率和脉动阵列利用率。

4

GQA+推测解码可以带来多大收益？

推测解码的优势可以进一步量化。

假设：

• 模型共有 128 个 Attention Head；
• 采用 GQA，将其划分为 G 组；
• 每组包含 128/G 个 Head；
• 推测解码一次处理 4 个 Token。

由于每个 Token 对应一个 Query，因此对于同一组 Key，现在能够同时处理：

4×128/G 个 Query

相比最初只有一个 Query 的情况，KV Cache 数据复用程度和脉动阵列利用率理论上可以提高：

4×128/G 倍

例如，如果：

G=8

那么：

4×128/8=64

即理论上可以获得约64 倍改善。

因此，虽然 Decode 仍然可能存在脉动阵列利用率较低、内存带宽需求较高的问题，但实际情况并没有最初看起来那么糟糕。

通过GQA+推测解码，原本极端不利的 Decode Attention 矩阵形状可以获得数量级上的改善。

不过，如果使用 256×256 脉动阵列，那么从最初 N=1 开始，理论上最好能够获得：

256 倍

提升，才能使对应维度达到 256，从而充分利用整个阵列。

而 GQA+推测解码可能只能达到例如 64 倍。

因此，这也说明了一个重要问题：

脉动阵列规模越大，在 Decode Attention 中实现高利用率就越困难。

尤其是在 Attention Decode 阶段，这一问题更加明显。

5

为什么不直接令 G=1？

既然 G 越小，同一套 Key 能够被越多 Attention Head 共享，那么为什么不直接令：

G=1？

如果共有 128 个 Attention Head，此时所有 128 个 Head 都共享同一套 Key，理论上能够获得极高的 Key 复用率和 Decode 效率。

问题在于，这会影响模型质量。

当 G=1 时，每个 Token 只有一套 Key 供所有 Attention Head 共同使用，不同 Head 无法通过独立 Key 充分表达不同的信息。

结果通常是模型能力明显下降。

因此，需要在两个目标之间进行权衡：

较小 G→更高 Decode 效率

较大 G→更好的模型质量

实际设计需要选择足够大的 G 以保持模型能力，同时又不能大到使 Decode 性能明显下降。

作者指出，这里只是对实际问题的简化描述，但能够体现基本思想。

6

为什么不一次推测 16 个 Token？

类似地，如果一次推测 4 个 Token 能够带来 4 倍收益，那么是否可以直接一次推测 16 个 Token，从而再提高 4 倍？

问题在于，用于生成这些候选 Token 的模型通常是一个更小、更快但能力较弱的模型。

之所以使用它，就是因为它必须比主模型快得多，否则就无法提前多个 Token 完成预测。

但由于这个模型能力较弱，它很难连续准确预测主模型未来 16 个 Token。

很可能在到达第 16 个 Token 之前，它就已经产生类似“pineapple”这样的错误预测。

一旦中间某个 Token 错误，其后的推测计算通常都会失效，从而造成大量计算浪费。

因此，推测 Token 数量同样需要进行权衡：

推测过少→并行度不足；

推测过多→预测错误导致大量无效计算。

实际系统可以采用更加复杂的方法，例如根据预测置信度动态调整推测 Token 数量。

7

进一步降低脉动阵列所需的 N

作者还提出第三种优化方法。

如果希望在前述 64 倍改善基础上再获得约 4 倍提升，从而达到：

64×4=256 倍

理论上就有可能使 256×256 脉动阵列也达到接近 100%的利用率。

实现方式是降低脉动阵列达到满利用率所要求的 N 维度，也就是前文“非方形脉动阵列”部分讨论的方法。

不过，这种优化主要适用于两个条件：

第一，没有专门为 Attention 配置较小的脉动阵列；

第二，当前系统并没有受到内存带宽限制。

8

Decode 并不一定天然是 Memory-bound

综合来看，虽然实现起来并不容易，但大型脉动阵列在 Decode 阶段仍然有可能获得较高利用率。

而且，前述很多技术不仅提高计算利用率，也能够降低内存带宽需求。

一般而言，任何能够减少 KV Cache 容量的方法，通常也能够相应降低 Decode 阶段的内存访问压力。

目前存在大量这样的优化技术，数量甚至多到作者无法在这里全部列举。

传统观点通常认为：

Decode 天然严重受内存带宽限制。

即使配置大量内存带宽，也很容易进入 Memory-bound 状态。

但作者认为，这一结论并不绝对。

Decode 最终是否受到内存带宽限制，很大程度上取决于软件和模型对各种优化技术的使用程度。

例如，稀疏注意力（Sparse Attention）或索引注意力（Indexed Attention）可能彻底改变 Decode 阶段的带宽需求。

如果相关研究进一步发展，作者甚至认为，Decode 的带宽需求理论上可能降低到SSD 带宽都足以支持的程度。

真正的问题在于，芯片厂商通常无法完全控制客户的软件和模型。

如果客户的软件没有充分采用这些优化，那么芯片仍然必须提供极高的内存带宽，否则客户就无法获得满意的性能。

NVIDIA GPU 和 Google TPU 配置巨大内存带宽，其中一个原因就在于此。

因此，“AI 芯片究竟应该配置多少内存带宽”并没有简单答案。

答案很大程度上取决于软件。

9

软件优化可能造成 10 倍甚至 100 倍的硬件需求差异

作者认为，这里已经可以观察到一个贯穿 AI 芯片设计的重要规律：

AI 芯片应该配置多少硬件资源，几乎所有问题最终都与软件有关。

根据软件栈和 AI 模型是否针对硬件进行了充分优化，同一个问题的最终答案可能相差：

10 倍甚至 100 倍。

因此，AI 芯片的硬件规格不能脱离软件单独讨论。

10

多芯片并行不会降低单位芯片的 Decode 带宽需求

对于存储容量而言，多芯片并行可以将 KV Cache 分散到不同芯片上。

但对于读取 KV Cache 所需要的Decode 内存带宽而言，并行化并不会自然降低每颗芯片所需要的带宽。

假设将计算并行到 C 颗芯片。

每颗芯片对于单个 Token 只需要读取原来的：

1/C

数据。

看起来每颗芯片的带宽需求降低了 C 倍。

但是，采用 C 颗芯片的目的通常是希望在相同时间内处理：

C 倍 Token

因此：

C×1/C=1

最终，每颗芯片所要求的内存带宽基本保持不变。

除非通过多芯片分布之后，KV Cache 能够全部放入 SRAM，否则仅依靠并行化并不能从根本上降低单位芯片的 Decode 内存带宽需求。

11

AI 芯片必须同时提供配套软件

如果一家厂商准备销售 AI 芯片，就必须确保 GQA、推测解码以及其他相关优化能够在其硬件上高效运行，同时还必须提供足够完善的软件，使客户能够真正使用这些能力。

否则，同竞争对手相比，实际性能差距可能达到例如：

64 倍

作者强调，本节描述实际上已经进行了高度简化。

真正实现这些优化所需要的软件远比这里介绍的复杂，而且这仅仅是 AI 芯片软件栈需要解决的众多问题之一。

因此，AI 芯片产品真正困难的地方并不只有硬件，软件同样具有极高的实现门槛。

12

多 Agent 带来的进一步复用

作者进一步讨论了一种额外的并行来源，即Agent Multiplication（多 Agent 并行）。

例如，一些 AI 助手可能同时运行多个 Chain-of-Thought Agent。这些 Agent 不仅能够关注此前 Token 对应的 Key，也能够关注其他 Agent 生成的 Token。

这种机制可以在：

推测解码×GQA×非方形脉动阵列

之外，再增加一个新的乘数。

例如，如果同时运行 4 条 Chain of Thought，那么 Attention 阶段的 Key 可以进一步获得：

4 倍复用

同时，对于单个用户而言，也相当于额外获得了约4 倍可利用 Batch。

严格来说，多 Agent 并行并不能直接视为最终回答生成速度提高 4 倍，因为它主要加速的是并行 Chain-of-Thought 部分，而不是最终向用户输出回答的部分。

但作者指出，大规模 Chain of Thought 恰恰是 Token 消耗非常大的阶段，因此会显著提高单用户 Token 吞吐需求。

所以，即使只是提高这一阶段的吞吐率，也具有较高价值。

13

那么 Batch=1 芯片还有必要吗？

假设同时采用：

4 倍推测解码×4 倍 Agent 并行=16 倍单用户有效 Batch

那么即使只有一个用户，也可能获得相当于 Batch=16 的并行度。

这似乎意味着专门针对 Batch=1 设计的芯片不再必要。

但作者认为，这并不能完全否定 Low Batch 芯片的价值。

Groq LPU 主要面向高价格、极低延迟 Token。

Low Batch 最主要的价值之一，就是获得极高的单用户 Token 生成速度。

现在，通过推测解码和多 Agent 并行，普通硬件似乎也能够获得更高的单用户吞吐率，但这并不意味着 LPU 失去优势。

如果单用户已经能够提供 16 倍并行度，那么理论上也可以购买 16 倍数量的 LPU，并进一步获得更高的单用户 Token 吞吐率。

因此，在不计成本、追求极高单用户吞吐率的场景中，LPU 仍然可能具有优势。

不过，这些技术确实会降低 LPU 方案的吸引力。

首先，为了进一步利用这种并行度，需要购买更多 LPU，而运行大型模型本身就已经需要大量 LPU。

其次，在实际应用中，单用户 Token 生成速度通常存在一个“足够快”的阈值。

如果推测解码和多 Agent 等技术使更加经济的硬件方案达到这一阈值，那么就没有必要为了进一步降低延迟而购买价格更高的 LPU。

此外，随着单用户有效 Batch 增加，其他硬件方案的经济效率也会提高，因为其对极端高带宽设计的依赖下降。

这样一来，LPU 所生成的低延迟 Token 在相对意义上会变得更加昂贵。

14

对端侧 AI 同样适用

这一结论同样适用于端侧和嵌入式 AI。

端侧设备通常只有一个用户，因此人们经常认为这类设备天然需要 Batch=1。

这与数据中心不同。数据中心可以同时服务大量用户，因此很容易通过多个用户形成较大的 Batch。

但结合本节介绍的技术可以发现：

单用户并不一定意味着有效 Batch 只能等于 1。

通过推测解码等方法，即使只有一个用户，也可能形成一定程度的 Token 级并行。

对于端侧文本 Decode 而言，总体 Batch 可能仍然比较小，但即使能够从：

Batch=1 提高到 Batch=4

其硬件效率也会明显改善。

最后需要注意，本节讨论的主要是文本 Decode。

对于图像生成，通常使用 Diffusion 模型。这类模型天然具有更多可利用的并行度，因此 Low Batch 问题没有文本自回归 Decode 那么突出。

目前也存在用于文本生成的 Diffusion 模型，尽管其整体效果通常还不如主流自回归模型。

如果未来文本 Diffusion 模型能够达到足够好的效果，那么它也可能成为解决文本 Decode 低 Batch 问题的另一条技术路径。