在MI355X上運行Kimi K3：每美元性能優於B300

過去幾個月，開源模型的能力快速提升。DeepSeek V4-Pro 和 GLM5.2 已接近 Opus 級智能，開源逐漸成為閉源之外具有成本效益的選擇。而 Kimi K3 的出現更進一步，宣稱達到 Fable/Sol 級別的智能，標誌着開源模型的新階段。

不過，更強的能力也意味着更大的規模。GLM5.2 有 7530 億參數，DeepSeek V4-Pro 有 1.6 萬億，Kimi K3 更是達到 2.8 萬億參數，僅權重就需要超過 1.5TB 顯存，還沒算上 100 萬 token 上下文所需的 KV 緩存。即使是 8 卡 B200 節點也放不下，只能選擇 B300 節點或兩個 B200 節點組成的 TP16 配置。

此時 AMD MI355X 成為另一個選擇：它同樣擁有 288GB 顯存，價格卻比 B300 平均低約 2.4 倍，比 B200 低約 1.7 倍。AMD 的短板一直是軟件生態，但 Wafer 認為 AI agent 正在加速縮小這一差距，而且 AMD 為 Kimi K3 提供了 day-0 支持，大部分工作已經完成。

實際測試中，在 1024 token 輸入、400 token 輸出的基準下，MI355X 達到每節點 952 tok/s 的聚合吞吐和 118 tok/s 的單流解碼。相比之下，TP16 B200 配置的聚合吞吐為 498 tok/s（兩個節點的總和，約每節點 249），單流僅 90 tok/s。B300 的聚合吞吐為 1,568 tok/s，仍比 MI355X 高約 1.65 倍，但價格是後者的 2.4 倍，因此在性能/美元上 MI355X 明顯勝出。按 MI355X 每 GPU 小時 2.50 美元、B300 6.00 美元、B200 4.25 美元計算，每美元吞吐分別為 48、33 和 7 tok/s。

B200 的成績其實有點吃虧，因為它需要跨節點 all-reduce（RoCE v2，約 195 Gb/s），成為唯一跨兩個節點的配置。這恰好説明，Kimi K3 的規模讓 MI355X 的高 HBM 容量優勢第一次轉化為實際可測的領先。

如何做到的？Kimi K3 雖然開箱即用，但要達到這個吞吐還需要一些工程修復。主要槓桿是投機解碼。K3 沒有自帶草稿張量，唯一的外部捷徑是 RadixArk 的 Kimi-K3-DSpark 塊擴散草稿。CUDA 上可以直接運行，但在 ROCm 上會報錯 NameError: name 'top_k_renorm_prob' is not defined。原因是 sglang 的 accept-sampling 驗證器有兩條路徑：密集路徑調用 top_k_renorm_prob，稀疏快速路徑直接用 torch.topk。CUDA 構建從 sgl_kernel 引入了該函數，而 ROCm 構建只引入了 Triton top-p 內核，沒有 gfx950 可用的 top-k renorm 內核，導致變量未定義。

修復只需要一個簡單的 PyTorch 函數：對概率向量做 top-k 截斷、其餘置零並重新歸一化，也就是排序、masked_fill 和除法，直接放入 ROCm 採樣分支即可。無需自定義內核——這次的問題是缺少定義，而不是缺少內核。修復後單流性能提升約 2.2 倍，中等負載下每流約 1.7 倍，峯值聚合提升 18%，而且峯值出現在更高併發（c64 vs c24）。

預填充優化同樣關鍵。MI355X 在 TTFT 上原本很掙扎：同樣的 172k token 冷預填充，MI355X 需要約 51 秒，B300 只要約 23 秒。對 100 萬上下文模型來説，大量工作負載會有巨大的預填充，GPU 空轉幾分鐘會讓整個節點陣列失效。

差距幾乎全部來自一個內核。K3 在 ROCm 上回退到慢速的通用 Triton attention，因為 AITER 的快速 MLA 預填充內核無法加載。原因是形狀不匹配：K3 在 TP8 下每個 rank 有 12 個 attention 頭，而 AITER 的 MLA 路徑只支持 4、8 或 16 的倍數。修復方法很簡單：把頭數從 12 零填充到 16，運行快速內核，再從輸出中取出真正的 12 個頭。結果同一預填充負載下，AITER MLA 內核達到約 13k tok/s，而 Triton 回退只有約 4–7k，預填充速度提升約 2–3 倍。這直接影響 TTFT，不影響聚合吞吐，但用户等待首個 token 的時間明顯縮短。

總的來説，在 MI355X 上獲得最佳性能/美元基本是開箱即得，雖然也遇到一些框架相關 bug，但比 GLM5.2 少，而且完全不需要自定義內核。AMD 上達到 SOTA 指日可待——CUDA 的護城河是否正在消失？