这可能是 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。价格更低,显存更大,性能够用,软件也不再需要折腾几个月。
对此你怎么看?
参考链接:
https://x.com/wafer_ai/status/2083628389903315406
https://x.com/ChiragAsarpota/status/2083864019870634151
本文来自微信公众号 “机器之心”(ID:almosthuman2014),作者:关注LLM的
