《Designing AI Chip Software and Hardware》中文整理版04- 芯片为什么这么贵

大家好，我是拓驰猎头CEO。基于20年AI硬科技猎头经验及百人规模猎头团队，我们在传统猎头模式的基础上，还发展了股东合伙人模式、OPC业务合伙人模式、SOHO猎头模式等创新性多元化业务模式，欢迎大家交流探讨。我的公众号也会持续分享行业信息及专业洞察，欢迎关注。

原文标题：Designing AI Chip Software and Hardware

原作者：Bjarke Hammersholt Roune

授权状态：原文公开免费阅读，但未见明确的开放许可或中文翻译发布授权。本文为非商业学习交流翻译；中文分篇标题为译者拟题；本文为第 4 / 9 部分。

激活和权重的压缩

深度学习革命一开始是每个数字使用 32 bit，后来降到 16 bit，再后来是 8 bit，现在甚至有时会使用 4 bit，而且准确率损失很小——也就是说，并没有让 AI 明显变笨。这是对权重（weights）和激活（activations）的有损压缩（lossy compression），也是 AI 中的一种标准技术。

不过，无损压缩（lossless compression）又怎么样呢？这听起来可能是个奇怪的问题，因为我们已经在使用有损压缩了，而且有损压缩不是比无损压缩更强大吗？这是一个误解。有损压缩的工作方式是移除不必要的精度。无损压缩的工作方式是利用数据中的非均匀模式。这是两件不同的事情。这就是为什么 JPEG 图像格式同时使用有损压缩和无损压缩。尤其是，JPEG 会在数据已经经过有损压缩之后，再对数据应用 Huffman 编码（Huffman coding）。为什么我们不也在 AI 里这样做呢？我不知道为什么不这样做。

关于这个主题已经有一些研究，但并不多。Nvidia GPU 对纹理压缩 / 解压缩（texture compression/decompression）有硬件支持，但那是图形相关的东西，不是 AI 相关的东西。较新的 Nvidia GPU 中也有一个更通用的解压缩功能，但它完全跟不上内存速度，大约只有内存速度的 1/10。所以，如果你在 HBM 中压缩数据，使用它实际上会让你得到比不使用它更低的带宽（bw）；不过，如果你只压缩部分数据而不是全部数据，也许仍然能得到一些提升。AMD 最近宣布它们未来 GPU 会有 “Universal compression”，听起来很有意思，但到目前为止还没有细节，所以目前还不知道它到底是什么。

你可以在这篇论文中看到一些把 Huffman 编码应用到 AI 上的结果。不过，在那篇论文里，他们做了一些专门针对 BF16 的处理，并把它拆成 5-bit 片段，这不太可能是最优的。我更想看到 8-bit 整数或者 FP8 上的结果。另一篇论文没有包含太多信息，但声称在真实世界 AI 工作负载上可以减少 20% 的数据大小。这里使用的只是无损压缩。在那篇论文中，他们希望在一个 cycle 内完成解压缩，这看起来并不必要，因为由于软件流水线（software pipelining）的存在，cycle 级延迟对 AI 来说并没有那么关键。当然，cycle 级延迟并不是完全无关，但也不是关键重要。

阻碍 AI 使用无损压缩的一个问题是：压缩结果是可变的——并不是所有向量都能被同样好地压缩。这会带来一些麻烦，不过可以通过高优先级 token 和低优先级 token 来解决。如果某个 batch 的 token 花费了比平时更多的时间或空间，那就直接丢掉一些低优先级 token，之后再重新运行它们。用这种技术，像压缩带来的这种轻微变化就不再是问题。

在这一节中，我想提醒大家注意：Huffman 编码表和解码表对于有损压缩也很有用。这类表比乍看起来更有用。我认为，在未来的 AI 加速器中，值得考虑把它们做成能够以内存和网络速度运行的功能。

举例来说，要对一个 8-bit 整数进行 Huffman 编码，你会创建一个包含 2^8 = 256 个条目的表。每个条目都有一个长度可变的 bit pattern。于是，你拿到一个 8-bit 整数，在表里查它，然后输出对应 bit pattern 中的那些 bit。就是这么简单。这个想法是：更常见的值获得更短的 bit pattern，而罕见的值获得更长的 bit pattern，这样净效果就是压缩一个向量或矩阵。

假设支持的最长 bit pattern 是 10 bit。为了解压缩，你也会有一个表，现在这个表有 2^10 = 1024 个条目。表中的值是编码成该 bit pattern 的原始 8-bit pattern，以及编码长度，也就是你在解码时需要向前跳过多远。因此，编码和解码由两个互为逆映射的表组成。由于解码表比编码表更大，所以解码表中并不是所有条目都会被使用。

这怎么也能用于有损压缩呢？当你创建编码查找表时，可以把多个值映射到同一个 bit pattern。在这种情况下，这个表就在无损 Huffman 编码之上，又编码了一个有损压缩。所有 FP4 值乘以 2（为了避免 half，也就是 0.5 这种半值）之后，都可以编码为 8-bit 整数。所以，确实存在一个编码表，可以把 Int8 编码成 FP4，再编码回来。这里乘以 2 只是创建了一个定点表示（fixed-point representation）。

这并不会改变脉动阵列的内部结构或吞吐量；脉动阵列仍然可以保持例如 8-bit。不过这意味着，通过改变编码表和解码表，你可以动态构造类似 FP4 的数值格式，或者任何定制的专用格式，用于压缩，从而节省内存存储和带宽。而且，在有损压缩之上，你还可以获得无损压缩的收益。所以，对于每一层，你都可以运行一个自动化优化过程，为这一层找到有损压缩和无损压缩的最优组合。

在上面提到的论文中，压缩只应用于权重，并且是在脉动阵列处使用。这是一种选择。不过，压缩也可以应用于激活，也就是 AI 内部的非权重向量。这意味着它不仅可以用于权重，也可以用于例如 KV 缓存（KV cache）。这会带来轻微的可变性，但可以按上面提到的方式解决。

这类压缩需要硬件支持，才能在编码和解码时都以内存和网络速度运行。Huffman 编码是串行的，不过可以通过在压缩文件格式内部设置多个独立数据流来并行化。如果把这个功能加入芯片的 DMA 功能中，而软件又使用 DMA 来访问数据，那么它对软件来说可以相当透明——问题只是为每个 DMA 配置是否执行压缩 / 解压缩。这需要高吞吐量，但在软件已经做了软件流水线的情况下，不一定需要低延迟——而软件本来就应该总是这样做。反正 DMA 本身已经可能有比较高的延迟。

我还没有从经验上量化这种优化相对于面积和功耗成本的影响，所以一家 AI 芯片初创公司在决定是否为此花费芯片面积之前，当然应该做这个量化。据我所知，目前没有公开研究对这个问题做过彻底调查。上面两篇论文是我能找到的最接近的内容，但它们并没有给出做决策所需的完整图景。

我相信，前面章节描述过的 Int7+1 结构化稀疏格式，如果和这种有损 + 无损压缩方法结合起来，会相当有吸引力。这种压缩意味着你可以支持 7 bit，但如果某一层适合更小的非整数 bit 宽，比如 4.3 bit，也可以无缝地缩放到那种更小的 bit 宽。

