Giới thiệu mKernel: Thư viện nhân hợp nhất đa GPU, đa nút cho giao tiếp dựa trên GPU

Chi phí truyền thông GPU là một nút thắt cổ chai có thể đo lường được trong các tác vụ AI sản xuất. Theo dữ liệu được dự án mKernel trích dẫn, truyền thông có thể tiêu tốn 43,6% thời gian chuyển tiếp (forward pass) và 32% tổng thời gian huấn luyện (end-to-end training). Trong các mô hình Mixture-of-Experts (MoE) phổ biến, truyền thông giữa các thiết bị có thể chiếm tới 47% tổng thời gian thực thi. Các nhà nghiên cứu từ dự án UCCL của Đại học UC Berkeley đã phát hành mKernel, một thư viện các nhân CUDA (CUDA kernels) bền vững, tích hợp truyền thông NVLink nội bộ nút, RDMA liên nút và tính toán vào một nhân duy nhất. Vấn đề: Truyền thông do máy chủ điều khiển Mô hình tiêu chuẩn cho truyền thông đa GPU là do máy chủ điều khiển: CPU chạy đường dẫn điều khiển và gọi vào một thư viện như NCCL hoặc NVSHMEM. Thư viện này thực hiện thao tác tập thể — một AllReduce, một AllGather, v.v. — trên các GPU. Tính toán và truyền thông chạy trên các luồng CUDA riêng biệt và chồng chéo tại các ranh giới nhân. Nhóm nghiên cứu xác định hai vấn đề với phương pháp này: (1) CPU không mở rộng quy mô cùng với tính toán GPU. Một giá đỡ GB300 NVL72 tích hợp 72 GPU Blackwell Ultra và 36 CPU Grace, cung cấp hiệu suất Tensor Core 720 PFLOP/s FP8/FP6, 1,44 EFLOP/s FP4 và băng thông NVLink toàn bộ 130 TB/s trong cùng một giá đỡ. Với tốc độ đó, chi phí điều phối máy chủ ở quy mô micro giây — một lệnh gọi cudaLaunchKernel, một kiểm tra "tất cả các ghi đã hoàn tất" phía CPU, một sự kiện giữa các luồng — xuất hiện trực tiếp dưới dạng các bong bóng đường ống. (2) Các hệ thống do máy chủ điều khiển chồng chéo tính toán và truyền thông tại các ranh giới nhân thô. Việc chồng chéo chi tiết hơn ở cấp độ ô hoặc khối không thể thực hiện được từ phía máy chủ. Giải pháp thay thế là truyền thông do GPU điều khiển: bản thân GPU kích hoạt các truyền tải, với truyền thông được tích hợp vào cùng một nhân với tính toán. Hầu hết các thư viện nhân tích hợp hiện có hoạt động trong một nút duy nhất hoặc một GPU duy nhất. mKernel nhắm mục tiêu đến trường hợp đa nút. mKernel hoạt động như thế nào mKernel là một thư viện các nhân CUDA bền vững. Mỗi nhân tích hợp truyền thông NVLink nội bộ nút, RDMA liên nút và tính toán dày đặc vào một nhân duy nhất. Đa GPU + đa nút, trong một nhân: Cả NVLink nội bộ nút và RDMA liên nút đều nằm trong cùng một nhân bền vững. Chồng chéo nội bộ nhân chi tiết: Tính toán và truyền thông chồng chéo ở mức độ hạt ô/khối, bao gồm cả truyền thông GPU nội bộ nút và liên nút. Nhân bền vững với chuyên môn hóa SM: Các CTA tự gán vai trò: tính toán, truyền thông nội bộ, gửi liên nút, giảm liên nút. Số lượng SM dành cho mỗi vai trò có thể điều chỉnh theo hình dạng. Mạng do GPU điều khiển được xây dựng trên libibverbs: mKernel sử dụng các ghi RDMA do GPU khởi tạo mà không phụ thuộc vào NCCL hoặc NVSHMEM. Phần phụ trợ truyền thông được viết lại từ đầu để tối đa hóa hiệu suất và hỗ trợ các thiết bị mạng không đồng nhất. Năm nhân tích hợp Kernel (nhân) | Chức năng kết hợp | Mô tả ---|---|--- AllGather + GEMM | AllGather → GEMM | Mỗi thứ hạng (rank) giữ một phần của A. Trong khi các thứ hạng thu thập các phần của các thứ hạng ngang hàng qua NVLink/RDMA, GEMM cục bộ tiêu thụ các ô (tile) ngay khi chúng đến. GEMM + AllReduce | GEMM → AllReduce | Tính toán C = A @ B và giảm các đầu ra cục bộ trên tất cả các thứ hạng trong một lần khởi chạy. Các ô đầu ra được đẩy vào cây giảm ngay khi chúng được tạo ra. MoE Dispatch + GEMM | Điều phối All-to-All → GEMM nhóm | Định tuyến các token MoE đến các thứ hạng chuyên gia của chúng (NVLink nội bộ nút + All-to-All liên nút) và chạy GEMM nhóm cho từng chuyên gia trong cùng một kernel. Các token được xử lý ngay khi chúng đến – không có chuyến đi khứ hồi bộ đệm tạm thời. Ring Attention | Trao đổi KV vòng → FlashAttention | Chú ý song song theo chuỗi trên các thứ hạng. Mỗi bước xoay một khối KV quanh vòng trong khi FlashAttention cục bộ tiêu thụ khối đã nhận trước đó. Tính toán và gửi/nhận vòng chạy đồng thời bên trong một kernel liên tục duy nhất. GEMM + ReduceScatter | GEMM → ReduceScatter | Tính toán C = A @ B và giảm-phân tán đầu ra. Mỗi ô đầu ra được giảm và chuyển tiếp đến thứ hạng sở hữu nó ngay khi nó được tạo ra. Thiết lập đánh giá Nhóm nghiên cứu đã đánh giá mKernel trên hai cụm 2 nút × 8 H200, chỉ khác nhau về cấu trúc liên kết giữa các nút: Hệ thống thử nghiệm | Nút × GPU | Kết nối nội bộ nút | Giao thức liên nút ---|---|---|--- NIC AWS EFA | 2 × 8 H200 | NVLink | AWS EFA / SRD 16 × 200 Gb/s EFA mỗi nút ConnectX-7 | 2 × 8 H200 | NVLink | InfiniBand 8 × 400 Gb/s NVIDIA ConnectX-7 mỗi nút mKernel đã được đánh giá so với NCCL, Triton-distributed, Flux, Mercury, MagiAttention, Transformer-Engine và ring-flash-attention. Nhóm nghiên cứu lưu ý rằng việc đánh giá thêm ở quy mô lớn hơn vẫn đang được tiến hành. Các phần phụ trợ và yêu cầu mKernel hỗ trợ hai phần phụ trợ mạng: Phần phụ trợ | Macro | Giao thức | Nơi chạy ---|---|---|--- CX7-D | INTERNODE_BACKEND_IBVERBS | libibverbs RC | ConnectX-7 / InfiniBand / RoCE EFA-D | INTERNODE_BACKEND_EFA | libibverbs + efadv (SRD) | AWS p5/p5e (H200, EFA) Cả hai phần phụ trợ đều chia sẻ cùng một API phía máy chủ và cùng một kernel trên GPU. Chỉ có việc triển khai proxy/phiên là khác nhau (session.h cho CX7, session_efa.h cho EFA). Yêu cầu: GPU NVIDIA Hopper (bản dựng mặc định nhắm mục tiêu sm_90a), CUDA 12.9, Python với PyTorch. Phần phụ trợ CX7 yêu cầu các tiêu đề và thư viện phát triển libibverbs. Phần phụ trợ EFA yêu cầu cài đặt AWS EFA với libfabric, libibverbs, efadv và các tiêu đề EFA trong EFA_HOME=/opt/amazon/efa theo mặc định. Giải thích trực quan của Marktechpost UCCL Hướng dẫn mKernel — Các Kernel kết hợp đa GPU, đa nút 01 / 07 — Tổng quan mKernel là gì? mKernel là một thư viện mã nguồn mở gồm các kernel CUDA liên tục từ dự án UCCL của UC Berkeley. Nó kết hợp giao tiếp NVLink nội bộ nút, RDMA liên nút và tính toán dày đặc vào một kernel duy nhất. Hầu hết các thư viện kernel kết hợp hiện có hoạt động trong một nút hoặc một GPU duy nhất. mKernel được thiết kế ngay từ đầu để mở rộng qua các ranh giới nút. 43,6% thời gian chuyển tiếp bị tiêu thụ bởi giao tiếp trong sản xuất 47% tổng thời gian thực thi trong các mô hình MoE phổ biến 32% tổng thời gian huấn luyện bị tiêu thụ bởi giao tiếp 02 / 07 — Vấn đề Tại sao giao tiếp do máy chủ điều khiển không hiệu quả Mô hình tiêu chuẩn là do máy chủ điều khiển: CPU gọi NCCL hoặc NV