/AI 芯片的软件和硬件设计(15)——权重与 KV Cache 为什么需要大容量 HBM/

翻译自 Bjarke Hammersholt Roune 的博文

/昂贵的 HBM 存储容量/

AI 芯片通常配置大量昂贵的 HBM。那么，为什么需要如此大的 HBM 容量？这些容量真的有必要吗？本章主要讨论这一问题。

1

权重与 KV Cache 为什么需要大容量 HBM

HBM 是一种非常昂贵的 RAM。由于厂商通常不会公开实际采购价格，因此不同估算存在差异，但 HBM 单位容量的成本可能达到普通 RAM 的约 3 倍，而大容量普通 RAM 本身也并不便宜。

HBM 的主要优势是高带宽。如果按照单位带宽成本，即$/GB/s 衡量，HBM 实际上并没有那么昂贵。真正昂贵的是 HBM 的单位存储容量。

因此，一个看似合理的设计思路是：

只配置较小容量的 HBM，同时保持较高带宽。

但目前的 AI 芯片并没有朝这一方向发展。相反，现代 AI 芯片配置的 HBM 容量越来越大。例如，NVIDIA 曾宣布 Blackwell Ultra GPU 将配置高达 288GB 的 HBM3e。

对于客户而言，288GB 显存看起来很有吸引力，但这些昂贵的 HBM 最终都需要由客户承担成本。

那么，为什么 AI 芯片仍然需要如此大的 HBM 容量？

主要原因有两个：

模型权重（Weights）和 KV Cache。

2

KV Cache 为什么会占用大量存储空间

在前面的 Decode 和 Prefill 部分中，已经介绍了大量 Key 向量。

这些 Key 被保存在一个称为KV Cache（Key-Value Cache）的数据结构中。

KV Cache 中的 V 表示 Value。这里暂时不需要深入讨论 Value 的具体计算作用，只需要知道：每一个 Key 通常都对应一个 Value，因此 Key 和 Value 都需要存储。

前文在 Decode 部分已经讨论过 KV Cache 的一个主要问题：

读取 KV Cache 需要大量内存带宽。

除此之外，KV Cache 还存在另一个密切相关的问题：

存储 KV Cache 本身需要大量内存容量。

因此，需要进一步分析到底需要保存多少 Key 和 Value，以及这些数据最终需要占用多少存储空间。

KV Cache 容量可以近似表示为：

2×Batch×Layers×Seq_length×Attention_head_groups×Vector_width×Element_size×Idle_magnification×Pipeline_factor

即：

KV Cache 容量=2×B×L×N×G×W×S×M×P

后面将逐一解释这些参数。按照实际模型配置，这个公式最终可能得到一个非常大的数值，从而直接推动 AI 芯片配置大容量、昂贵的 HBM。

3

Value 向量：×2

公式最前面的：

2×

用于考虑 Value 向量。

原因是，通常情况下，每个 Key 都有一个对应的 Value，而且 Value 向量与 Key 向量占用的空间相同。

因此，如果只计算 Key 所占空间，还需要再乘以 2，才能得到完整 KV Cache 的容量。

GQA 也可以采用 Key 和 Value 数量不同的实现方式，但通常情况下，可以认为二者数量相同。

如果用一个简单的类比来理解，可以把 KV Cache 想象成一本电话簿：

Key≈姓名

Value≈电话号码

姓名用于找到对应的信息，而电话号码才是查找到的实际数据。一本电话簿必须同时保存姓名和电话号码才有意义。

类似地，KV Cache 保存的是过去 Token 对应的、可以通过 Key 进行查找的信息 Value。

4

Transformer 层数：L=32

Transformer 通常包含多个 Layer。

例如，可以假设：

L=32

非常大的模型可能拥有远多于 32 层的结构，例如接近 100 层。不过，为了方便后续计算，这里采用 32 层作为示例。

每一个 Transformer 层都拥有自己的 Key，因此 KV Cache 容量需要乘以层数：

×L=32

即使 32 层对于现代大模型而言并不算特别大，这一参数也已经会明显放大 KV Cache 容量。

5

序列长度：N=1,000,000

AI 能够向前“记住”多少 Token，通常由序列长度（Sequence Length）或Attention Window Length决定。

这里将其记为：

N

对于超长上下文模型，N 可以达到：

N=1,000,000

也就是 100 万个 Token。

为了观察超长上下文对 KV Cache 容量的影响，作者后续计算采用这一数值。

