Đây có lẽ là khoảnh khắc mà AMD đã chờ đợi từ lâu.
Gần đây, Wafer AI đã triển khai mô hình Kimi K3 trên AMD MI355X. Kết quả là, một mô hình vốn cần đến 16 card NVIDIA B200 chạy trải dài qua hai máy chủ, thì giờ đây có thể được triển khai trong một máy chủ AMD với chỉ 8 card MI355X.
Điều quan trọng hơn là, nó không chỉ đơn thuần là nhét mô hình vào trong.
Trong bài kiểm tra với đầu vào 1024 Token và đầu ra 400 Token, MI355X đạt được tổng thông lượng 952 Token/s, tốc độ sinh token đơn luồng đạt 118 Token/s.
Tính trên một nút (node), thông lượng của nó cao gấp khoảng 3.8 lần so với phương án sử dụng 16 card B200. Hiệu quả chi phí cũng vượt trội hơn B200 và B300.
Và điều bất ngờ nhất là, lần này ROCm không gây ra quá nhiều rắc rối.
Mô hình quá lớn, bộ nhớ GPU bắt đầu quan trọng hơn sức mạnh tính toán
Kimi K3 sở hữu 2.8 nghìn tỷ tham số, chỉ riêng trọng số mô hình đã cần hơn 1.5 TB bộ nhớ GPU, chưa kể đến KV Cache cần thiết cho ngữ cảnh (context) hàng triệu Token.
Một máy chủ 8 card B200, mỗi card có 192 GB bộ nhớ GPU, tổng dung lượng khoảng 1.5 TB. Nghĩa là, ngay cả việc chứa đủ trọng số mô hình cũng khó khăn, chứ đừng nói đến chỗ cho KV Cache. Vì vậy, B200 buộc phải sử dụng hai máy chủ với tổng cộng 16 GPU.
B300 mỗi card sở hữu 288 GB bộ nhớ GPU, có thể chứa mô hình trong một nút duy nhất. Trùng hợp thay, AMD MI355X cũng sở hữu 288 GB bộ nhớ GPU, 8 card MI355X cộng lại khoảng 2.3 TB, một máy chủ là đủ.
Đây không chỉ là việc sử dụng ít hơn một máy. Khi mô hình chạy xuyên qua nhiều nút, mỗi lần sinh ra một Token đều có thể cần đồng bộ dữ liệu qua mạng. Ngay cả khi sử dụng mạng RoCE v2 tốc độ khoảng 195 Gb/s, việc truyền thông xuyên nút vẫn làm chậm quá trình giải mã.
MI355X nhờ bộ nhớ GPU lớn hơn, đã giữ toàn bộ mô hình trong một nút duy nhất.
Nhìn vào kết quả cuối cùng, tổng thông lượng đỉnh của 8 card MI355X đạt 952 Token/s, tốc độ sinh đơn luồng là 118 Token/s.
Để so sánh, tổng thông lượng của triển khai hai nút với 16 card B200 là 498 Token/s, quy đổi thành một nút thì khoảng 249 Token/s.
Nghĩa là, thông lượng một nút của MI355X cao gấp khoảng 3.8 lần so với thông lượng trung bình một nút của triển khai hai nút B200. Về tốc độ sinh token đơn người dùng, 118 Token/s của MI355X cũng cao hơn 90 Token/s của B200.
B300 vẫn là phương án có hiệu năng tuyệt đối cao nhất. Tổng thông lượng của một nút 8 card B300 đạt 1568 Token/s, tốc độ sinh đơn luồng là 172 Token/s, tổng thông lượng gấp khoảng 1.65 lần MI355X.
Nhưng giá cả đã thay đổi kết luận. Wafer tính toán dựa trên giả định giá mỗi card MI355X là 2.5 USD/giờ, B200 là 4.25 USD/giờ, B300 là 6 USD/giờ.
Với giả định giá này, MI355X có thể cung cấp khoảng 48 Token/s thông lượng đỉnh trên mỗi đô la; B200 khoảng 7 Token/s; B300 khoảng 33 Token/s.
B300 nhanh hơn, nhưng hiệu quả chi phí trên mỗi đơn vị của MI355X lại cao hơn. Đối với các trung tâm dữ liệu cần vận hành quy mô lớn các mô hình mở, điều này có thể quan trọng hơn việc chỉ tranh giành ngôi vương về hiệu năng.
Bất ngờ hơn nữa, ROCm về cơ bản có thể dùng trực tiếp
Trong thời gian dài, vấn đề lớn nhất của GPU dành cho trung tâm dữ liệu của AMD thường không phải là phần cứng, mà là phần mềm.
Cùng một mô hình chạy trực tiếp được trên CUDA, nhưng khi chuyển sang ROCm, có thể phải sửa khung làm việc, bổ sung toán tử, thậm chí viết lại nhân (kernel) cấp thấp.
Nhưng tình hình với Kimi K3 lại khác.
AMD đã cung cấp hỗ trợ gần như đồng bộ ngay từ đầu. Wafer cho biết, mô hình về cơ bản có thể chạy trực tiếp trên MI355X, công việc tiếp theo chủ yếu tập trung vào một số vấn đề tương thích nhỏ và tối ưu hóa hiệu năng.
Một vấn đề xuất hiện ở khâu giải mã suy đoán (speculative decoding). Kimi K3 bản thân không cung cấp tham số mô hình nháp (draft model) cần thiết cho MTP hoặc EAGLE, vì vậy Wafer đã sử dụng một mô hình nháp khuếch tán khối (block diffusion draft model) bên ngoài.
Giải pháp này có thể chạy trực tiếp trên CUDA, nhưng trong môi trường ROCm, yêu cầu thực tế đầu tiên đã khiến bộ lập lịch (scheduler) báo lỗi. Nguyên nhân là trong nhánh ROCm thiếu định nghĩa một hàm có tên là `top_k_renorm_prob`.
Hàm này làm việc không quá phức tạp: chọn ra k giá trị cao nhất từ phân bố xác suất, đưa các xác suất khác về 0, sau đó chuẩn hóa lại các xác suất được giữ lại.
Wafer cuối cùng đã sử dụng một hàm PyTorch thông thường để bù vào đoạn logic này, không cần viết kernel GPU thủ công, cũng không cần thiết kế lại hệ thống giải mã suy đoán.
Sau khi sửa chữa, giải mã suy đoán đã cải thiện hiệu năng đơn luồng khoảng 2.2 lần, hiệu năng đơn luồng ở mức độ song song trung bình tăng khoảng 1.7 lần, tổng thông lượng đỉnh tăng khoảng 18%.
Quan trọng hơn, hệ thống có thể đạt được thông lượng đỉnh ở mức độ song song cao hơn.
Chữ đầu tiên quá chậm, cuối cùng chỉ bổ sung bốn số không
Tất nhiên, thông lượng không phải là toàn bộ dịch vụ suy luận (inference). Đối với người dùng thực tế, một chỉ số khác trực tiếp ảnh hưởng đến trải nghiệm là TTFT, tức là thời gian chờ từ khi gửi yêu cầu đến khi nhìn thấy Token đầu tiên.
Ở mục này, hiệu suất ban đầu của MI355X không tốt. Đối với một tác vụ điền sẵn (prefill) khởi động lạnh (cold start) với khoảng 172 nghìn Token, MI355X cần khoảng 51 giây, trong khi B300 chỉ cần khoảng 23 giây.
Trong các mô hình hỗ trợ ngữ cảnh hàng triệu Token, tác vụ điền sẵn có thể rất lớn. Nếu khi xử lý ngữ cảnh dài, người dùng mỗi lần đều phải chờ hàng chục giây hoặc lâu hơn, thì tốc độ giải mã có cao đến mấy cũng khó bù đắp được vấn đề về trải nghiệm.
Wafer cuối cùng phát hiện ra, sự chênh lệch hiệu năng gần như hoàn toàn đến từ một kernel chú ý (attention kernel). Kimi K3 trong cấu hình song song tensor 8 luồng, mỗi GPU sẽ nhận được 12 đầu chú ý (attention head). Trong khi đó, kernel điền sẵn MLA tốc độ cao trong AMD AITER chỉ hỗ trợ các kích thước như bội số của 4, 8 hoặc 16.
12 đầu không khớp, vì vậy hệ thống đã quay về sử dụng kernel tổng quát Triton chậm hơn.
Giải pháp rất đơn giản: bổ sung thêm 4 số 0 vào 12 đầu chú ý để thành 16 đầu, gọi kernel tốc độ cao hiện có, sau khi tính toán xong thì lấy lại 12 đầu thực sự cần thiết. Không sửa cấu trúc mô hình, cũng không viết kernel assembly mới, chỉ là bổ sung bốn số không.
Sau khi tối ưu, tốc độ điền sẵn ổn định của kernel AITER MLA đạt khoảng 13 nghìn Token/s, trong khi đường dẫn quay về Triton ban đầu chỉ đạt khoảng 4000~7000 Token/s, thời gian điền sẵn khởi động lạnh vì vậy được rút ngắn khoảng hai đến ba lần.
Tối ưu hóa này sẽ không thay đổi thông lượng giải mã cuối cùng, nhưng sẽ giảm đáng kể thời gian người dùng chờ đợi chữ đầu tiên xuất hiện.
Điều này cũng cho thấy, khoảng cách phần mềm giữa AMD và NVIDIA nhìn có vẻ lớn, đôi khi không phải do năng lực cốt lõi không đủ, mà chỉ là các kernel tốc độ cao hiện có tạm thời chưa bao phủ một hình dạng mô hình mới nào đó.
Hào sâu CUDA vẫn còn, nhưng lỗ hổng đã xuất hiện
Một lần thử nghiệm tất nhiên không thể chứng minh AMD đã bắt kịp NVIDIA toàn diện.
B200 vì bộ nhớ GPU không đủ, buộc phải chạy xuyên nút; Hiệu năng tuyệt đối của B300 vẫn dẫn đầu; Chuỗi công cụ, hỗ trợ framework và hệ sinh thái nhà phát triển của ROCm cũng vẫn kém hơn CUDA.
Nhưng các mô hình mở đang nhanh chóng bước vào kỷ nguyên nghìn tỷ tham số. Khi mô hình lớn đến mức một máy chủ không chứa nổi, dung lượng bộ nhớ GPU không còn chỉ là con số trên bảng thông số, mà sẽ trực tiếp ảnh hưởng đến chi phí truyền thông, độ phức tạp triển khai và thông lượng cuối cùng.
Chiến lược của AMD là trang bị nhiều HBM hơn cho mỗi card, đang trở thành một lợi thế hệ thống thực tế.
Nếu AMD có thể tiếp tục nâng cao độ ổn định của ROCm, mở rộng hỗ trợ hình dạng cho các kernel tốc độ cao, và cung cấp hỗ trợ thích ứng ngay từ ngày đầu cho các mô hình mới, thì các trung tâm dữ liệu sẽ phải cân nhắc nghiêm túc đến những GPU này. Giá thấp hơn, bộ nhớ GPU lớn hơn, hiệu năng đủ dùng, phần mềm cũng không cần mất hàng tháng trời để "vật lộn".
Bạn nghĩ gì về điều này?
Tài liệu tham khảo:
https://x.com/wafer_ai/status/2083628389903315406
https://x.com/ChiragAsarpota/status/2083864019870634151
Bài viết này từ tài khoản công chúng WeChat "机器之心" (ID:almosthuman2014), tác giả: 关注LLM的