给硬核读者看的 TurboQuant 怎么样？

TurboQuant 是 Google 的一种 KV 缓存压缩技术。据说它可以在不损失准确率的情况下，让 KV 缓存大小减少约 8 倍。这篇已有论文是通过一篇 Google 博客文章向公众介绍的，当时也被媒体广泛报道。博客文章和媒体报道都非常具有误导性。

那个约 8 倍的基线是每个元素 32 bit，而在这里这并不是一个严肃的比较，所以这个数字没有意义。在这种情况下，8 倍意味着降到 4 bit，而这以前已经有人做过了，所以这并不是新东西。换句话说，媒体报道关注的东西甚至都不是被改进的部分！

如果这篇博客文章的其余内容可以按字面相信，那么 TurboQuant 真正的进展在于：它能在不重新训练、不明显损失准确率，甚至没有校准数据（calibration data）的情况下做到这一点。这确实是一个重要进展。所以它不一定是一个糟糕结果，但真实故事和媒体报道几乎不是一回事。

我对此持怀疑态度，但也许 TurboQuant 就是你应该用来压缩 KV 缓存的方法。你必须自己去弄清楚。如果确实如此，你就需要确保你的芯片能够以内存速度执行 TurboQuant 中并不简单的额外开销。

AI 工作负载的类型——训练、解码/生成、预填充

本章解释 AI 芯片必须支持的三类主要 AI 负载：训练（training）、预填充（prefill）和解码/生成（decode）。从 AI 硬件和软件的需求来看，它们三者都显著不同；因此，制造 AI 芯片的人需要清楚：自己的芯片到底面向这些工作负载中的哪一种。

训练（Training） 这是 AI 模型中的权重被更新、从而让 AI 变得更聪明的阶段。它是 AI 上学。AI 在这里并没有做任何有用的事，就像学生也不做任何有用的事一样；这里的重点是学习如何去做，而不是正在做。就芯片所需功能以及实现训练所需软件而言，训练是最具挑战性的工作负载。构建一颗优秀的训练芯片，并配套相应软件，是最具挑战性的 AI 芯片项目。如果你的硬件和软件做得正确，那么在训练期间，你的芯片应该是计算受限（compute-bound）的，也就是说，瓶颈应该是脉动阵列，而不是内存带宽或其他东西。

推理（Inference） 这是 AI 从训练学校毕业，并可能开始做有用工作的阶段。在推理期间，权重不再更新。从硬件和软件复杂度角度看，为推理制造芯片没有训练那么难。对于 Transformer 模型来说，推理有两种类型：预填充和解码。

推理：预填充（Inference: Prefill） 这是 Transformer 推理的一种特定形式。在这个阶段，AI 读取提供给它的文本或其他数据，但它不会对此说什么；它只是在注意并记住所写的内容。预填充在硬件和软件实现上没有训练那么复杂。如果你的硬件和软件做得正确，那么在预填充期间，你的芯片也应该是计算受限的，也就是说，瓶颈应该是脉动阵列，而不是内存带宽或其他东西。

推理：解码/生成（Inference: Decode） 这是 Transformer 推理的另一种形式。在这个阶段，AI 会“说话”，也就是生成文本、音频、图像或视频等输出。解码通常比预填充更复杂，但没有训练那么复杂。

取决于你的实现方式，解码比预填充和训练需要更多内存带宽。如果你的软件不够优化，而且你的芯片没有太多内存带宽，那么你在解码时会变成内存带宽受限（memory-bound）。你也应该努力让解码变成计算受限，但你未必能做到。即使内存带宽不是问题，解码也可能产生不均衡的矩阵乘法，导致脉动阵列无法获得足够高的利用率。因此，在当前市场上，购买一个解码 token，也就是输出 token 的价格，要比购买一个预填充 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 更重要——它们可能有更高的摊销后利润率。

训练芯片需要哪些推理芯片不需要的东西？

高吞吐量转置（High-throughput transposition） 训练的反向传播过程需要转置权重，而你需要硬件来做这件事，否则它会太慢。推理通常不需要这一点。你可以使用矩阵乘法来做转置，所以它可以在现有的推理脉动阵列上完成，但这并不高效——转置单元（transposition units）所需的功耗和面积远低于一个大型脉动阵列。

看起来推理期间似乎也需要转置硬件，因为注意力公式涉及计算 Q*K^T。但你可以绕开这一点，让你的脉动阵列直接执行这个操作，而不是执行 Q*K。也就是说，总是把行和行相乘，而不是行乘以列。你仍然可以完成 Transformer 以及大多数模型中的所有其他矩阵乘法，因为那些是激活乘以权重，也就是 A*W；而你可以在推理之前离线转置权重，利用的是 A*W = A*(W^T)^T。

有没有脉动阵列做 A*B^T 的先例？有，TPU 上就是这样。TPU 的脉动阵列可以选择执行 A*B 或 A*B^T 中的一种。可以读一下 TPUv2 的硬件论文。

你也许还能在推理中避免使用置换单元（permutation units），但在多头注意力（multihead attention）内部的数据移动上要小心这一点——如果每个 head 的特征维度不是你的向量维度的倍数，比如不等于你的向量维度，那么你可能需要某种方式在这里做置换。卷积（convolutions）如果有一定置换能力以及跨 lane 求和能力，也可能更容易在你的硬件上实现。卷积很难实现好，比矩阵乘法更难。你可能需要，也可能不需要实现卷积。

你也许更愿意通过为推理也配置显著的转置和置换能力，来确保无论模型是什么都不会出问题。但如果你只保持标准 Transformer 模型，并且张量维度都是向量维度的倍数，那么这也许并不是严格必要的。我认为我还是建议配一个置换单元。至于转置单元，我不确定。不过我知道，如果你有它，你会睡得更踏实。

额外的向量计算（Additional vector compute） 训练需要支持各种优化器（optimizers），比如 Adam。这类优化器涉及大量向量计算。这是在推理期间已有工作之外的额外需求，所以与推理相比，训练需要更高比例的向量处理能力。

昂贵的数值格式（Expensive numerics） 所有推理都可以在脉动阵列内部使用 8 bit 或更少 bit 的乘法完成。加法需要更多 bit，但加法更便宜。这会让脉动阵列更便宜。即便如此，如果你的产品也支持 16-bit 数值格式，即使只是用于推理，也会更容易销售，因为有些推理客户就是懒得把模型降到 8 bit。因此，作为一家做推理芯片的公司，你可以在 8 bit 和 16 bit 之间做选择。

但对于训练来说，如果你的芯片只有 8-bit 数值格式，你不太可能把它卖给任何人，所以它必须支持 16 bit。虽然我实际上认为 8 bit 可以、也将会足够用于训练，但这不是行业标准，而且你不能指望客户现在知道怎么处理这一点。

额外的训练算子和融合（Additional training ops and fusions） 为 AI 推理芯片编写软件的大量工作，是实现推理中使用的大量各种算子（ops）。这已经比大多数供应商愿意做的工作更多了，但它是必要的。训练中使用的算子更多，而且它们往往比推理算子更复杂。

