Mixture-of-Kittens: Nhân kernel MoE nguồn mở tối ưu cho hệ thống NVL72

Hôm nay, chúng tôi chính thức mở mã nguồn Mixture-of-Kittens (MoK), megakernel đào tạo MoE (Mixture-of-Experts) chuyên dụng của chúng tôi dành cho các hệ thống NVL72.

Trong quá trình mở rộng quy mô đào tạo và suy luận cho Composer, mô hình lập trình dựa trên tác nhân (agentic coding model) của chúng tôi, lớp mixture-of-experts luôn là điểm nghẽn chính. Tùy thuộc vào khối lượng công việc và cấu hình đào tạo, nó có thể chiếm hơn một nửa thời gian đào tạo từ đầu đến cuối (end-to-end).

MoK giải quyết điểm nghẽn đó bằng cách hợp nhất toàn bộ quá trình truyền tải và tính toán của MoE vào một kernel duy nhất, hoàn toàn tất định (deterministic). Hiện tại, nó đang vận hành quá trình đào tạo Composer trên hàng chục nghìn GPU.

Bạn có thể dùng thử MoK và khám phá mã nguồn trên GitHub. Chúng tôi rất mong nhận được phản hồi và đóng góp từ các bạn.

MoK ra đời từ một vài nỗ lực trước đó nhằm tăng tốc lớp MoE. Trong năm qua, chúng tôi đã tự viết các kernel đào tạo MXFP8 và NVFP4, đồng thời phát triển phương pháp "warp decode" cho suy luận MoE.

Tuy nhiên, các kỹ thuật đó chỉ tối ưu hóa phần tính toán của lớp và giả định rằng việc truyền tải giữa các GPU sẽ được xử lý riêng biệt. Trong các khối lượng công việc thực tế của chúng tôi, truyền tải đã trở thành yếu tố giới hạn. Điều đó dẫn đến việc chúng tôi phải thiết kế lại toàn bộ lớp MoE từ những nguyên lý cơ bản, với việc tích hợp trực tiếp truyền tải vào kernel.

Ngoài ra, việc chuyển sang sử dụng GB300 NVL72 đã thay đổi bài toán theo hai cách quan trọng. Thứ nhất, NVL72 là một rack đa nút nằm trong một miền NVLink duy nhất, cho phép chồng lấp (overlap) nhanh và chi tiết giữa tính toán và truyền tải trên tất cả 72 GPU.

Thứ hai, các CPU Grace tích hợp (chữ "G" trong GB300) thường chậm hơn so với GPU. Chúng tôi nhận thấy các luồng GPU dễ dàng bắt kịp công việc phía CPU, khiến GPU hoàn toàn nhàn rỗi trong thời gian đó. Vì vậy, chúng tôi phải giảm thiểu tối đa công việc của CPU và sự đồng bộ hóa CPU-GPU.

Giải pháp của chúng tôi cho tập hợp các thách thức này là Mixture-of-Kittens (MoK), một megakernel đào tạo MoE được tối ưu hóa cao, xây dựng từ những nguyên lý cơ bản cho NVL72. MoK hợp nhất toàn bộ truyền tải và tính toán MoE vào một kernel duy nhất, hoàn toàn tất định và đạt hiệu suất hàng đầu so với các triển khai công khai hiện có.

Trong stack đào tạo thực tế của chúng tôi trên nhiều rack NVL72, MoK đã tăng số lượng token mỗi giây (tokens per second) lên 1,41 lần.

Phần còn lại của bài viết này giải thích các ý tưởng chính đằng sau MoK, bao gồm cách chúng tôi chọn hướng truyền tải phù hợp, cấu trúc sự chồng lấp giữa tính toán và truyền tải, cũng như loại bỏ sự đồng bộ hóa CPU-GPU bằng bộ đệm vòng (ring token buffers). Chúng tôi cũng đề cập đến thiết kế megakernel, tính tất định, hỗ trợ MXFP8 và một số chi tiết triển khai khác.

Chồng lấp tính toán và truyền tải trong MoE

MoK nhắm đến các lớp MoE theo phong cách DeepSeek-V3 (DSV3), vốn được sử dụng rộng rãi trong các mô hình mở như GLM, Qwen, Kimi (lên đến K2.7) và chính DSV. Các lớp này kết hợp một chuyên gia dùng chung (shared expert) với nhiều chuyên gia được định tuyến (routed experts), thường là hàng trăm chuyên gia.

Đối với mỗi token đi vào lớp MoE, một phép chiếu bộ định tuyến (router projection) sẽ chọn top-k chuyên gia được định tuyến và gán trọng số bộ định tuyến cho từng chuyên gia. Sau đó, mỗi chuyên gia được chọn sẽ thực hiện tính toán mạng chuyển tiếp (feed-forward network) tiêu chuẩn, bao gồm các phép chiếu up và gate, hàm kích hoạt SwiGLU và phép chiếu down. Lớp này kết hợp đầu ra của các chuyên gia dùng chung và chuyên gia được định tuyến bằng cách sử dụng các trọng số bộ định tuyến.

Chúng tôi sử dụng ký hiệu sau:

Và lớp MoE tính toán:

trong đó