不过，并不是所有 Attention Head 都必须使用相同的 N。

例如，可以让部分 Attention Head 只关注较短的上下文，这称为局部注意力（Local Attention）。此时真正影响 KV Cache 容量的，是所有 Attention Head 对应的平均序列长度，而不是最大序列长度。

此外，如果部分或全部 Attention 采用完全不同的机制，例如线性注意力（Linear Attention），那么 KV Cache 容量的计算方式也会相应改变，不能继续直接使用这里的公式。

6

Batch：B=64

Batch 表示 AI 同时处理多少个相互独立的对话。

这里记为：

B

例如：

B=64

在 Decode 阶段，为了使 Transformer 中的 FF 层充分利用脉动阵列，通常需要一定大小的 Batch。

例如，对于一个：

256×256 脉动阵列

如果完全依靠 Batch 提供并行度，那么可能希望：

B=256

从而使脉动阵列保持较高利用率。

但是，如果采用推测解码，例如一次处理 4 个 Token，其中包含 3 个推测 Token，那么可以从推测解码获得约 4 倍额外并行度。

这样，Batch 就可以相应降低为：

B=256/4=64

因此，后续计算采用：

B=64

这意味着 KV Cache 中需要保存 64 个并发对话的数据，因此 Key 数量也需要乘以 64。

需要注意，在 Prefill 和训练阶段，可能完全不需要 Batch，或者只需要很小的 Batch。

原因是，在 Prefill 和训练中，序列长度本身就可以提供 Decode 阶段通常需要依靠 Batch 获得的并行度。

但是，一个只执行 Prefill 或者只执行训练的 AI 助手永远不会真正生成输出 Token。因此，完整 AI 服务最终仍然需要 Decode，也就需要考虑 Decode 对应的 Batch 和 KV Cache 容量。

7

Attention Head Group：G=8

大型模型可能包含例如：

128 个 Attention Head

如果不采用 GQA，每个 Attention Head 都拥有自己独立的一组 Key。

这样 KV Cache 容量就需要：

×128

但如果采用GQA（Grouped Query Attention），多个 Attention Head 可以共享同一组 Key。

例如，可以将 128 个 Attention Head 划分为：

G=8 组

这样只需要保存 8 组 Key，而不是 128 组。

因此，KV Cache 容量只需要：

×G=8

这也说明，GQA 除了能够改善前文讨论的 Decode 计算效率和内存带宽之外，还能够显著降低KV Cache 存储容量。

某些 GQA 实现仍然可以让不同 Head 拥有独立 Value，这会增加 KV Cache 容量。不过，为了简化后续计算，这里假设同一 Group 中的 Value 也进行共享。

8

向量宽度：W=128

接下来需要考虑每个 Key 向量包含多少个元素。

这里将 Key 向量宽度记为：

W

对于一个大型模型，可以假设：

W=128

也就是说，每个 Key 包含 128 个数值元素。

9

单元素大小：S=1Byte

接下来需要确定 Key 向量中的每一个元素占用多少存储空间，将其记为：

S

许多 LLM 的 KV Cache 仍然使用 16bit 格式，因为这种方式实现起来最简单。

不过，这里假设模型开发者或部署团队进行了进一步优化，将 KV Cache 降低到：

8bit

由于：

8bit=1Byte

因此：

S=1Byte

4bit KV Cache 同样是可行的，目前甚至已经存在低于 4bit 的相关研究，但进一步降低位宽可能影响模型精度。

作者认为，随着 AI 算法不断发展，未来 4bit 可能成为行业常见的 KV Cache 格式，但目前尚未达到这一阶段。

前文讨论的压缩技术同样可以用于进一步降低 KV Cache 容量。

10

空闲放大系数：M=1.0

作者进一步引入了一个自行定义的概念：

Idle Magnification（空闲放大系数）

记为：

M

它用于描述这样一种情况：HBM 中实际保存的对话数量，可能远多于芯片当前正在计算的对话数量。

例如，当用户正在阅读 AI 生成的回答时，AI 正在等待用户继续输入。

此时，这个对话并没有进行任何计算，但它对应的 KV Cache 仍然可能保存在 HBM 中。

这样做的原因是，一旦用户继续输入，AI 就能够立即恢复该对话。

如果 KV Cache 已经被转移到 SSD，那么用户重新输入后，系统还需要先将 KV Cache 从 SSD 重新加载到内存中，可能增加响应延迟。

