/AI 芯片的软件和硬件设计/

翻译自 Bjarke Hammersholt Roune 的博文

1

关于本文档

他原本计划创办一家 AI 芯片公司，而这份文档记录了他对于如何设计 AI 芯片的思考，涵盖软件和硬件两个方面。文档主要面向未来负责 AI 芯片项目或参与相关工作的人员，介绍了开展这类项目所需要了解的部分基础知识。

他认为，其他人可能也会对 AI 芯片领域感兴趣，因此文档的大部分内容都尽量写得通俗易懂，即使是业余爱好者也能够理解。其中部分内容技术性较强，需要更深入的背景知识，这些内容会以粗体标注为“极客内容（for nerds）”。

文档中还穿插了他在硅谷工作期间的一些轶事，希望这些经历既具有一定趣味性，也能够帮助读者理解相关问题。这些内容都有明确标记，如果读者对此不感兴趣，可以直接跳过。

他了解到，目前已有一些 AI 芯片公司将这份文档作为新员工入职培训材料的一部分。他认为这种做法很合适，因为文档中的许多内容在其他渠道并不容易系统获取。

2

作者为何有资格撰写本文

他曾在 Google 参与 TPU 相关工作，也曾在 Nvidia 参与 GPU 相关工作，并担任技术领导职位。此外，他还曾就职于 Amazon、Facebook 等从事 AI 软硬件研发的公司。

在 Google 期间，他曾担任 TPUv3 的软件技术负责人，并深度参与了当时 TPU 硬件功能的选择与设计。AI 软硬件系统正是他的专业领域。

3

作者为何撰写本文

最终，由于个人原因，他决定不再创办一家 AI 芯片公司。此前，他希望创办 AI 芯片公司的部分原因在于，他认为这个领域存在更好的做事方式，这不仅可能为投资者带来收益，也能够让客户从中受益。

如今，他希望通过撰写这份文档并将其免费公开，来实现一部分原本希望通过创办公司所产生的积极影响。

4

摘要

本文的核心观点是，

要基于这一思路设计出具有竞争力的 AI 芯片，还需要处理大量具体的架构与工程细节，而逐一讨论这些问题也是本文的重要内容。不过，本文并不局限于这一特定架构思路；无论采用何种设计方案，其中涉及的底层约束和关键影响因素基本是相同的。

本文讨论中最突出的矛盾，一方面是如何最大化单位成本所能获得的 token 吞吐量，另一方面则是是否需要为 AI 芯片配置大量昂贵且实际并非必需的功能和资源，只为了让客户获得一种与现有平台尽可能相似、符合既有使用习惯的产品。文中的建议建立在一个重要前提之上：芯片团队拥有一支能力非常强的软件团队，能够在不过度堆叠硬件资源的情况下，仍然让客户方便地使用芯片。但现实中，招聘到具备这种能力的软件团队绝非易事，因此这一点必须充分考虑。

对于从事 AI 软件开发的读者，作者还建议参考其关于 AI 编译器测试的文章。由于测试问题已经在该文章中进行了讨论，因此本文不再展开。此外，作者还提出过一些关于“可解释 Transformer”的想法，认为其中也有值得进一步探讨的内容。

5

硬件设计方案（概要）

下面是作者所提出硬件设计方案的概要。他建议读者最后再阅读这一部分，不过如果愿意，也可以先了解整体设计思路。该方案的主要特点是硬件实现相对简单，编程方式像 CPU 一样灵活，同时在单位计算性能和功耗成本方面具有较高的性价比。作者并不认为会有 AI 芯片公司完全照搬这一方案，但建议在评估其他设计时，以这一方案作为参照：对于每一项偏离，都应明确理解为什么这种偏离是合理且必要的。如果作者当初真正开始研发这款芯片，他本人也很可能会随着项目推进不断调整这一设计。

有人可能会质疑：“这不就是在 CPU 旁边加了一个脉动阵列吗？仅凭这样的设计，真的能够与复杂的 GPU 和 AI ASIC 竞争吗？”作者的答案是肯定的。其核心原因在于，芯片的大部分面积都用于缓存和脉动阵列，因此对于标量和向量计算，并没有必要脱离 CPU 式的设计思路。CPU 具有很强的可编程性，而且硬件和软件工程师已经非常熟悉如何设计和使用 CPU，这正是该方案的基本思想。

至于这种设计为何能够成立，需要结合全文才能完整理解。以下首先列出该设计的一些主要特点：

• 设计一个传统 CPU，在此基础上增加更多向量计算单元和向量寄存器，并为每个 CPU 核心连接一个脉动阵列。如果是 Intel 或 AMD，可以考虑采用 x86；否则，应选择一种性能较高、成本较低且能够自主实现或获得授权的处理器架构。