更高的软件和硬件灵活性（Increased software and hardware flexibility） 那些训练超大 LLM 的人，也就是你最想把训练芯片卖给的人，是 AI 研究人员。一整队一整队的 AI 研究人员每天都在提出新想法。他们会不断尝试新想法。这要求你提供给那些想做大规模训练的客户的软件和硬件具备很高灵活性。

推理客户不太可能对你的软硬件灵活性有同样高的需求，尽管他们可能仍然会喜欢其中一部分灵活性。另外，有些 AI 研究人员不一定特别关心如何使用你的产品背后的计算机科学。即使他们用一些奇怪方式使用它，它也需要直接能工作。

大型推理客户往往会配备一些工程师，他们更关心计算机科学，而不是 AI 研究；并且他们对“不那么易用但能给公司带来切实收益”的软件有更高容忍度，比如能提高单位成本 token 产出（tokens per dollar）的软件。

你当然不是不能把训练芯片作为第一颗芯片来做。但你真的想这样做吗？我的建议是大概率不要。所以，出于这个原因，我在这份文档里不会对训练展开太多细节。

解码/生成

解码/生成（decode）指的是模型必须“说”点什么的时候——写文本、生成图像，或者任何存在所谓自回归输出（autoregressive output）的场景。解码是更难支持的那类推理。它天然并不太适合脉动阵列，但只要有足够的 AI 研究和软件英雄主义，你仍然可以让它跑起来——所以这就是你必须做的事。

回忆一下，Transformer 包含两类层：注意力（attention）和 FF。这两者在解码时都有问题，但在二者之中，解码注意力（decode attention）问题更大。

解码的问题在于，如果用直接方式实现，它一次只生成一个输出 token。这类似于你写作时一次只写一个词。这叫作自回归（autoregression），意思是模型在决定前一个 token 之前，不能开始处理下一个 token。“token” 可以理解成“词”，虽然 token 并不总是一个词。所以，它必须先看到，也就是回归到自己的输出上，才能继续。这里的 “auto” 表示 self，也就是自己。为了理解为什么这是个问题，可以考虑自回归解码期间注意力是如何工作的：

图示说明：已有的历史 token 会提供 key，新 token 会产生 query，然后这个 query 会和所有历史 key 做比较。

每个之前的 token 都会发布一个 key，这个 key 是一个向量，用来描述那个词是关于什么的。这件事之前已经完成了，所以我们会从所谓的 KV 缓存（KV cache，也就是 key-value cache）中取出这些 key。这里先不用管 value 是什么。

当我们需要处理一个新的 token，比如看起来它会是 “jumped” 时，我们会创建一个 query。query 也是一个向量，它会和所有 key 中的每一个进行比较，也就是做点积。如果你想了解更多 Transformer 注意力的细节，可以去网上读一篇正式的 Transformer 入门文章。这里的问题是，这个过程对应一次矩阵乘法；而由于我们只处理一个 token，所以只有一个 query。这意味着这个矩阵乘法有一个维度不仅很小，实际上等于 1。这会导致脉动阵列几乎没有利用率。所以，如果用朴素方式实现，解码的脉动阵列利用率会非常差，低于 1%。

而且，key 的数量可能很多。假设你已经和你最喜欢的 AI 聊了过去两年。如果你使用的是所谓长上下文窗口（long context window）的 AI，并且这个 AI 正在跟踪你的全部聊天历史，那么它可能必须拉取多达一百万个 key，只为了把它们分别和一个 query 向量进行比较——即使这些历史记录很少真正相关。这会需要大量内存带宽；但与此同时，对从内存读进来的数据要做的事情又非常少，于是整颗芯片都只是坐在那里，等待内存把一百万个 key 送过来。

更糟糕的是，Transformer 有多层，比如 32 层。所以你可怜的 AI 可能必须把这整套事情做 32 次，只为了弄清楚它接下来应该说 “jumped”。然后整个过程又从头开始，用来弄清楚下一个词应该是 “over”。如此继续，一次一个 token。

等等，你可能会想。我好像在哪里听说过 Transformer 会使用一种叫作 batch 的东西，也就是说一颗芯片会同时处理多个对话，而这应该能帮助脉动阵列利用率。不幸的是，batch 对 FF 层很重要，尤其是在解码期间；但它对注意力帮助不大，所以在这里帮不上我们。batch 也许能让你在多个对话之间复用某些权重，但在这里我们讨论的注意力问题中，加载权重并不是问题——加载 key 才是问题，而不同对话有不同的 key，所以无法在这里复用 key。

也许你还听说过 Transformer 使用多个注意力头（attention heads）。这意味着整个过程不是每层发生一次，而是每层最多发生 128 次，每个独立注意力头发生一次，并且每个注意力头都有独立的 key 向量和 query 向量，用来反映这些 token 的不同属性。仅靠这一点本身，并不能改善低效率问题。

不过，如果结合一种叫作分组查询注意力（Grouped Query Attention，GQA）的技术，它就会很有帮助。GQA 的想法是：在 G 个分组中，让多个注意力头复用同一组 key 向量。假设总共有 128 个注意力头，那么每个分组大小是 128/G。每个注意力头只有一个 query 向量，但现在我们把这些 head 合并成了 G 个分组，每组包含 128/G 个注意力头，而且组内使用相同的 key。于是，突然之间，我们就有了 128/G 个 query 在查看同一组 key。

所以，这意味着我们正在看的矩阵乘法维度从 1 增加到了 128/G。这会带来高得多的脉动阵列利用率。它还把内存带宽消耗降低了 128/G 倍，因为我们只需要读取 G 组 key，而不是 128 组 key。

除此之外，由于我们现在会把 key 复用 128/G 次，增大 key 向量和 query 向量的尺寸也变得合理，因为这样可以改善它们的质量，也就是让它们包含更多信息。结果表明，这对应于增大矩阵乘法的 K 维度；你可能还记得，前面章节说过，这也是提高注意力期间脉动阵列利用率的另一个因素。只要我们增加 key 向量尺寸的倍数小于 128/G，那么在内存带宽需求上，我们仍然获得了净收益。所以，GQA 是一个非常重要的优化。

我们还可以做得更好。由于一种叫作推测解码（speculative decoding）的技术，事实证明我们并不真的需要一次只处理 1 个 token。这个想法是：对接下来几个 token 会是什么，做一个便宜但有根据的猜测。例如，你可以尝试猜接下来的 3 个 token。在上面的例子里，我们对接下来内容的猜测可能是 “jumped over the”。既然模型现在已经有了关于接下来词语的猜测，它就可以开始处理 “the” 之后会发生什么，同时也在处理并确认 “jumped”、“over” 和 “the” 是否真的是它想说的内容。

这一切可以在一个步骤中完成，所以现在我们一次处理 4 个输出 token，而不是只处理 1 个。这会让 key 的内存带宽需求和脉动阵列利用率再额外改善 4 倍。