作者认为，如果这确实成为问题，更合理的方法可能是采用更快的 SSD。不过，将空闲对话继续保存在 HBM 中，确实也是 AI 服务可能采用的一种方案。

例如，如果 HBM 中只有一半对话正在进行计算，而另一半处于等待状态，那么：

M=2

即实际 KV Cache 容量需求被放大 2 倍。

为了简化计算，这里假设：

M=1

也就是不考虑额外的空闲 KV Cache。

作者还特别说明，当 AI 助手在回答前等待一分钟并显示“思考中”时，这通常并不是因为系统正在从 SSD 加载 KV Cache。

这种较长等待更可能来自模型自身生成内部 Token，即 Decode，以及读取网页等外部信息，即 Prefill。

此外，HBM 中的**内存碎片（Fragmentation）**也可能降低实际可用容量。从效果上看，也可以将其视为一种 Idle Magnification，只不过很难给出统一的具体数值。

11

流水线系数：P=1.0

完整流水线可能使 KV Cache 容量需求最高增加约：

4 倍

这一问题将在并行计算章节中进一步讨论。

为了简化当前计算，作者假设没有采用这种完整流水线，因此：

P=1

不过，作者实际上认为应该采用完整流水线。

虽然它会增加 KV Cache 容量需求，但能够显著改善：

• 计算资源利用率；
• 网络带宽需求。

因此，从整体系统角度看，增加 KV Cache 容量可能是值得付出的代价。

12

一个 KV Cache 究竟需要多大容量？

将前述参数代入公式：

KV Cache 容量=2×B×L×N×G×W×S×M×P

取：

B=64

L=32

N=1,000,000

G=8

W=128

S=1Byte

M=1

P=1

则：

2×64×32×1,000,000×8×128×1Byte

得到：

4,194GB

也就是大约：

4.2TB KV Cache

这是一个极其庞大的 HBM 容量需求。

因此，大规模 KV Cache 正是现代 AI 芯片不断配置更多昂贵 HBM 的重要原因之一。

需要说明的是，100 万个 Token 的上下文长度确实非常大，但这并非纯粹为了举例而虚构的参数，目前 AI 服务厂商已经能够提供这种级别的长上下文产品。因此，这一计算所反映的容量问题具有实际意义。

13

单颗芯片显然无法容纳 4.2TB KV Cache

近期 NVIDIA GPU 和 Google TPU 通常配置约 100～200GB 量级的 HBM。

其中一个重要原因就是需要保存大容量 KV Cache。

但即使一颗芯片拥有 200GB HBM，也显然无法保存：

4,194GB

的 KV Cache。

解决方法是：

将 LLM 并行部署到多颗 AI 芯片，并将 KV Cache 分布到多颗芯片上。

因此，真正重要的并不是单颗芯片拥有多少 HBM，而是整个多芯片系统总共拥有多少 HBM 容量。

例如，如果将一个大型 LLM 的层内计算并行分布到 32 颗芯片，那么每颗芯片需要保存的 KV Cache 约为：

4,194GB÷32≈131GB/芯片

这样就可以理解为什么现代 AI 芯片可能需要配置接近：

200GB HBM/芯片

虽然这一计算进行了简化，但能够说明其基本原理。

14

那么模型权重呢？

大型 LLM 同样拥有大量权重，因此权重显然也需要大量存储空间。

例如，假设模型权重总容量达到：

1000GB

但与 KV Cache 不同，权重可以更加充分地利用多种并行方式进行分布。

假设模型采用：

32 路层内并行

同时又对：

32 个 Transformer 层进行流水线并行

那么总共可以使用：

32×32=1024 颗芯片

如果将 1000GB 权重均匀分布到 1024 颗芯片上，那么平均每颗芯片只需要保存：

1000GB÷1024≈1GB 权重

此时，单颗芯片的权重存储压力实际上并不大。

因此，权重容量主要在并行规模较小时成为问题。

这里存在一个非常重要的区别：

KV Cache 容量可以通过层内并行降低，但不能通过层间流水线并行进一步降低；

而：

权重容量既可以通过层内并行降低，也可以通过层间流水线并行降低。

因此，在超大规模多芯片系统中，KV Cache 可能比模型权重更容易成为单颗芯片 HBM 容量的决定因素。

关于不同并行策略如何改变每颗芯片的 KV Cache 和权重容量需求，作者将在后续并行计算章节中进一步讨论。