Cerebras 如何用 59 片 WSE-3T 晶圆部署 3 万亿参数 Kimi K3

我们最常被问到的问题之一，就是如何在 Cerebras 晶圆级芯片上部署大模型。单块 Cerebras WSE-3T 晶圆拥有 44 GB 片上内存，而万亿参数的前沿模型仅权重就需要数 TB 的空间。

本文将解释我们如何将大型前沿模型拆分到多块 WSE-3T 晶圆上，从而以每秒 1000-2000+ token 的速度运行数万亿参数的模型。

我们以开源模型 Kimi K3（约 3 万亿参数）作为示例进行说明。

读完后，你将了解：

• 我们如何将完整模型装入片上内存
• 我们如何将模型映射到多块芯片上
• 我们如何处理大量并发请求

以 Kimi K3 为例拆解

Kimi K3 是一个开源的 3T 参数模型，包含 69 个线性注意力层和 24 个稠密层。

在本示例中，我们采用：

• 假设平均请求大小为 256k
• 4 位权重量化（类似 Hugging Face 实现）
• 8 位 KV 缓存

#1：每一层的大小

为了计算支持 Kimi K3 所需的 Cerebras WSE-3T 晶圆总数，我们先计算每个模型层的大小。

Kimi K3 拥有 3 万亿参数。如果全部 3T 参数均以 4 位存储，则原始参数负载为：

将该负载除以 93 层，得到：

#2：晶圆数量

现在我们已知每层约需 16 GB，就可以算出每块 44 GB 晶圆能容纳多少层。答案取决于层的类型：

• 稠密层需要额外空间用于 KV 缓存，因为每个稠密（全注意力）层都要为每个活跃请求中的每个 token 存储键和值。因此，我们每块晶圆放置一个稠密层，并将剩余 SRAM 留给其 KV 缓存。
• 线性注意力层不需要额外空间用于 KV 缓存。线性注意力层不为每个 token 存储键和值，而是为每个请求维护一个小的、固定大小的状态，与上下文长度无关。因此每块晶圆可容纳两个线性层（2 × 16 GB = 32 GB）。

因此，对于 Kimi K3 的 69 个线性层和 24 个稠密层，我们得到：

这样晶圆总数为 35 + 24 = 59 块晶圆。

#3：将晶圆连接起来

Cerebras 以流水线方式执行模型。Cerebras 编译器将模型层映射到各晶圆区域后，运行时负责规划 token 在晶圆之间流动的路径。为实现超低延迟，晶圆之间只需传递激活值，这可以通过基于以太网的 RDMA 完成。每跳大约 2 微秒。

需要强调的是，由于每块 Cerebras 晶圆拥有 44 GB 片上 SRAM，而每层仅需 16 GB，我们通常能够将模型的完整层直接放在芯片上。

这意味着在大多数情况下，模型权重无需移动，层也不需要跨晶圆拆分。只有激活值会从一块晶圆传到下一块。

#4：如何服务单个请求

现在晶圆已连接在一起、模型也已布局完毕，我们可以把每块晶圆看作流水线上的一个工位：每个工位的模型权重保持不动，请求的中间结果依次流经各层工位。

例如，晶圆 1 运行模型第 1–2 层，并将激活值发送给晶圆 2。晶圆 2 运行第 3–4 层，并将结果传给晶圆 3，依此类推，直到最后一块晶圆生成下一个 token。

#5：如何服务大量请求

这引出了一个常见的误解。如果一个 token 必须依次穿过每个晶圆，那么每个晶圆在等待轮到自己时，不都是空闲的吗？实际上，系统不就是在以 batch size 为 1 运行吗？

不！如果我们一次只沿流水线发送一个请求，并等它完成后再开始下一个请求，那将是对流水线的低效利用。相反，我们让整条流水线始终充满不同的请求。由于模型已经被拆分为多个流水线阶段，每个阶段都可以同时处理不同的请求。

当阶段 3 为请求 A 处理一个 token 时，阶段 2 正在为请求 B 处理一个，阶段 1 正在为请求 C 处理一个。在我们的示例中，每个阶段是一层。在生产环境中，Cerebras 使用粒度细得多的阶段，而且阶段数量也多得多。

正在加载播放器...

每个阶段还可以处理一个微批次，即一小组请求，而不是单个请求。因此，在途请求的数量大致等于阶段数乘以微批次大小。

正在加载播放器...

要理解这种权衡：较小的微批次会让每个请求更快，而较大的微批次会提高总吞吐量。

#6：我们究竟能服务多少个请求？

分配给 KV cache 的片上内存量决定了最大并发请求数。每个在途请求都需要在每个 dense 层晶圆上拥有自己的 KV cache，因此片上内存的大小设定了并发请求的上限。如上所述，我们只需要计算 dense 层的 KV cache 需求。

对于每块已经包含一个 dense 层的 44 GB 晶圆，我们可以用以下方式计算 KV cache 的剩余空间：

晶圆总存储量减去一个 dense 层的大小。在本例中：

计算并发请求数

要找出任意特定模型映射所支持的并发请求数，我们首先计算单个 256k 请求占用多少缓存：

然后计算此配置支持的总并发请求数，我们进行除法：

总并发请求数表示在任意时刻流经流水线的同时请求数量。与 GPU 的 batch 不同，这些请求不会互相等待完成，而是持续流经流水线。

当然，并发数和 batch size 都不能确切告诉你推理系统的运行效率。在评估推理系统的吞吐量时，更有意义的指标是在相同输出生成速度下测量的单位功耗总 token 吞吐量（例如每秒每兆瓦的 token 数）。

这如何扩展到 3T 参数以上

同样的方法可扩展到更大的模型：计算权重和请求状态所需的内存，将模型分布到多个晶圆上，并将它们连接成流水线。每增加一块晶圆都会贡献算力和片上内存，使我们能够容纳超出任何单芯片容量的模型。

感谢与我们一同学习！有兴趣在 Cerebras 上部署模型或使用我们的推理云进行构建吗？

• 通过我们的免费 API 开始使用：cloud.cerebras.ai
• 联系我们的销售团队：sales@cerebras.ai

设计：Halley Chang。感谢 Atharva Talpade、Zhenwei Gao、Angela Yeung、Jessica Liu、Joyce Er、Emma Call 和 Rohan Gupta 的审阅与反馈。