• 配置大容量、由软件管理的 SRAM 缓存，容量达到数 MB。当某个程序未使用全部 SRAM 时，剩余部分还可以作为传统的硬件自动管理缓存使用。

• 支持所有存储器与网络之间的 DMA 传输，包括 SRAM→网络→SRAM 的直接 DMA。同时，为 DMA 提供硬件 8-bit Huffman 压缩与解压缩能力，其处理速度应达到 HBM 或互连链路的带宽水平。

• 无需额外配置主机，因为该芯片本身已经是一颗通用 CPU。应确保芯片能够运行 Linux 以及 TensorFlow、PyTorch 等软件框架。此外，还应配置一些不带脉动阵列的标量 CPU 核心。

• 不需要为单颗芯片配置大量 HBM。模型权重可以通过多芯片协同方式容纳。

• 为芯片配套 SSD，并结合稀疏注意力机制，使 KV Cache 能够存放在 SSD 中。

• 脉动阵列仅支持 7-bit 乘法，而不是 16-bit 乘法；累加结果采用更高位宽，并支持饱和运算。

• 脉动阵列针对右操作数（RHS）支持 1-of-2 结构化稀疏。该方案实现简单，并可将计算吞吐量提升至 2 倍。

• 主要互连网络采用基于铜走线的 1D 或 2D 环面（Torus）拓扑。使用成本较低的定制主板，各主板能够沿一个或两个维度直接拼接，并通过板间铜走线完成连接，从而避免使用昂贵的高速线缆。

• 为该芯片开发 XLA 后端。大多数用户将通过 XLA 使用这款设备。同时，应支持使用标准 C++编写自定义 Kernel 和自定义 XLA 算子，并通过 intrinsics 访问底层硬件能力。

• 该设计同样适用于训练，但初期应优先面向推理，因为推理场景的实现相对更简单。

这种设计避开了 AI 芯片中一些成本高昂但实际并非必需的部分，硬件可以较为直接地通过 RTL 实现，同时其编程方式与普通 CPU 类似。作者不太理解为什么类似的方案没有成为 AI 芯片行业的主流标准，因为目前行业似乎更倾向于探索各种新型的非 CPU 架构。

不过，要充分发挥这款芯片的能力，需要较高水平的软件支持。同时，还需要开展一定的 AI 算法研究，使稀疏索引注意力、1:2 结构化稀疏等技术能够取得良好效果。此外，在本文建议的低 HBM 配置下，如果要运行大型 LLM，系统必须达到较大的最小部署规模。

当然，也可以通过修改设计、增加部分硬件资源来降低上述使用门槛，但这样会牺牲一定的成本效率。作者认为，即使采用这种更高成本的设计，整体方案仍然具有竞争力。至于哪种选择能够带来更高利润，他无法预判；但从大规模部署条件下最大化单位成本 token 吞吐量，同时保持较高可编程性的角度来看，本文提出的方案是他认为更优的选择。

6

关于 Groq LPU 的技术说明（极客内容）

Groq 的 LPU 设计不需要 HBM。那么，它与本文提出的方案是否相同？在模型权重存储方面，两者的思路确实类似：都通过将权重分布到多颗芯片上，避免单颗芯片配置大容量存储，因此都要求较大的最小部署规模。

但在 KV Cache 的处理方式上，两者存在明显区别。Groq 在解码阶段采用低 Batch 方案，因此能够获得较低的 token 延迟，同时显著降低 KV Cache 的容量需求，这也是其无需使用 HBM 存储 KV Cache 的主要原因。不过，低 Batch 意味着 Groq 芯片无法使用需要较高 Batch 才能充分利用的大规模脉动阵列，因此在作者看来，这会降低芯片的功耗效率和面积效率。

Groq 的方案同时会带来很高的存储带宽需求，而 LPU 通过使用 SRAM 代替 HBM 来提供所需的高带宽。相比之下，作者提出的方案采用索引注意力（indexed attention），以显著降低注意力计算的存储带宽需求，从而能够将 KV Cache 卸载到 SSD。这样仍然可以采用较大的 Batch，进而充分利用大规模脉动阵列。

不过，这一方案的前提是稀疏索引注意力能够取得良好的实际效果。如果芯片厂商及其客户无法解决这一问题，就必须增加更多 HBM 容量，这也是目前大多数 AI 芯片所采用的方式。

简而言之，作者认为 Groq 的低 Batch 方案更适合提供成本较高的“高端”低延迟 token 服务，而本文提出的方案则更侧重于提供成本较低、中等延迟的 token 服务。他认可 Groq 方案在特定市场中的价值，但认为其更适合规模较小、对成本不敏感且高度重视低延迟的细分市场；相比之下，本文提出的方案更面向广泛用户和大多数客户的实际需求。