ARTICLE · 1042771
AI 芯片的软件和硬件设计(16)——降低 AI 芯片对大容量内存的需求
/AI 芯片的软件和硬件设计(16)——降低 AI 芯片对大容量内存的需求/
翻译自 Bjarke Hammersholt Roune 的博文
1
降低 AI 芯片对大容量内存的需求
前一节指出，AI 推理芯片之所以需要配置大容量 HBM，一方面是为了存储庞大的 KV Cache，另一方面是为了保存大型 LLM 的权重矩阵。
作者首先对前一节计算中采用：
N=1,000,000
这一参数提出疑问。
支持 100 万 Token Attention Window 的 AI 助手确实已经成为实际产品，因此 AI 系统需要支持这种场景。但问题在于：是否存在成本更低的方法来实现超长上下文，而不是简单地将全部历史 Token 对应的 KV Cache 长期保存在昂贵的 HBM 中？
作者认为，实际上存在很多方法可以在模型精度损失很小的情况下解决这一问题，其基本思想是：
将 Attention 的一部分转化为 Retrieval。
例如，可以对最近的：
1024 个 Token
继续采用传统的 Dense Attention，因为最近的上下文通常更加重要。
而对于更久远的历史信息，则不再让每个 Token 都参与完整 Attention，而是使用Sparse Retrieval，只检索与当前计算真正相关的历史内容。
这样，更早的 Attention 历史就不一定需要全部保存在 HBM 中，而可以存放在：
SSD
甚至数据中心中距离计算芯片更远的存储设备中。
采用这种方法后，实际需要长期保存在高速内存中的上下文长度可以从：
N=1,000,000
降低到例如：
N=1,000
相当于降低：
1000 倍
虽然系统仍然需要购买较大的 SSD，但 SSD 单位容量的成本远低于 HBM。
按照前一节的示例计算，KV Cache 容量就可以从：
4,194GB
降低到大约：
4GB
而且，这还只是通过降低有效 N 获得的收益。
如果进一步结合其他 KV Cache 优化方法，那么容量还可以继续下降。更重要的是，这一计算甚至还没有考虑通过层内并行（Within-layer Parallelism）将 KV Cache 进一步分布到多颗芯片。
因此，作者提出一个问题：
既然存在这些方法，为什么现代 AI 芯片仍然配置如此庞大的 HBM？
2
是否真的需要如此大的 HBM？
作者认为，一些公司可能已经在使用 Sparse Retrieval 等方法。
但从当前产品来看，近期 NVIDIA GPU 和 Google TPU 仍然配置了非常大的 HBM，而业界对此的一个重要解释就是需要存储大型 KV Cache。
因此，作者认为，AI 行业可能并不真正需要目前如此庞大的昂贵 HBM 容量。
换言之，如果软件和模型算法能够更加充分地优化 KV Cache，那么 AI 芯片的 HBM 容量可能存在明显下降空间。
3
一颗只有 4GB 内存的 AI 芯片？
对于 AI 芯片初创公司而言，作者认为这是一个值得认真研究的问题。
作者甚至提出了一种相当激进的设想：
设计一颗 32bit 芯片，仅配置 4GB 内存。
之所以是 4GB，是因为 32bit 地址可以寻址：
2³²Byte=4GB
相比之下，Google 和 NVIDIA 目前的高端 AI 芯片通常配置约 100～300GB 级别的高带宽内存。
作者推测，如果软件能够做到足够优化，同时系统采用大规模并行，那么这种低成本、小容量内存芯片理论上仍可能具备运行超大模型的能力。
其中需要满足两个重要条件。
第一，对于较久远的 Attention 历史，采用逐层 Retrieval（Per-layer Retrieval），而不是将全部长上下文 KV Cache 保存在高速内存中。
第二，使用足够大规模的并行，将大型 LLM 的权重分布到大量芯片上。
不过，目前一代 HBM 似乎无法以每个 Stack 仅 4GB 这样的小容量采购。至少从公开资料来看如此，实际是否可以通过定制获得则并不确定。
此外，如果考虑 Interposer 等相关成本，HBM 按照：
GB/s/$
衡量时，也不一定始终是最经济的内存技术。
因此，放弃 HBM 并不一定像表面上看起来那样不可行。
不过，这种方案可能更适合：
大量较小芯片+高速板级网络
的系统架构。
而采用这种路线，就需要对单芯片成本进行更加严格的控制。
4
问题一：最低芯片数量会明显增加
作者同时指出，这种设计存在两个明显问题。
第一个问题是：运行一个大型模型所需要的最低芯片数量会非常高。
例如，一个：
1TB 模型
如果每颗芯片只有：
4GB 内存
那么至少需要：
1TB÷4GB≈256 颗芯片
才能仅仅将模型装入系统。
这显然会增加部署的不便，尤其是在每颗芯片本身仍然很昂贵的情况下。
不过，作者认为，1TB 模型本身已经属于非常大的模型。对于大型 AI 助手等高性能应用，即使没有容量限制，为了降低 Token 延迟，通常也会采用较高程度的层内并行。
同时，大型 AI 服务本身需要非常高的吞吐量，因此最终部署规模原本就会很大。
因此，对于这类应用而言，较高的最低并行度未必是无法接受的问题。
而且，4GB 只是作者为了说明设计空间而提出的极端例子。实际产品完全可以采用：
16GB 或 32GB/芯片
作为更折中的选择。
5
问题二：对软件能力要求极高
第二个问题更加现实。
前述设计实际上隐含了一个重要前提：
必须有人能够把整个软件系统精确地优化好。
包括：
正确使用 SSD 保存远期上下文； 正确执行 Retrieval； 正确配置模型； 正确进行大规模并行； 确保有限的高速内存不会成为性能瓶颈。
对于拥有较强 LLM 部署能力的客户，这些优化可能可以自行完成。
但现实中，很多购买 AI 芯片的客户未必拥有如此强的软件和系统能力，即使他们愿意投入数百万美元购买芯片。
因此，芯片厂商最终可能不得不通过增加昂贵的硬件资源，使产品对软件优化不足的客户也能够较容易使用。
从纯技术角度看，这些客户本来并不一定需要如此昂贵的硬件；但如果他们实际上不会进行相应的软件优化，那么这种“理论上不需要”对产品设计就没有实际意义。
作者推测，这可能是近期 Google TPU 和 NVIDIA GPU 配置大容量昂贵 HBM 的一个重要原因：通过增加硬件资源降低客户的软件部署难度。
作者进一步推测，高利润率也可能使这种策略更加容易接受。对于高利润 AI 芯片而言，增加一些制造成本对最终利润的影响可能相对有限，但如果因此能够吸引更多客户，则可能具有商业价值。
6
芯片公司是否应该更多参与客户的软件优化？
由此，作者认为可能存在一种新的商业机会：
芯片厂商不仅销售硬件，还更加深入地帮助客户建立标准化、高效率的软件部署方案。
如果能够针对所有客户形成统一的软件优化方法，就可能通过软件降低客户对昂贵硬件资源的需求。
另一条路线则更加激进：
AI 芯片初创公司直接提供低成本 AI 助手服务。
这样，芯片公司自己就是自己硬件的客户，可以完全控制软件栈，并利用自身的软件能力尽可能降低系统成本。
作者认为，Groq 在 LPU 方面采用的策略至少部分具有这一特点，即使用自己的硬件提供相应服务。
这两条路线都具有较高难度，但作者认为都值得关注。
7
HBM 成本真的值得如此关注吗？
这里还存在一个反对意见。
虽然 HBM 非常昂贵，但 Google 和 NVIDIA 销售或出租的 AI 芯片本身同样非常昂贵。
假设 HBM 只占最终产品成本的：
20%
那么为了节省这 20%的制造成本，却设计一个更难使用、HBM 容量更小的产品，可能并不划算。
作者说明，这里的 20%只是示例，因为大型企业实际采购 HBM 的价格并不公开；一些外部估算甚至认为 HBM 可能占到相关硬件成本的 50%。
因此，一个合理的问题是：
如果 HBM 只是整个昂贵系统中的一部分，是否真的值得重点优化？
8
“每一项都只有 20%”可能导致整体成本膨胀
作者认为，上述逻辑可能存在一个不容易察觉的问题。
原因在于，一款产品中可能同时存在很多类似的成本冗余。
单独观察其中任何一项，都可能得到：
“只占总成本 20%，不值得优化。”
但是，如果多个 20%的冗余同时存在，并且全部被消除，最终成本下降就可能非常明显。
问题在于，每一种冗余都可以用其他冗余作为理由保留下来：
这一项只占整个高成本系统的一小部分，因此没有必要优化；
而其他部分也使用同样的理由。
作者将其概括为：
“冗余保护冗余（Bloat protects bloat）。”
作者举出的另一个例子是 Host。
TPU 和 GPU 通常需要连接到独立 Host 计算机上。
作者提出疑问：为什么不能直接在 TPU 中集成一些 CPU Core，使 TPU 本身成为一台完整计算机？
Google 本身已经能够设计高性能 CPU，因此从技术能力角度看，这并非完全不可想象。
Host 计算机本身也需要相当高的成本。
但如果 Host 只占整个系统成本的例如 10%，设计者又可能认为：
“只有 10%，为什么还要专门改变架构？”
如果 HBM、Host 以及其他资源都采用类似逻辑，那么最终整个系统中就可能积累大量成本。
9
从 TCO 和 Tokens/$重新考虑 AI 硬件
作者推测，如果 Google 未来从更加激进的角度降低 AI 硬件的：
TCO（Total Cost of Ownership，总拥有成本）
并重点优化实际：
Tokens per Dollar
同时继续提供高质量软件并降低价格，那么可能显著增强其 AI 硬件的市场竞争力。
不过，作者认为 Google 目前未必具有足够强的动力这样做，因为 TPU 需求本身已经很高；作者同样不认为 NVIDIA 会主动采取这种极端降成本策略。
因此，作者认为，这可能为其他 AI 芯片厂商留下市场机会。
10
比较 HBM 经济性时必须考虑单芯片性能
作者特别强调，比较两颗芯片的 HBM 经济性时，不能只比较每颗芯片配置了多少 GB HBM。
更加合理的指标应该类似于：
单芯片 HBM 容量/单芯片实际性能
例如，假设设计了一颗芯片：
性能只有 Blackwell 的 1/2
同时：
HBM 容量也是 Blackwell 的 1/2
那么是否意味着 HBM 成本效率提高了一倍？
并不是。
因为为了获得与一颗 Blackwell 相同的系统性能，需要购买：
2 颗这样的芯片
最终购买的 HBM 总容量仍然与一颗 Blackwell 相同。
因此，从系统层面看，HBM 经济性没有任何改善。
这一逻辑不仅适用于 HBM，也适用于芯片上的其他资源。
作者认为，这可能也是 AI 加速器不断采用更大规模芯片的原因之一。
不过，如果一颗芯片内部开始包含多个彼此独立的内存接口，并形成高度非一致的内存空间，那么情况会变得更加复杂。
此时，从某种意义上说，可能只是将多个独立组件封装到一起，然后统一称为“一颗芯片”。
当然，如果这些组件之间具有足够高速的互连，那么这种整合又不仅仅只是名称变化。因此，其真实经济效果需要更加复杂的系统级分析。
/轶事：作者对 Google TPU 产品策略的回顾/
以下内容主要是作者根据其在 Google TPU 团队工作期间的个人经历，对当时 TPU 项目发展方式作出的回顾和评价。
作者表示，在其参与 TPUv1～TPUv4 工作的时期，其长期感受到一种相对保守的产品发展氛围。
从外部看，Google 很早就投入资源开发 TPU，似乎具有较强的前瞻性。
但作者认为，从内部实际情况看，TPU 项目最初更多是为了解决 Google 自身的计算成本问题，而不是以建立一个大型外部 AI 硬件业务为目标。作者无法确认这种情况在其离职之后是否仍然持续，但其根据外部观察认为，这种倾向可能仍然存在。
11
TPU 最初主要解决 Google 内部需求
按照作者的回忆，TPUv1 最初的出发点非常具体：
如果未来 Google 需要大规模使用语音合成与用户进行交互，那么使用当时的硬件运行这些任务成本过高，因此需要开发一种更便宜的专用硬件。
TPUv1 最初并没有在 Google 内部进行非常广泛的部署，不过最终取得的效果明显超过最初目标。
TPUv2 的目标同样比较具体，最初主要针对 Google 内部与：
搜索和广告
相关的两个训练应用，并没有明确追求更广泛的用途。
后来 TPU 逐渐能够被 Google 内部其他团队通用使用，但作者称，当时仍然没有明确目标将 TPU 提供给 Google 之外的客户。
再后来 TPU 进入 Google Cloud，但作者认为，当时为外部客户提供 TPU 所投入的产品化资源仍然有限。
例如，虽然 XLA 在某些情况下能够让程序只需将：
GPU→TPU
修改很少的代码即可迁移，但在 Google Cloud 中，启动和使用 TPU 的方式与 GPU 仍存在差异。
因此，客户实际从 GPU 迁移到 TPU 时，仍然需要进行额外的系统层修改。
作者认为，如果当时 Google 的目标是积极推动客户采用 TPU，那么更合理的方法应该是尽量将 TPU 设计成 GPU 的直接替代方案，从而降低迁移成本。
12
作者认为 TPU 的商业潜力没有被充分利用
作者认为，从今天 AI 硬件市场的规模来看，Google 当时拥有非常巨大的 TPU 商业机会。
但在作者看来，Google 长期主要将 TPU 视为：
降低内部 AI 硬件成本的工具
而不是：
独立的大规模利润来源。
作者认为，TPU 在降低 Google 内部 AI 计算成本方面已经取得成功，但 Google 并没有像 NVIDIA 发展 GPU 业务那样，全力推动 TPU 商业化。
至于 Google 为什么采取这种策略，作者表示自己并不知道，也没有在 Google 工作期间得到明确解释。
作者进一步认为，如果未来 TPU 没有成为 Google 利润的重要组成部分，那么其中至少包含管理层战略选择的影响。
至于 Google 如果采取不同策略是否能够达到 NVIDIA 当前的市场地位，则属于作者个人的反事实判断，并不能通过现有事实直接验证。
13
TPU 成功在一定程度上超出了最初预期
那么，在作者认为 Google 并没有全力推动 TPU 商业化的情况下，为什么 TPU 今天仍然取得了较大成功？
作者将其中一部分原因归结为 Google 拥有大量优秀工程师。
根据作者在公司内部的观察，TPU 团队长期交付了超过最初预期的成果，因此最终形成了一个比公司最初规划更加成功的平台。
作者认为，Google 最初并没有将自身定位为 AI 硬件公司，而主要希望通过 TPU 降低自身 AI 计算基础设施成本。
因此，当 TPU 逐渐展现出巨大的外部商业潜力之后，公司对这一机会的投入在作者看来仍然相对谨慎。
14
作者与管理层的一次交流
作者还回忆，其曾向汇报链中一位级别很高的管理者询问：
为什么 Google 不能像 NVIDIA 发展 GPU 那样，更加积极地发展 TPU？
对方最初回答：
“但 GPU 是 NVIDIA 的主要产品。”
作者则表示，这恰恰是其想表达的问题，即是否应该让 TPU 成为更加重要的业务。
随后，对方以比较强烈的态度终止了这一讨论，并回应：
“你是认真的吗？”
作者表示，这次交流进一步强化了其个人判断，即 TPU 相对保守的发展策略可能并不只是来自其附近的管理层，而可能来自更高层级。
不过，这仍然是作者根据个人经历形成的推断，而不是对 Google 内部决策过程的完整记录。
15
作者对其 TPU 工作经历的总结
作者表示，其曾是 TPU 项目中的重要软件工程师，负责 TPUv3 软件，并长期参与 TPU 软硬件协同设计。
但在其任职期间，从未有人询问其：
如何让 TPU 发展得更加激进，或者如何扩大 TPU 业务。
作者也没有遇到类似：
“如果团队实现某个技术目标，公司就会进一步扩大 TPU 投入。”
这样的明确目标。
根据其个人感受，当时组织内部并不存在非常强烈的 TPU 商业扩张意愿。
作者表示，即使从今天外部观察 Google，其个人仍然没有明显改变这一判断。
16
与 NVIDIA 的对比
作者还曾在 NVIDIA 工作，因此其将两家公司的内部体验进行了对比。
在作者看来，NVIDIA 对于 GPU 业务的推进明显更加积极，而 Google 对于 TPU 则没有表现出同等程度的商业扩张意愿。
不过作者也明确说明，其本人并不是 Google 的 Director 或 Vice President，也没有参与最高层级的相关决策会议。
因此，其无法确认这些会议中究竟发生了什么，只能描述自己从工程团队层面观察到的情况。
17
对 AI 硬件竞争的看法
作者认为，这种情况不仅可能使 Google 错失部分商业机会，也可能影响 AI 硬件市场的竞争程度。
在作者看来，更低成本、更高质量的 AI 服务需要建立在具有更好经济性的 AI 硬件之上。
因此，市场需要更多厂商在 AI 硬件领域与 NVIDIA 展开充分竞争。
作者认为 Google 目前并没有完全承担这一角色，因此其他 AI 硬件公司仍然存在发展机会。
作者也表示，撰写这份文档的一个目的，就是公开分享其对 AI 芯片设计的经验和观点，希望能够帮助更多厂商参与 AI 硬件竞争。
最后，作者表示，其曾就上述个人经历联系 Google 新闻部门寻求评论，但收到的自动回复说明，如果联系者不是媒体成员或 Google 员工，通常不会获得自动邮件之外的进一步回应。作者称，截至其撰写本文时并未收到其他回复。