这里的猜测也可能是 “jumped over pineapple”。在这种情况下，“pineapple” 是错的，这意味着我们浪费了处理 “pineapple” 的时间，也浪费了计算第四个 token 应该是什么的时间，因为那个计算假设前一个 token 会是 “pineapple”，但结果它是错的。所以这里是有风险的。只要这些快速猜测经常是正确的，它对内存带宽需求和脉动阵列利用率仍然是有益的。

为什么推测解码具体会有收益？因为我们现在一次处理 4 个 token，并且我们使用 G 个分组，每组有 128/G 个注意力头，假设总共有 128 个注意力头。于是，我们最终得到总共 4 * 128/G 个 query，它们查看同一组 key。所以，我们把内存消耗和脉动阵列利用率又额外改善了 4 倍，总改善倍数是 4 * 128/G。如果 G 是 8，那么总改善就是 64 倍。

结论是：解码仍然可能有很差的脉动阵列利用率，内存需求也仍然可能很高；但解码并不像乍看起来那么糟，因为通过结合推测解码和 GQA，你可以获得例如 64 倍的改善。如果你有一个 256 × 256 的脉动阵列，理想情况下你会希望一路达到 256 倍，因为我们一开始只有 1。但这样大的倍数未必可用，所以你可以看到，在注意力期间，尤其是在注意力解码期间，更大的脉动阵列更难获得良好利用率。

为什么不选择 G=1，从而得到极快的解码，假设总共有 128 个注意力头？问题在于，在这种情况下，每个 token 只有一组 key，而所有注意力头都必须共享这组 key。这对模型质量并不好。换句话说，它往往会让 AI 变笨太多，所以这样不行。因此，我们需要足够大的 G 来保留 AI 的智能，但又不能大到让解码变得太慢。这是一个简化描述，但它抓住了正确思路。

为什么不推测解码 16 个 token，而不只是 4 个？那不是又好 4 倍吗？原因是，我们用来做这些猜测的过程本身是一个很笨的 AI，它比真正的 AI 便宜得多、运行也快得多——否则，这个笨 AI 就不可能提前很多 token。它不太可能正确猜出那个聪明得多的真正 AI 接下来会说的那么多 token。这个笨 AI 很可能在到达 16 个 token 之前就脱口而出 “pineapple”（或者其他错误的东西），然后我们为那些 token 做的计算就浪费了。所以这里同样需要取得平衡。完整细节也非常复杂，比如可以根据猜测置信度改变推测 token 的数量，但这个简化视角给出了正确思路。

给硬核读者看的解码用脉动阵列细节

还有第三种优化，它可能让你获得最后那个 4 倍，从而达到 256 倍，并因此有可能让 256 × 256 脉动阵列也达到 100% 利用率。你可以通过降低所需的 N 维度来做到这一点——参见前面关于非方形脉动阵列的章节。只有在你没有采用“为注意力使用较小脉动阵列”这个想法，并且你也不是内存带宽受限时，这才相关。

所以，看起来虽然很困难，但我们大概仍然可以在解码期间让大型脉动阵列获得高利用率。许多这些技术也有助于降低内存带宽需求。一般来说，任何降低 KV 缓存存储需求的技术，都会对解码内存需求产生类似影响，而这类技术有很多，多到无法在这里列完。

这个领域的通常看法是：解码高度受限于内存带宽，即便你配置了很多内存带宽也是如此。但事实是，这是否成立取决于你把这些优化执行得有多好。例如，稀疏 / 索引式注意力（sparse/indexed attention）可以完全改变这个数学关系——甚至到这样一种程度：对它的研究也许会让 SSD 带宽都足以支持解码（！！！）。

问题在于，如果你不能控制客户的软件和模型，那么他们可能不会在这里做到最优。所以你就需要配置惊人的内存带宽来让他们满意，而 Nvidia GPU 和 Google TPU 就有这种内存带宽，这也是原因之一。因此，你需要配置多少内存带宽是个很难的问题。它取决于软件。

你应该已经注意到这里有一个趋势：关于如何配置 AI 芯片，几乎所有事情都取决于软件。根据软件和 AI 模型本身是否被设置得足够高效，答案可能相差 10 倍或者 100 倍。

和存储不同，并行化不会改变读取 KV 缓存所需的解码内存带宽，除非你可以因此把它全部放进 SRAM。跨 C 颗芯片并行，也许会让每颗芯片每个 token 只需要从内存中拉取 1/C 的字节数；但人们希望这样做会让系统在同样时间内处理 C 倍数量的 token，所以它们相互抵消，对每颗芯片所需内存带宽来说是中性的，因为 C * 1/C = 1。

如果你要销售一颗 AI 芯片，你必须确保所有这些技术都能很好地配合你的硬件，而且你有正确的软件来支持客户使用这些技术。否则，你的芯片和竞争对手相比可能会差出例如 64 倍。这里的描述和真正必须编写的软件相比已经非常简化了，而且这只是你必须提供软件支持的许多方面之一，只有这样人们才能使用你的 AI 硬件。AI 芯片产品的软件侧相当有挑战。

给硬核读者看的 agent 乘法

如果你试用 Grok 的 AI 助手，注意 Grok 和 Groq 是完全不同的 AI 公司，你可能会注意到它说的是 “agents thinking”，复数，而不是 “agent thinking”，单数。这里的想法是有多个 chain-of-thought agent，它们会“看到”（attend to）之前 token 的 key，也会看到彼此的 token。

这是另一个乘数，除了推测解码、GQA 和非方形脉动阵列之外，又多了一个乘法因子。例如，如果有 4 条这样的思维链，你就在注意力的 key 复用上又得到一个 4 倍。同时，你也获得了一个 4 倍的 batch，这个 batch 可以有收益地应用到单个用户上。

agent 乘法难道不是只适用于并行 chain-of-thought 吗？所以我们不能说它是一个直接加速，因为我们并没有加速 AI 给你发送回答的那部分？这是对的，但大型 chain-of-thought 才是真正需要大量 token 的地方，它推动了对更高单用户 token 吞吐的需求。所以，即使只是改善“那一部分”，也仍然很重要。

给硬核读者看的 batch=1 怎么样？

假设你使用 4 倍推测解码和 4 倍 agent 乘法。对单个用户来说，这就是 16 的 batch。这似乎说明我们永远不需要 batch=1 产品。不过，虽然这些因素确实会让更高 batch 的产品更有吸引力，但我不认为它们能完全消除 low batch 的理由。

回忆一下，Groq LPU 面向的是低 batch 下昂贵、低延迟的 token。低 token 延迟的主要意义是获得高单用户 token 吞吐。现在看起来，我们已经在没有 batch=1 操作开销的情况下获得了这种高单用户吞吐，或者至少获得了 16 倍。那么 LPU 就没用了？我听过有人提出这个论点，但它并不完全正确。如果你对单个用户有 16 倍 batch，那么你可以买 16 倍数量的 LPU，也同样让 LPU 获得 16 倍的单用户 token 吞吐。所以，对于高单用户吞吐下的昂贵 token，LPU 仍然领先。