Với tính song song chuyên gia (expert parallelism - EP), chúng tôi phân mảnh (shard) các chuyên gia được định tuyến và trải rộng trọng số của chúng trên nhiều GPU (hoặc các rank), và chúng tôi gọi số lượng rank cùng giữ tất cả trọng số chuyên gia là bậc EP (EP degree). Ví dụ, với 256 chuyên gia được định tuyến và bậc EP là 64, mỗi rank giữ 4 chuyên gia được định tuyến, cộng với chuyên gia dùng chung. Kết quả là, các token phải được chuyển qua các GPU trước và sau lớp MoE, tùy theo phép chiếu bộ định tuyến.

Triển khai đơn giản nhất của MoE phân tán là gửi từng token đến các rank giữ các chuyên gia được chỉ định (dispatch all-to-all), chạy FFN, trả kết quả về rank gốc của từng token (combine all-to-all) và lấy tổng có trọng số của các đầu ra chuyên gia. Nhưng vì việc truyền tải có thể mất thời gian bằng chính việc tính toán, nên việc chạy tuần tự cả hai là không hiệu quả.

Giải pháp tiêu chuẩn là chồng lấp việc truyền tải dispatch/combine với FFN của từng chuyên gia thông qua đường ống (pipelining): truyền một khối token, tính toán FFN trên đó trong khi chồng lấp việc truyền khối tiếp theo, và lặp lại. MoK là một biến thể của sơ đồ này, với một tập hợp các kỹ thuật mới, đặc thù giúp nó nhanh hơn các baseline hiện có.

Chọn hướng truyền tải phù hợp

Khi gửi token qua các GPU, người ta có thể chọn cơ chế dựa trên đẩy (push-based), nơi GPU sở hữu token chủ động lưu trữ chúng vào các GPU đích từ xa, hoặc cơ chế dựa trên kéo (pull-based), nơi GPU cần token sẽ tải chúng từ các GPU nguồn từ xa. Các phương pháp hiện tại thường dựa vào truyền tải push-based để phân tán và thu thập token qua các GPU (ví dụ: DeepEP).

Quan niệm phổ biến là việc đẩy (push) làm bão hòa các liên kết giữa các GPU tốt hơn, vì nó liên quan đến ít truyền tải giao thức hơn, do đó nó trở thành lựa chọn mặc định. Tuy nhiên, quan sát của chúng tôi là mỗi cơ chế đều có những đánh đổi riêng, và việc chọn đúng cơ chế cho từng toán tử truyền tải là rất quan trọng để tối đa hóa hiệu suất, vì ba lý do sau.

Lập lịch

Để dispatch và combine token nhanh nhất có thể, các điều kiện sau phải được thỏa mãn:

Với dispatch dựa trên push, chúng tôi cần tạo một bảng lịch trình với các cột {src_index, dst_rank, dst_index}, trong đó src_index là chỉ số vào bộ đệm kích hoạt đến cục bộ và dst_index là vị trí trong bộ nhớ của rank đích nơi token phải đến. Chỉ số hàng của bảng này sẽ quyết định thứ tự các token được gửi qua NVLink. Chúng tôi muốn các hàng của bảng này xoay vòng qua dst_rank theo kiểu round-robin, để mọi lane được kết nối luôn bận rộn.

Trong số các hàng nhắm đến một dst_rank nhất định, các giá trị dst_index phải được xen kẽ đều trên tất cả các GPU nguồn, và chúng cũng phải tuân theo thứ tự chuyên gia cục bộ tăng dần trên rank đích đó. Việc xây dựng bảng này liên quan đến nhiều phép sắp xếp, và lịch trình của mỗi rank phải tính đến lịch trình của mọi rank khác, vì không có hai rank nguồn nào có thể ghi vào cùng một dst_index.

Lịch trình dựa trên push: các mục liên tiếp phải xen kẽ rank đích để đạt độ bão hòa mạng lưới tối đa.

Lịch trình dựa trên pull: các mục liên tiếp phải xen kẽ rank nguồn để đạt độ bão hòa mạng lưới tối đa.

Với dispatch dựa trên pull, bảng lịch trình được đơn giản hóa thành hai cột, {src_rank, src_index}, và chỉ số hàng của bảng tương ứng trực tiếp với chỉ số token trong bộ đệm đích cục bộ.

Về lý thuyết, số lượng hàng trong bảng và trong bộ đệm đích sẽ khớp nhau (điều này đòi hỏi phân bổ bộ nhớ lớn; sẽ nói thêm ở phần sau). Không cần sắp xếp ở đây. Chúng tôi duyệt qua các kết quả phép chiếu bộ định tuyến, và bất cứ khi nào tìm thấy một token cần đến rank hiện tại, chúng tôi sẽ ghi rank nguồn và chỉ số của nó vào lịch trình của mình. Thuật toán trông như sau:

Đầu vào: tensor định tuyến, trong đó là chỉ số toàn cục của chuyên gia được gán cho tuyến thứ -th của token trên rank nguồn; rank cục bộ; số chuyên gia mỗi rank

Đầu ra: số lượng token mỗi chuyên gia cục bộ, số lượng token mỗi chuyên gia cục bộ và rank nguồn, độ lệch vùng, bảng lịch trình

cho đến, đến, đến thực hiện

nếu thì

kết thúc nếu

kết thúc cho

cho đến thực hiện

kết thúc cho