Tóm tắt
Wafer chạy Kimi K3 trên một node 8x AMD MI355X và lập luận rằng dung lượng HBM đang trở thành lợi thế thực tế khi phục vụ các open-weight model rất lớn. Bài viết có nhiều số đo hữu ích về throughput, speculative decode và prefill, nhưng kết luận về performance-per-dollar phụ thuộc mạnh vào giả định giá thuê GPU.
Kết quả chính
Kimi K3 là model 2.8T parameters, cần hơn 1.5TB VRAM trước khi tính KV cache cho context 1M token. Theo Wafer, cấu hình 8x MI355X có 288GB VRAM mỗi GPU, đủ để chạy Kimi K3 trên một node; một node 8x B200 thì không đủ, nên phải dùng hai node B200 hoặc chuyển sang B300.
Trong benchmark input 1.024 token và output 400 token, Wafer báo cáo:
- 8x MI355X đạt 118 tok/s cho single stream decode.
- Peak aggregate đạt 952 tok/s mỗi node.
- B300 vẫn nhanh hơn về aggregate throughput, đạt 1.568 tok/s, nhưng Wafer cho rằng MI355X thắng về throughput trên mỗi USD/GPU-giờ.
- So với triển khai hai node B200 của họ, MI355X có lợi thế vì tránh cross-node all-reduce trên đường decode.
Điểm dev nên để ý: đây là bài toán memory capacity và serving topology, không chỉ là so FLOPS. Khi model quá lớn để nằm gọn trong một node, interconnect và KV cache bắt đầu quyết định latency lẫn chi phí.
Cách Wafer tối ưu
Kimi K3 có thể chạy sẵn trên MI355X, nhưng Wafer vẫn phải sửa một số điểm trong stack ROCm/sglang.
Tối ưu đầu tiên là speculative decode. K3 không ship draft tensor kiểu MTP hay EAGLE, nên Wafer dùng draft model ngoài RadixArk Kimi-K3-DSpark. Trên ROCm, nhánh verifier của sglang bị lỗi vì top_k_renorm_prob không được định nghĩa. Wafer thay bằng một hàm PyTorch đơn giản: lấy top-k probability, zero các token còn lại, rồi renormalize tổng về 1. Sau khi sửa và harden speculative decode, họ báo single-stream nhanh hơn khoảng 2.2x và peak aggregate tăng 18%.
Tối ưu thứ hai nằm ở prefill. Wafer nhấn mạnh decode tok/s dễ được khoe quá mức, trong khi time-to-first-token mới là độ trễ người dùng cảm nhận. Với cold prefill 172k token, MI355X ban đầu mất khoảng 51 giây, còn B300 mất khoảng 23 giây. Nguyên nhân chính là AITER MLA prefill kernel không load vì shape mismatch: K3 ở TP8 có 12 attention heads mỗi rank, trong khi kernel hỗ trợ 4, 8 hoặc bội số của 16. Cách sửa là pad 12 head lên 16, chạy kernel nhanh, rồi lấy lại 12 head thật. Wafer nói steady-state prefill tăng từ khoảng 4-7k tok/s lên khoảng 13k tok/s.
Điều dev nên kiểm chứng
Bài viết đáng đọc nếu bạn vận hành inference cho model rất lớn, đặc biệt khi đang so AMD MI355X, NVIDIA B200 và B300. Tuy vậy, đừng bê nguyên kết luận performance-per-dollar vào kế hoạch mua hoặc thuê GPU nếu chưa kiểm tra lại các điều kiện sau:
- Giá thuê trong bài là snapshot từ marketplace, không phải TCO sở hữu phần cứng.
- Availability của MI355X/B300 theo vùng và nhà cung cấp có thể làm đảo ngược phép tính.
- Bài viết tập trung vào throughput và prefill, chưa trình bày benchmark correctness sau khi tối ưu speculative decode.
- Kimi K3 là open-weight model; gọi nó là open source dễ làm mờ câu hỏi về training data và tooling.
Kết luận ngắn: memory đang là lợi thế thật của MI355X ở lớp model khổng lồ, nhưng “CUDA moat đã chết” vẫn là câu hỏi mở. Số đo này nên đi vào spreadsheet capacity planning, không nên đi thẳng vào slide chiến thắng.