原文：Designing AI Chip Hardware and Software，作者 Bjarke Hammersholt Roune，2026 年
原文由作者公开发布并声明 freely available，鼓励传播。本文为中文翻译版。
设计 AI 芯片的硬件与软件#
Bjarke Hammersholt Roune
2026
这篇文档是什么？#
我本来打算创办一家 AI 芯片公司，这篇文档包含了我关于如何制造 AI 芯片的思考，涵盖软件和硬件两方面。它是写给那些将要负责或参与 AI 芯片项目的人的——这是你需要了解的一部分内容。
我想其他人也许也会对了解 AI 芯片这个领域感兴趣，因此大部分内容应该在业余爱好者水平上就能容易理解。有些部分更为技术化，需要更深的背景知识——这些部分以粗体标记为"技术细节（for nerds）"。
这篇文档还包含了我在硅谷工作期间的一些轶事，我希望它们既有趣又有启发性。它们都有明确标记，如果你对此不感兴趣可以跳过。
我知道这篇文档目前已经被一些公司用于 AI 芯片公司的员工入职培训，我认为这是个好主意，因为文档中的许多内容在其他地方不容易找到。
我是谁，为什么有资格写这个#
我曾在 Google 从事 TPU 工作，在 Nvidia 从事 GPU 工作，担任技术领导职位，也曾在其他生产 AI 软件和硬件的公司工作，如 Amazon 和 Facebook。我曾是 Google TPUv3 的软件技术负责人，并在当时深度参与了 TPU 硬件特性的选择和设计。这是我的专业。
执行摘要#
我在这篇文档中的主要论点是，AI 加速器的最终未来将是 AI CPU——带有大缓存和脉动阵列（systolic array）的传统式 CPU，其中缓存和脉动阵列占据了芯片的大部分面积和功率。要从这个想法中得到一个有竞争力的 AI 芯片，有很多细节需要做对，逐一讨论这些细节是本文档的主要内容。
你在这篇文档中会发现的最大张力在于，一方面是最大化每美元的 token 数，另一方面是用大量昂贵且不必要的功能和容量严重过度配置你的 AI 芯片，以便客户收到的东西看起来像他们已经在使用和习惯的东西。我在这篇文档中给出的建议假设你拥有一个最出色的软件团队，能够让这一切运作起来，这样客户即使没有严重的过度配置也能轻松使用你的芯片。你绝对不会轻松地招到能做到这一点的软件团队。所以请记住这一点。
如果你从事 AI 软件工作，你可能还想阅读这篇关于 AI 编译器测试的帖子。由于我已经在那篇文档中讨论了测试，我不会在这篇文档中深入测试话题。无关的：我确实有一些关于可解释 transformer 的想法，我认为很有趣。
我会为你的公司工作吗？#
我已经退休了，不在找工作。我不做付费咨询。我不会管理你的初创公司，我对接受资金也不感兴趣。我没法帮你找联合创始人。如果你想问问题而我也愿意回答的话，我会免费回答问题，而我很可能愿意回答。
硬件设计提案（摘要）#
这是我提议的设计摘要。我建议最后再读这一节，但如果你愿意也可以先读。这个设计的定义特征是，它在硬件上容易实现，像 CPU 一样容易编程，并且在计算/瓦特方面非常便宜。我怀疑没有人会简单地照搬我的计划，但我建议仔细审视这个设计，并理解对于每一个偏离，你到底为什么认为偏离是个好主意。如果我真的开始构建这个，我肯定也会随着时间修改设计。
“等等，这就是一个带脉动阵列的 CPU，你的意思是你能用这个击败复杂的 GPU 和 AI ASIC？“是的，这就是我相信的。这个芯片的大部分是缓存和脉动阵列，所以你不需要在标量和向量计算上偏离 CPU 方式。CPU 是最可编程的东西，硬件和软件工程师都已经知道如何构建和使用它们。这就是这个想法。你需要阅读整个文档才能理解为什么这能行得通，但以下是设计的一些亮点：
- 设计一个传统 CPU，带有额外的向量计算（和寄存器），并为每个核心附加一个脉动阵列。如果你叫 Intel 或 AMD，可能用 x86，否则找一些你能制造或授权的快速且便宜的东西。
- 拥有一个大型软件管理的 SRAM 缓存（MB 级），在给定程序未使用其部分时可以兼作自动缓存。
- 支持所有内存和网络之间的 DMA。允许 SRAM-网络-SRAM 的 DMA。提供硬件 8 位 Huffman 压缩/解压缩，速度匹配 HBM/链路速度的 DMA。
- 不需要主机计算机——你的芯片本身就是通用 CPU。确保你的芯片能运行 Linux 和 TensorFlow/PyTorch。添加没有脉动阵列的标量核心。
- 不要为你的芯片提供大量 HBM。通过使用多个芯片来容纳权重。
- 为你的芯片配备 SSD。通过使用稀疏注意力使 KV cache 能够存储在这里。
- 脉动阵列只支持 7 位乘法，不支持 16 位。求和更宽并饱和。
- 脉动阵列支持 RHS 的 1-of-2 稀疏性——简单且允许 2 倍吞吐量。
- 主网络是铜走线的 1D 或 2D torus。使用（便宜的）定制主板，可以在 1D 或 2D 中相互扣合，连接走线。没有昂贵的线缆。
- 为你的芯片编写 XLA 后端。大多数用户将通过 XLA 使用这个设备。也支持用标准 C++（带 intrinsics）编写的自定义 kernel 和自定义 XLA op。
- 这个设计也适用于训练，但首先专注于推理——那更容易。
这避免了 AI 芯片中你并不真正需要的昂贵部分，用 RTL 制造这个硬件很直接，你可以像编程任何其他 CPU 一样编程它。我不知道为什么类似的东西不是 AI 芯片的行业标准。行业似乎更倾向于追求新颖的非 CPU 架构。
这个芯片需要相当的软件复杂度才能用好，你必须做一些 AI 研究来让稀疏（索引）注意力与 1:2 稀疏性良好工作，并且在这里建议的低 HBM 配置中，它需要一个较大的最小安装规模来运行大型 LLM。可以选择修改设计使其不那么经济以绕过这些因素，我认为这仍然是一个好设计，但那时它会更贵。我无法预测哪种选择会优化你的利润，但这里建议的是在大量使用时在最大化每美元 token 数方面最好的。
[[[ 关于 Groq LPU 的技术说明（for nerds）：Groq 的 LPU 设计不需要任何 HBM。这和这里说的一样吗？在分布式权重以避免每个芯片需要大内存容量来存储权重方面是一样的，导致对大安装规模的相同要求。然而，在处理 KV cache 需求方面不一样。Groq 使用低批量（因此低 token 延迟）的方式来做 decode（输出 token），这大幅降低了 KV cache 需求——这就是他们不需要任何 HBM 来存储 KV cache 的原因。这意味着 Groq 的芯片不使用（也不能使用）大型脉动阵列（那需要更高的批量），这必然导致更差的功耗和面积效率。Groq 的方法也导致非常高的带宽需求，LPU 拥有这些带宽因为它使用 SRAM 代替 HBM。我在这里提议的设计则使用索引注意力（导致更低的带宽需求）来实现将 KV cache 卸载到 SSD，这样大批量因此大脉动阵列仍然可以使用。但这样你就必须让稀疏（索引）注意力工作——如果你（和你的客户）搞不定这个，你就得添加更多 HBM，这确实是大多数 AI 芯片所做的。]]]
脉动阵列——AI 硬件的基础#
脉动阵列（systolic array）是 AI 硬件（即深度学习硬件）的基础。为什么？
AI 中几乎所有的数学运算（OP）都发生在矩阵乘法内部。大型脉动阵列是在硬件中实现矩阵乘法最高效的方式。所以，脉动阵列是 AI 硬件的基础。
像 Nvidia GPU（如 Blackwell）或 Google TPU 这样的 AI 加速器芯片就是访问脉动阵列的一种方式。AI 芯片上有比脉动阵列更多的元素，但脉动阵列是最重要的元素，就像汽车的引擎一样。
这篇文档不是解释什么是脉动阵列或它们如何用于矩阵乘法的——我假设你已经知道了，但如果你以前没听说过，可以让 Gemini 这样的 AI 助手来解释。
如果明天开发出一种不依赖矩阵乘法的更好的 AI，那么本文档中的结论可能都需要改变。这似乎不是目前的发展方向，因为矩阵乘法已经统治了很长时间，但对于任何 AI 硬件芯片生产商来说，他们当前的 AI 芯片明天可能没用是一个风险，取决于 AI 研究的发展。没有办法绕过这种情况。十年前是这样，今天也是这样。
像 Nvidia GPU 和 Google TPU 这样的 AI 芯片上有多个大型脉动阵列。如果 AI 芯片做得得当，那么 AI 消耗的大部分电力都用于驱动这些脉动阵列。所以当你听到 AI 使用大量电力时，实际上是脉动阵列使用大量电力，至少对于设计得当的 AI 硬件来说是这样。这并不是因为脉动阵列功耗低效——没有脉动阵列的话电力需求会高得多。
脉动阵列定义了 AI 芯片速度和功耗效率的上限。事实上，当我们谈论 AI 芯片的利用率时，理想情况下应该接近 100%，我们主要谈论的是脉动阵列有多忙。芯片的其余部分是为了让脉动阵列保持忙碌。
在 Nvidia GPU 上，脉动阵列位于 Nvidia 称为 TensorCore 的内部。Google 使用术语矩阵乘法单元（Matrix Multiply Unit, MXU）来指代他们的脉动阵列（你可能注意到"multiply"中没有 X，但 x 有时用来表示乘法，如 2 x 3 = 6，所以他们用带 X 的 MXU 来称呼它）。AMD 将他们的脉动阵列称为 Matrix Cores。Intel 将其脉动阵列的接口称为高级矩阵扩展（Advanced Matrix Extensions, AMX）。在 Amazon Inferentia/Trainium 上，脉动阵列被称为 NeuronCores。听起来似乎有许多不同的 AI 芯片方式，当然确实存在差异，但这都是脉动阵列加上其他东西把它们粘在一起。全都是脉动阵列。品牌名称是为了市场营销。有些 AI 芯片不包含脉动阵列，比如很老的 Nvidia GPU。与包含脉动阵列的芯片相比，这些芯片不适合大多数 AI 应用，因为功耗和面积效率差。
所以下次当你听到带有"北极光速度爱因斯坦天才核心"或"牛顿苹果落地尤里卡核心"或任何他们将要叫的名字的 AI 芯片时，记住他们说的是脉动阵列。全都是脉动阵列。供你参考，脉动阵列是 1978 年发明的。我们最近看到带有大型脉动阵列的 AI 芯片出现，是因为 AI 创造了对这种东西的需求。这不是什么先进的现代事物，这是旧技术，突然之间有了大量的客户基础。如果深度学习在 1978 年就被发明了，那时也会有一波脉动阵列芯片的热潮，我说的是 AI 芯片。那将是同样的事情。它并不是什么复杂的现代东西。
你可能会说，高内存和网络带宽对 AI 也很重要，所以不仅仅是脉动阵列。是的，但我们为什么需要高带宽来做 AI？那是因为脉动阵列如此高效，你需要高带宽来喂它们——否则它们会闲置等待更多工作。没有脉动阵列，芯片会慢得多，你也不需要每个芯片那么多带宽。AI 革命是建立在脉动阵列的成本效益之上的。
所以现在你知道了。AI 芯片的引擎是脉动阵列，每家公司都给他们特定的脉动阵列芯片取了不同的名字。
在本文档的其余部分，我将专注于 Nvidia GPU 和 Google TPU，因为我在 Google 和 Nvidia 都参与了这两者的工作，所以我很了解它们，而且截至 2026 年，这是业界排名前两位的 AI 芯片产品。当然还有许多其他脉动阵列芯片，我是说 AI 芯片。
[[[ 例外（for nerds）：有一些例外。Groq 的 LPU 没有大型脉动阵列，也许根本没有脉动阵列。LPU 是为了低 token 延迟，不是为了经济实惠的 AI。Nvidia 授权了 Groq 的技术，这正是 Jensen（Nvidia 的 CEO）在 GTC 2026 上说的——它是为了昂贵和低延迟的 token，不是一般的经济实惠 AI。还有模拟芯片，使用光或电，试图以其他方式实现经济性——这些尚未被证明，但如果被证明的话，可能是未来重要的东西。]]]
更大的脉动阵列更高效#
这一节是本文档中唯一包含一点数学的几节之一。所以，为了弥补这一点，它也有一些（画得很差的）图片。假设你想将两个 N × N 矩阵 A 和 B 相乘，得到矩阵积 C = A * B。那么你可以这样可视化一个脉动阵列：
A 的行 → [脉动阵列] → C 的行
↑
B 的行在每个周期（cycle），这个脉动阵列从左边接收一个向量，从顶部接收一个向量，并向右输出一个向量。你可以看出这个图是简化的，因为它需要知道整个矩阵 B 才能产生 C 的一行，所以在看到所有 B 之前不能产生 C 的任何行。这就是为什么在现实中，整个矩阵 B 被送入脉动阵列并存储在那里，然后才开始处理 A 的行。然而，如果你连续做多个独立的矩阵乘法，你可以在处理当前 A 矩阵的同时送入下一个 B 矩阵，并输出前一个 C 矩阵，这样脉动阵列的左、上和右在每个周期都可以忙碌，这叫做 100% 利用率。100% 是好的。脉动阵列还使用一种叫做偏斜/交错（skew/stagger）的东西，使图片更加复杂，但我会在讨论中省略它，因为它与本节无关。
N × N 脉动阵列的向量宽度是 N。当我们将向量宽度从 N 加倍到 2N，创建一个 2N × 2N 的脉动阵列而不是 N × N 的时候会发生什么？向量宽度增加了一倍，脉动阵列大了 4 倍，所以它每个周期可以做 4 倍的数学运算。要看到我们通过加倍向量宽度得到了什么，我们需要将图片扩展到包括与每个脉动阵列相关的 CPU（或类似 GPU 的 CPU 结构），它告诉脉动阵列该做什么并处理来自 A、B 和 C 的行（即向量）。
那么我们使用这个更大的脉动阵列得到了什么？我们更大的 2N × 2N 脉动阵列每个周期做的数学运算是之前的 4 倍，所以使用较小 N × N 脉动阵列的等效吞吐量结构将需要 4 个脉动阵列。它还需要 4 个额外的 CPU 核心来告诉 4 个脉动阵列该做什么并处理它们的向量。所以通过使用 2N × 2N 脉动阵列而不是 N × N 的，我们将 CPU 核心/工作的数量减少了 4 倍。然而，右边的 CPU 必须处理宽度是两倍的向量，这使它更贵，但不到两倍那么贵。所以通过从左图到右图，我们在芯片其余部分的效率上获得了超过 2 倍的改进。但我们在做这个时没有损失任何矩阵乘法能力。所以我们通过加倍向量宽度在芯片其余部分获得了 2 倍的效率增益。再做一次，你得到 4 倍。再做一次，你得到 8 倍，以此类推。
难道不能使用同一个 CPU 核心来驱动左图的所有 4 个脉动阵列吗？可以，但那时它需要 4 倍的能力/速度，所以这不会改变总体数学。
这些图解释了为什么我们看到 AI 的脉动阵列越来越大。Nvidia GPU 在 Volta（2017）中获得了 4 × 4 的脉动阵列，在 Turing（2018）中获得了 16 × 16 的脉动阵列，在 Blackwell（2025）中获得了 64 × 64 的脉动阵列。Google TPU 一直使用 128 × 128 或 256 × 256 的脉动阵列大小。
你可能认为这个图不适用于 GPU，因为我在图上放了脉动阵列和 CPU 核心，但图上没有叫 GPU 的东西。就本特定观点而言，GPU 只是 CPU 的一种特殊形式，一种有大量向量计算的形式，所以这个图完全适用于 GPU。在图中看到"CPU"的地方读作"GPU”，你会得到正确的想法。
**这里有个谜题可以思考：**Nvidia 早在 2017 年就已经是一家非常大且盈利的公司，GPU 是且一直是其主要产品。与此同时，虽然 Google 也是一家非常大的公司，但 Google 内部的 TPUv2 团队很小，硬件和软件合计大约 30 人（XLA 团队只有 5 人）。然而从那时起 Google TPU 持续与 Nvidia GPU 竞争得很好。这怎么可能？答案在很大程度上与脉动阵列的大小有关，我相信。
让我们比较 Volta 和 TPUv3 的脉动阵列大小。Volta 有 4×4 的脉动阵列，而 TPUv3 有 128×128 的脉动阵列。要看这有什么影响，让我们分别看标量计算和向量计算。
**标量计算：**标量计算包括告诉 AI 加速器其余部分应该做什么的处理。我们看到向量宽度加倍使标量工作减少 4 倍。从 4 到 128 是 5 次向量宽度加倍，所以从 4 × 4 脉动阵列到 128 × 128 脉动阵列的标量计算收益是：
4^5 = 4 * 4 * 4 * 4 * 4 = 1024x
AI 芯片还需要使用标量计算来做矩阵乘法以外的事情，所以这不是全部故事，但它是故事的很大一部分。TPU 只需要做 1/1024 的标量工作来驱动脉动阵列进行矩阵乘法（实际上甚至更多，因为 TPU 向量寄存器是 8x128 而不只是 128，所以是 8 * 1024 = 8192x）。标量工作不是 AI 芯片上的主要成本，但它不是完全微不足道的，1024 倍的差异不能忽视。
**向量计算：**我们看到向量宽度加倍使脉动阵列外的向量工作减少 2 倍。通过 5 次从 4 到 128 的加倍，变成了：
2^5 = 2 * 2 * 2 * 2 * 2 = 32x
这非常显著。你可能会说脉动阵列内部发生的事情本身就是一种向量工作，所以我们并没有消除这么多向量工作，我们只是把它从脉动阵列外面移到了里面。技术上这是对的，但脉动阵列比向量处理器高效得多，包括检索输入和存储结果的工作，所以实现的效果是把这些工作从效率较低的地方移到了效率高得多的地方。所以净效果不是 32 倍那么大，但仍然非常显著。
如果你问一个出色的硬件工程师更完整地解释这一点，他们会给你一个涉及寄存器文件、皮焦耳和芯片内数据传输效率的长篇解释，也许还包括关于各种替代设计的题外话，但总体结论是更大的脉动阵列效率高得多。
所以你开始看到一个小小的 Google 团队如何能跟上 Nvidia 的强大力量。还有其他原因，但这是一个重要原因。Nvidia GPU 现在有更大的脉动阵列了，但它们仍然不如 Google TPU 上的大。
这里的显然问题是：那为什么不是每个人都使用更大的脉动阵列，比如 1024 × 1024？那应该更高效，对吧？这是下一节的主题。
限制脉动阵列大小的因素#
从到目前为止的讨论来看，似乎最高效的 AI 芯片将完全由一个非常大的脉动阵列组成（这几乎就是 TPUv1）。芯片面积中由一个（或多个）大型脉动阵列组成的部分越多，你的芯片就越高效。在其他条件相同的情况下。那为什么不把整个芯片做成一个大脉动阵列？
最明显的原因是芯片的其余部分需要做一些标量和向量工作，你还需要一些空间用于缓存、内存控制器等，所以整个芯片不能只是一个大脉动阵列。事实上这些其他元素占 AI 芯片的很大一部分，尽管随着脉动阵列越来越大这种情况会越来越少。
[[[ 技术细节（for nerds）：TPUv1 基本上就是一个大脉动阵列。这需要从一个单独的外部 CPU 驱动它，因为芯片上没有类似 CPU 的元素。你可以这样做，但现在你必须买一个 CPU 来驱动脉动阵列，所以可以说你并没有节省芯片面积，你只是把那些面积放到了另一个你仍然需要购买的芯片上。]]]
对于 AI 来说，过大的脉动阵列还有另一个更具限制性的问题。N × N 的脉动阵列无法高效地执行小于 N × N 的矩阵乘法。如果你的 AI 模型中的向量是 64 个元素宽，而你有一个 128 × 128 的脉动阵列，你的脉动阵列将不会被高效使用。这里有个图来可视化。在图中，向量是 128 个元素宽，但只有前 64 个元素（红色）实际上在被使用：
你可以看到只有 1/4 = 25% 的脉动阵列在被使用。这叫 25% 利用率，非常糟糕。幸好有一种方法可以改善这个情况，尽管它需要在驱动脉动阵列的软件中增加一些复杂性：
这个想法是脉动阵列的左上和右下子集可以独立使用，这样 50% 的脉动阵列可以被使用。这是矩阵乘法数学的结果，所以你不需要硬件特性来启用这个优化。这样，脉动阵列可以达到 50% 的利用率。但 50% 的利用率仍然不好。我们想要 100%。
所以我们可以看到，每次脉动阵列是它运行的 AI 模型的向量宽度的两倍时，效率就会降低 2 倍（或 4 倍，如果软件没有使用上面提到的优化）。这个问题是 AI 芯片通常包含许多较小的脉动阵列而不是一个巨大的脉动阵列的主要原因。
所以结论是，AI 模型中矩阵乘法的维度限制了 AI 硬件中脉动阵列的大小。如果你的客户在他们的 AI 模型中使用小的向量宽度，你不能卖给他们一个大向量宽度的芯片，除非他们改变他们的 AI。
[[[ 三个维度（for nerds）：矩阵乘法涉及 3 个维度而不是 1 个。上面的图显示的是中间维度小于脉动阵列时发生的情况。如果左或右维度太小也会发生类似的问题，尽管对于那些维度通常有一些巧妙的方法可以找到更多数据来乘法，这样你实际上增加了那些维度以达到 100% 利用率。例如，AI 模型中的批量（batch）维度就服务于这个目的。对于我们在这里关注的中间维度，“找到更多数据"没有帮助：如果矩阵乘法中的点积只对这么多项求和，那么这是一个硬性限制，你对此无能为力。除了上面展示的优化。所以这就是为什么在本节中我关注中间维度，它对应于向量宽度——它是更根本的限制。]]]
AI 模型的向量宽度有其存在的原因。如果你的模型有一定数量的参数，那么有一个特定的向量宽度能最好地利用这些参数来创造聪明的 AI。你可以使用比那更高的向量宽度，但那时你的 AI 在给定相同数量的参数时会更笨。
按照行业经验法则，更大的 LLM 有更大的向量宽度，所以只要 LLM 继续增长，向量宽度因此脉动阵列也将继续增长。这似乎是一个将继续的趋势。
让我们关注 transformer 架构，它是当前 LLM AI 革命的驱动力。Transformer 有两种不同的矩阵乘法：注意力（attention）和前馈（feed-forward, FF）。我不会在这里解释 transformer，但如果你想知道 AI 中发生了什么，请研究它们。注意力矩阵乘法不幸地不是那么大。FF 矩阵乘法可以非常大。所以是 transformer 的注意力部分最限制 AI 芯片上脉动阵列的大小。否则我们现在可能已经看到 1024 × 1024 的脉动阵列了，而不是 64 × 64（Nvidia）和 256 × 256（Google）。
[[[ 技术细节（for nerds）：这里的完整图景比这个简化视图复杂得多——注意力不是经常被报告为等待内存多于等待脉动阵列吗？所以这听起来好像在注意力期间脉动阵列利用率低可能是可以接受的，因为在做注意力时它不需要非常高效地运行就能跟上内存。这有时是对的，但仍然有限制。这也涉及到训练、预填充和解码之间的区别，训练和预填充期间的内存带宽需求远低于解码期间，所以对于这些你确实需要脉动阵列的高利用率——详见后面关于这个的部分。像推测解码和 GQA 这样的技术使得甚至解码也大大减少了内存密集性（详见解码部分）。此外，如果你做得对，你可以重叠注意力和 FF（需要更深的流水线，或者更高效地，需要模型修改使它们并行进行），在这种情况下你确实总是想要脉动阵列的高利用率。稀疏注意力，特别是索引注意力，也可以大幅降低注意力的带宽需求，而且在我看来，所有注意力都应该是稀疏（索引）注意力。还有许多使 KV cache 更小的技术，因此降低解码带宽需求，我在这篇文档中甚至没有提到。所有这些的简化结论如上所述：FF 可以使用非常大的脉动阵列，而注意力，如果做得对，使用较小的脉动阵列效果更好。所以真正应该是 transformer 注意力矩阵乘法在限制 AI 芯片上脉动阵列的大小，尽管确实在实践中，如果你或你的客户不付出努力来实现这一点，你有时可能被其他东西限制。]]]
听到注意力矩阵乘法不是那么大可能会让人惊讶，因为 LLM 花了很多时间做注意力。如果注意力只使用小的矩阵乘法，那注意力不应该很快完成吗？解释是到目前为止的术语有点简化。矩阵乘法是一个 N × K 矩阵乘以一个 K × M 矩阵，产生一个 N × M 的结果。所以有三个独立的维度，不只是一个 N。注意力矩阵乘法在 K 维度小的意义上是"小"的，但其他两个维度通常很大，所以总计算量也很大。K 维度定义了多少个数字被加在一起来创建结果输出矩阵的一个元素。如果 K 小，那么一个被设计为对比这更多的数字求和的脉动阵列就无法获得满利用率，这在脉动阵列越大相比 K 时越糟糕。尽管 K 对注意力来说很小，N 和 M 可以很大，而且注意力需要做很多矩阵乘法（如果有很多注意力头），所以这就是注意力即使 K 很小也可以花很长时间的原因。在 FF 中，所有三个维度 N、K 和 M 通常都足够大以填满脉动阵列。
[[[ 关于图像和视频的细节（for nerds）：如果你的 AI 在处理图像或视频，注意每个像素只有 3 个值：红、绿、蓝。这可能是个问题，因为这意味着处理图像的矩阵乘法的 K 只有 3。这极其小，比注意力小得多。然而，这并不像看起来那么糟糕。一旦 AI 开始查看数据，我们可以且确实立即创建一个大得多的 K，所以只有在最初的处理中你有 K=3。此外，那个初始处理可能可以在 CPU 上完成，在你把数据交给 AI 芯片之前，因为它是轻量级的（由于 K=3），在这种情况下即使对于大型脉动阵列，在图像和视频处理期间你仍然可能获得 100% 的 AI 芯片利用率。]]]
[[[ 关于可分解脉动阵列的细节（for nerds）：Etched 的 Sohu 芯片使用可分解的脉动阵列，可以在硬件中配置为用于注意力的许多独立脉动阵列或用于 FF 的一个大型脉动阵列。这不是看起来那么完整的解决方案。鉴于芯片能够在许多分解的小脉动阵列模式下运行，服务该模式所需的昂贵片上带宽和标量/向量计算必须存在于芯片上。所以你并没有获得大型脉动阵列通常会给你的芯片面积节省，实际上这种方法反而花费了额外的面积，因为现在你还需要芯片面积来实现组合脉动阵列的额外功能。不过从这种方法中可能仍然有功耗节省。]]]
[[[ 关于批量、token 延迟和 KV cache 的细节（for nerds）：更大的脉动阵列也可能导致需要更大的批量维度，这与本节指出的其他问题是分开的。更大的批量降低内存带宽需求但导致增加的 token 延迟和 KV cache 存储。详见关于非方形脉动阵列的小节。]]]
数值格式，或如何制造便宜的脉动阵列#
并非所有脉动阵列都是平等的。我们已经看到更大的脉动阵列更高效。还有什么重要的？数值格式（numerics）。数值格式很重要。但那意味着什么？
矩阵乘法由点积组成。点积是数的乘积之和，例如向量 (a, b, c) 和 (x, y, z) 的点积是 ax + by + c*z。所以点积涉及 N 次乘法和 N - 1 次加法，但是什么的加法和乘法？嗯，是数的，但在计算机上，一个数不只是一个数。有不同种类的数。这就是"数值格式"所指的。那数的种类是什么？
**FP32：**FP32 就是 C++ 中所谓的"float”。它是一个 32 位浮点数。最初的突破性深度学习网络 AlexNet 对加法和乘法都使用 FP32，使用的是 Nvidia GPU。然而，FP32 的加法，尤其是乘法，在芯片面积和功率方面实现起来非常昂贵。
**FP16：**FP16 是一个 16 位浮点数。这是 Nvidia GPU 上的一个创新，其中乘法用 16 位完成，而加法仍然用 32 位完成（加法需要比乘法更高的精度）。事实证明你可以对乘法使用较低的精度而不会降低 AI 的质量，所以这很有效。使用 FP16 的脉动阵列所需的面积和功率不到 FP32 的一半。此外，如果你使用 FP16，你只需要一半的内存带宽和内存容量。所以从 FP32 到 FP16 是一件大事。
FP16 的一个缺点是它的动态范围比 FP32 低得多，即它无法表示那么大范围的数字。这意味着在用 FP16 训练 AI 模型时，有必要使用损失缩放（loss scaling）。损失缩放完全不难做，但它确实需要稍微修改模型，所以因此 FP16 对于训练来说不是 FP32 的直接替代品。对于推理来说它是直接替代品，因为推理不需要损失缩放（AI 推理期间不计算损失）。但这发生在 AI 研究很多但 AI 使用很少的时期，所以当时推理不那么重要。
**BF16：**BF16 也是一个 16 位浮点数，但它以较低精度为代价保留了 FP32 的动态范围。事实证明 BF16 对 AI 的效果和 FP16 一样好，但它不需要损失缩放，所以它是 FP16 的直接替代品。此外，BF16 在硬件中实现所需的功率和面积比 FP16 更少。
BF16 是由 Google 在 TPUv2 中引入的，这在当时相对于使用 FP16 的 Nvidia GPU 来说有点抢先。Nvidia 后来在他们的 GPU 上添加了对 BF16 的支持，由于 BF16 的优势，似乎缩小了差距。但差距真的缩小了吗？不完全是，因为 Nvidia GPU 仍然支持 FP16。这意味着 Nvidia GPU 要为实现 FP16 支付更高的芯片面积代价，而 Google TPU 至今不需要这样做，因为它们不支持也不需要支持 FP16。
**FP8：**FP8 指的是两种不同的 8 位浮点格式，即 E4M3 和 E5M2。FP8 在训练或推理期间都不是 FP32 的直接替代品，尽管它通常可以在推理期间相当容易地使用。FP8 在脉动阵列中实现乘法比 BF16 便宜得多。所以 Nvidia GPU 以及现在的 Google TPU 在使用 FP8 时实现了脉动阵列吞吐量的 2 倍提升，尽管两者仍然支持 BF16，后者以半速运行。此外，FP8 当然只占 FP16 一半的位数，所以你在带宽上也获得了 2 倍的收益（或与 FP32 相比为 4 倍）。
BF16 标志着 Google 在数值格式上相对于 Nvidia 的创新优势期，而 FP8 标志着这一时期的结束。FP8 是 Nvidia、Arm 和 Intel 的联合工作发明的，Nvidia 在 Google 之前数年就在他们的 GPU 中引入了它。Google 确实在稍早时候对 TPU 有了 int8 支持，但仍然落后。
**FP4：**FP4 是，你此时可能已经猜到了，一个 4 位浮点格式。Nvidia GPU 支持它，可能下一版 Google TPU 也支持。FP4 脉动阵列比 FP8 脉动阵列实现起来更便宜，尽管由于加法仍然需要在更高精度下完成，而且脉动阵列确实有其他开销，使用更少位数的芯片面积和功率优势越来越少——这种收益有递减回报。不过，Nvidia GPU 为 FP4 提供的脉动阵列吞吐量是 FP8 的两倍。FP4 的另一个大优势，不受递减回报影响的，是有效内存带宽的又一次翻倍，因为传输的位数是 FP8 的一半。并非所有 AI 推理都能使用 FP4，但有些可以，在这种情况下它可以是一个重大优化。FP4 用于权重是有挑战性的。FP4 用于激活是非常有挑战性的。
[[[ FP4 的整数求和（for nerds）：乘法通常比求和昂贵得多。然而，随着我们在脉动阵列中对乘法使用越来越少的位数，求和开始成为越来越大比例的开销，因为求和位数不像乘法位数那样减少得多。所以出于这个原因，你可能会遇到像 FP4 乘积被求和到 FP32 中的情况，就像在 Nvidia GPU 上做的那样——这时求和可能比乘积更贵。整数类型的一个好处是整数求和特别便宜，而且你不需要像浮点求和那样多的位数来做整数求和。两个 FP4 的乘积可以用 int8 使用定点表示。所以用整数值而不是浮点来求和 FP4 乘积实际上可能更便宜也更精确。不过这不适用于 MXFP4，因为它包含额外的乘数使 FP32 求和更合适。]]]
**Int8：**Google TPU 和 Nvidia GPU 也提供 int8 支持，这是一个 8 位有符号整数。浮点有一些整数可以避免的复杂性。Int8 适用于所有合理的推理需求，尽管 FP8 也是如此，所以如果你已经为 FP8 支付了硬件代价，不清楚你为什么还需要 Int8。尽管如此，这种支持是存在的。
**MXFP8、MXFP6、MXFP4、NVFP4：**这些是 Nvidia GPU 支持但 Google TPU 不支持的额外格式。Nvidia 在他们的脉动阵列中支持许多不同的格式有点疯狂了。这些是复杂的格式，在一定程度上改善了数值精度，使更多推理工作负载可以使用较低精度。每个数字有适度的位成本，例如 MXFP4 每个数字使用 4.25 位而不是 FP4 使用的 4 位。MXFP4 和 NVFP4 是为那些几乎在 4 位下工作但不完全行的模型——它们可能用这些格式就行。MXFP6 是如果那不行你可以尝试的——也许 6 位就够了。MXFP8 是下一个 nerd 框的主题。Nvidia 从支持这些格式中获得一些客户锁定的好处，因为这些格式在 Nvidia GPU 之外没有那么广泛的支持。
[[[ MXFP8 的奇特案例（for nerds）：MXFP8 是一种旨在改进 FP8 的复杂格式。然而，如果部署（量化）做得正确，你永远不需要高于 8 位的精度来做推理（累加器除外，但那不是我们这里讨论的）。事实上，8 位已经超过你推理所需要的了。那 MXFP8 为什么存在？有两个原因。高质量量化需要一些努力，所以人们有时使用可能需要 MXFP8 才能工作的较差方法。此外，人们有时会尝试 8 位训练。你可以用 8 位做训练，但这不容易，MXFP8 对此有一点帮助。]]]
所以 Google TPU 和 Nvidia GPU 都在它们的脉动阵列中支持许多不同的格式。这有灵活性的好处，但也导致脉动阵列内部硬件复杂性的显著增加。这在一定程度上削弱了脉动阵列作为简单高效硬件结构的观点。并非所有脉动阵列都是平等的。Nvidia 在数值格式创新方面领先于 Google，他们从中获得了巨大优势。然而，Nvidia 现在支持如此多不同的格式，他们肯定要为此付出 Google 不需要那么多付出的巨大代价。即使 Google 也在为 TPU 支持的许多格式付出可能显著的代价。
TPUv1 是一个纯 Int8 脉动阵列（用于乘法）。它在 Google 内部的采用有限，部分原因是软件难以使用，部分原因是当时的 AI 从业者通常不知道如何准备 AI 模型以在 8 位下运行。那是很久以前的事了，现在这是一件标准且容易的事情。也许市场现在已经准备好重新审视 TPUv1 的这种方法了，这将带来比不必要地灵活的 Google TPU 和 Nvidia GPU 更便宜的推理硬件——在大型脉动阵列内部支持的数值格式方面不必要地灵活。
为什么不做一个只做 FP4 的芯片？因为 FP4 并不总是脉动阵列中所有 AI 需求的足够精度。有时是的，但不总是。未来也许改进的 AI 方法将实现纯 FP4 的方式，但我们目前还没到那一步（也许永远不会）。所以目前对于严肃的 AI 芯片你必须支持超过 4 位。
FP8 如果做得对总是足够好的，所以你可以做一个只有那个的芯片。然而，Int8 更简单也够好，所以纯 Int8 芯片在我看来是一个值得构建的真正竞争者。这里的问题是说服你的客户。他们总是可以使用 8 位，但他们可能不知道怎么做，尽管现在这相当容易。
最后我要指出，脉动阵列的数值格式不需要与标量和向量数值格式相同。脉动阵列是正确的 AI 芯片上最大、最昂贵的结构。所以你希望它们使用便宜的数值格式。你芯片上的向量和标量结构没有那么大，所以你可能支持例如 FP32 向量和标量数值格式，即使在脉动阵列只支持 Int8 的芯片上。有一些 AI op 如归一化和 softmax 不是矩阵乘法（因此不会在脉动阵列上运行），它们可能更容易以这种方式实现。你其实任何东西都不真正需要 FP32，我相信，但那是最容易的。
脉动阵列的结构化稀疏性#
结构化稀疏性（structured sparsity）是 Nvidia 在脉动阵列上领先于 Google 创新的另一个领域。
稀疏矩阵是其中许多数字为零的矩阵。稀疏表示通常是只指示非零条目的位置然后说明这些数字是什么的表示。这种表示不适合脉动阵列，因为脉动阵列是一个需要规则数据模式的大型规则硬件模式。稀疏表示可能是高度不规则的。它没有结构。所以它不工作。
然而，结构化稀疏性对脉动阵列确实工作得很好。在结构化 A-of-B 稀疏性中，写作 A:B 稀疏性，每 B 个条目中最多可以有 A 个非零条目。所以在 2:4 稀疏性中，前 4 个条目中可以有 2 个非零条目，接下来 4 个条目中又有 2 个非零条目，以此类推。所以 (2, 0, 0, 3, 4, 0, 1, 0) 是 2:4 稀疏的，而 (2, 4, 0, 3, 0, 0, 1, 0) 不是。事实证明这种稀疏性非常适合脉动阵列。自 Ampere 以来的 Nvidia GPU 支持 2:4 稀疏性。这种技术可以在仅需要非常小的芯片面积和功率增加（以及每操作功率的大幅减少）的情况下使脉动阵列吞吐量翻倍。2:4 的结构化稀疏性也降低了内存带宽和存储需求，因为零条目不需要存储和传输。带宽和存储收益接近 2 倍，但不完全等于 2 倍，因为有必要存储非零条目的位置，每 4 个条目需要 3 位。内存带宽的减少通常也只适用于权重，不适用于激活（即 AI 中不是权重的向量），因为你通常不应该稀疏化激活，只稀疏化权重。截至 2025 年，Google TPU 不支持这种稀疏性。我不知道为什么不支持。
Nvidia 建议使用 2:4 稀疏性的方式是先密集地（即非稀疏地）训练模型，然后再将其稀疏化以符合 2:4 稀疏模式。他们报告说这导致 AI 精度没有或几乎没有降低，即 AI 不会因此变笨。
结构化稀疏性总是好主意吗？如果它导致精度降低（即 AI 变笨），即使降低只是轻微的呢？也许这种权衡有时不值得？我要声称如果你使用 Nvidia GPU，你通常应该总是使用结构化稀疏性。如果可能的轻微精度降低让你担心，那就训练或下载一个大 50% 的模型，然后稀疏化那个。结果将同时增加速度（更快的 AI）和精度（更聪明的 AI）。我不知道什么情况下你不应该在 Nvidia GPU 上使用结构化稀疏性，除非你无法修改和重新训练模型（这是稀疏化所需的）。如果你向他人提供模型，考虑给他们提供稀疏化模型的选项，这样他们就不会有这个问题。
看起来 Nvidia 在这里相对 Google TPU 获得了 2 倍的优势，而 Google 到目前为止不打算在 TPU 中缩小结构化稀疏性的差距。并非所有脉动阵列都是平等的。Google 最初在脉动阵列创新上领先于 Nvidia，将大型脉动阵列和 BF16 远早于 Nvidia 推向市场。结构化稀疏性是另一个例子，表明 Google 在脉动阵列创新上领先于 Nvidia 的时期已经过去了。Google 在 AI 的其他领域强力创新，但不在脉动阵列上。与此同时，Nvidia 在这个话题上全力以赴。
2:4 稀疏性是最佳方法吗？我不确定。如果我今天做一个芯片，我会大力研究 1:2 稀疏性与 7 位整数运算结合的可能性。这允许一种特别简单和吸引人的数据格式：一种 8 位格式，其中第一位指示非零条目的位置（在接下来的两个条目中），剩余 7 位是该整数条目的位。所以你会有一个 7 位乘法器脉动阵列，原生支持两种格式：7 位稀疏或密集整数（定点和整数是等价的）。密集格式只需将高位第 8 位设为零（这有助于压缩性）。为了使格式适合单个向量宽度（而不是半个，这对 SIMD 来说很麻烦），你可以将两个相邻的行打包到一行中。如果你的向量寄存器是 32 位的，你可以这样将 8 行打包到一个 32 位向量寄存器中。如果客户要求矩阵的行数不是 8 的倍数，我建议用 defenestration（把他们扔出窗外）作为解决方法。
我们通常谈论 4 位和 8 位数值，其中 8 位对推理来说总是足够的而不损失模型精度（如果量化做得对），而 4 位勉强可用但有时够好。我们关注 4 或 8 位是因为它们是 2 的幂，而我们行业习惯于偏好这个（有充分理由？我不太确定）。事实是 5 位通常就够了，比 4 位好用得多。所以很不幸 5 不是 2 的幂。有了 7 位，你肯定够了。由于我们其实不需要第 8 位，这允许将那一位专用于稀疏性而不是运算，所以我在这里提议的格式使用 7 位运算，不过如果你愿意，你可以考虑 5 或 6 位。
你可能需要一些仔细的 AI 工作来将 1:2 稀疏化而不损失太多模型精度。你需要在这方面有一些优秀的 AI 人才。如果这奏效了，我认为这种简单的 Int7+1 稀疏格式可能会成为 AI 推理数值格式多年来的终态，如果与后面章节描述的有损和无损压缩想法结合的话。我确实建议支持没有稀疏性的密集操作，但你可以让它以半速运行，就像 Nvidia 对他们的结构化稀疏性支持所做的那样。
你可能选择为密集格式支持 int8。严格来说这不是必需的，因为如果量化做得好 int7 可以做 int8 能做的几乎任何模型，但许多模型被量化到 int8 而不是 int7，因为 int7 不是一个流行的数据类型，所以如果你支持 int8 你的客户将更容易找到现成的量化模型。
请注意，虽然稀疏性对权重效果很好，但对激活不那么有吸引力。这意味着 2 倍的计算收益不适用于注意力中将两个激活矩阵相乘的矩阵乘法。它也不减少 KV cache 存储和带宽需求——除非你确实选择对激活应用稀疏性。
[[[ 1:2 稀疏性的想法：实现 1:2 稀疏性的简单方法是配对相邻列，并将每对条目中较小的条目设为 0。你可以用同样的方式实现 2:4 稀疏性，Nvidia 推荐这样做，然后进行蒸馏训练来恢复精度。Nvidia 还建议将量化与这种变换结合，同时做两件事。肯定有许多改进可以做到这里，也许可以允许使用 1:2 稀疏性而不损失精度或损失很小。例如，你可能可以通过建立列的匹配问题来获得更好的初始误差，使共同误差最小的列被匹配，而不是随机做。另一种你可以尝试的方法是对每对中较小的条目施加逐渐增强的 L1 损失，而不是直接设为零。过一段时间后，收集最好的，比如最小 50% 的列对，它们现在有最小误差，保留那些列匹配，并重新匹配其他列，因为我们现在有证据表明那些不是那么好的匹配（训练过程反推得太强了），然后重试。你可能这样最终有一个困难的残余，几列就是找不到好的匹配。你可能如果必要的话，将那些最后几列困难的列与一个 0 列匹配，这实际上通过为这些列禁用优化来完全保留最后几个困难的列。还有一个选择是对蒸馏模型（训练成本不那么高）使用重复的彩票假设训练——这应该是最强大的技术。你还可以注意到权重中的误差总和或误差平方和不完全是理想的匹配/误差度量，而是你应该测量在一组代表性输入向量上点积值的差异。肯定还有许多这样的想法可以用来推进结构化稀疏性，我怀疑（但不知道），一直到使 1:2 稀疏性成为总是可以使用的技术。如果这里提议的研究项目失败了，这是可能的，可以改为采用 Nvidia（和 AMD）的方法，使用 2:4 稀疏性，这已被证明可以工作，但这种格式更复杂，在权重存储方面不那么位效率，对 SIMD 不那么友好，而且在硬件中需要更多复杂性。注意 1:2 可以无损转换为 2:4，但反过来不行。另一个你可以考虑的选择是 Int6+2 稀疏 2:4。这在硬件上可能与 Int7+1 稀疏 1:2 相似或稍便宜。这一切都取决于通过出色的量化和训练配方证明这些方法总是有效。如果你认为你的大多数客户太简单了，即使给了简单的脚本也不会遵循配方，那么这个特性就不那么有吸引力了。如果你想帮助人类的话，那就要靠你来想出一种情况，你能让客户轻松做到这一点——轻松到他们会做并且做得好。一种方法是自己提供量化和稀疏化的模型，或成为自己的客户。]]]
[[[ 关于 TPU SparseCore 的细节（for nerds）：你可能认为 Google TPU 支持稀疏性，因为它们包含所谓的 SparseCore。然而，SparseCore 与脉动阵列中的结构化稀疏性无关。SparseCore 不包含任何脉动阵列，也不做矩阵乘法。它在以更低粒度访问 TPU 上昂贵的 HBM 内存的意义上是"稀疏"的，所以它可以比主 TPU 核心更高效地以更小的块检索数据，使它在这方面更类似于传统 CPU。它应该用于嵌入表查找，而不是矩阵乘法。TPU 也可以做块稀疏矩阵乘法，但那远不如结构化稀疏性好（它不是一种流行的技术），块稀疏不是硬件特性——它是你在软件中实现的东西。XLA:GPU 可以做结构化 2:4 稀疏性，但这在 XLA:TPU 上不行。TPU 确实缺少一个硬件特性。
如果你在考虑在你的 AI 芯片上复制 SparseCore，那我建议在那里放一个普通的 CPU 核心。那将起到"SparseCore"的作用，而且更通用。]]]
[[[ 训练的结构化稀疏性（for nerds）：完全可以在训练中也使用结构化稀疏性，但训练的反向传播需要转置矩阵使情况变得复杂。推理不涉及转置矩阵，但训练涉及。脉动阵列需要一些额外的针对训练的特殊功能才能在训练期间充分利用结构化稀疏性。此外，有必要选择一个即使在转置后仍保持 A:B 结构化稀疏的稀疏模式——不是大问题，但必须做。Nvidia 当前的结构化稀疏硬件支持是针对推理的，不是训练的，所以它不是目前用于推理的技术。随着硬件的一些改进，我相信每个人都应该在训练中也使用它，我预测这最终会发生，但我不预测"最终"会很快。人们仍在习惯推理的结构化稀疏性。]]]
单一大小的脉动阵列是不平衡的#
回顾一下，transformer 有两种矩阵乘法：注意力和 FF。这两者对 transformer 模型的性能都很重要。设计 AI 硬件的一个复杂因素是注意力内部的矩阵乘法很小（K 维度小），目前通常在 16 到 128 个元素的量级。FF 内部的矩阵乘法大得多，由 d_model 决定，对于大模型可以例如是 8,192 个元素。
Google TPU 有大型脉动阵列，现在是 256 × 256。这意味着它们对 FF 矩阵乘法非常高效，但在注意力上利用率较低，因为注意力不太适合像 256 这样大的向量大小。Nvidia GPU 有较小的脉动阵列（Blackwell 中现在是 64 × 64，是 256 × 256 的 1/16），这意味着它们在注意力上有可能获得更高的利用率（取决于内存带宽，这也是注意力的一个重要因素），但对 FF 来说硬件效率低下，因为与较大的脉动阵列相比，小脉动阵列本质上是低效的。所以 Google TPU 和 Nvidia GPU 在这里都有问题。
[[[ 技术性 AI 评论（for nerds）：你可以做使用更高向量宽度的注意力，从而提高 TPU 上的利用率。这很容易，假设你在训练自己的模型——只需在模型定义中为注意力向量宽度放入一个更大的数字。然而，注意力本质上在向量不如 FF 最优的较大时最具成本效益。所以通过增加向量宽度来提高利用率仍然是低效的，只是以不同于显示低利用率的方式低效。感兴趣的读者可以阅读关于分组查询注意力（Grouped-Query Attention, GQA）等想法，它通过使更大的向量宽度成为最优来减少这个问题。虽然它减少了问题，但并没有完全消除它。更多内容在后面的章节中。]]]
作为芯片设计师，你应该怎么做？FF 使用宽向量宽度效果最好。注意力使用较小的向量宽度效果最好。我提出的解决方案基于注意到 Google TPU 和 Nvidia GPU 都包含多个独立核心。给定 GPU 上的所有核心有相同的向量宽度，给定 TPU 上的所有核心也有相同的向量宽度。如果你拿半个 Nvidia GPU 来做注意力和半个 Google TPU 来做 FF，并将它们组合在一个芯片上呢？那硬件总是在做它最优化的事情。Google 和 Nvidia 不太可能同意生产这样一个芯片，但我们可以在不混合不同公司产品的情况下实现这个解决方案。我的建议是有一个带有巨大脉动阵列的 FF 核心，用于 FF，以及几个较小的注意力核心，每个有一个较小的脉动阵列，用于注意力。都在同一个芯片上。这样注意力可以有高利用率（像在 Nvidia GPU 上），FF 可以在硬件中高效实现（像在 TPU 上）。
这个解决方案解决了注意力和 FF 之间最优向量大小的不匹配。这个想法有一些挑战，但我会论证它们不是严重的问题。
**问题：又多了一件事要理解。**这是为你的芯片编写软件的人需要理解和设计的又一件事。这是一个真正的缺点。然而，至少对于 Google TPU，编程这些东西的人不是天真的用户，这是他们必须管理的最不复杂的事情。对于 Nvidia，这可能是像 Blackwell 这样的昂贵产品的技术，不一定适合预算型 GPU。
**问题：找到正确的平衡。**首先，FF 核心可以做注意力，注意力核心也可以做 FF，只是效率没那么高（否则 Nvidia GPU 根本不能做 FF，TPU 也不能做注意力——但它们当然可以）。所以当你选择在芯片上为 FF 和注意力核心分配面积时，你是在选择 AI 从业者在使用你的硬件时应该瞄准的 FF 操作与注意力操作的比例，以保持所有核心做它们最擅长的事情。所以 AI 从业者现在必须配合你作为芯片设计师选择的注意力和 FF 之间的比例，否则就会降低效率。无论你选择什么比例，在某些场景下都会不是最优的。这是一个问题，是的，但即使你选择了一个糟糕的比例它仍然是一个改进：目前，使用这里使用的术语，TPU 是 100% FF 核心，GPU 是 100% 注意力核心。100% 偏向任一边是你能做到的最不平衡的。所以，如果你明智地选择任何不是 100% 偏向任一边的比例，无论场景如何这都将是一个改进（没有 AI transformer 只做注意力或只做 FF——两者都是繁重的计算且都需要）。所以如果你被选择平衡的想法吓到了，就按 OP 做 50%/50%（你不能只是半数注意力核心半数 FF 核心，因为一个注意力核心会比一个 FF 核心小得多，那就不平衡了）。它会比 Nvidia GPU 或 Google TPU 上的更平衡，因为它们由于使用单一大小的 MMU 而完全不平衡。
[[[ 非常技术性的实现细节（for nerds）：如果你的实现中同一个脉动阵列先做注意力再做 FF，这里有个问题。为了让注意力核心能只做注意力，FF 核心能只做 FF（这是理想的），你需要对两者进行流水线化。如果你之前没有这样做，现在你可能需要流水线中两倍的阶段来划分工作，增加 KV cache 内存需求和增加延迟。最简单的解决方案是仍然使用所有核心来做两种用途，即使不是最优的，完全解决这个问题，同时仍然提供比单一大小 MMU 架构更高的吞吐量。然而，真正的解决方案，也能将你的延迟减半（而不是增加），是修改模型使注意力和 FF 并行运行。请注意，重叠 FF 和注意力也有好处，因为注意力比 FF 需要更多的内存带宽。所以如果你不重叠它们，你可能在注意力期间被内存带宽限制，然后在 FF 期间有闲置的内存带宽被浪费。净结果是如果你不重叠这两者你需要更多的总内存带宽，导致更昂贵的芯片。这是另一个例子，说明由于软件实现差，今天的 AI 芯片被过度配置了。]]]
**问题：创建软件和硬件。**乍看起来如果你的系统上有两种不同的核心，那你需要在软件和硬件中都实现一切两次：一次用于小向量宽度（注意力核心），一次用于大向量宽度（FF 核心）。这里要注意的重要事情是核心只在向量宽度上不同，其他什么都一样。所以，如果在你的软件和硬件中，你从不提到向量宽度，除了一个名为 VECTOR_WIDTH 的常量，你可以对两种核心使用完全相同的代码。你甚至可以运行相同的二进制文件，未修改的，因为 ISA 不需要根据向量宽度而改变。
如果你使缓存大小与向量宽度成正比（你可能应该这样做），甚至缓存层次优化也可以在不同向量宽度的核心之间保持不变，因为你能在缓存中放入的向量数量是相同的。
除了能在同一芯片上有不同向量宽度之外，还有其他原因避免在代码中依赖向量宽度（在硬件和软件中都是）：
**理由 1：向量宽度已经在同一硬件的代际之间变化。**Nvidia 和 Google 都在它们设备的代际之间改变了向量宽度。Nvidia 已经做了好几次。所以对向量宽度有依赖是不好的业务。它会导致代际之间不必要的软件和硬件（RTL）重写。
**理由 2：不这样做会大大恶化你的软件测试情况。**当你为你的 AI 芯片测试 AI 软件时，你需要为每个测试提供输入和预期输出，以查看软件和硬件是否正确计算。你当然需要测试有多个 tile 的情况，即大于脉动阵列大小的输入。所以如果你的 tile 大小是 256 × 256，你必须编写包含超过 256 * 256 = 65536 个数字的测试。你可能需要数百甚至数千个测试（甚至更多——TPUv2 测试系统按 flops 计算当时会是世界上第五大的计算机，尽管这主要是因为 TPUv2 相对其他计算机的速度有多快）。理想情况下你必须手动理解和验证每个测试中的每个数字。所有 65536 个。每个测试。那很糟糕。如果你的软件参数化了向量宽度，你可以为 2x2 的 tile 编写测试，并且可以确信它也会为 256 × 256 的 tile 工作（也许在那个大小有几个测试）。这听起来可能像一个细节，但实际上这是一件大事。行业中有多少人今天使用这个想法？据我所知，没有人。这是一个好主意。即使你正在开发的硬件有单一大小的脉动阵列，你仍然应该编写你的软件测试就好像向量宽度是可变的（你需要一个支持改变向量宽度的模拟器）。
你可能认为测试听起来不够重要到能部分驱动这样的决策。我建议你可能低估了测试对 AI 软件的重要性。
**理由 3：不这样做也会恶化你的硬件测试情况。**硬件测试在这里有同样的好处。此外，硬件设计通常在现场可编程门阵列（FPGA——如果你不知道这是什么就去查）上测试。许多硬件设计无法完整放入 FPGA，这在 FPGA 上测试设计时导致许多麻烦。如果你绝对聪明到将你的硬件设计也参数化了向量宽度，你可以 FPGA 测试一个 2x2 的脉动阵列和 2 宽的向量宽度，然后也许它能放入你的 FPGA。所以就这么做。即使你的最终产品不支持改变向量宽度，开发你的芯片时就好像它会支持仍然对你有利。所以不，你不需要实现两个不同的芯片，只需在你的硬件设计中也参数化向量宽度。芯片上的给定核心必须有特定的向量宽度，当然，但你的 RTL/Verilog/whatever 定义的核心设计可以将其保留为符号常量。测试和调试 2 宽的芯片比调试 256 宽的芯片容易。在软件中是这样，在硬件中也是这样。
总结是在其核心之间有不同向量宽度的芯片最终迫使软件和硬件团队做一些他们本来就应该做的事情。如果你在创建 AI 芯片，这个想法不是竞争力的必需品，但我认为值得仔细考虑。
更大脉动阵列的芯片更容易制造#
更大的脉动阵列对 AI 芯片产品的竞争格局有影响。脉动阵列可以越大，实现有竞争力的 AI 芯片硬件就越容易。为什么？因为 AI 芯片上不断占据更多面积的部分是脉动阵列和缓存。其余部分在面积百分比上正在缩小。或者应该是，如果你做得对的话。
脉动阵列和缓存与例如优化的传统 CPU 的惊人复杂性相比并不难制造。脉动阵列周围的一切才变得复杂，因为脉动阵列周围的部分类似于传统 CPU——或者实际上就是传统 CPU，我认为这是理想的。
如果传统 CPU 是硬件工程师制造过的最复杂的东西（这是事实），那我为什么说制造传统 CPU 比制造某种定制且不灵活的、不太像 CPU 的 AI 芯片更容易呢？传统 CPU 在硬件复杂性方面的优势在于，虽然它们很复杂，但所有受过教育的硬件工程师都已经知道如何制造它们。硬件工程学位几乎就是多年来关于如何制造传统 CPU 的教育。他们都知道怎么做。所以如果你要在硬件中制造复杂的东西，如果你制造的是传统 CPU，你就已经遥遥领先了。CPU 不简单，但即使如此，对于拥有硬件工程学位的人来说它们是容易制造的（在制造 CPU 上超越 Intel 或 AMD 当然不容易，但只是制造一个还行的 CPU 对硬件工程师来说不难）。它们也是对拥有软件工程学位的人来说最容易编程的东西。所以你可以看到这里有好处。
如果脉动阵列非常大，那么你能否非常高效地在面积和功率方面做周围的类 CPU 部分就不那么重要了，因为随着脉动阵列越来越大那就是芯片越来越小的一部分了。因此，创建有竞争力的 AI 硬件所需的最低复杂度随着脉动阵列越来越大而逐渐降低。更大的 AI 硬件复杂性有好处，但随着脉动阵列越大，收获这种好处的需要越来越低。
不管脉动阵列的大小如何，软件都不会更容易做（也不会更难，一旦超过了一定的大小），所以最低所需的硬件复杂度随着脉动阵列大小的增加而降低，但软件复杂度无论如何都是一样的。以高利用率使用 AI 芯片所需的软件复杂度非常高，即使本文档中稍后的想法确实有助于显著降低它。因此，为了节省硬件而创建一个围绕脉动阵列的难以编程的硬件结构是一个糟糕的主意。那在逐渐变小的硬件部分（非脉动阵列和非缓存部分）上节省了硬件成本，同时在已经太高的软件复杂度上付出了代价。不是好的权衡。这个论点对小脉动阵列是错误的，随着脉动阵列越来越大变得越来越正确。
因此我预测 AI 硬件的最终未来就是标准 CPU，像任何标准 CPU 一样完全灵活可编程，驱动非常大的脉动阵列。这里的"标准 CPU"仍然包括高内存带宽、快速网络、高向量计算等。
没有比 CPU 更容易编程的东西了。GPU 通过 CUDA 可编程，但 CUDA 不是对 C++（CUDA 基于它）的改进。CUDA 的惊人之处在于它使 GPU 普遍可编程——这确实令人惊叹。但 GPU 不如 CPU 容易编程，因为 C++ 比 CUDA 容易。CPU 是最可编程的东西，这就是为什么 CPU 是 AI 的最终未来，而不是 GPU——因为 CPU 可以成为比 GPU 更方便和更高效的脉动阵列驱动器。
但是等等，使用运行在 GPU 上的 Nvidia 软件似乎对许多 AI 从业者来说工作得很好？嗯，是的，那是因为 Nvidia 在 CUDA 软件上进行了大量投资，整个 AI 社区也是如此。为 Nvidia GPU 构建易用的软件绝对是可能的——Nvidia 在每一代新 GPU 上都这样做。我说的是为 CPU 构建同样的软件会更容易。这没有发生是因为目前可用的 AI CPU 没有足够的竞争力。
Jensen 是 Nvidia 的 CEO，如果你看他的一些采访，我注意到的一件事是，别人谈论 CUDA 护城河的时候，Jensen 更多地谈论 CUDA 生态系统的好处。我认为 Jensen 关注那里是对的。远不是 CUDA 是一个好处，我会说 Nvidia 陷在了 CUDA 沼泽里，但第四座城堡屹立不倒。不是 CUDA 令人惊叹，而是 Nvidia 和其他人在它之上构建并继续构建的东西。如果不必使用 CUDA 这些会更容易做到，但尽管如此，它是用 CUDA 做的，现在不幸地我们有了一个大的 CUDA 生态系统/沼泽城堡。
为什么目前可用的 AI CPU 不能与 Nvidia GPU 和 Google TPU 竞争？我相信这纯粹是因为 Nvidia 和 Google 在执行他们的 AI 愿景方面比 AMD 和 Intel 做得更好——到目前为止。AI CPU 的潜力很大，但顶级 CPU 公司就是还没有实现这种潜力，这是执行力的问题。
Nvidia 让 GPU 看起来不仅适合而且实际上最适合 AI，因为他们的 GPU 非常出色，他们的执行力非凡，他们的软件帝国庞大，但在最终，AI 的未来将是 CPU，而不是 GPU。Nvidia 不能永远跑赢这个未来，但如果他们保持目前出色的执行力，他们也许能跑赢很长时间。此外，Nvidia 当然有可能创建 AI CPU（Nvidia 目前的 CPU 是用来支持 AI GPU 的，但那些 CPU 本身不是带有大型脉动阵列的 AI 设备）。所以我不是在预测 Nvidia 的衰落，但我确实在预测，最终我们不会用 GPU 来做 AI。
高级章节：非方形脉动阵列#
方形脉动阵列在软件中最容易处理。然而，非方形脉动阵列有一些优势，本节讨论这些。本节更偏数学，使用了许多在文档后面解释的概念，如 decode 工作负载、token 延迟与吞吐量、KV cache、推测解码和批量，所以如果你还不知道这些概念可能需要推迟阅读本节，或者如果你想要轻松阅读可能更愿意不读本节。
假设我们正在做如下的 tile 矩阵乘法：
- NxK：左手边（LHS）激活，乘以
- KxM：右手边（RHS）权重，产生
- NxM：结果
在本节中我们将仔细研究每个维度 N、K 和 M。
N 维度（LHS 行）#
首先要注意的是 LHS NxK tile 的一行产生输出的对应一行。LHS 行之间不交互。所以如果 LHS 行逐一传入，一个 KxM 的脉动阵列实际上在做一系列 1xK 乘 KxM 的矩阵乘法，产生大小为 1xM 的单独行。
所以看起来 N 并不是脉动阵列的真正属性，它只是关于我们向脉动阵列传入多少次行。然而，这不是我建议的看法。问题是需要多少 LHS 行 N 才能从脉动阵列获得 100% 利用率。从某种意义上说这是脉动阵列的真正 N，即使 KxM 的脉动阵列没有物理 N 维度。
什么决定了 N？简而言之，取决于加载 RHS tile/矩阵所需的时间。将 KxM 的 RHS tile 加载到脉动阵列中需要时间，为了 100% 利用率，我们需要有足够的 LHS 行来处理，使得这花费的时间与将下一个 tile 加载到 RHS 中一样长。最明显的脉动阵列实现方式是使 LHS 和 RHS 的带宽完全相同（1:1）。在这种情况下，获得满利用率的最低所需 N 要求 LHS（NxK）和 RHS（KxM）大小相同，这在 N=M 时实现。在这种情况下，128x128（KxM）的脉动阵列将需要 N=128 行才能以 100% 利用率运行。在这种情况下，例如 N=1 只得到 1/128 = 1% 的利用率。所以即使 N 不是脉动阵列的明显物理属性，它也很重要。
[[[ 术语细节（for nerds）：你可以同时向脉动阵列传入多行以提高效率。例如 TPUv3 向量寄存器是 8x128，你一次向脉动阵列传入整个寄存器，导致你可能称之为 8x128 乘 128x128 的脉动阵列，你可以认为这意味着 N=8。那也是定义 N 的一种合理方式，实际上如果你谈论脉动阵列的 N 作为第三维度，这可能是人们假设你意思的，但这与我在这里讨论的不同。]]]
假设我们为脉动阵列提供 RHS 的带宽是 LHS 的两倍。那么我们只需要 N=M/2。我们可以将此推广为 100% 利用率的脉动阵列所需最低 N 的公式：
N = M * bw_ratio = M * (LHS 带宽) / (RHS 带宽)
允许更小的 N 有好处。在 decode 注意力中，N 维度对应查询向量的数量，这可能很低。所以支持更低的 N 并保持高效率可以增加 decode 注意力期间的脉动阵列利用率。
在 decode FF 中，N 维度对应批量乘以推测 token 的数量。所以允许更低的 N 可以让你降低批量，这直接导致更低的 token 延迟（不改变网络需求）和整体 decode 的 KV cache 存储需求——两个都是重要指标。但请注意降低批量完全不会改善吞吐量，所以它不有助于每美元 token 数。
然而，情况并不像增加进入脉动阵列的 RHS 带宽以支持更低 N 那么简单。RHS 权重（或注意力中的 key 激活）通常存储在 HBM 中，然后在 SRAM（缓存）中，然后才加载到脉动阵列中。所以你不仅需要改善进入脉动阵列的带宽，还需要增加 HBM 和 SRAM 带宽以跟上。这是一个昂贵的优化，但在某些情况下可能值得以降低 token 延迟。
为什么不为 N=1 配置带宽？有些芯片确实这样做！它们被称为"batch=1"或"低批量"芯片。但这非常昂贵。此时你拥有的已经不真正是脉动阵列了，它肯定不能捕获你从脉动阵列期望的效率。如果你这样做你需要每操作 M 倍的带宽，其中例如 M 可以是 128。所以你需要配置 128 倍的带宽。昂贵！
[[[ 关于低批量 Groq LPU 的细节（for nerds）：Groq LPU 正是这种低批量芯片，有着微小的 N，也许小到 N=1——我找不到具体数字，他们可能有一个非常小但不到 N=1 的 N。所以 LPU 不使用 HBM 是合理的，因为 HBM 对这种带宽需求极高的极端情况来说太慢了。相反，LPU 将所有权重存储在 SRAM 中，缓存必须提供异常高的带宽。每个 LPU3 芯片有高达 500 MB 的 SRAM。如此大的缓存又必须有比普通 SRAM 缓存更高的带宽，在芯片功率和面积方面都需要大量开销。此外，500 MB 的缓存虽然作为缓存很大，但与模型大小相比仍然很小，所以例如一个 100B FP8 模型至少需要 200 个 LPU3 芯片才能存储权重。
我在其他地方写过 LPU 不使用大型脉动阵列。人们想知道：鉴于我无法接触其芯片内部，我怎么能知道？我这么说是因为如此微小的 N 实际上意味着你拥有的不是真正的脉动阵列，即使你试图让它看起来像一个（你可以的——我不知道他们是否这样做了）。你没有捕获人们期望从大型脉动阵列获得的面积和功率特性。
大型脉动阵列在每美元 token 方面非常经济，而 LPU 低批量方法放弃了这些好处。那为什么 Nvidia 会在这种情况下授权 Groq 技术？因为所有这些开销带来了一个有价值的好处：非常低的 token 延迟。64 的批量意味着 64 倍的 token 延迟，在其他条件相同的情况下（实际上它们不同！）。所以你可以看到低批量如何在这方面有帮助。
Nvidia GPU 使用大型脉动阵列，已经获得了出色的 token 延迟，但它们确实需要一定大小的批量才能在 decode FF 上良好运行。所以 Groq LPU 可以提供比 Nvidia GPU 更低的 token 延迟，但价格更高。所以 Nvidia 使用 Groq 技术来服务昂贵但低延迟 token 的利基市场。Jensen 在 GTC 2026 主题演讲中解释了这一点，LPU 是为了昂贵的（“高端的”）低延迟 token，尽管从事件的报道来看，许多评论者并没有传达他说的这一部分。
N 维度在注意力期间不对应批量，所以你可以在大型脉动阵列上以大 N 做低批量注意力，就像在 Nvidia GPU 上，仍然获得高利用率。所以如果你买 LPU，Nvidia 仍然建议使用 GPU 来做注意力（和预填充），将 LPU 保留给 decode FF。GPU 每 token 更经济，有大 HBM 来存储 KV cache，所以这种混合策略是有道理的。许多评论者确实传达了这种混合策略的想法，但省略了这只是针对高价低延迟 token 的。具有出色但不极端低延迟需求的常规 token 用纯 GPU 方法更经济。换句话说，用大型脉动阵列方法。
Groq 的营销也大力关注他们所谓的静态调度的优点。对他们来说把这当作与 GPU warp 的混乱的对比来引起注意是合理的，尽管这对 Groq 芯片来说不是独特的。TPU 当然也允许仔细定时，甚至在 GPU 上，用特殊技术（通常不使用）你可以大致预测哪个 warp 什么时候做什么。Groq 围绕静态调度的营销按他们陈述的来说是正确的，但他们夸大了影响以及这与其他芯片上发生的事情有多大不同。]]]
K 维度（归约维度，又称"输入特征”）#
这是之前章节关注的维度。在之前的章节中，已经指出注意力有低 K 而 FF 有高 K。所以注意力限制了你在脉动阵列上想要多大的 K。之前也指出你可能更喜欢在芯片上有用于注意力的核心（较低的 K，例如 64）和用于 FF 的核心（较高的 K，例如 256）。在之前的章节中，假设我们讨论的是 N=K=M 的方形脉动阵列。但如果我们允许 K 大于其他维度呢？
更大的 K 对注意力没那么有帮助，但对 FF 有帮助。在 FF 期间，你可以有超过 16000 的特征维度（即 K 维度）。所以这里可以使用 K=1024 甚至更大。注意这个特征维度也用于 FF 的层内并行化跨多个芯片，所以非常大的脉动阵列 K 维度确实减少了你可以提取的最大并行度。例如模型 K=16384 和脉动阵列 K=1024，你最多可以并行 16 路，因为 1024 * 16 = 16384。
脉动阵列的 K 翻倍大致相当于拥有 2 个之前大小的脉动阵列并使用附着在脉动阵列上的通用标量/向量核心的向量计算来将结果相加。具有更大 K 的单个脉动阵列的好处是，这在向量/脉动阵列带宽、控制/标量逻辑和向量核心上的向量计算方面有更少的开销。
请注意方形脉动阵列涉及所有相同向量宽度的向量非常方便。如果你引入 K 不等于 M 的脉动阵列，那么向量核心上的向量宽度将与 K 或 M（或两者）不同。你会因此在硬件和软件中遇到并发症。
参见 Wu et al. 了解更多关于他们的高 K 方法。
M 维度（又称"输出特征"）#
也许将此维度翻倍也有意义？确实如此，在某种程度上，但它大致对应于拥有 2 个共享 LHS 输入的脉动阵列。所以你可能更愿意只有 2 个脉动阵列，尽管共享 LHS 输入确实在片上带宽和控制逻辑上节省了一些，所以有一些好处。我不推荐这个，但如果你没有仔细考虑过这个，你就没有做彻底的架构探索。
如果你想增加 M，你可能更愿意同时增加 K 和向量核心的向量宽度，因为那样所有向量大小相同，这更简单。这只是对应于制造更大的方形脉动阵列，所以这是之前一直假设的标准想法。你可能选择增加 M 和 K 而不增加 N——这只是一个具有更低 N 的方形脉动阵列，我们在之前关于 N 的小节中讨论过。
复制维度#
这通常不被认为是脉动阵列的一个维度，但它有与其他维度类似的效果，所以我为了完整性将其包含在这里。我已经提到了这个想法几次：你也可以通过在芯片上拥有多个脉动阵列来扩展脉动阵列计算。要么在单独的核心上，要么每个核心有多个脉动阵列，或两者兼有。Google TPU 和 Nvidia GPU 都这两种都做过。
一般来说，N 倍复制将所有成本和收益乘以 N。有一些例外，某些东西可以共享。例如，如果你在同一个核心上放置两个脉动阵列，你可能能够重用一些共享的 SRAM（缓存）容量。
拥有许多脉动阵列是正常的，因为脉动阵列的大小受到我们在之前章节中讨论的各种因素的限制，但一个脉动阵列即使是 256 × 256 也远不足以用完现代 AI 芯片上的所有面积。所以人们使用多个。那些在一个芯片上的脉动阵列可以共享一些缓存，并且可以以高于片外合理提供的带宽相互连接。
周期的暴政与脉动阵列 DMA#
使脉动阵列难以编程的事情之一是它们需要驱动它们的核心大量关注。本节讨论这个问题以及在硬件中可以做些什么。
假设你有一个脉动阵列，每个周期接收一个 LHS 向量（行）、一个 RHS 向量（列）并输出一个结果向量（行）。那么控制脉动阵列的核心必须在每个周期处理这 3 个向量来保持脉动阵列忙碌。你可能还需要两个向量操作来从某处读取这两个输入向量，你可能还需要对输出做其他向量操作如 softmax，这可以涉及许多向量操作。所以我们达到了每个周期大约 10 个向量操作，这些向量可能宽达 256 个元素。这是要在硬件中配置的大量向量计算能力，而且在软件中这种设置几乎没有给做其他事情留下空间，你必须设置好以便在每个周期都准备好做你需要做的事情而不会失败。这很有挑战性。你需要一个出色的编译器来让这一切工作。我称这个问题为周期暴政（cycle tyranny）——核心必须在每个周期一直可用以应对脉动阵列。
[[[ 周期暴政灾难（for nerds）：为什么我们生活在周期暴政之下？因为替代方案是脉动阵列闲置。我个人了解一个芯片项目（不是 TPU）因为没有认真对待周期暴政的问题而遭遇灾难——他们制造了一个无论你怎么编程都永远不能以满脉动阵列容量运行的芯片，因为控制核心跟不上暴政。服务脉动阵列的周期暴政程度悄悄地压上来。你真的必须仔细考虑如何处理这个问题。]]]
[[[ 通过展开缓解周期暴政（for nerds）：前面的故事是关于在周期暴政上做得太少。这个是关于做得太多——这好得多！
当为 TPUv2 配置标量计算时，工程师们查看了可能在 TPUv2 上运行的示例内循环（TPUv2 与 TPUv1 完全不同），意识到这些循环需要高速率的标量计算来计算内存地址以跟上脉动阵列。实际上，他们观察到的地址计算量如此之高，以至于他们添加了一个专门的硬件单元来计算循环的地址。这个单元提供了非常高的地址计算能力，但使用起来有些困难。所以我们从未使用过它。
循环展开（loop unrolling）是一种古老而强大的编译器优化。事实证明，至少稍微展开循环使生活在周期暴政下容易得多，所以对于 TPU 循环总是至少稍微展开一些。任何编译器开发者都能告诉你，展开循环可以大大优化地址计算，因为许多地址变成常量偏移，计算起来很便宜，这就是为什么我们从不需要地址计算单元的原因。用于推导地址计算加速器需求的示例内循环根本没有展开，所以这就是为什么这个单元看起来是必要的，即使它不是。]]]
周期暴政的问题当然不会因为你在每个核心有多个脉动阵列而改善。训练也因为使用 Adam 等优化器更新权重需要额外的向量工作而在这方面更糟。
有多种方法可以减少周期暴政。一种是 2D 向量寄存器，例如 TPUv2 向量寄存器是 8x128。这样你可以提供一个向量寄存器作为脉动阵列的输入，那将提供 8 整行（行），所以在脉动阵列处理完这 8 行之前你不必再担心。这已经将周期暴政减少了 8 倍。你也可以有更多行。但要注意引入 2D tiling 可能导致非常显著的软件和硬件复杂性。如果要处理的数据量不是 8 的倍数（或你的向量寄存器有多少行），你可能特别会遇到问题。
另一种方法是在脉动阵列之前和之后放置硬件队列，这样你可以快速连续提供一堆行，然后让向量核心做一会儿其他事情。这有一些帮助，尽管深队列在硬件中不是免费的。
你也可以以较低的时钟速率运行脉动阵列，使驱动它的标量/向量核心更容易跟上——在硬件和软件中都是如此。这种方法在早期 TPU 中使用过。不过如果你在脉动阵列中使用 8 位或更低的算术，你将需要大量这样的脉动阵列来用计算填满合理大小的芯片。如果你使用慢时钟的脉动阵列，那你就需要更多。这个选择可能限制芯片之间可用的并行度和/或增加 token 延迟。
这里有一个我没见任何地方使用过的不同想法（尽管它开始类似于数据流设计）：脉动阵列 DMA。这是一个从某处向脉动阵列（或从脉动阵列）传送可能数千行/列的 DMA——可能来自另一个脉动阵列（也许带有向量加法）、SRAM、HBM 或通过网络从另一个芯片。这样你大大减少了周期暴政——在数据只是直接从某处传输或传输到某处的情况下，核心不必非常频繁地服务脉动阵列。在推理期间，权重总是只是被传输，KV cache 通常在你读取它时也是如此。除了（解）压缩，但在那个章节中我建议将此作为 DMA 的特性添加。如果以这种方式传输的数据不会被重用，那你也避免了将这些数据读入和写入寄存器或 SRAM 然后再进入（或离开）脉动阵列的带宽开销。这也允许你将数据从你的脉动阵列直接流入（或来自）某种网络操作（如 all-reduce）或流到另一个核心。这是否是个好主意取决于（也取决于）在硬件中实现这个的开销，但对我来说这似乎是一个有前途的想法。
脉动阵列 DMA 的主要问题是你可能需要在脉动阵列之前或之后以某种自定义方式转换数据。在这种情况下，数据必须以通常的方式通过标量/向量核心。如果你在脉动阵列 DMA 可用的论据上减少了标量/向量核心的吞吐量，你可能在这种情况下无法达到 100% 脉动阵列利用率，被标量/向量计算瓶颈限制。所以如果你走这条路，你将需要记住这种场景的可能性。
脉动阵列 DMA 也与我建议的多核心很配，其中一些核心没有附加脉动阵列。这些可以用于其他目的，如输入流水线处理、运行 all-reduce 的网络算法、做检索（用于 RAG 或层内）或运行我也推荐在设备上运行的 OS。所以这些核心可以做各种其他工作，而不浪费脉动阵列的容量。在直接附加到脉动阵列的核心上的向量或标量容量被超过的情况下，你可能能够使用脉动阵列 DMA 将一些多余的工作以对主要脉动阵列标量/向量核心特别轻量的方式卸载到这些额外核心上。
轶事：脉动阵列很热#
本节包含一个关于管理 TPUv3 中意外脉动阵列热量产生的长篇故事。
脉动阵列包含如此令人难以置信的计算密度，以至于芯片中包含脉动阵列的区域可能比芯片的其他区域热得多。这是一个问题，因为它可能损坏（熔化）芯片上的那个点——这就是为什么热量最好在芯片上均匀产生，而不是都在一个点上。有办法解决这个问题，如使脉动阵列以比芯片其余部分更慢的时钟运行更大的脉动阵列、将脉动阵列彼此远离放置在芯片上，或在极端情况下，也许（也许不会）甚至将一个脉动阵列的部分分布在芯片上，在部分之间放置其他（不那么热的）逻辑。这是一个真正的关注点，但可以管理，不应该限制脉动阵列的大小。
在 TPUv3 中，芯片上脉动阵列的热量产生在设计阶段被低估了，这导致当芯片生产出来并到达 Google 总部时出现了一个令人不快的惊喜。这种事情如果你修改芯片并做全新的生产批次可以纠正，但这会延迟 TPUv3，而那是不可接受的。那怎么办？
硬件团队做了大量工作，实施了一长串缓解措施，这意味着我们能够在不重新制作芯片的情况下发布 TPUv3。这个问题的一部分也在软件中得到纠正，使用基于编译器的静态分析技术来避免以会产生危险热量的模式驱动脉动阵列，基于一个经验验证的芯片内部热量产生数学模型。结果是 TPUv3 可以安全地以 25% 提高的时钟速率运行，这得益于这个软件解决方案，显著减少了问题的影响。我个人构想并编写了这个软件解决方案。我是 TPUv3 软件的技术负责人。
在内部，我称这个功能为"编译器节流"（compiler throttling），后来导致我上司的上司礼貌地告知我，我在市场营销方面很糟糕。编译器节流听起来像是让芯片变慢的东西，不是吗？但这个功能的实际结果是产生快达 25% 的 TPUv3，一些重要的工作负载确实实现了完整的 25% 加速。所以如果它让事情更快，为什么不叫它 Turbo Mode？现在那是个好名字。他一告诉我这个，我就知道他是对的，但那时改名已经太晚了。:(
编译器节流，呃，我是说"Turbo Mode"，对我来说是一个有趣的（也非常具有挑战性的）项目，但如果你在生产包含脉动阵列的芯片，确实要仔细考虑脉动阵列的热量产生。你特别可能想要考虑包含一种方式让芯片在硬件中注意到它正在使用太多功率并快速自我节流。TPUv3 确实有硬件节流机制，但它基于事后测量温度，而不是在任何给定时刻有多少电流流入芯片。如果脉动阵列内部可能发生的热量峰值高于预期，那正常情况下这没问题，但这种响应太慢了。换句话说，你可能在更高的温度到达芯片外部并导致节流发生之前就在内部损坏芯片了。Turbo Mode 是一种软件解决方案，确保编译器不能产生可能导致那种结果的指令模式，这使得提高最大时钟速率成为安全的。
你可能会问为什么不直接创建一个可以处理最坏情况而不需要节流的冷却方案。为什么芯片需要节流？只要让它不用节流就能工作。我当时一开始也在想这个问题。原因是最坏情况的数据模式和指令模式产生的热量远远超过平均情况，也远远超过芯片实际使用时会发生的。真正糟糕的数据和指令模式不可能偶然发生，它们必须是故意的才会发生——即有人在搞你的芯片，不是在尝试用它做任何有用的事情。所以你需要节流，因为在不节流的情况下处理最坏情况太昂贵了，不太可能对真正的客户有任何好处。这对所有高性能计算芯片都是如此，不仅仅是 TPU。这就是为什么它们都有一个在高温时降低时钟速率的硬件节流功能。
即使节流在某些情况下只在糟糕的工作负载下才发生，你也不能忽视最坏情况，假设或希望没有人会愚蠢到触发它，因为一个好奇的实验者、你计算机上的病毒或云中的恶意用户可能会在瞬间摧毁你所有昂贵的 AI 硬件。特别是，在云环境中可能发生的是，恶意程序会摧毁它运行在上面的硬件，然后由于数据中心自动检测到硬件故障而在新硬件上重新启动（“哦，抱歉，让我们马上给你一台新机器”）。所以，没有节流的话，恶意用户可以以这种方式摧毁 Google 全球数据中心舰队中可能所有的空闲 TPU 容量。所以 AI 芯片，和一般的芯片，确实需要节流功能。你也许从这个故事中还可以得出结论……Turbo Mode……必须以非常谨慎的方式开发，获得批准发布到 Google 数据中心可能需要付出很大的努力。是的，确实如此。如果你做对了，大赢。如果你做错了，完全不可接受。
为什么 25% 是一个大赢？嗯，如果你在 TPU 上投资了 $X，你从这项技术中免费获得了 25% 更多的 flops，你可以说实现了 $X 的 25% 的经济价值。所以如果 X 是一个大数字，而它确实是，那就是很多钱。远超过我年薪的 100 倍。我认为这是一个相关的计算，因为如果我当时不在 Google，这不会发生。不幸的是，尽管有这个出色的论证，我没有得到基于这个计算的佣金。哦好吧。
你可能听说硅谷高级工程师的薪水很高，确实如此。即使如此，我刚告诉你 Google 本可以多雇用超过 100 个我这样的人来换取他们从我做的这一个项目中得到的东西，而那只是因为我在那里才发生的。有时候就是这样。所以这个故事可能开始解释为什么大公司如此热衷于即使以很高的价格雇用优秀工程师。
激活和权重的压缩#
深度学习革命从每个数字 32 位开始，然后降到 16 位，然后 8 位，现在甚至有时使用 4 位而几乎不损失精度——即不会让 AI 变得太笨。这是对权重和激活的有损压缩（lossy compression），是 AI 中的标准技术。
然而，无损压缩呢？这可能看起来是个奇怪的问题，因为我们已经在使用有损压缩了，无损压缩不是比有损压缩更弱吗？这是一个误解。有损压缩通过去除不必要的精度来工作。无损压缩通过利用数据中的非均匀模式来工作。这是不同的事情。这就是为什么 JPEG 图像格式同时使用有损压缩和无损压缩。特别是，JPEG 在数据经过有损压缩后对其应用 Huffman 编码。为什么我们不对 AI 也这样做？我不知道为什么不。
关于这个话题的调查不多。Nvidia GPU 有纹理（解）压缩的硬件支持，但那是图形方面的，不是 AI 方面的。更新的 Nvidia GPU 中还有更通用的解压缩功能，但它完全跟不上内存速度（~1/10），所以如果你在 HBM 中压缩数据，使用它实际上比不使用获得的带宽少得多——尽管如果你压缩部分数据但不是全部，你可能仍然获得一些提升。AMD 最近宣布了他们未来 GPU 的"通用压缩"，听起来很有趣，但到目前为止没有详细信息，所以目前还不知道它是什么。
我在这一节想提高人们对 Huffman 编码和解码表在有损压缩中的有用性的认识。这样的表比乍看起来更有用。我认为值得考虑在未来的 AI 加速器中以内存和网络速度包含它。
要对例如 8 位整数进行 Huffman 编码，你创建一个有 2^8 = 256 个条目的表。每个条目有一个某种可变长度的位模式。所以你取一个 8 位整数，在表中查找它，输出那么多位的那个位模式。就这么简单。这个想法是更常见的值收到更短的位模式，而罕见的位模式收到更长的位模式，这样净效果是压缩一个向量或矩阵。假设支持的最长位模式是 10 位长。要解压，你也有一个表，现在有 2^10=1024 个条目。表中的值是编码到那个位模式的原始 8 位模式，以及编码长度（这样你知道解码时要向前跳多远）。所以编码和解码由两个互为逆的表组成。由于解码表比编码表大，解码表中不是所有条目都会被使用。
这如何也能用于有损压缩？当你创建编码查找表时，你可以将多个值映射到相同的位模式。在这种情况下，该表在无损 Huffman 编码之上编码了一个有损压缩。所有 FP4 值乘以 2（避免分数）可以编码为 8 位整数，所以确实存在一个编码表将 Int8 编码为 FP4 然后再解码回来（乘以 2 只是创建了一个定点表示）。这不会改变脉动阵列的内部或吞吐量，它将保持在例如 8 位，但这意味着，通过改变编码和解码表，你可以动态地制作像 FP4 或任何定制格式这样的数值格式来压缩，从而节省内存存储和带宽。你还可以在有损压缩之上获得无损压缩的好处。所以对于每一层，你可以运行一个自动化优化来找到该层的有损和无损压缩的最优组合。
这种压缩需要硬件支持才能以内存和网络的速度运行，用于编码和解码。Huffman 编码是串行的，但可以通过在压缩文件格式中拥有多个独立的数据流来并行化。如果这个功能被添加到芯片的 DMA 功能中，而软件使用 DMA 来访问数据，那它对软件来说可以相当透明——只是配置每个 DMA 做（解）压缩的问题。这需要高吞吐量但不一定需要低延迟，鉴于软件是软件流水线化的——它应该总是如此。DMA 本来就可以有相当高的延迟。
我没有经验性地量化这种优化相对于面积和功率成本的影响，所以 AI 芯片初创公司在决定是否在此上花费芯片面积之前肯定应该这样做。据我所知，没有公开研究彻底调查这个问题。
我相信之前章节描述的 Int7+1 结构化稀疏格式，如果与这种有损 + 无损压缩方法结合，可能相当有吸引力。压缩意味着你可以支持 7 位，但如果对给定层来说合适的话，可以无缝地缩小到更小的非整数位宽（例如 4.3 位）。
[[[ 关于 TurboQuant（for nerds）：TurboQuant 是 Google 的 KV cache 压缩技术。它据说在不损失精度的情况下将 KV cache 大小减少 ~8 倍。这篇论文由 Google 博客文章向公众介绍，并在当时被新闻界广泛报道。博客文章和新闻报道都高度误导。该 ~8 倍的基线是每元素 32 位，这不是一个严肃的比较，所以这个数字毫无意义。在这种情况下 8 倍意味着降到 4 位，这以前就做过了，所以那不是新东西。换句话说，新闻报道关注的东西甚至没有改进！实际的 TurboQuant 进步，如果博客文章的其余部分可以信的话，是在不重新训练、不太损失精度甚至不需要校准数据的情况下做到这一点。那确实是一个重大进步。所以这不一定是个坏结果，但真实故事与新闻报道的几乎没有相似之处。我持怀疑态度，但也许 TurboQuant 是你应该压缩 KV cache 的方式。你得自己去发现。如果是的话，你需要确保你的芯片能够以内存速度执行 TurboQuant 中的非平凡开销。]]]
AI 工作负载的种类——训练、解码、预填充#
本章解释 AI 芯片必须支持的三种主要 AI 负载：训练（training）、预填充（prefill）和解码（decode）。这三种在 AI 硬件和软件需求方面有显著不同，制造 AI 芯片的人需要清楚自己的芯片针对的是哪种工作负载。
**训练（Training）：**这是 AI 模型中的权重被更新以使 AI 更聪明的过程。这是 AI 的学校。AI 在这里不做任何有用的事情，就像学生不做任何有用的事情一样，这是关于学习如何做，而不是做。训练在芯片所需的功能方面和实现训练所需的软件方面都是最具挑战性的工作负载。构建一个好的训练芯片及配套软件是最具挑战性的 AI 芯片项目。如果你的硬件和软件做对了，你的芯片在训练期间应该是计算受限的（compute-bound），即瓶颈是脉动阵列，而不是内存带宽或其他任何东西。
**推理（Inference）：**这是 AI 从训练学校毕业并可能开始做有用工作的阶段。推理期间权重不再被更新。在硬件和软件复杂度方面，推理芯片比训练芯片更容易制造。对于 transformer 模型，有两种推理：预填充和解码。
**推理：预填充（Prefill）：**这是一种特定的 transformer 推理，AI 阅读提供给它的文本（或其他数据），但它不对此说任何东西，它只是记录和记住所写的内容。预填充在硬件和软件上比训练更简单。如果你的硬件和软件做对了，你的芯片在预填充期间应该是计算受限的。
**推理：解码（Decode）：**这是另一种 transformer 推理，AI 在创建输出（如文本、音频、图像或视频）的意义上"说话"。解码通常比预填充更复杂，但不如训练复杂。
取决于你如何实现，解码比预填充和训练需要更多的内存带宽。如果你的软件不够优化且芯片上没有太多内存带宽，你在解码上将是内存受限的（memory-bound）。你应该争取在解码上也是计算受限的，但你可能无法做到。即使内存带宽不是问题，解码也会创建不平衡的矩阵乘法，无法在脉动阵列上获得那么高的利用率。因此，目前市场上购买一个解码（输出）token 的价格远高于预填充（输入）token，尽管预填充和解码每个 token 的计算量几乎相同。
你可能已经注意到 AI 助手对你说的话比你对它们说的多得多，所以 AI 助手平均而言在解码上负担很重而在预填充上很轻。对于在回复前自己"思考"的推理型 AI 来说更是如此——生成那些"思考"涉及自言自语，这就是解码。然而，AI 助手正在做越来越多的预填充。这是因为 AI 助手现在正在搜索网络（和其他数据库）并阅读大量信息以确定如何回复，这种方法叫做检索增强生成（Retrieval Augmented Generation, RAG）。由于 RAG 不会消失，未来大多数 token 可能会是预填充 token 而不是解码 token。这对如何构建 AI 芯片有重大影响。这意味着脉动阵列一如既往地重要。
所以 AI 芯片不仅仅是 AI 芯片。训练芯片需要快速的脉动阵列和推理不需要的高级功能。预填充芯片需要快速的脉动阵列。而解码芯片需要较小的脉动阵列和比训练和预填充更高的内存带宽（至少目前如此——我相信稀疏（索引）注意力将大幅降低解码的内存带宽需求）。
你有时会看到专门针对训练或推理的芯片，如 Amazon Inferentia（用于推理）和 Amazon Trainium（用于训练）。对于推理，截至 2026 年的行业标准是在同一芯片上做预填充和解码，芯片不专门针对其中任何一个来营销。这意味着脉动阵列在解码期间通常大量闲置，而内存带宽在预填充期间大量闲置。
训练#
如果你创办 AI 芯片初创公司，你的第一个芯片应该是训练芯片吗？我建议不要。至少一开始，让其他公司占据这个市场。
训练现在看起来可能是一个更大的市场，反映了目前用于 AI 服务的 AI 投资美元（流向训练）多于客户美元（流向推理）。随着该领域的成熟，这显然必须反转。推理会有更多的钱，但也有更多的竞争——其他人也都从推理芯片开始。
训练是一个大市场，因为训练非常大的 LLM 需要大量硬件。许多资金充足的 AI 应用初创公司购买大量训练芯片来开发他们的 AI。他们一开始可能不会买很多推理芯片，因为他们一开始不会有任何客户让他们的 AI 与之交谈。所以 AI 开发的早期阶段涉及购买大量训练芯片而不买推理芯片。训练芯片也可以做推理，只是不那么经济，所以如果你卖训练芯片，你可能能让客户在你的训练芯片上一直待到他们（希望）过渡到有客户因此需要做推理。听起来做训练芯片公司很聪明，对吧？Google TPU 和 Nvidia GPU 都能做训练和推理。所以跟随行业领导者做训练芯片，对吧？
这不是 Google 的做法。TPUv1 是推理芯片，所以 Google 也不是从训练芯片开始的。
问题是训练在软件和硬件上比推理需要显著更多的复杂性。训练芯片每芯片的生产成本也更高，因为它们需要推理不需要的额外硬件功能。如果你是一个可能需要在价格上竞争否则就破产的初创公司，这对你来说比对 Nvidia 和 Google 更重要。
训练芯片需要而推理芯片不需要什么？
**高吞吐量转置：**训练的反向传播需要转置权重，你需要硬件来做这个，否则太慢。推理通常不需要这个。你可以用矩阵乘法来做转置，所以它可以在现有的推理脉动阵列上完成，但这不高效。
看起来推理期间也需要转置硬件，因为注意力公式涉及计算 QK^T，但你可以通过让脉动阵列直接做这个操作而不是 QK 来绕过这个，即总是将行乘以行而不是行乘以列。你仍然可以做 transformer 中的所有其他矩阵乘法，因为那些是激活 * 权重（AW），你可以在推理之前离线转置权重，利用 AW = A*(W^T)^T。
TPU 上有做 AB^T 的先例吗？有的，TPU 上就是这样，脉动阵列可以选择做 AB 或 A*B^T（请读 TPUv2 硬件论文）。
**额外的向量计算：**训练需要支持各种优化器如 Adam。这些优化器涉及大量的向量计算。这是推理之外的额外工作，所以训练比推理需要更高速率的向量处理。
**昂贵的数值格式：**所有推理都可以用 8 位或更少的位数在脉动阵列内部做乘法（加法需要更多但更便宜），这使脉动阵列更便宜。即使如此，如果你的产品也支持 16 位数值格式会更容易卖，即使对于推理也是如此，因为一些推理客户就是不愿意降到 8 位。对于训练，如果你的芯片只有 8 位数值格式，你不太可能把它卖给任何人，所以必须是 16 位（尽管我实际上认为 8 位对训练来说可以也将会足够，但那不是行业标准，你不能期望你的客户目前知道如何处理这个）。
**额外的训练算子和融合：**为 AI 推理芯片编写软件的大部分工作是实现大量用于推理的算子。这已经比大多数供应商想做的更多工作了，但这是必要的。训练使用的算子更多，而且往往比推理的更复杂。
**增加的软件和硬件灵活性：**训练非常大的 LLM 的人，也就是你最想卖训练芯片给的人，是 AI 研究人员。整个团队的 AI 研究人员每天都在想出新想法。他们将不断尝试新想法。这需要你提供给大规模训练客户的软件和硬件有很大的灵活性。推理客户不太可能对你的软件和硬件的灵活性有那么大的需求。此外，一些 AI 研究人员不一定对使用你产品的计算机科学非常感兴趣。它只需要工作就行，即使他们以奇怪的方式使用它。大型推理客户倾向于雇用对计算机科学更感兴趣（而不是 AI 研究）的工程师，他们对不那么容易使用但能为公司在每美元 token 数等方面提供实际好处的软件有更高的容忍度。
你不是不能把训练芯片作为你的第一个芯片。但你想要吗？我建议可能不要。因此，我不会在本文档中过多详细讨论训练。
解码#
解码是模型必须说些什么的时候——写文本、生成图像，任何有所谓自回归输出的东西。解码是更难支持的推理类型。它天然不太适合脉动阵列，但通过足够的 AI 研究和软件的英雄般努力，你可以让它工作——所以这就是你必须要做的。
回顾一下 transformer 包含两种层：注意力和 FF。两者对解码来说都有问题，但两者中解码注意力更成问题。
解码的麻烦在于，如果以直接的方式做，它一次只生成一个输出 token。这类似于你写东西时一次只写一个词。这叫做自回归（autoregression），意思是模型在决定了前一个 token 是什么之前不能开始下一个 token。所以它必须看到（“回归”）自己的（“auto”=自身）输出才能继续。要看到为什么这是个问题，考虑一下自回归解码期间注意力如何工作：
之前的 token: The quick brown fox
↓ ↓ ↓ ↓
keys keys keys keys
↑
新 token 的 query（"jumped"?）每个之前的 token 发布一个 key，这是一个描述那个词是关于什么的向量。这已经在之前做过了，所以我们从所谓的 KV cache（“key-value” cache）中拉取这些 key。当我们需要处理一个新 token（看起来将是"jumped"）时，我们创建一个 query，也是一个向量，与每个 key 进行比较（点积）。这里的问题是这个过程对应一个矩阵乘法，因为我们只处理一个 token，所以只有一个 query，这意味着矩阵乘法有一个维度不仅仅是小，实际上等于 1。这导致脉动阵列几乎没有利用率。所以天真地做解码，脉动阵列利用率非常差，不到 1%。
而且可能有很多 key。假设你已经和你最喜欢的 AI 聊了 2 年。如果你使用的 AI 有所谓的长上下文窗口，而且 AI 在跟踪你的整个聊天历史，它可能必须拉入多达一百万个 key，只是为了将每个与一个 query 向量比较——即使所有那些历史很少是相关的。这将需要大量的内存带宽，但同时从内存读入的数据几乎没什么要做的，整个芯片就坐在那里等内存传递一百万个 key。更糟的是，transformer 有多层，例如 32 层，所以你可怜的 AI 可能必须做所有这些 32 次才能确定下一个词应该是"jumped"。然后整个过程从头开始确定下一个词将是"over"。如此反复，一次一个 token。
等一下，你可能会想。我在某处听说 transformer 使用叫做"batch"的东西，即芯片同时处理多个对话，这应该有助于脉动阵列利用率。不幸的是，batch 对 FF 层很重要，特别是在解码期间，但它对注意力没什么帮助。Batch 可能让你在对话之间重用一些权重，但加载权重不是我们这里看注意力时的问题——加载 key 才是，而不同的对话有不同的 key，所以没有 key 的重用。
也许你听说 transformer 使用多个注意力头。这意味着这整个过程不是每层发生一次，实际上每层最多发生 128 次，每个单独的注意力头一次。单独来看这对低效率的情况完全没有帮助。然而，与一种叫做**分组查询注意力（Grouped Query Attention, GQA）**的技术结合时它可以帮助很多。GQA 的想法是在大小为 128/G 的 G 个组中重用相同的 key 向量（假设有 128 个注意力头）。每个注意力头只有一个 query 向量，但现在我们将头组合成 G 组，每组 128/G 个注意力头，在每组内使用相同的 key，我们突然有了 128/G 个 query 查看相同的 key。这意味着我们关注的矩阵乘法维度从 1 增加到 128/G。这导致脉动阵列利用率大大提高。这也将内存带宽消耗减少了 128/G 倍。此外，由于我们现在重用 key 128/G 次，增加 key（和 query）向量的大小来提高它们的质量变得合理，这对应于增加矩阵乘法的 K 维度，这又是另一个会增加注意力期间脉动阵列利用率的东西。所以 GQA 是一个非常重要的优化。
我们可以做得更好。事实证明我们不真的需要一次只做 1 个 token，这得益于一种叫做**推测解码（speculative decoding）**的技术。这个想法是对接下来几个 token 做一个便宜且有根据的猜测，例如你可能尝试猜测下 3 个 token。在上面的例子中，我们对接下来内容的猜测可能是"jumped over the"。由于模型现在有了词将是什么的猜测，它可以开始处理"the"之后会是什么，同时它也在处理和确保"jumped"、“over"和"the"确实是它想说的。这一切可以在一个步骤中完成，所以现在我们一次处理 4 个输出 token 而不是 1 个。这将 key 内存带宽需求和脉动阵列利用率额外提高了 4 倍。
推测解码为什么正好是一个好处？因为我们现在一次处理 4 个 token，并且使用 G 组每组 128/G 个注意力头（假设总共 128 个注意力头），我们最终有总共 4 * 128/G 个 query 查看相同的 key。所以我们总共改善了 4 * 128/G 倍。如果 G 是 8，那这就是总共 64 倍的改善。
结论是解码的脉动阵列利用率仍然可能不好，内存需求仍然可能很高，但解码并不像乍看起来那么糟糕，因为你可以通过组合推测解码和 GQA 获得例如 64 倍的改善。如果你有 256 × 256 的脉动阵列，理想情况下你会一路到 256 倍，因为我们从 1 开始，但这么大的倍数可能不可用，所以你可以看到更大的脉动阵列在注意力期间更难获得好的利用率，特别是在注意力解码期间。
为什么不选 G=1 来获得超快解码？麻烦在于这种情况下每个 token 只有一组 key，所有注意力头必须共享，这对模型质量不太好。换句话说，它往往使 AI 变得太笨了。所以我们需要足够的 G 来保持 AI 的智能但不要太多以至于解码变得太慢。
为什么不推测解码 16 个 token 而不是 4 个？因为做这些猜测的过程是一个非常笨的 AI，比正式的 AI 便宜和快得多——否则笨 AI 不够快到领先很多 token。它不太可能正确猜出聪明得多的正式 AI 要说的那么多 token。笨 AI 会在我们到达 16 个 token 之前很久就说出"pineapple”（或其他错误的东西），然后我们为那些 token 做的计算就浪费了。所以再次需要取得平衡。
[[[ 关于解码脉动阵列的细节（for nerds）：还有第三种优化可能让你获得最后的 4 倍以达到 256 倍，因此甚至 256 × 256 脉动阵列也可能达到 100% 利用率。你可以通过降低所需的 N 维度来做到——参见关于非方形脉动阵列的章节。这只在你没有使用更小的注意力脉动阵列且也不是内存受限时才相关。]]]
看起来通过努力，我们可能可以在解码期间也获得大型脉动阵列的高利用率。许多这些技术也有助于内存带宽需求。一般来说，任何降低 KV cache 存储需求的东西都对解码内存需求有类似的效果。该领域的普遍看法是解码是高度内存带宽受限的，即使你配置了大量内存带宽，但事实是这可能是也可能不是这样，取决于你执行所有这些优化的好坏。例如稀疏/索引注意力可以完全改变这个数学——以至于对此的研究可能允许甚至 SSD 带宽对解码来说也足够（!!!）。麻烦在于如果你不控制客户的软件和模型，他们可能没有在这里做到最优。所以你将需要难以置信的内存带宽量来让他们满意，这就是 Nvidia GPU 和 Google TPU 拥有的以及这样做的原因之一。
你应该注意到这里的一个趋势——关于如何配置 AI 芯片的几乎一切都取决于软件，根据软件和 AI 模型本身被设置为高效运行的好坏，你最终可能得到相差 10 倍或 100 倍的答案。
[[[ Agent multiplication（for nerds）：如果你尝试使用 Grok 的 AI 助手（注意 Grok 和 Groq 是完全不同的 AI 公司！），你可能注意到它说"agents thinking"，复数，不是"agent thinking"，单数。这个想法是让多个思维链 agent “看到”（attend to）之前 token 和彼此 token 的 key。这是除了推测解码、GQA 和非方形脉动阵列之外的进一步乘数。例如 4 个这样的思维链，你在注意力的 key 重用上又获得了 4 倍。你也在可以有利地应用于单个用户的 batch 上获得了 4 倍。]]]
[[[ 关于 batch=1（for nerds）：假设你使用 4 倍推测解码和 4 倍 agent multiplication。那是单个用户 16 batch。这可能暗示人们永远不需要 batch=1 产品。然而虽然这些因素确实使更高 batch 的产品更有吸引力，我不认为这完全消除了低 batch 的需求。
回顾 Groq LPU 是针对低 batch 的昂贵低延迟 token。低 token 延迟的主要目的是获得高的每用户 token 吞吐量。现在看来我们无论如何都获得了那个高的每用户吞吐量，或至少 16 倍，而不需要 batch=1 操作的开销。那 LPU 没用了？我听过这个论点，但它不完全对。如果你对单个用户有 16 倍 batch，你可以买 16 倍多的 LPU 来获得 LPU 上的 16 倍每用户 token 吞吐量。所以 LPU 对高每用户吞吐量的昂贵 token 仍然领先。
这些观点也适用于设备端/嵌入式用例，人们经常引用需要 batch=1，因为嵌入式设备往往一次只有 1 个用户。鉴于本节的技术，我们可以看到即使对于嵌入式用例你可能也比你想的有更多 batch 可用。嵌入式对文本解码来说仍然可能是低 batch 的，但即使只是到 batch=4 也比 batch=1 高效得多。
另一个要记住的点是这是针对文本解码的。对于图像，通常使用扩散模型（diffusion model），它有充足的 batch。这里我不会解释扩散模型，但你可以去查。也有文本的扩散模型，尽管它们往往不太好用。不过如果你发现你可以使用文本扩散模型，那是另一种解决低 batch 文本解码困境的方式。]]]
预填充#
预填充本质上和解码做同样的事情，除了我们读取的是已经写好的 token。所以 AI 在理解已经存在的文本（或其他数据）（"（预）填充 KV cache"），它不是在生成文本。这很重要，因为这意味着我们可以同时处理文本中的所有词（上限是芯片内存的限制），而不是逐词处理：
The quick brown fox
↓ ↓ ↓ ↓
keys keys keys keys
↕ ↕ ↕ ↕
queries queries queries queries之前解码时我们只有一个 query 向量要处理，现在我们可以同时处理所有 token，给我们与 token 数量一样多的 query 向量。这很像推测解码，只是我们已经确切知道下一个 token 会是什么，所以没有任何推测性。输入可以是包含数千个 token 的整个网页。这大幅改善了内存带宽需求，因为每次从内存读入（或首次生成）的每个 key 现在每次读入时可能被使用数千次。所以我们没有解码期间 query 向量太少的大问题。
注意同样的充足 query 向量的好处也适用于训练，因为训练也是在已经生成的文本（或其他数据）上进行的。解码的问题是推理特有的。
我们在预填充期间仍然可以从 GQA 中获益。对于预填充，GQA 不会帮助我们最担心的解码时矩阵乘法的 N 维度，因为该维度已经足够大了（我们有足够的 query 向量）。GQA 确实有助于预填充期间的内存带宽需求，因为使用 GQA 需要读取的 key 更少，但预填充不太可能像解码那样需要内存带宽改善。然而，注意力还有一个 K 维度，该维度也需要足够大以填满脉动阵列。这是 GQA 在预填充期间帮助最大的地方。解码章节提到 GQA 经常与增加 key 和 query 向量的宽度结合。这是另一种说法：我们在直接增加 K 维度，帮助从预填充注意力中获得高（可能 100%）的脉动阵列利用率。所以 GQA 对预填充也很重要。
解码-预填充（管理式）聚合 vs 分离#
**默认聚合：**解码和预填充的默认方式是做显而易见的事情：在同一芯片上用同样的 kernel 做预填充和解码。数学上预填充和解码几乎没有区别，它们只是每次每个对话处理多少 token 有着极其不同的数字（例如 1-4 vs 256），所以它们可以组合并同时做。这样解码与预填充的比率是可变且未管理的。
**分离（Disaggregation）：**解码-预填充分离是预填充和解码在不同芯片上做。甚至可以（但不必）是不同种类的芯片，因为解码可能需要更高的内存带宽。这需要将 KV cache 从预填充的地方传输到解码使用的地方。好处是可以使用专门为预填充或解码定制的软件甚至硬件。
分离可以帮助的另一个问题是混合预填充和解码可能导致芯片有时主要在做预填充，因为预填充的对话/序列每步处理时间可能长得多（有更多 token 要一次处理）。这可能导致解码的 token 延迟增加。
**管理式聚合（Managed Aggregation）：**管理式聚合是以管理的比率混合解码和预填充。当解码是内存受限的（等待内存）而预填充是计算受限的（等待脉动阵列）时，这是有益的。如果你有意地以正确的比率将两种工作负载组合在一起，你可以达到一个完美平衡的场景，其中内存带宽和脉动阵列都被充分利用，两者都不在等待对方。这比先做一个再做另一个在延迟（对优先 token）和吞吐量（对所有 token）方面都有改善。
管理式聚合是我刚刚编造的术语，尽管在实现上这与显然被称为"分块预填充"（chunked prefill）的东西非常相似或甚至相同。管理式聚合的含义是如果你期望使用这种技术，你的芯片不需要那么多内存带宽。它不仅仅是为了防止芯片被预填充工作淹没。
**达到平衡：**系统中解码与预填充的总体比率可能无法在所有芯片上达到完美平衡。然而管理式聚合仍然可以避免一个芯片主要做预填充而旁边的芯片主要做解码的情况。公司可以通过定价来大致实现平衡——许多公司已经将预填充（输入 token）定价远低于解码（输出 token）。此外如果你有自己生成的无限解码或预填充工作源，你可以用它来平衡你的系统。另一个选择是除了预填充之外还将解码与训练结合，因为训练也倾向于计算受限。要记住的重要事情是即使无法达到完美平衡，管理式聚合也是有益的。
你可以通过重叠注意力和 FF 使比率要求更宽容——FF 可能有一些空闲的内存带宽可以通过重叠提供给注意力。要实现这种重叠，你需要流水线化（注意力和 FF 在单独的阶段）或模型修改使注意力和 FF 并行发生。后一种技术也将 token 延迟减半。
如果你构建了一个旨在这样做管理式聚合的芯片，你甚至可以为它配置比其他情况下更少的内存带宽，降低生产成本。这是一个大胆的（读作：冒险的）举动，因为也许不是所有客户都会使用管理式聚合。这又是一个例子，说明不成熟的客户需要的产品与成熟客户最适合的产品不同且更昂贵。
昂贵的 HBM 存储容量#
AI 芯片附带大量昂贵的 RAM，称为 HBM。为什么？这真的必要吗？这是本章的主题。
权重和 KV cache 作为大型昂贵 HBM 内存的原因#
HBM 是一种非常昂贵的 RAM。估计各不相同（公司不会告诉你他们为 HBM 支付了多少），但 HBM 每 GB 存储可能比普通 RAM 贵 3 倍。HBM 提供高带宽，在 $/GB/s 带宽方面并不那么贵。**昂贵的是一 GB 的 HBM 存储容量。**所以你可能认为 AI 芯片会附带小的 HBM 容量但仍提供高带宽。今天在发生吗？完全没有。相反，现代 AI 芯片上巨大昂贵的 HBM 存储容量越来越大。Nvidia 宣布他们的 Blackwell Ultra GPU 将以惊人的 288 GB 特别昂贵的 HBM（即 HBM3e——“e"不代表"expensive”，但你可以那样想）推出。
如果昂贵的 HBM 这么贵，为什么 AI 芯片附带这么多？主要是为了存储大型 KV cache 和模型权重。
KV cache 需要多少存储？这里有一个公式：
KV cache 大小 = 2 × batch × layers × seq_length × attention_head_groups
× vector_width × element_size × idle_magnification × pipeline_factor各参数：
- ×2：Value 向量和 Key 向量各一份（Value 向量是 key 用来查找的数据，如果 KV cache 是电话簿，key 是人名，value 是电话号码）
- Layers L=32：每层有自己的 key
- Sequence length N=1,000,000：AI 能记住多远？这可以高达一百万
- Batch B=64：AI 同时处理 B 个独立对话。256×256 脉动阵列在 decode FF 期间需要 batch=256，但推测解码 4 个 token 可将其降至 B=64
- Attention head groups G=8：用 GQA，128 个注意力头分成 8 组共享 key
- Vector width W=128：key 向量有多宽
- Element size S=1 byte：每个元素多大（8 位 = 1 字节）
- Idle magnification M=1.0：内存中有多少空闲对话的 KV cache（等用户回复时）
- Pipeline factor P=1.0：完全流水线化可使 KV cache 用量最多增加 4 倍
用这些数字算出来：
2 × 64 × 32 × 1,000,000 × 8 × 128 × 1 = 4,194 GB
这是大量的昂贵 HBM！N=1,000,000 的序列长度确实很大，但这是 AI 助手公司提供的真实产品。
最近的 Nvidia GPU 和 Google TPU 有大约 100 到 200 GB 的 HBM。我们无法在 200 GB 中存储 4,194 GB，那这怎么工作？答案是 LLM 被并行化到许多 AI 芯片上，KV cache 分布在许多芯片上。例如在层内跨 32 个芯片分布，你需要 4194 GB / 32 = 131 GB/芯片。这就是你可能最终需要每芯片 200 GB 昂贵 HBM 的原因。
权重呢？如果有 1000 GB 的权重，你可以分布在所有芯片上。层内并行 32 路加上 32 层流水线并行 = 1024 芯片。1000 GB / 1024 = 每芯片不到 1 GB。所以权重存储不是大问题——只要你做了足够的并行化。注意：KV cache 存储受益于层内并行但不受益于层间并行（流水线）。相比之下，权重存储受益于两种并行。
减少 AI 芯片对大内存的需求#
我在这里要抱怨的是我们在计算中使用了 N=1,000,000。有没有更经济的方式？确实有许多方法可以在最小损失模型精度的情况下解决这个问题——将注意力变成检索。我最终相信我们需要的解决方案是：对最近过去（例如最近 1024 个 token）使用密集注意力，然后对其余注意力历史转为稀疏检索，这些可以存储在附加的 SSD 上，甚至更远在数据中心某处。这样你可以将内存需求从 N=1,000,000 减少到 N=1,000，改善 1000 倍。SSD 比 HBM 每 GB 便宜得多。所以我们从 4,194 GB 降到 4 GB。
如果你是 AI 芯片初创公司，我强烈怀疑你可以做一个 32 位芯片只有 4 GB 内存（32 位指针可以寻址 2^32 B = 4 GB），而不是 Google 和 Nvidia 目前卖的 100-300 GB。如果你的软件做到完全正确并大规模并行化，你的便宜芯片将和这些超贵的行业领先芯片一样有能力。“你只需要做的"是对注意力使用逐层检索（对过去近期历史）并足够并行化以分布式存储大型 LLM 的权重。
这里有两个问题。第一，这样加载大模型所需的最小芯片数会很大。1 TB 模型需要至少 256 个芯片才能运行如果每个只有 4 GB 内存。第二，以上假设是像我这样的人在编写软件和部署模型。许多客户在部署 AI 模型方面不够成熟，但他们可能仍然愿意支付数百万美元买你的芯片。所以你最终在硬件上做了非常昂贵的事情来服务那些如果他们在软件上做得更高效就不需要所有这些开销的客户。我怀疑这是最近 Google TPU 和 Nvidia GPU 上巨大昂贵 HBM 的真正解释——是为不那么成熟的客户（内部或外部）让事情变得更容易的努力，即使这不是完全必要的。
有人可能反对说 HBM 可能只是总成本的 20%，所以不值得为了节省 20% 的生产成本而制造一个不太方便的产品。我认为这种论证经常以非显然的方式出错。因为我相信你会在这些产品的许多地方发现这种臃肿，如果你同时去除所有臃肿，影响远大于区区 20%，但如果你逐个去除每种臃肿，每个看起来都"只有"20%。**这就是所有臃肿保持原样的原因——你不能决定去除一种臃肿，因为它被证明只是所有其他臃肿的一小部分，而你也不能因为同样的原因去除那些。臃肿保护臃肿。**例如，为什么 TPU（或 GPU）需要主机，一台你连接 TPU/GPU 的单独计算机？为什么 TPU 不直接有一些片上 CPU 核心然后 TPU 就是计算机？
我怀疑 Google 可以在一两年内通过在实现的每美元 token 基础上极端和创造性地降低 TCO（总拥有成本），同时提供优秀的软件，然后也降低价格，在 AI 硬件市场取得巨大进展。我不认为他们会这样做，部分因为他们不必这样做，我也绝对不认为 Nvidia 会，所以这为市场上其他玩家留下了成为世界上最大的公司而不是 Nvidia 的机会。
[[[ 每芯片性能对 HBM 经济性的影响（for nerds）：不应该直接比较两个芯片之间的 HBM 容量来理解 HBM 的成本效益，而应该比较每芯片 HBM 容量与每芯片实现性能的比率。假设你做了一个速度是 Blackwell 一半、HBM 也是一半的芯片。你改善了 HBM 的经济性吗？完全没有，因为你需要买两倍多的芯片才能得到等价于 Blackwell 的系统，所以你最终购买了与 Blackwell 完全相同量的 HBM。]]]
[[[ 警示故事：价值万亿美元的态度（for nerds）：
当我在 Google 从事 TPU 工作时（TPU v1-v4 时代），公司内部始终存在一种我称之为"因缺乏雄心和非股东利益导向的忽视"的微妙但清晰的氛围。
TPUv1 的动机是"如果 Google 要和你说话，作为助手或其他方式，在现有硬件上运行语音合成太贵了，所以让我们做些什么来降低成本”。TPUv2 旨在服务两个与搜索和广告相关的内部训练应用，明确没有超出此范围的雄心。当 TPU 进入 Google Cloud 后，投入非常少，结果如果你在意 Google 股东利益的话是难以置信的差——没有雄心让 TPU 成为 GPU 的即插即用替代品。
与此同时，Google 显然坐在一个万亿美元的金矿上（现在可能仍然如此）。Google 作为公司将 TPU 视为在 AI 硬件上为自己省钱的方式。Nvidia 按市值是世界第一的公司，如果有那个意愿的话，那很容易是 Google。
我曾向我汇报链中非常高层的管理者追问这个问题，问为什么我们不像 Nvidia 对 GPU 那样全力以赴做 TPU。最初的回应是"但那是 Nvidia 的主要产品"，我说"是的，这就是我的意思"，然后我得到了一个以极其敌意的方式传达的令人惊讶的回应：“are you for real with this question?“关闭了对话。这是原话——我永远不会忘记那句话。
我是 TPU 上非常突出的软件工程师，我领导了 Google 的 TPUv3 软件，深度参与了 TPU 硬件/软件协同设计很长时间，但从来没有人问过我一个关于如何对 TPU 更有雄心的问题。
我仍然认为这一切很不幸，因为 Google 在这方面的缺乏雄心不仅对 Google 股东不利，而且对人类不利。我们需要负担得起的高质量 AI，这必须从经济性比市场上现有的好得多的 AI 硬件开始。]]]
AI 助手的并行化#
本章给出了如何并行化 AI 助手（即 LLM）的一些全局想法。“并行化"指的是将 LLM 的计算分布在多个芯片和/或芯片内的多个核心上。并行化是每芯片性能本身无关紧要的原因，尽管大多数 AI 芯片广告完全关注这个无关紧要的指标。重要的是你的系统能以什么价格交付什么，而不是单个芯片能交付什么。
对于 AI 芯片初创公司，理解并行化的精确影响至关重要，这在我看来是一个极其棘手且反直觉的主题。关于 AI 并行化的三个事实：
- **网络带宽是层内并行化的限制。**这就是为什么 AI 推理需要快速网络。层间并行化不需要那么多网络带宽。
- 如果你将网络带宽翻倍，你可能能做两倍的层内并行化，这最多可以将 token 延迟减半。所以大型脉动阵列是推动更快网络的力量。
- **层间并行化（流水线）有助于权重存储需求，但不有助于 KV cache 存储需求。**相比之下，层内并行化确实有助于 KV cache 存储需求。所以网络带宽与你需要多少 HBM 相关。
如果你使用 MoE，这一切变得更加复杂。另一个事实：如果你将芯片速度翻倍，你也需要两倍快的网络来跟上。
并行化概念#
**吞吐量（tokens/s）：**并行化最明显的好处。一个芯片 1000 tokens/s，两个芯片 2000 tokens/s。
首 token 延迟（seconds to first token）：AI 在开始考虑你的请求后生成第一个 token 所需的时间。这不是 token 延迟。
Token 延迟（ms/token）：稳态下生成一个 token 的时间。Token 延迟和吞吐量的关系不像看起来那么密切：
- 吞吐量 1000 tokens/s 不意味着你能在 1 秒内得到 1000 token 的回复——吞吐量是跨对话计算的
- 即使只有你一个用户，你也不太可能获得全部吞吐量——AI 可能只有在独立对话中才能达到
- token 延迟 1ms 意味着吞吐量至少 1000 tokens/s，但可能更高（多对话并行时）
- 结论：高吞吐量不一定导致低延迟，低延迟也不一定导致高吞吐量。事实上往往相反。
**变异性（Variability）：**吞吐量和延迟都是平均值。MoE 是变异性的主要原因。
**存储：**并行化可以根据方式不同大幅增加或减少每芯片存储需求。
**本地 vs 全局需求：**1 TB 模型需要 1 TB HBM。单芯片 100 GB 放不下？但 20 个芯片的全局容量是 20 × 100 GB = 2 TB——放得下。
**最小安装规模：**你的 AI 能运行之前必须拥有的最少芯片数。大的最小安装规模不方便。
**可靠性与故障域：**1024 芯片协作时，一个芯片失败可能导致全部不可用。在早期 TPU 中正是这个问题——连接线缆偶尔失败导致整个 1024 芯片 pod 不可用。Sameer Kumar 修改了并行化软件使其即使有几条失败线缆仍能良好运行，代价是将 torus 有效网络带宽减半。但 TPUv3 torus 网络本来就过度配置了——这是过度配置带来灵活性的例子。
通过复制进行并行化#
最简单的方式：多个独立实例各处理一部分对话。
| D 路复制 | |
|---|---|
| 吞吐量 | ×D |
| 延迟 | 不变 |
| 每芯片 HBM | 不变 |
| 可靠性 | 改善（一个实例失败其他仍可用） |
| 网络带宽 | 变化小 |
| 脉动阵列利用率 | 不变（仍可 100%） |
复制是我们讨论的唯一改善可靠性的并行形式。
使用流水线的层间并行化（Decode）#
如果模型有 32 层，你可以将它们分布在 32 个芯片上。每层在自己的芯片上，芯片间通过网络交换向量。Layer 2 必须等 Layer 1 完成。解决方案是流水线化：有 3 批对话 X1, X2, X3 依次在各芯片上流动，使所有芯片同时忙碌。
流水线最明显的效果是将吞吐量乘以 S（S 个阶段）。但流水线不减少 token 延迟，因为每个对话/token 一次只被一个芯片处理。
为什么用流水线而不是复制？流水线有减少权重存储需求的强大好处。每个芯片/阶段只需存储自己层的权重。32 层 = 每芯片权重存储减少 32 倍。复制没有这个好处。
KV cache 存储呢？流水线对减少 KV cache 存储毫无帮助。原因非显然：虽然每个芯片只存储自己层的 key，但流水线导致同时处理 32 倍的对话批次。每芯片每对话减少 32 倍 × 对话增加 32 倍 = 两者抵消。
完整流水线（Full pipelining）：将流水线阶段数翻倍来隐藏传输时间。一半对话在被处理，另一半在网络传输中。每步交换。这确保 100% 脉动阵列利用率，但将 token 延迟翻倍、KV cache 需求翻倍。好处是网络不再需要那么快。
| Decode 部分流水线 ×S | Decode 完整流水线 ×S | |
|---|---|---|
| 吞吐量 | ×略小于S | ×S |
| Token 延迟 | 增加（取决于网络） | 翻倍 |
| 权重存储/芯片 | ÷S | ÷S |
| KV cache/芯片 | 不变 | 翻倍 |
| 可靠性 | 更差（故障域×S） | 更差（故障域×S） |
| 网络带宽 | 关键路径上 | 只要够快就不影响 |
| 利用率 | 降低（永远<100%） | 不变（可100%） |
使用流水线的层间并行化（Prefill）#
流水线对预填充比解码更有吸引力。原因是预填充不需要引入额外对话批次来做流水线。我们查看的是同一文档/对话的跨度/子集，而不是不同对话的批次。由于只看同一文档，完整流水线不会显著增加 KV cache 存储需求。
| Prefill 完整流水线 ×S | |
|---|---|
| 吞吐量 | ×S |
| Token 延迟 | 翻倍 |
| 权重存储/芯片 | ÷S |
| KV cache/芯片 | ÷S |
| 网络带宽 | 只要够快就不影响 |
| 利用率 | 不变（可100%） |
层内并行化#
在流水线中我们假设每层运行在单个芯片上。也可以将单层并行化到多个芯片上。这比流水线更可取，因为它减少 KV cache 需求和 token 延迟。大代价是需要更多网络带宽。
注意力很容易并行化。128 个注意力头可以分布到多个芯片。用 GQA 分 8 组 = 8 路并行。还可以将 KV cache 拆分到多个芯片。
FF 也并行化得很好。FF 主要是两个大矩阵乘法，中间向量宽度增大（通常×4）。可以让每个芯片负责中间宽度的一个子集。
两者都需要跨芯片求和（all-reduce）来合并结果。每芯片层内网络带宽需求在并行化芯片数翻倍时略多于翻倍。
[[[ 量化"略多于翻倍”（for nerds）：做分布式求和（all-reduce）时，如果有 P 个芯片，每个芯片传输的总数据量为 2X × (1 - 1/P)，其中 X 是每个芯片上数据的大小。1-1/P 因子在 P 小时显著（P=2 时为 50%，P=4 时为 75%），P 大时接近 1。
虽然 2X × (1-1/P) 对大 P 大致等于 2X 且不依赖于 P，但步骤速率翻倍了，所以你需要在相同时间内做两倍多的分布式求和。这是 AI 芯片大量带宽需求的主要驱动力。]]]
[[[ 每芯片层间带宽需求对层内并行化中立（for nerds）：P 路层内并行化时，现在 P 倍快地处理 P 倍多的批次。但如果你让 P 个芯片中的每个只发送 1/P 的数据到下一层，这就改善了 P 倍，与处理 P 倍快相抵消。所以层间带宽在每芯片基础上是中立的。这意味着层内需要比层间多得多的带宽，所以你可以把层间传输放在慢得多的网络上。]]]
与流水线类似，我们可以选择：在需要做层内网络传输时让芯片闲置等待，或使用 2 阶段流水线将对话数翻倍隐藏传输时间。后者将 KV cache 和 token 延迟增加 2 倍但降低网络带宽需求。
注意注意力和 FF 之间有通信步骤，所以完整流水线每个 transformer 层有 4 个阶段。建议做模型修改使注意力和 FF 独立、并行运行——这将 token 延迟和阶段数减半。代价是收敛模型精度略有下降。
如果在层内和层间都做完整流水线，权重存储收益相乘，非常显著。这是你如何实现跨 1024 芯片的并行化并将每芯片权重存储降到极小的方式。
[[[ 完整流水线好处的例子：假设你在层内并行化 8 路而不用完整流水线，观察到额外 25% 的时间在等网络。你不想进一步并行化。你也不想用完整流水线因为 token 延迟会翻倍。
发现思维错误了吗？如果网络是 25% 额外开销，那如果你做完整流水线，你可以并行化 4 倍而完全不等网络，总共 32 路。你确实将 token 延迟翻倍了，但你也将其除以 4，所以总的来说你将 token 延迟减半了，而不是翻倍。而且比这更好，因为你可以减去原来的 25% 开销。]]]
混合专家模型（MoE）#
MoE 是 FF 的一种变体，不是一个大矩阵而是 N 个不同的较小矩阵称为专家（expert）。每个 token 只被路由到相关知识的专家。这是加速，因为 token 不必考虑不相关知识的权重。MoE 是加速 transformer 模型的强大技术，但如果你想要高利用率，它在部署方面造成很多复杂性。
专家数量在模型之间差异很大（Mixtral 8×7B 有 8 个，Llama 4 有 128 个）。
MoE 的大问题是专家不一定被均匀使用。如果专家 A 知道英语名词，专家 B 知道公元前 405 年蒙古戴的帽子，你可以看到专家 A 会被咨询得比 B 多得多。结果是你需要为专家 A 配置比 B 多得多的芯片。这种动态变异性是之前没有讨论过的。
如果专家太少token 无法填满脉动阵列，利用率就低。120% 负载需要复制到 2 个专家各 60% 利用率。但 1020% 负载可以分配 11 个芯片各 93% 利用率。安装规模越大，MoE 利用率越高。
不让一些 token 等待就无法获得 100% 利用率。最简单的解决方案是接受较低利用率。更高级的系统是有高优先级 token（有客户等待的 AI 助手 token）和低优先级 token（批量离线工作），专家优先服务高优先级 token，低优先级 token 在队列中等待来填满专家。
**作者的预测：**LLM 在权重中包含大量知识，其中大部分大多数时候无关。这就是为什么有很多专家的 MoE 如此有效——它在丢弃大量无关权重。问题是我们首先在权重中存储了太多知识。我们应该使用层内检索从向量数据库（不同于 RAG）来存储大部分知识，大幅减少权重数量同时提高模型精度。这也将大幅降低 MoE 的积极影响，尽管我认为我们仍然会使用它，但只用少量专家如 4 或 8 个。
AI 网络拓扑（MoE 世界）#
本章解释路由器和 torus 拓扑以及 AI 工作负载需要的算法。
网络概念#
**节点（Node）：**芯片是节点，路由器也是节点。路由器是节点这一事实引入了 2 倍的路由开销——这通常被忽视。
**链路带宽：**当有人说"10 gigs 带宽"时要问：bits 还是 bytes？双向的？Giga 还是 Gibi？Goodput 还是线速？理论还是实际？你可能需要问几个问题才能弄清楚别人说的带宽数字到底什么意思。
**Mesh 拓扑：**1D mesh 是线，2D mesh 是网格。Tesla 的已停产 Dojo 芯片用了一个非常大的 2D mesh。不推荐使用 mesh 因为 torus 更好。
**Torus 拓扑：**1D torus 是环。2D torus 添加另一个维度。torus 像 mesh 但有额外的长链路环绕，在每个维度创建环形连接。早期 Nvidia GPU 通过点对点 NVLink 以环形拓扑连接。
**折叠 torus（Folded torus）：**torus 中的长环绕链路是个问题（更长 = 更贵）。TPUv3 用了非常长的光缆。折叠 torus 保持所有链路短——所有链路最多 2 单位长，使 torus 能使用只适用于短距离的链路技术如铜走线。
**路由器拓扑：**路由器是有许多端口的节点，可以灵活地在端口之间传输数据。
链路类型（按距离从短到长，带宽/美元从高到低）#
- 同一芯片上连接 — 最高带宽、最低延迟、最便宜。导致越来越大的芯片。
- 硅桥芯片到芯片连接 — Blackwell 的多芯片设计和 HBM 连接方式。
- 芯片引脚连接 — 连接内存、网络和其他组件。
- 被动铜走线 — PCB 上的铜线，比铜缆便宜且带宽更高。可以在同一 PCB 上紧密放置的芯片之间创建非常快的网络。
- 主动铜走线 — 较长距离需要中继器加强信号。
- PCB 到 PCB 直接连接 — “可拼接"主板。
- 背板连接 — 多个 PCB 连接到背板 PCB。
- 铜棒连接 — 刚性铜棒而非柔性线缆，更便宜更低功耗。
- 柔性铜缆 — 以太网线缆类型。
- 光缆 — 长距离最优。Google 和 Nvidia 都使用。最昂贵的方案。
为什么 Nvidia 和 Google 使用最贵的方案？因为最容易使用，且两个产品都没有为成本优化。
多跳浪费的幽灵#
网络的平均每芯片有用带宽：
平均链路利用率 × 所有芯片链路带宽 / 平均跳数对于单路由器拓扑，平均跳数始终为 2（芯片 → 路由器 → 芯片）。如果只有邻居通信，torus 的平均跳数为 1。所以路由器最多只有 50% 效率，而 torus 可以达到 100% 效率。
路由器不是有用的目的地，所以路由器上的每个端口和所有带宽都被浪费了。路由器上的每个端口与芯片上的端口 1:1 匹配，所以正好一半的带宽被浪费。
一旦单个路由器端口用完，你需要多级路由器，涉及长电缆。路由带宽是昂贵的带宽。
作者的网络方案提议#
作者不推荐 Google 的方法（宽 3D torus + 昂贵光连接，不适合 MoE）也不推荐 Nvidia 的方法（昂贵的路由拓扑）。
**提议：便宜的低带宽路由器 + 便宜的高带宽窄 torus 的混合方案。**没有昂贵线缆，仍然适合 MoE 工作负载。用铜走线 torus 取代光缆，用可拼接主板实现物理连接。
协同设计（Co-design）#
Co-design 是硬件和软件团队一起设计的过程，使硬件和软件相互适应。作者分享了他在 Google 和 Nvidia 的 co-design 经验。
关键洞察：
- 智力质量和良好判断力比资历重要。Nvidia 之所以能从 BF16 的落后中恢复并超越 Google，是因为他们能快速适应并掌控局面。
- Co-design 需要扎根于产品——不能是脱离实际的理论讨论。价格信息是 co-design 的关键组成部分。
- 组合搜索：AI 芯片设计涉及大量相互关联的决策。标准的"爬山法”（hill climbing）co-design 过程会陷入局部最优。臃肿保护臃肿——每种臃肿单独来看"只有"20%，但组合起来消除所有臃肿的效果是超线性的（$1.01 的两个 $0.50 节省分别是 2x，但合计是 100x，因为 1/$0.01 vs 1/$1.01）。
[[[ 故事：等等，这人是谁？Google 曾雇用一个团队做某种 TPU co-design 相关的事情，但与实际的 TPUv2+ 项目联系有限。最初软件团队没人知道这个团队存在。负责这个团队的人获得了图灵奖。当我作为 TPU v2+ co-design 的主要软件联系人加入那些会议时，从来没有人问过我关于那个的问题。我提出了一些我们不确定的话题——这完全不受欢迎。那不是他们在做的事情。
Gemini 这样总结：“Patterson 于 2016 年加入 Google 担任杰出工程师，成为 TPU 设计理念的公共代言人和主要学术权威。“结果是在公众心目中将 Patterson 与 TPU 联系起来，为 Google 品牌整体产生了图灵奖光环。]]]
对 GPU AI 特性的批判性审视#
Nvidia GPU 包含一些 AI 不需要的昂贵功能。如果你移除这些功能，剩下的大致是一个带脉动阵列和高内存带宽的 CPU。所以本节的结论是 Nvidia GPU 在 AI 上成功是因为 Nvidia 在构建 GPU 和编写软件方面非常出色，不是因为 CUDA 编程模型最适合 AI。
Kernel 和 warp 特化#
GPU 程序被分成"kernel”，它们是一个接一个运行的独立子程序。这源于 GPU 运行图形 shader 的古老方式。
在 CPU 上，程序只是一个大二进制文件运行多久就运行多久。CPU 模型更优越。当你想同时做多件事（如重叠计算和网络操作）时问题就出现了。在 GPU 上你不能在单个 SM 上复用两个 kernel。解决方案是megakernel——整个程序就是一个大 kernel。所有 CPU 二进制文件（和 TPU 二进制文件）已经是这种方式了。
Path divergence（向量子集）#
GPU 可以操作向量的子集——用于 AI 维度不是向量大小倍数时。这是方便但不必要的功能。AI 从业者应该选择是向量大小倍数的维度——“2 的幂是对的，否则你的代码就是灾难”。
软件管理缓存#
CPU 通常有自动缓存。GPU 有自动缓存和手动缓存（“shared memory”）。对于高性能 AI，完全手动控制缓存更方便。AI CPU 应该做 Nvidia 做的：有手动和自动缓存。理想情况下整个缓存默认自动，按需将部分转为手动。
Warp 调度（内存访问流水线）#
warp 访问内存时如果不在缓存中就要等待。GPU 快速切换到另一个不在等内存的 warp。这使硬件（理想情况下）从不闲置。这是一个昂贵的功能，可以说是使 GPU 成为 GPU 而不是 CPU 的最大特性。
**Warp 调度用于 AI 不是因为 AI 需要它，而是因为 CUDA 编程模型需要它。**软件流水线可以为 AI 实现相同效果但更快更便宜。
需要主机计算机#
GPU 需要连接到主机计算机，增加了成本和不必要的开销。如果 GPU 或 TPU 本身就是计算机，主机的成本就没有了。我的提议是完全去除主机计算机。
许多线程的需要和 kernel 切换开销#
CUDA 涉及许多线程的额外开销：每个线程需要计算自己是谁和该做什么；kernel 开始时脉动阵列闲置；kernel 结束时要等最慢的线程。CPU 和 TPU 在这些问题上更好——它们不需要那么多独立线程，TPU 也相当确定性。
AI kernel 的软件流水线#
软件流水线对 AI 软件是必需的，也是 AI 软件中最难的事情——但如果你按我建议的方式做，就不会那么难。
什么是（软件）流水线#
假设你需要在循环中做三个操作 A、B、C。如果顺序执行，每次只有一个单元忙碌，利用率只有 33%。
软件流水线通过重排操作使多个单元同时忙碌：
步骤 1: A(x_1)
步骤 2: A(x_2) 和 B(a_1)
步骤 3: A(x_3) 和 B(a_2) 和 C(b_1) ← 稳态，100% 利用率
步骤 4: A(x_4) 和 B(a_3) 和 C(b_2)
步骤 5: B(a_4) 和 C(b_3)
步骤 6: C(b_4)前两步是序幕（prologue）（~50% 利用率），中间是稳态（steady state）（100%），最后是收尾（epilogue）（~50%）。
手动软件流水线很难#
实际代码中 A、B、C 不是单行的东西，实际的手动软件流水线代码往往远比示例复杂。再加上分 tile——如果数组太大放不进缓存，必须细分为固定大小的 tile，tile 大小要对齐向量宽度、要适合缓存、维度要整除…选择 tile 大小是一个充满竞争性关注的困难业务。
编写 tile 化软件流水线库（带 op 融合）#
既然手动做太难，不要那样做。
作者建议创建一个 C++ 模板库，你声明阶段 A、B、C 及它们的输入输出，库产生流水线化的实现。最重要的用途：预取内存——先预取一个 tile 的数据，然后处理数据，然后写输出，重复多次。
作者确实为 Google TPU 编写了这样一个库，Google TPU 上的所有操作都是使用这个库完成的。这就是为什么 TPUv2 能由 5 人 XLA 团队发布的原因。Jax 团队后来将类似接口以 Pallas 的形式暴露出来。
建议的最小方案：支持 XLA 和 C++（带 intrinsics）。XLA 是最经过实战检验的 AI 编译器且是开源的。C++ 让客户能做你能做的一切。如果你的芯片是类 CPU 的，它可以用 C++ 高性能编程——Google TPU 可以但没有提供这个选项，作者认为这是个错误。
[[[ 关于用 AI 自动写 kernel（for nerds）：一个想法是用 AI 编码模型，在大量如何为你的硬件编写代码的示例上重新训练/微调，然后将生成的 AI 编码器作为你产品的一部分提供给客户。AI 可能还能生成形式化证明来证明正确性。]]]
为什么不用 X、Y 或 Z？#
有许多选择：Triton、Jax-Pallas、Mojo、TVM、Halide、自己的语言等。但 XLA 是最经过实战检验的，C++ 不比它更直接。这两个是最小的合理方案。
Nvidia GPU 上的软件流水线 = Warp 调度#
CUDA 中不写循环而是写函数，GPU 同时执行多个实例。效果等同于软件流水线但通过硬件实现。如果使用好的 C++ 流水线库，我认为用那种方式写高性能 kernel 实际上比用 CUDA 的 warp 调度更容易。warp 调度（加上 warp 的所有派生成本）是一个非常昂贵的硬件功能。我预测这最终将成为 Nvidia 的竞争问题。
异议与回答#
作者收到了读者反馈，将其改写为异议并逐一回答。
OS 中断脉动阵列#
异议（Sean Silva）： 在 AI 加速器上运行 Linux，OS 中断不会导致脉动阵列出现不可预测的空闲周期吗？
回答： 提案是每个芯片有多个核心，只有部分附有脉动阵列。中断和 OS 杂务可以限制在没有脉动阵列的核心上。
很少有人能获得所需的 CPU#
异议（Sean Silva）： 没有一个真心想做 AI CPU 加速器的公司拥有可修改的、性能足够高且能运行 Linux 的 CPU IP。拥有的那些（AMD、Intel、Apple）走的是"稍好一点的 CPU 推理"路线。
回答： 同意这些公司有条件但到目前为止没有尝试。Google 也可以基于 Axion 做未来 TPU 但也没发生。获得现有高端 CPU 是加分项但不是必需——从现有 CPU 开始甚至可能是弊端如果它不适合需求且你无法修改。这些 AI 设备不需要那么多标量工作。TPU 使用 VLIW 指令集来降低每个标量核心的面积和功率成本，如果你做自己的 CPU 也值得考虑。
缺乏科学验证#
异议（Sean Silva）： 一些建议也许最好改写为要进行的实验。
回答： 同意这些是有根据的观点，制造 AI 芯片的人应该自己验证而不是仅仅相信我的话。
这个设计可能不够通用#
异议（Sean Silva）： 似乎不太可能任何一个特定设计能扩展到所有用例。
回答： 如果 AI 将继续基于大型矩阵乘法，这个设计通常适用于大多数 AI 应用。主要例外是非常低功耗的边缘 AI。这个设计可以同时服务 prefill 和 decode——GQA + 推测解码的强部署减少了两者之间的差异。如果想做专用芯片，只需调整 HBM 容量、带宽和脉动阵列尺寸等常数。
每美元 token 数不是一切#
异议（Rama Govindaraju）： 还有其他重要指标，例如 token 质量。
回答： 假设硬件 A 的总拥有成本是硬件 B 的一半。选择 A 的公司（Cod）相比选择 B 的竞争对手（Bass）有三个选项：
- 以更低价格提供相同 token
- 用两倍计算产生每个 token，以相同价格提供更高质量的 token
- 用相同投资支持两倍客户
所以更好的 token 经济性可以转化为相同价格下更高质量的 token。更经济的 AI 硬件让公众获得的不仅是负担得起的，而且是更高质量的 AI 服务。这正是我写这篇文档的原因——没人付钱让我这样做。
AI 和 AI 硬件行业仍然不成熟。即使不改进芯片光刻，AI 硬件和配套软件中仍有远超 2 倍的改善空间。
模拟计算呢？#
异议： 模拟计算或光子计算呢？
回答： 希望那些都能成功！到目前为止这些方法在市场上还没被证明优于传统脉动阵列。脉动阵列的堡垒在我看来相当坚固。
为什么要更多 FLOPs？#
异议（Yiren Zhao）： 更大的脉动阵列不一定有效——如果你不是计算受限的，增加更多计算没有帮助。
回答： 完全正确：不是计算受限时增加更多计算没有帮助。但考虑一个降频的更大脉动阵列使其与较小的有相同吞吐量——更大的仍然更高效（有上限）。所以更大脉动阵列的效率独立于是否需要更多 FLOPs。在 co-design 搜索中应包含更大降频脉动阵列的可能性。
AI 芯片需要是平衡的设备——co-design 不能只最大化计算，需要仔细考虑脉动阵列计算、向量计算、标量计算、内存带宽、内存容量、网络带宽、功耗/散热等之间的复杂依赖。
结构化稀疏性很烂#
异议（Jonas M. Kübler）： 没听说有人在 Nvidia GPU 上大规模使用结构化稀疏性。精度下降太大，性能增益不够。
回答： 本文档不是对当前部署实践的调查，而是关于应该做什么。最好的部署将使用 2:4——Nvidia 在 MLPerf 中用了它，腾讯也在使用。缺乏结构化稀疏性是"臃肿保护臃肿"的例子——浪费但熟悉且更容易部署。
2:4 结构化稀疏性的简单配方对人类的价值至少是两位数十亿美元（可能超过 $100B）。可以尝试的改进方向：
- 重排列以获得更好的匹配 — 列顺序任意，可无运行时开销地重排
- 保留几个异常值列为密集 — 通过添加零列实现
- 保留几个异常值行为密集
- 使用渐进稀疏化方法（除了或补充近端梯度法）
- 使用彩票假设训练 — 最昂贵但可能最有效
- 修改原始训练的超参数和正则化 — 对较小元素放置稍高的正则化损失
- 训练期间小概率地将较小元素设为零 — 本质上是 elementwise dropout 用于不同目的
- 将稀疏性限制在 FF 矩阵 — 不理想但合理，因为注意力中大部分 op 是激活-激活矩阵乘法
脉动阵列不适合 depthwise convolution#
异议（Trung Ngo）： 混合/SSM 模型中的因果 depthwise 1D 卷积对 TPU 大型脉动阵列是性能杀手。
回答： 确实如此。使用大量 depthwise convolution 的模型在现代 GPU 或 TPU 上运行不好。Mamba 模型是例外——其 depthwise convolution 非常小，有望用向量计算能力来做，同时在脉动阵列上做更适合的事情。
Depthwise convolution 的问题：算术强度低、求和数字极少（低 K 维度）。3x3 depthwise convolution 做的是 Ax9 乘 9xB 的矩阵乘法——9 对大型脉动阵列来说太小了。
路由器是否优于 torus？#
异议（Aygul Galimova）： NVSwitch 已实现了 Clos 式交叉开关来解决 MoE 路由问题，这是否改变了你的分析？
回答： 我将原来的长回应扩展成了一整个新章节——“MoE 世界中 AI 的网络拓扑”。
现有的 AI CPU 怎么样？#
异议（Jan-Eric Schäfrich）： ARM 的 SME 使用外积计算做矩阵乘法而不是脉动阵列，似乎更灵活但功耗和面积效率为代价。
回答： 当我说 AI CPU 时，指的是专门加速 AI 的设备。你需要最高效的矩阵乘法——那是脉动阵列，不是外积引擎。
这些芯片缺乏顶级 AI 加速器的内存和网络带宽，矩阵乘法单元在 32 × 32 或更小。Nvidia GPU 中等（64 × 64），Google TPU 大（256 × 256）。2X 边长 = 4X 大小，所以 256 × 256 比 16 × 16 大了 256 倍。
这些设备是拿普通 CPU 加一点"AI 东西"得到的。我设想的是从零开始专为顶端经济 AI 设计的 CPU 芯片。如果 AI 加速器是 F1 赛车，这些 CPU 只是快的普通汽车。我设想的是能赢 F1 的 AI CPU。
ARM 的"AI CPU”（Arm AGI CPU）只是一个普通 CPU 被营销为与 AI 加速器配合使用——完全不是我在这篇文档中说的。实际上我提议的是完全去掉主机计算机。