不过，这仍然会改变 LPU 的吸引力。首先，你现在需要买更多 LPU 才能收获这个收益，而你本来就已经需要买很多 LPU 才能运行一个模型。另外，很多时候会存在某个“足够快”的单用户 token 生成速率。这里的 16 倍有可能把一种更经济的方法推到“足够快”的范围内；一旦到了这个范围，购买 LPU 就变得没必要了。此外，随着更高的单用户 batch，替代方案的经济效率会提高，因为你不再需要那么多带宽；这样，那些昂贵的 LPU token 相对来说会变得更加昂贵。

这些观点也适用于端侧 / 嵌入式（on-device / embedded）使用场景。人们经常说这些场景需要 batch=1，因为一个嵌入式设备通常一次只有 1 个用户。这不同于数据中心应用，数据中心同时服务许多用户，因此可以通过这种方式获得 batch。根据本节中的技术，我们可以看到，即使在嵌入式用例中，你可用的 batch 也可能比你想象得更多。对于文本解码来说，嵌入式可能仍然是 low batch，但即使只是达到 batch=4，也比 batch=1 高效得多。

另一个需要记住的点是：这里讨论的是文本解码。对于图像，通常使用扩散模型（diffusion models），而扩散模型有很多 batch。我不会在这里解释扩散模型，你可以自己查。也有用于文本的扩散模型，不过它们往往效果没那么好。尽管如此，如果你发现可以使用文本扩散模型，那么那也是解决 low-batch 文本解码困难的另一种方式。

预填充

预填充（prefill）本质上和解码做的是同一件事，只不过我们读取的是已经写好的 token。所以，AI 是在理解已经存在的文本（或其他数据），也就是“预先填充 KV 缓存（(pre)filling the KV cache）”，而不是在生成文本。

这一点很重要，因为这意味着我们可以同时处理文本中的所有词，最多到芯片内存限制为止，而不是一个词一个词地处理。因此，图像会变成大概这样的情况：

图示说明：和解码时只有一个 query 向量不同，预填充时可以同时处理大量已有 token，因此会有许多 query 向量一起参与计算。

这里你可以看到，之前我们只有一个 query 向量需要处理，而现在我们可以同时处理所有 token，从而得到和 token 数量一样多的 query 向量。这很像推测解码，只不过我们已经确定知道接下来的 token 是什么，所以没有任何推测成分。输入可以是一整页网页，包含数千个 token。

因此，这会显著改善内存带宽需求，因为每次从内存中读取一个 key，或者第一次生成一个 key，这个 key 现在都可能被使用数千次。所以，我们不再有解码期间“query 向量太少”的大问题。

请注意，充足 query 向量带来的同样收益也适用于训练，因为训练同样发生在已经生成好的文本（或其他数据）上。解码的问题是推理中特有的问题。

预填充期间，我们仍然可以从 GQA 中受益。对于预填充来说，GQA 不会帮助我们最担心的那个矩阵乘法 N 维度，因为这个维度已经足够大了——我们有足够多的 query 向量，能够在预填充期间让这个维度上的脉动阵列达到满利用率。

GQA 在预填充期间也会帮助降低内存带宽需求，因为使用 GQA 时要读取的 key 更少；但和解码相比，预填充不太可能那么需要改善内存带宽，所以这一点对预填充没有对解码那么重要。

不过，注意力还有一个 K 维度，前面章节提到过，这个维度也必须足够大，才能填满脉动阵列并获得满利用率。这就是 GQA 在预填充期间最有帮助的地方。我们在解码章节中提到，GQA 经常会和 key 向量、query 向量宽度的增加结合使用。换句话说，我们正在增加注意力操作的向量宽度，也就是直接增加 K 维度，从而帮助在预填充注意力期间获得较高的、甚至可能是 100% 的脉动阵列利用率。所以，GQA 对预填充也很重要。

解码-预填充的（管理式）聚合与解耦

默认聚合（Default aggregation） 解码和预填充的默认方式，就是做最显然的事情：在同一颗芯片上，使用相同 kernel 同时做预填充和解码。从数学上看，预填充和解码之间差别很小；它们只是每个对话一次处理的 token 数量有巨大差异，比如 1–4 对比 256。所以，它们可以被组合起来同时执行。这样一来，解码与预填充的比例就是可变且不受管理的——一颗芯片可能任意地主要在做预填充，而另一颗芯片可能主要在做解码。

解耦（Disaggregation） 解码-预填充解耦（decode-prefill disaggregation）指的是在不同芯片上分别做预填充和解码。它甚至可以是不同类型的芯片，当然不一定必须如此，因为解码可能需要更高的内存带宽。这要求把 KV 缓存从它被预填充的地方，传输到它将被用于解码的地方。这样做的好处是，可以使用专门针对预填充或专门针对解码的软件，甚至硬件。

解耦还能帮助解决另一个问题：如果混合预填充和解码，有时芯片可能会主要在做预填充，因为一个预填充对话 / 序列在每一步中可能要处理多得多的 token，所以它每一步处理时间可能长得多。这可能会增加解码的 token 延迟，因为那个单独的解码 token 可能要等待成千上万个预填充 token 完成，才能再次运行。我见过有人把这称为解耦的一个好处，不过如果把预填充和解码放进 batch 的系统在创建 batch 时能够智能地考虑这一点，那么这其实才会真正成为问题。

管理式聚合（Managed aggregation） 解码-预填充管理式聚合（decode-prefill managed aggregation）指的是按一个受管理的比例混合解码和预填充。当解码是内存受限，也就是等待内存，而预填充是计算受限，也就是等待脉动阵列时，这样做是有益的。如果你有意把这两种工作负载按正确比例组合起来，就可以达到一个完全平衡的场景：内存带宽和脉动阵列都被充分利用，并且二者都不需要等待对方。相比于先只做其中一种、再只做另一种，这会同时改善延迟和吞吐量：优先 token 先执行时延迟更好，所有 token 的总体吞吐量也更好。

“管理式聚合”是我刚刚编出来的一个术语，不过在实现上，这和一种显然叫作 “chunked prefill” 的东西非常相似，甚至可能是同一个东西。管理式聚合的含义是：如果你预期会使用这种技术，那么你的芯片不需要配置那么多内存带宽。它并不只是为了防止芯片被预填充工作淹没。

达到平衡（Reaching balance） 在你的 AI 系统中，解码和预填充的总体比例未必能达到完美平衡。在这种情况下，不可能让所有芯片都获得完全平衡的工作负载。不过，管理式聚合仍然可以避免一些情况，比如一颗芯片主要做预填充（向一个方向不均衡），而旁边那颗芯片主要做解码（向另一个方向不均衡）。

公司可以通过定价方式大致实现平衡，至少对离线 token 是这样。这里的离线指的是“答案什么时候给你就什么时候给你，不一定非得现在立刻给”。许多公司已经把预填充，也就是输入 token 的价格，定得远低于解码，也就是输出 token 的价格。

另外，如果你有一个无限来源的解码或预填充工作，而且这些工作是你自己生成的，自动生成的工作可以包括例如挖掘你的 AI 中不一致的回答，比如问题取反时答案也应该翻转，那么你可以用它来平衡系统。根据当前时刻的平衡情况，多做或少做其中一种工作。这也有助于确保你的系统始终被充分利用。

