Inferact, công ty do đội ngũ cốt cán của vLLM sáng lập, vừa công bố một loạt bài kiểm tra suy luận: chạy mô hình Kimi K3 của Moonshot AI trên 16 chip Google TPU v7 Ironwood, tốc độ giải mã đơn luồng đạt 709 token mỗi giây, trong khi nhóm đối chứng dùng 16 chip Nvidia GB200 chỉ đạt 452 token mỗi giây, nhanh hơn khoảng 57%.
Báo cáo kiểm tra và mã nguồn được công bố cùng lúc vào ngày 23/9, kho mã inferact/tpu-megakernels được mở theo giấy phép Apache 2.0. Inferact là startup do đội ngũ vLLM lập nên, trước đó đã gọi được 150 triệu USD vòng hạt giống, định giá 800 triệu USD, do a16z và Lightspeed dẫn dắt.
Trên giấy tờ, GB200 vẫn nhỉnh hơn
Trước hết là các thông số trên giấy. Theo số liệu trong báo cáo, một chip TPU v7 có hiệu năng đỉnh BF16 là 2,31 PFLOPS, FP8 là 4,61 PFLOPS, đi kèm 206GB HBM với băng thông 7380 GB/s; một chip GB200 tương ứng là 2,5 PFLOPS, 5 PFLOPS, 186GB HBM, băng thông 8000 GB/s. Về hiệu năng tính toán và băng thông, GB200 đều nhỉnh hơn một khoảng, TPU chỉ hơn ở dung lượng bộ nhớ, nhiều hơn 20GB.
Kimi K3 là mô hình MoE 92 lớp, phần attention kết hợp cả Kimi Delta Attention và Multi-head Latent Attention (MLA). Bài kiểm tra dùng song song tensor 4 chiều, song song chuyên gia 8 chiều, chia trên 16 chip.
So sánh tốc độ thô khi chưa bật giải mã suy đoán:
| Kích thước batch | TPU v7 (megakernel) | GB200 (vLLM) |
|---|---|---|
| 1 | 249 | 127 |
| 2 | 392 | 227 |
| 4 | 515 | 373 |
| 8 | 865 | 636 |
Đơn vị là số token mỗi giây. Khi kích thước batch bằng 1, TPU gần như nhanh gấp đôi; khi batch tăng lên 8, khoảng cách thu hẹp còn hơn ba mươi phần trăm. Sau khi bật giải mã suy đoán DSpark do DeepSeek đề xuất, tốc độ đơn luồng trở thành 709 so với 452, mỗi bước giải mã mất khoảng 8,5 mili giây.
Nhanh nhờ gộp 92 lớp thành một chương trình
Cách làm của Inferact gọi là megakernel. Trong các framework suy luận thông thường, phép tính của mỗi lớp được tách thành nhiều nhân (kernel) độc lập, khởi chạy lần lượt, giữa các kernel có khoảng trống, trọng số cũng chỉ có thể bắt đầu chuyển từ bộ nhớ lên chip khi kernel tương ứng bắt đầu chạy.
Inferact dùng Pallas viết toàn bộ 92 lớp của một bước giải mã thành một chương trình duy nhất. Theo báo cáo, một megakernel xóa bỏ ranh giới giữa các kernel, việc nạp trọng số không còn bị ràng buộc với kernel sử dụng nó, bộ nhớ trên chip chứa được bao nhiêu thì có thể lấy trước bấy nhiêu. Khi giải mã đơn luồng, phần lớn thời gian chip đang chờ dữ liệu, việc lấy trước dữ liệu bù đúng vào chỗ đó; khi batch tăng lên, tỷ trọng tính toán tăng theo, lợi ích của cách này mỏng dần, điều này khớp với xu hướng trong bảng số liệu.
Một lợi ích đi kèm là thời gian biên dịch: trước đây biên dịch bằng XLA mất hơn 30 phút, còn megakernel chưa đến 90 giây. Để chứng minh tốc độ tăng không đánh đổi bằng chất lượng, báo cáo kèm theo điểm số của Kimi K3 chạy trên TPU: GPQA-Diamond đạt 94,4%, GSM8K đạt 97,2%.
Ước tính sơ bộ hiệu suất trên mỗi đơn vị năng lực tính toán
Dùng kết quả giải mã suy đoán đơn luồng để ước tính: hiệu năng đỉnh BF16 của TPU v7 bằng khoảng 92% của GB200, nhưng tốc độ token tạo ra lại gấp 1,57 lần đối phương, quy đổi trên mỗi đơn vị hiệu năng đỉnh, TPU đạt khoảng 1,7 lần GB200. Quy đổi theo băng thông, kết quả cũng quanh mức 1,7 lần.
Con số này cần chiết khấu một chút. Phía GB200 dùng công thức Kimi K3 do vLLM công bố công khai, thuộc lộ trình nhiều kernel thông thường, chưa được cải tạo theo hướng megakernel với mức độ tương đương. Báo cáo không có đối chứng tối ưu tương xứng trên GB200, nên chưa thể tách bạch khoảng cách này đến từ phần cứng hay từ công sức tối ưu phần mềm. Phía Nvidia cũng chưa phản hồi về những con số này.
Đối với Moonshot AI, bài kiểm tra này cho thấy một hướng triển khai khả thi cho Kimi K3 trên phần cứng ngoài Nvidia. Kimi K3 là mô hình trọng số mở với 2,8 nghìn tỷ tham số, quy mô này khiến ngưỡng tự triển khai khá cao; có thêm nhân TPU mã nguồn mở, việc thuê TPU trên nền tảng đám mây để chạy K3 có thêm một lựa chọn.
Nguồn tham khảo: blog kỹ thuật Inferact, CocoLoop, kho mã inferact/tpu-megakernels, QbitAI; số liệu thông lượng theo cách đo 16 chip đối 16 chip trong báo cáo của Inferact, thông tin gọi vốn theo các nguồn công khai.