更小、更快、更安全:大规模运行Kimi与GLM
Smaller, faster, safer: running Kimi and GLM at scale
Workers AI 在 Cloudflare 数据中心 GPU 上运行 Moonshot Kimi K 系列和 Z.ai GLM 等开放模型,采用三项技术优化内存:量化 KV 缓存(FP8 替代 BF16)使 Kimi K2.6 上下文容量从约 686,000 token 翻倍至 137 万,decode 吞吐提升 41%;压缩模型权重(INT4 替代 FP8)使 GLM 5.2 检查点从 705 GB 减至 421 GB,decode 加速且精度无差异;构建 KV 缓存完整性检查,成本低于 1%。实验使用 SGLang 框架,与团队合作上游化补丁。
Workers AI 在 Cloudflare 数据中心内的 GPU 上,为全球一些最优秀的开放模型运行推理,这些数据中心靠近你的用户。其中能力最强、要求也最高的两个模型是 Moonshot 的 Kimi K 系列和 Z.ai 的 GLM。它们是大规模、长上下文、混合专家模型,使用体验极佳,但由于内存限制,高效服务它们也颇具挑战。我们之前写过如何在 Workers AI 上服务大型模型,以及如何分离推理的 prefill 和 decode 阶段以充分利用每块 GPU。本文探讨了在此之上叠加的三项技术,以将这些模型装入内存并保持高速:量化 KV 缓存、压缩模型权重,以及——由于这两项技术都会在共享硬件上承载更多请求——保护这些请求共享的缓存。这些优化使我们能够以更低成本支持更多客户,且模型精度不受影响。我们所有的实验和生产流量均使用 SGLang(一个开源推理服务框架)运行并进行基准测试。我们发现 SGLang 在市场上提供了最佳性能,并且我们与 SGLang 团队紧密合作,将补丁和新功能上游化,使我们的工作惠及开源社区。
量化 KV 缓存
模型生成文本时,会将已处理每个 token 的注意力键(K)和值(V)存储在一个称为 KV 缓存的结构中。该缓存使模型能够扩展长对话,而无需在每个新 token 上重新读取整个上下文。对于长上下文模型,它会迅速增长,通常最先填满 GPU 内存的是 KV 缓存,而非模型权重。默认情况下,缓存以 16 位精度(BF16)存储。我们改为使用 8 位浮点格式(FP8,e4m3),将其大小减半。在 Kimi K2.6 上,这使内存中可容纳的上下文量从约 686,000 个 token 提升至约 137 万,翻了一倍。
值得精确说明收益来源,因为它并非原始速度。量化缓存会在每个 token 上增加少量工作,因为 FP8 注意力核在读取值时需要转换。它改变的是我们一次能驻留的请求数量。以下测量针对 Kimi K2.6 在分离式 H200 部署上的解码,直接比较注意力核:
在任何单一并发级别下,BF16 每个 token 快几个百分点。但 BF16 在 32 个并发请求时耗尽缓存,无法接纳第 33 个,而 FP8 可继续至 64 个,达到每秒 2,192 个 token,比 BF16 峰值高出约 41%,且每个 token 成本降低约 30%。由于我们将 prefill 和 decode 作为独立池运行,可以在最有益处的地方应用此技术:prefill 受计算限制而非内存限制,因此我们保留 BF16 缓存以维持其略高的吞吐量。
如果这改变了模型的答案,一切就毫无意义,因此我们进行了验证。在我们的评估套件中,FP8 和 BF16 缓存表现无差异:
压缩模型权重
KV 缓存是对 GPU 内存的一种需求;模型权重则是另一种。对于 GLM 5.2,我们将权重从 8 位浮点压缩至 4 位整数(INT4),且精度无损失。检查点从 705 GB 缩减至 421 GB,约减少 40%,在 8 路张量并行部署中,每 GPU 内存从约 88 GB 降至 52 GB,为同一硬件上约 118 万 token 的 KV 缓存留出空间。在我们的评估套件中,INT4 和 FP8 权重表现无差异:
更小的权重使 decode 阶段更快,原因明确:生成每个 token 意味着将模型权重从 GPU 内存中流出,因此 decode 速度受内存带宽限制。移动更少数据,每个 token 更快到达。效果在低并发下最大,此时每请求延迟最为关键:
Prefill 行为不同。它受计算限制,INT4 权重必须扩展回原状才能进行乘法,因此额外步骤使 prefill 变慢而非变快,GLM 在 FP8 下维持约每秒 10,160 个 token 的 prefill,而 INT4 为 8,660。与 KV 缓存一样,分离式设计将这一权衡转化为选择而非妥协:我们在 decode 中使用 INT4,它在此胜出;在 prefill 中使用 FP8,它在此胜出。在我们运行的每个基准测试中,模型精度保持在 FP8 模型的 0.8 个点以内,质量无差异。
保护共享 KV 缓存
上述两项技术效果相同:它们让更多请求同时共享一块 GPU 的内存。这种效率是核心目的,但也意味着数百个请求在读写同一物理 KV 缓存的页面。使这变得快速的机制——分页注意力、连续批处理、缓存复用——都依赖于精确无误的簿记,而在我们的请求量级下,即使是十亿分之一错误也会频繁出现。因此,我们构建了 KV 缓存完整性检查作为防御层。思路简单:每个物理缓存页面在重新分配时都会获得一个标签,服务器记录每个请求预期使用的页面和标签。在支持的 decode 操作从缓存读取前,会检查这些映射。如果任何不匹配,受影响的请求将被中止,而不是允许从错误页面返回数据。
决定安全检查是否上线的关键是其成本。我们在一个中等规模生产模型上进行了测量,配置为两个 prefill、两个 decode,输入 8,192 个 token,输出 1,000 个 token:
吞吐量和尾延迟的成本均低于 1%,即使 95% 置信区间的上限也保持在 1% 附近。我们通过将验证作为独立批处理检查而非融合到注意力核中来保持计算成本低廉,后者会在 GPU 线程组之间引入竞争。它按部署启用,默认路径使用无操作跟踪器,无额外开销,因此不需要它的部署无需付出任何代价。
下一步
高效服务前沿模型是一个动态目标,这是其背后的持续工作。我们正在将 FP8 KV 缓存扩展到更多集群,在 Blackwell(NVIDIA 的 GPU 架构)上验证 NVFP4 权重,并努力使完整性检查能在任何地方以可忽略的成本保持开启。这些优化将使我们能够以更低成本和相同精度继续支持更多客户。如果你觉得将最优秀的开放模型压缩到 GPU 上并为数百万开发者提供服务听起来像你感兴趣的问题,欢迎加入我们。