另一个选择是，除了预填充之外，还把解码和训练结合起来，因为训练通常也倾向于计算受限，虽然不一定像预填充那样明显。需要记住的重点是：即使不可能达到完美平衡，管理式聚合仍然是有益的。

管理式聚合会带来一个最低所需的预填充与解码比例，以确保你始终是计算受限，而不是内存受限。你还可以通过把注意力和 FF 重叠执行，让这个比例要求更宽松一些。FF 很可能有一些空闲内存带宽；通过把它们重叠起来，你可以把这些带宽提供给注意力使用。

如果每颗芯片上有 2 个或更多核心，这种方式效果最好。在 GPU 上，这些核心由于某种原因被称为 SM。如果只有 1 个核心，而你试图把注意力和 FF 同时复用到它上面，那么这两者会竞争缓存空间，在这种情况下未必有好处。为了启用这种重叠，你要么需要把注意力和 FF 放在流水线的不同阶段中进行流水线化，要么需要修改模型，让注意力和 FF 并行发生，而不是一个接一个发生。后一种技术还会把 token 延迟减半。

混合系统也是可能的。在这种情况下，你的主服务器集群会被设置为接近完美平衡，但你还会有一小组只做预填充或只做解码的芯片，取决于你有时哪一种会更多，用来处理不均衡的溢出工作。

如果你构建的芯片本来就是面向这种管理式聚合的，那么你甚至可以给它配置比其他情况下更少的内存带宽，从而降低生产成本。这是一个大胆的决定，也就是有风险的决定，因为也许不是所有客户都会使用管理式聚合；如果他们不用，就可能因为你的这个选择而看到更差性能。这又是一个例子：不够成熟的客户所需要的产品，和更成熟客户真正适合的产品不同，而且前者更昂贵。

昂贵的 HBM 存储容量

AI 芯片会配备大量昂贵的 RAM，这种 RAM 叫作 HBM。为什么会这样？这真的有必要吗？这就是本章的主题。

权重和 KV 缓存为什么导致需要大容量、昂贵的 HBM 内存

HBM 是一种非常昂贵的 RAM，也就是计算机内存。估算数字各不相同，因为公司不会告诉你它们为 HBM 付了多少钱，但按每 GB 存储容量计算，HBM 可能比普通 RAM 贵 3 倍。普通 RAM 本身在大规模使用时就已经很贵了。

HBM 提供高带宽，而且如果按每 GB/s 带宽的美元成本来算，HBM 并没有那么贵。真正非常昂贵的是 1 GB 的 HBM 存储容量。所以你可能会以为，AI 芯片会配备较小容量的 HBM，同时仍然提供高带宽。今天是这样吗？完全不是。相反，现代 AI 芯片上巨大而昂贵的 HBM 存储容量正在变得越来越大。

Nvidia 已经宣布，它们的 Blackwell Ultra GPU 将配备惊人的 288 GB 特别昂贵的 HBM，也就是 HBM3e。这里的 “e” 并不代表 “expensive”，但你这样理解也没问题。[^2] 作为客户，你也许能理解为什么 288 GB 内存会让人兴奋——直到你想起来，每一个 GB 都是你自己在付钱。

如果昂贵的 HBM 这么贵，为什么 AI 芯片还要配这么多？主要是为了存储大型 KV 缓存和模型权重。什么是 KV 缓存？

在前面关于解码和预填充的章节里，我们看到了很多 key 向量。这些 key 向量被存储在一种叫作 KV 缓存（KV cache）的数据结构中，KV 是 key-value cache 的缩写。先不用管 value 是什么——value 向量是和每个 key 向量关联的某个向量，用于一些我们这里不会展开的目的。我们会继续重点关注 key。

KV 缓存的一个麻烦之处是，正如我们在解码章节中看到的，检查它，也就是把它从内存中取出来，可能需要大量内存带宽。KV 缓存还有另一个相关问题：存储它可能需要大量内存容量。为什么会这样？

问题在于，我们需要存储多少 key 和 value，以及它们会占多少空间。下面是 KV 缓存占用空间的公式，公式中的每个量会在后面解释：

2 * batch * layers * seq_length * attention_head_groups *vector_width * element_size * idle_magnification * pipeline_factor

结果会表明，这可能是非常大的存储量，从而推动对大容量、昂贵 HBM 内存的需求。

Value 向量 ×2 这个公式一开始乘以 2，是为了计入 value 向量。这是因为 value 向量和 key 向量数量一样多，而且 value 向量占用的空间和 key 向量一样大。你可以用某种 GQA 做法让这一点不成立，但通常情况就是这样。我没有解释 value 向量是什么，而这对我们当前目的并不重要；但它们确实必须被存储，所以我们必须乘以 2。

好吧，我也可以稍微解释一下 value 向量：value 向量是 key 用来查找的数据。所以，如果 KV 缓存是一本电话簿——虽然它不是——那么 key 向量就是人的名字，而 value 向量就是电话号码。一本电话簿只有同时有这两者才有用。所以，KV 缓存包含的是关于过去 token 的、带 key 的、可搜索的信息，也就是 value。

层数 L=32 你可能还记得，一个 Transformer 有多层 L，例如 L=32。非常大的模型可以有明显多于 32 的层数，比如约 100 层，但这次计算我们使用 32。32 层已经不是一个很小的层数。每一层都有自己的 key，所以我们需要把 key 的数量乘以 L=32。

序列长度 N=1,000,000 我们的 AI 能记住多远以前的内容？这叫作序列长度 N，或者注意力窗口长度 N。它可以高达 100 万。所以，为了理解 KV 值存储的影响，我们就假设 N = 1,000,000。

请注意，可以让某些注意力头使用小得多的 N，这被称为局部注意力（local attention）。在这种情况下，重要的是平均序列长度 N，而不是最大值。有时，人们也会对部分或全部注意力使用完全不同的注意力方法，比如线性注意力（linear attention）。在那种情况下，数学关系就会和这里的公式不同。

Batch B=64 一个 batch 为 B，意思是 AI 同时处理 B 个独立对话。为了在 Transformer 的 FF 层解码期间充分利用脉动阵列，需要 batch。我们这里一直重点讨论注意力层，而不是 FF，但 FF 也是需要的。对于一个 256 × 256 的脉动阵列来说，你会希望有 batch=256，才能在解码期间让脉动阵列保持忙碌。

不过，如果你使用推测解码，比如一次处理 4 个 token，也就是推测 3 个 token，那么你可以把 batch 降到 B=64，把那 4 倍收益从推测解码中获得，而不是从 batch 中获得。所以，我们需要把 key 的数量乘以例如 B=64。

请注意，在预填充和训练期间，我们可能不需要使用 batch，或者只需要很小的 batch，因为在预填充和训练期间，序列长度可以承担 batch 对 FF 层起到的全部作用。但在解码期间不是这样。然而，一个只做预填充或只做训练的 AI 助手永远不会做任何事或说任何话，也就是不会输出 token。所以我们需要解码，因此也需要 batch。训练和预填充也可能需要 batch，但只有在序列长度很小时才需要。

