Chạy Kimi K3 trên AMD MI355X: Hiệu năng vượt trội so với Nvidia B300

Liệu bộ nhớ có phải là "con hào" bảo vệ?

Với tốc độ vận hành Kimi K3 đạt ~952 tok/s/node, AMD tiếp tục chứng minh vị thế dẫn đầu về hiệu năng trên mỗi đô la chi phí.

Trong vài tháng qua, chúng ta đã chứng kiến sự bùng nổ về khả năng của các mô hình mã nguồn mở. Với việc DeepSeek V4-Pro và GLM5.2 đạt trình độ thông minh gần tương đương Opus, mã nguồn mở đã nổi lên như một giải pháp thay thế thực tế, tiết kiệm chi phí cho các mô hình đóng mà chúng ta vốn đã phụ thuộc từ lâu.

Nhưng chúng ta vẫn chưa thấy mô hình nào như Kimi K3. Hứa hẹn mang lại trí thông minh ở cấp độ Fable/Sol, Kimi K3 đánh dấu sự khởi đầu của một kỷ nguyên mới cho mã nguồn mở.

Tuy nhiên, mô hình thông minh hơn đồng nghĩa với mô hình lớn hơn — và kích thước của chúng đang tăng nhanh không kém gì khả năng. GLM5.2 có 753 tỷ tham số, DeepSeek V4-Pro có 1,6 nghìn tỷ và Kimi K3 nặng tới 2,8 nghìn tỷ (!!) tham số. Đó là hơn 1,5TB VRAM trước khi phân bổ KV cache cho 1 triệu token ngữ cảnh. Ngay cả một node B200 (8 GPU) cũng không thể chứa nổi Kimi K3. Điều đó khiến bạn chỉ còn vài lựa chọn hạn hẹp: phục vụ trên một node B300 (vốn có 288GB VRAM mỗi GPU), hoặc sử dụng hai node B200 (TP16) để vận hành Kimi.

Nhưng hãy đoán xem GPU nào khác không phải của NVIDIA có 288GB VRAM? Đó là MI355X của AMD. Bạn đã thấy chúng tôi ưa chuộng các con chip này chưa? Với chi phí trung bình rẻ hơn khoảng ~2,4 lần mỗi GPU so với B300 và rẻ hơn ~1,7 lần so với B200, MI355X là giải pháp thay thế tiết kiệm chi phí cho Blackwell với thông số phần cứng tương đương. Vấn đề duy nhất với AMD là hỗ trợ phần mềm — các kernel chậm hơn và ít hỗ trợ "day-0" trên các framework suy luận khiến việc phục vụ các mô hình tiên phong trên AMD trở thành một nỗ lực kỹ thuật thực sự. Tại Wafer, chúng tôi khẳng định rằng các tác nhân (agents) đang cải thiện việc tối ưu hóa kernel và mô hình, dần thu hẹp khoảng cách này ngay lúc chúng ta đang nói chuyện. Nhưng với việc AMD cung cấp hỗ trợ day-0 cho Kimi K3, phần lớn công việc đã được thực hiện sẵn cho chúng tôi.

Kết quả rất khả quan: trên benchmark đầu vào 1.024 token / đầu ra 400 token, MI355X đạt 952 tok/s/node và 118 tok/s ở luồng đơn — cao hơn 3,8 lần tổng thông lượng mỗi node và hơn 1,3 lần tốc độ giải mã luồng đơn so với triển khai TP16 B200 của chúng tôi (vốn đạt 498 tok/s cho tổng cộng 16 GPU, 2 node — khoảng 249/node). Các node B300 vẫn thắng khoảng 1,65 lần về tổng thông lượng so với MI355X, nhưng với mức giá đắt hơn 2,4 lần, MI355X đã đánh bại B300 về hiệu năng trên mỗi đô la.

Hiệu năng/đô la ở mức 2,50 USD/GPU-giờ cho MI355X, 6,00 USD cho B300 và 4,25 USD cho B200.

Để bảo vệ B200, các con số của nó phần nào bị giảm sút do phải thực hiện cross-node all-reduce trên đường dẫn giải mã quan trọng (RoCE v2 ở mức ~195 Gb/s) — đây là cấu hình duy nhất ở đây trải dài trên hai node, vì Kimi K3 không thể chứa vừa trọng số cộng với 1 triệu token KV pool trên một node 8×192GB duy nhất. Nhưng đó chính là điểm mấu chốt: Kimi K3 với kích thước hiện tại là một trong những mô hình đầu tiên mà chúng tôi thấy sự tập trung của MI355X vào dung lượng HBM mang lại lợi thế thực tế, có thể đo lường được so với B200.

Cách chúng tôi thực hiện

Mặc dù Kimi K3 có thể vận hành ngay lập tức (out of the box), nhưng vẫn cần phải làm thêm một số việc để đạt được con số thông lượng hiện tại.

Đòn bẩy chính là giải mã suy đoán (speculative decode). K3 không đi kèm các draft tensor — không MTP, không EAGLE — vì vậy con đường suy đoán duy nhất là một draft block-diffusion bên ngoài: Kimi-K3-DSpark của RadixArk. Trên CUDA, nó chạy ngay lập tức. Trên ROCm, yêu cầu thực tế đầu tiên của chúng tôi làm hỏng bộ lập lịch với lỗi này:

