機器之心編輯部
(相關資料圖)
這可能是 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