注意力头分组 G=8 有多少个注意力头？对于一个大模型来说，例如可以有 128 个。如果没有 GQA，每个 head 都会有自己的一组 key，所以我们需要乘以 128。不过，有了 GQA，这些 head 会被分成共享 key 的若干组。组数可能少到 8 组，所以我们只需要乘以 G=8。所以你可以看到，GQA 对于降低 KV 缓存所需存储量也很重要。

GQA 可以被实现成 value 仍然各自不同，这会增加 KV 缓存存储量。对于这里的计算，我们假设 value 在一个组内是共享的。

向量宽度 W=128 一个 key 向量有多宽，也就是它有多少个元素，或者说向量中有多少个单独的数字？对于一个大模型来说，可能有 W=128。我们就用这个数。

字节数 S=1 存储 key 向量中的一个元素，也就是一个数字，需要多少字节 S？许多 LLM 使用 16-bit KV 缓存，因为这是最容易的。但我们假设这个 LLM 的创建者和/或部署团队付出了努力，把它降到了 8 bit——任何严肃场景都确实应该这么做。所以每个元素是 S=1 byte，因为 1 byte 有 8 bit。

使用 4 bit 也是可能的，也有研究在探索使用更少 bit，但这可能会降低模型准确率，也就是让 AI 变笨。也许随着 AI 方法不断改进，最终整个行业会把这里降到 4 bit 并作为标准做法，但今天还不是这样。不过可以参见前面关于压缩的章节——压缩在这里会有帮助。

空闲放大 M=1.0 空闲放大（idle magnification）是我刚刚发明的一个术语，我不知道有没有标准术语。它考虑的是这样一个事实：昂贵的 HBM 内存中可能存放着比当前正在计算的对话多得多的对话。所以它们当前是空闲的。

例如，当你花时间阅读 AI 助手的回复时，就可能发生这种情况。在这段时间里，AI 正在等待你的回应。等待期间，这个对话没有被处理，因为没有任何事情正在发生。不过，一旦你回应了，你会希望 AI 立刻或者至少很快给出回应。如果这个对话和它的 KV 缓存已经在内存中准备好了，那么 AI 就可以很快继续处理它。如果它存储在某个 SSD 上，那么在 AI 开始之前，需要花一点时间把它读回内存，这会让速度变慢。

如果这真的是问题，我建议换更快的 SSD。但把空闲对话存放在 HBM 中，确实是一种 AI 助手可能采用的技术。如果内存中只有一半对话当前正在被处理，那么这就代表 KV 缓存有一个 M=2 的空闲放大因子。我们这里假设 M=1，也就是没有放大。

说清楚一点，当你的 AI 助手等了一分钟才回复，并说它正在“思考”时，这不太可能是因为它正在把 KV 缓存读入内存，因为那不应该花那么久。它真正做的事情，是在对自己说话，也就是生成 token，即解码；同时还会查看网站等数据，也就是预填充，以弄清楚如何最好地回答你的问题。

由于 HBM 碎片化（fragmentation），你也可能出现有效容量降低。你也可以把这看成一种空闲放大，尽管很难给它一个具体数字。

完整流水线 P=1.0 完整流水线（full pipelining）可能会让 KV 缓存内存使用量最多增加 4 倍——参见后面关于并行化的章节。对于这个计算，我们假设没有这样做；不过你很可能应该这么做，因为它对 KV 缓存存储需求不好，但对计算利用率和网络带宽需求非常有利。

那么，在这些数字下，KV 缓存有多大？

2 * B * L * N * G * W * S * M * P= 2 * 64 * 32 * 1,000,000 * 8 * 128 * 1 byte= 4,194 GB

这是大量昂贵的 HBM！像这样的大型 KV 缓存，是 AI 芯片上放置大量昂贵 HBM 内存的重要驱动因素。公平地说，N=1,000,000 的序列长度确实很大，但这是 AI 助手公司正在提供的真实产品，所以这个计算得到的是一个真实相关的数字。

近年的 Nvidia GPU 和 Google TPU 大约有 100 到 200 GB 的昂贵 HBM 内存。它们有这么多昂贵 HBM 内存，部分原因就是需要存储大型 KV 缓存。不过我们可以看到，4,194 GB 不能被放进只有 200 GB 的内存里，那这件事要怎么工作？

答案是，LLM 会跨许多 AI 芯片并行化，KV 缓存也会分布在许多芯片上。所以问题不在于单颗芯片上能找到多少存储，而在于许多芯片合起来总共有多少昂贵 HBM 内存。

例如，你可以把一个大 LLM 的层内部分分布到 32 颗芯片上。在这种情况下，KV 缓存每颗芯片需要：

4194 GB / 32 = 131 GB

所以你可以看到，为什么最终可能会需要每颗芯片 200 GB 的昂贵 HBM。这是一个简化视角，但总体上给出了正确思路。

那么权重呢？LLM 不是有很多权重吗？所以这大概也需要很多存储吧？确实需要，但如果有 1000 GB 的权重，你可以把它分布到所有芯片上。所以，如果你在层内并行化 32 倍，并且跨全部 32 层做流水线并行化，那么你使用的是：

32 * 32 = 1024

颗芯片。如果你把 1000 GB 权重分布到 1024 颗芯片上，那么每颗芯片只需要存储大约 1 GB 权重。所以这并不多。只有在你没有做太多并行化时，权重存储才会成为问题；否则不是问题。

注意，KV 缓存存储需求可以通过层内并行化改善，但不能通过层间并行化，也就是流水线并行化改善。相反，权重存储需求可以通过这两种并行化方式同时降低。关于并行化如何改变每颗芯片上的 KV 缓存和权重存储需求，更多细节见后面的并行化章节。

降低 AI 芯片对大容量内存的需求

Reducing the need for large memories on AI chips

上一节得出的结论是：AI 推理芯片需要大型 HBM 内存，这是因为需要存储大型 KV 缓存，也因为需要为大型 LLM 存储大型权重矩阵。

我想在这里抱怨的一点是，我们在计算中使用了 N=1,000,000。拥有这么长注意力窗口的 AI 助手确实是正在销售的真实产品，所以它需要被支持；但这件事难道不能用更经济的方式实现吗？

事实上，有很多方法可以在模型准确率损失很小的情况下解决这个问题，也就是把注意力转变成检索（retrieval）。这里有一篇相关论文，你也可以找到其他论文。我也在 LinkedIn 上发过关于这件事的内容。最终，我认为我们需要这样的方案：最近的历史，例如过去 1024 个 token，可能仍然使用稠密注意力；但对更早的注意力历史，则转向稀疏检索（sparse retrieval）。这些更早的历史可以存储在一个连接的 SSD 上，甚至也许可以存储在数据中心中更远的地方。

通过这种方式，你可以把内存需求从 N=1,000,000 降到 N=1,000，也就是提升 1000 倍。你仍然需要买一个大 SSD，但按每 GB 计算，SSD 比 HBM 便宜得多。所以，通过把 N 降到 1000，我们可以从 4,194 GB 降到 4 GB。还有其他优化 KV 缓存的方法，所以如果你也应用这些思路，那么甚至可以得到比这更小的 KV 缓存。而且这还是在通过层内并行分布 KV 缓存之前。那为什么这件事没有发生呢？

