机器之心Pro 6小时前
16张B200才能跑的Kimi K3,8张AMD就装下了
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_keji1.html

 

这可能是 AMD 等了很久的时刻。

最近,Wafer AI 在 AMD MI355X 上部署了 Kimi K3。结果是,原本需要 16 张 NVIDIA B200、横跨两台服务器运行的模型,在一台配备 8 张 MI355X 的 AMD 服务器中就能完成部署。

更关键的是,它不只是把模型装进去了。

在输入 1024 Token、输出 400 Token 的测试中,MI355X 跑出了 952 Token/s 的总吞吐量,单用户生成速度达到 118 Token/s。

按照单节点计算,它的吞吐量约为 16 卡 B200 方案的 3.8 倍,性价比也超过了 B200 和 B300。

而最让人意外的是,ROCm 这次居然没有特别折腾。

模型太大,显存开始比算力更重要

Kimi K3 拥有 2.8 万亿参数,仅模型权重就需要超过 1.5 TB 显存,还没算百万 Token 上下文所需的 KV Cache。

一台 8 卡 B200 服务器,每张卡有 192 GB 显存,总容量约为 1.5 TB。也就是说,模型权重都很难完整放下,更别说给 KV Cache 留空间。因此,B200 必须使用两台服务器、16 张 GPU。

B300 每张卡拥有 288 GB 显存,可以在单节点内装下模型。巧的是,AMD MI355X 同样拥有 288 GB 显存,8 张 MI355X 合计约 2.3 TB,一台服务器就够了。

这不只是少用一台机器。模型跨节点运行后,每生成一个 Token,都可能需要通过网络同步数据。即使使用约 195 Gb/s 的 RoCE v2 网络,跨节点通信仍然会拖慢解码。

MI355X 靠更大的显存,把整个模型留在了一个节点里。

从最终结果看,8 张 MI355X 的峰值总吞吐量达到 952 Token/s,单路生成速度为 118 Token/s。

作为对比,16 张 B200 的双节点部署总吞吐量为 498 Token/s,换算到单节点约为 249 Token/s。

也就是说,MI355X 的单节点吞吐量约为 B200 双节点部署平均单节点吞吐量的 3.8 倍。单用户生成速度方面,MI355X 的 118 Token/s 也高于 B200 的 90 Token/s。

B300 依然是绝对性能最高的方案。8 张 B300 的节点总吞吐量达到 1568 Token/s,单路生成速度为 172 Token/s,整体吞吐量大约是 MI355X 的 1.65 倍。

但价格改变了结论。Wafer 按照 MI355X 每卡每小时 2.5 美元、B200 为 4.25 美元、B300 为 6 美元进行计算。

在这一价格假设下,MI355X 每美元可以提供约 48 Token/s 的峰值吞吐量;B200 约为 7 Token/s;B300 约为 33 Token/s。

B300 更快,但 MI355X 的单位成本效率更高。对于需要大规模运行开放模型的数据中心来说,这可能比单纯争夺性能冠军更加重要。

更意外的是,ROCm 基本可以直接用

长期以来,AMD 数据中心 GPU 最大的问题往往不是硬件,而是软件。

同一个模型在 CUDA 上可以直接运行,到了 ROCm 上,可能就要改框架、补算子,甚至重写底层内核。

但 Kimi K3 的情况有所不同。

AMD 为它提供了接近首发同步的支持。Wafer 表示,模型基本可以直接在 MI355X 上运行,后续工作主要集中在少量兼容性问题和性能优化上。

其中一个问题出现在推测解码环节。Kimi K3 本身没有提供 MTP 或 EAGLE 所需的草稿模型参数,因此 Wafer 使用了一个外部的块扩散草稿模型。

这套方案在 CUDA 上可以直接运行,但在 ROCm 环境中,第一个真实请求就导致调度器报错。原因是 ROCm 分支里少定义了一个名为 top_k_renorm_prob 的函数。

这个函数所做的事情并不复杂:从概率分布中选出最高的 k 个值,将其他概率归零,再把保留下来的概率重新归一化。

Wafer 最终使用一个普通的 PyTorch 函数补上了这段逻辑,不需要手写 GPU 内核,也不需要重新设计推测解码系统。

修复后,推测解码让单路性能提高约 2.2 倍,中等并发下的单流性能提高约 1.7 倍,峰值总吞吐量则提升了约 18%。

更重要的是,系统能够在更高并发下达到峰值吞吐量。

首字太慢,最后只补了四个零

当然,吞吐量并不是推理服务的全部。对于真实用户来说,另一个直接影响体验的指标是 TTFT,也就是从发送请求到看到第一个 Token 之间的等待时间。

在这一项上,MI355X 最初的表现并不好。面对一段约 17.2 万 Token 的冷启动预填充任务,MI355X 需要约 51 秒,而 B300 只需要约 23 秒。

在支持百万 Token 上下文的模型中,预填充任务可能非常庞大。如果处理长上下文时,用户每次都要先等几十秒甚至更久,再高的解码速度也很难弥补体验上的问题。

Wafer 最终发现,性能差距几乎全部来自一个注意力内核。Kimi K3 在 8 路张量并行配置下,每张 GPU 会分到 12 个注意力头。而 AMD AITER 中速度较快的 MLA 预填充内核,只支持 4、8 或 16 的倍数等形状。

12 个头无法匹配,于是系统退回了速度较慢的通用 Triton 实现。

解决办法很朴素:把 12 个注意力头补零到 16 个,调用现有的高速内核,计算完成后再取回真正需要的 12 个头。没有修改模型结构,也没有编写新的汇编内核,只是补了四个零。

优化后,AITER MLA 内核的稳定预填充速度达到约 1.3 万 Token/s,而原来的 Triton 回退路径大约只有 4000~7000 Token/s,冷预填充时间因此缩短了约两到三倍。

这项优化不会改变最终的解码吞吐量,但会显著减少用户等待第一个字出现的时间。

这也说明,AMD 和 NVIDIA 之间看起来很大的软件差距,有时并不是底层能力不足,只是现有高速内核暂时没有覆盖某种新模型形状。

CUDA 的护城河还在,但缺口已经出现

一次测试当然不能证明 AMD 已经全面追上 NVIDIA。

B200 因为显存不足,被迫跨节点运行;B300 的绝对性能依然领先;ROCm 的工具链、框架支持和开发者生态,也仍然不如 CUDA。

但开放模型正在快速进入万亿参数时代。当模型大到一台服务器装不下时,显存容量就不再只是参数表上的数字,而会直接影响通信成本、部署复杂度和最终吞吐量。

AMD 给单卡配置更多 HBM 的策略,正在变成一种实际的系统优势。

如果 AMD 能继续提升 ROCm 的稳定性、扩大高速内核的形状支持,并为新模型提供更及时的首日适配,那么数据中心就必须认真考虑这些 GPU。价格更低,显存更大,性能够用,软件也不再需要折腾几个月。

宙世代

宙世代

ZAKER旗下Web3.0元宇宙平台

一起剪

一起剪

ZAKER旗下免费视频剪辑工具

相关标签

吞吐量 amd gpu
相关文章
评论
没有更多评论了
取消

登录后才可以发布评论哦

打开小程序可以发布评论哦

12 我来说两句…
打开 ZAKER 参与讨论