Cerebras 在官方博客中介绍了将万亿参数模型部署到多块 WSE-3T 晶圆上的方法,并以约 3 万亿参数的开源模型 Kimi K3 为例,给出使用 59 块晶圆的部署计算。单块晶圆只有 44 GB 片上内存,而这类模型仅权重就需要数 TB,如何分配权重、缓存和请求,是这套方案要解决的问题。Cerebras 称,多晶圆方案可让多万亿参数模型达到每秒 1000 至 2000 以上 token 的速度;这是发布方的性能表述。
据 Cerebras 介绍,Kimi K3 包含 69 层线性注意力层和 24 层稠密(全注意力)层,共 93 层。博客示例假设平均请求长度为 256k,权重采用 4 位量化,KV 缓存采用 8 位存储。将量化后的参数总量按 93 层平均计算,每层权重约占 16 GB。
两类层的内存需求不同。稠密层需要为每个活跃请求中的每个 token 保存键和值,因此每块晶圆只放一层,剩余 SRAM 留给 KV 缓存。线性注意力层则为每个请求保留一个固定大小的小状态,不随上下文长度增长,所以一块晶圆可容纳两层,权重合计约 32 GB。
按这一分配方式,69 层线性注意力层需要 35 块晶圆,24 层稠密层需要 24 块,合计 59 块。这是博客在上述假设下给出的示例配置。
Cerebras 将模型组织成流水线:编译器把模型层映射到晶圆上,运行时安排数据在晶圆之间流动。每个阶段的权重留在原处,请求的中间结果依次传向后续阶段,直至生成下一个 token。
据 Cerebras 介绍,晶圆之间只需传递激活值,通过以太网 RDMA 连接,每跳耗时约 2 微秒。由于多数情况下整层模型能放进单块晶圆的片上 SRAM,层无需跨晶圆拆分,权重也无需移动。
流水线可让不同阶段同时处理不同请求,每个阶段还可处理一小组请求,即微批。Cerebras 给出的关系是:同时在途的请求数,约等于阶段数乘以微批大小。较小的微批有利于提高单个请求的速度,较大的微批则提高总吞吐量。
并发也受内存限制。每个在途请求都需要在各块稠密层晶圆上保存自己的 KV 缓存,因此,分给缓存的片上内存决定了并发请求数的上限。
Cerebras 认为,仅看并发量或批大小不足以判断推理效率;更有意义的比较,是在相同输出生成速度下,衡量单位功率对应的总 token 吞吐量,例如每兆瓦每秒输出的 token 数。