嗯，我为什么认为这件事没有发生？难道公司不会使用这些用于注意力的稀疏检索思路吗？也许有些公司正在用，但事实是，最近的 Nvidia GPU 和 Google TPU 正在配备极大的 HBM，而给出的说法是：这是为了存储大型 KV 缓存。所以我认为，AI 行业正在不必要地造成当前全球 RAM 短缺。毕竟我们其实并不真的需要所有那些昂贵 HBM。

如果你是一家 AI 芯片初创公司，这似乎是一个值得弄清楚的相关问题。我的强烈怀疑是：你可以做一颗只有 4 GB 内存的 32-bit 芯片，因为 32-bit 指针可以寻址 2^32 B = 4 GB 的值，而不是像 Google 和 Nvidia 现在销售的那样配备 100–300 GB 内存。如果你把软件做得完全正确，并且进行大规模并行化，那么你的便宜芯片会和这些超级昂贵的行业领先芯片一样有能力。

“你所需要做的全部事情”就是：对注意力使用逐层检索，也就是只对最近历史之外的部分做检索；并且进行足够多的并行化，让你能够以分布式方式存储一个非常大的 LLM 的权重。

当前一代 HBM 不能以每个堆栈只有 4 GB 这样小的数量购买，至少从公开信息看起来是这样，当然如果你去问，也许有可能。一旦把 interposer 和其他成本考虑进去，从每美元 GB/s 的角度看，HBM 不一定是最经济的内存技术。是的，这一点带感叹号。所以这件事未必像你想的那样是个大问题。不过，离开 HBM 可能更适合那些通过快速板上网络连接的小芯片，这也要求更严格地最小化每颗芯片的成本。

我确实看到这里有两个问题。

第一个问题是，这种方式会让加载大模型所需的最小芯片数量变得很大。一个 1 TB 模型——这是一个非常大的模型——如果每颗芯片只有 4 GB 内存，那么至少需要跨 256 颗芯片并行化才能运行。你可以看到这有多不方便，尤其是如果这 256 颗芯片中每一颗本身都很贵的话。不过，如果你遵循我的想法，它们就不会很贵。

但一个 1 TB 模型确实非常大，而且无论如何，为了速度，也就是 token 延迟，你都会希望做层内并行化。对于严肃用途，比如知名 AI 助手，大的最小并行化要求并不是什么大问题，因为知名 AI 助手无论如何都需要大量吞吐量，所以最终安装规模本来就会很大。当然，你也可以不用我这种极端的 4 GB，而是每颗芯片使用 16 GB 或 32 GB——我把这个想法推到极限，是为了突出什么是可能的。

还有第二个问题。上面我默默假设的是：写软件、部署模型、确保 SSD 被使用、并把所有东西都精确设置到即使 HBM 很小也能刚好工作的人，是我，或者像我这样的人。那些对 LLM 部署非常成熟的客户，也许可以自己做到这一点。

不过，你的许多客户很可能并没有那么擅长部署 AI 模型，但他们仍然可能愿意花几百万美元购买你的芯片。所以你最终会在硬件上做非常昂贵的事情，以服务那些其实如果能在软件中更高效地做事，就不真正需要这些昂贵东西的客户。但他们不会在软件中更高效地做事，所以如果他们做了会怎样，其实并不重要。

我怀疑，这就是最近 Google TPU 和 Nvidia GPU 上巨大而昂贵 HBM 的真正解释。我怀疑，这是为了让不那么成熟的客户，无论是内部客户还是外部客户，更容易使用这些产品，尽管这并不是完全必要的。而我认为，这反过来也部分是因为这些产品的利润率很高。让一颗高利润 AI 芯片的生产成本变得更贵，对最终利润没有那么大影响；但在这种情况下，获取更多客户确实有影响。

所以你可以看到，也许有一个机会：一家芯片公司可以更深入地帮助客户，以一种标准化方式得到高效方案，并把这种方案推广到所有客户身上。也许还有另一个机会：一家芯片初创公司可以转向直接以低价格提供 AI 助手服务，这样它们就可以成为自己芯片的客户，并用自己优秀的软件降低成本。这似乎是 Groq 正在用它们的 LPU 芯片追求的策略，至少“使用自己的硬件”这一部分是这样。这两条路看起来都很难，但我觉得都很有意思。

这里可能有一个反对意见：虽然 HBM 很贵，但 Google 和 Nvidia 正在销售/出租的芯片本身也贵得离谱。非常、非常贵。所以，如果 HBM 只是可怜客户最终支付成本中的例如 20%，这只是猜测，大公司支付的昂贵 HBM 价格并不公开，虽然我见过估算说昂贵 HBM 可能占成本的 50%，那么为了节省 20% 生产成本而制造一个更不方便、HBM 更少的产品，也许没有太大意义。所以也许 HBM 成本没那么重要，也不是当前应该重点关注的地方？[^3]

我认为，这类论证经常会以一种不明显的方式出错。为什么？因为我相信你会在这些产品的很多地方都发现这种膨胀（bloat）。所以，如果你同时把所有这些膨胀都去掉，影响会远大于区区 20%；但如果你单独去掉每一种膨胀，它看起来就“只有”20%。

这就是为什么所有膨胀都会留在原地：你不能决定去掉某一种膨胀，因为它会被解释成只占其他所有膨胀中的一个不大比例（比如 20%），而那些其他膨胀也会因为同样理由无法被去掉。膨胀会保护膨胀。

例如，为什么 TPU（或 GPU）需要一台 host，也就是一台单独连接 TPU/GPU 的计算机？为什么 TPU 不能直接有一些片上 CPU 核心（Google 已经在做优秀的 CPU），然后 TPU 本身就是那台计算机？我碰巧知道这些 host computer 相当昂贵。但也许 host 只占总体成本的 10%（我不知道具体数字），那为什么要改呢？诸如此类。

我怀疑，如果 Google 以实际单位成本 token 产出为目标，极端且创造性地降低 TCO（total cost of ownership），同时提供优秀软件并降低价格，那么它们可能在一两年内就在 AI 硬件市场取得巨大进展。我不认为它们会这样做，部分原因是它们没有必要；我当然也不认为 Nvidia 会这样做。所以，市场上就留下了机会，让其他玩家成为取代 Nvidia 的全球最大公司。

[^2]: 原文脚注：HBM3e 代表 “High Bandwidth Memory 3 Extended”。也许它是 “extended price tag”，因为需要更多数字来写价格？

[^3]: 原文脚注：这个可怜客户不是得很有钱，才付得起这些价格吗？嗯，是的，这个客户原本很有钱，但付完这些价格之后也许就不再有钱了。所以也许应该写成 “now poor”，也就是“现在变穷的客户”，而不是 “poor”。

持续创作高质量文章，欢迎关注公众号。

• 拓驰猎头2010年成立，总部在深圳，全职猎头80人+，全球Office10个+，专注在芯片半导体、AI大模型、智能硬件、新能源、汽车5大赛道30个+细分行业。

• Jason，拓驰猎头CEO，20年猎头经验。

  长期招募合伙人、猎头顾问、SOHO顾问等，欢迎+V：hunterjason