Bộ xác thực accept-sampling của sglang có hai cách để xây dựng phân phối mục tiêu: đường dẫn dày (dense path) gọi top_k_renorm_prob, và đường dẫn nhanh thưa (sparse fast path) định tuyến trực tiếp qua torch.topk. Bản dựng CUDA nhập top_k_renorm_prob từ sgl_kernel; bản dựng ROCm chỉ alias một kernel Triton top-p và để top_k_renorm_prob không xác định — không có kernel top-k renorm cho gfx950 để alias. Vì vậy, ngay khi một yêu cầu rơi vào đường dẫn dày, bộ xác thực gặp lỗi NameError đó và làm sập cả bộ lập lịch.

Bản sửa lỗi chỉ là một hàm PyTorch duy nhất. Top-k renorm là một thao tác nhỏ: lấy vector xác suất của mô hình, giữ lại k mục cao nhất, đặt các mục còn lại bằng 0 và chia tỷ lệ phần còn lại để tổng bằng 1. Một lệnh sắp xếp, một lệnh masked_fill, một lệnh chia — được đưa thẳng vào nhánh lấy mẫu ROCm của sglang, cùng một phép tính mà bản dựng CUDA nhận được từ sgl_kernel. Không cần kernel tùy chỉnh: phản xạ trên ROCm là giả định bạn cần một kernel, nhưng ở đây chỉ là thiếu định nghĩa, không phải thiếu kernel.

Với việc spec dec đã được sửa và củng cố, chúng tôi đạt được hiệu năng luồng đơn tăng ~2,2 lần, tăng ~1,7 lần mỗi luồng ở mức tải trung bình và tăng +18% tổng thông lượng đỉnh. Quan trọng hơn, thông lượng đỉnh tổng thể của chúng tôi đạt được ở mức đồng thời cao hơn nhiều (c64 so với c24 khi không dùng spec).

Tối ưu hóa Prefill

Các cuộc thảo luận về hiệu năng mô hình thường nhấn mạnh vào số token giải mã mỗi giây. Nhưng trong nhiều trường hợp, decode tok/s là "vàng giả" — giải mã được tôn vinh quá mức, trong khi thời gian đến token đầu tiên (time-to-first-token), con số mà người dùng cảm nhận rõ nhất, lại bị bỏ qua.

MI355X gặp khó khăn ở đây: một lần prefill nguội (cold prefill) 172k-token giống hệt nhau mất ~51 giây trên MI355X so với ~23 giây trên B300. Trên một mô hình có ngữ cảnh 1 triệu token, rất nhiều khối lượng công việc có prefill khổng lồ (đôi khi là nguội), và việc để các GPU quay cuồng với prefill trong nhiều phút có thể khiến toàn bộ các node trở nên vô dụng.

Khoảng cách này gần như hoàn toàn do một kernel. K3 trên ROCm đã quay lại sử dụng attention Triton chung chậm chạp vì kernel prefill AITER MLA nhanh không thể tải được. Vấn đề là sự không khớp về hình dạng (shape mismatch), không phải thiếu kernel — K3 tại TP8 cung cấp 12 attention head mỗi rank, và đường dẫn MLA của AITER được xây dựng cho 4, 8 hoặc bội số của 16. Bản sửa lỗi cực kỳ đơn giản: thêm zero-pad vào số lượng head từ 12 lên 16, chạy kernel nhanh và trích xuất 12 head thực từ đầu ra.

Kết quả: trên cùng một lần prefill nguội 172k, ASM prefill AITER MLA chạy ở mức ~13k tok/s ổn định so với ~4–7k của Triton fallback, tăng tốc prefill lên ~2–3 lần. Đó là đòn bẩy cho TTFT, không phải cho tổng thông lượng — việc giải mã không thay đổi, vì vậy nó không làm thay đổi các con số ở trên; nó làm thay đổi con số mà người dùng phải chờ đợi trước khi token đầu tiên xuất hiện.

Kết luận

Việc đạt được tỷ lệ hiệu năng trên mỗi đô la tốt nhất trên MI355X tương đối dễ dàng ngay từ đầu. Có một số lỗi liên quan đến framework như dự kiến — nhưng ít hơn so với GLM5.2, và lần này chắc chắn không cần đến các kernel tùy chỉnh.

SOTA trên AMD sắp trở thành hiện thực. Liệu "con hào" CUDA đã chết?

Bài viết liên quan

Tích hợp Wafer với TrueFoundry AI Gateway

Cách máy chủ suy luận serverless nhanh, tương thích với OpenAI của Wafer tích hợp với TrueFoundry AI Gateway để định tuyến thống nhất, khả năng quan sát và không lưu giữ dữ liệu.

Hiệu năng trên mỗi đô la đang ngày càng nhanh hơn và rẻ hơn

Cách chúng tôi phục vụ GLM5.2 trên AMD MI355X ở mức 2626 tok/s/node và 213 tok/s luồng đơn với chi phí thấp hơn gấp 2 lần so với Blackwell.