Giao tiếp mixture-of-experts hiệu quả và có tính di động

Chúng tôi giới thiệu một thư viện mã nguồn mở, có tính di động và hiệu suất cao dành cho giao tiếp Mixture-of-Experts (MoE), đạt hiệu suất nhanh gấp 10 lần so với các nguyên hàm giao tiếp All-to-All tiêu chuẩn.

Việc triển khai của chúng tôi có một số cải tiến kỹ thuật chính mang lại hiệu quả giao tiếp MoE vượt trội:

Giao tiếp do GPU khởi tạo (IBGDA): Hỗ trợ giao tiếp trực tiếp từ GPU đến NIC, giúp giảm đáng kể độ trễ bằng cách bỏ qua sự tham gia của CPU

Chồng lấp giao tiếp và tính toán: Kiến trúc kernel phân chia với các giai đoạn gửi và nhận riêng biệt cho phép quá trình tính toán diễn ra đồng thời với việc truyền dữ liệu qua mạng

Hiệu suất đơn nút nhanh nhất: Độ trễ thấp hơn 2,5 lần so với việc triển khai nhanh nhất trước đó trên cấu hình đơn nút

Hiệu suất đa nút hiệu quả và có tính di động: Quá trình triển khai của chúng tôi đạt tốc độ nhanh hơn tới 10 lần so với giao tiếp all-to-all tiêu chuẩn. Mặc dù chậm hơn khoảng 2 lần so với các cách triển khai mang tính chuyên biệt cao, phương pháp của chúng tôi mang lại tính di động tốt hơn trên các phiên bản NVSHMEM và môi trường mạng (NVLink, CX-7 và EFA)

Thư viện này hoàn toàn là mã nguồn mở và có sẵn tại https://github.com/ppl-ai/pplx-kernels.

Trong bài viết này, chúng tôi khám phá những thách thức của tính song song chuyên gia trong các mô hình MoE quy mô lớn và mô tả cách tiếp cận của mình đối với việc điều phối và kết hợp token hiệu quả trên các môi trường GPU phân tán.

Giới thiệu

Các mô hình Mixture-of-Experts (MoE), chẳng hạn như DeepSeek R1 và Mixtral 8x7B, cải thiện các mô hình dày đặc bằng cách giới hạn số lượng trọng số được kích hoạt cho mỗi token. Thay vì đạt được số lượng tham số cao bằng cách tăng số lượng lớp hoặc tăng kích thước của các lớp tuyến tính trong Multi-Layer Perceptron (MLP) của mỗi lớp bộ giải mã, các mô hình MoE thay thế MLP truyền thống bằng nhiều chuyên gia (expert) và một bộ định tuyến (router). Ví dụ, trong tổng số 671B tham số của DeepSeek R1, chỉ có 37B tham số được nhân với một token cho trước trong quá trình suy luận. Điều này làm giảm đáng kể lượng tính toán cần thiết để giải mã một token và cải thiện độ trễ so với mô hình dày đặc.

Các mô hình MoE đặt ra một số thách thức bổ sung cho quá trình suy luận so với các mô hình dày đặc. Mặc dù bản thân các chuyên gia là các lớp MLP nhỏ, mỗi lớp bộ giải mã lại bao gồm một bộ định tuyến để quyết định token được chuyển đến chuyên gia nào, trong đó mỗi token được chuyển đến nhiều chuyên gia. Bộ định tuyến thường là một lớp tuyến tính nhỏ tạo ra phân phối xác suất. Thông thường, các chuyên gia có điểm số top-K sẽ được chọn và hoạt động kích hoạt cuối cùng được tính toán dưới dạng trọng số trung bình, tổng hợp các đầu ra của chuyên gia nhân với trọng số có nguồn gốc từ phân phối xác suất của bộ định tuyến.

Để giảm thiểu độ trễ trong các hệ thống MoE phân tán, tính song song có thể được khai thác trên nhiều thiết bị, nhưng điều này lại mang đến những thách thức về giao tiếp. Các mô hình như Mixtral 8x7B hoặc Llama-70B vừa vặn trong 8 thiết bị trên một nút duy nhất, hưởng lợi từ các kết nối liên lạc NVLink nhanh và có độ trễ thấp (lên tới 900Gbps). Tuy nhiên, các mô hình lớn hơn đòi hỏi phải phân đoạn (shard) các chuyên gia trên nhiều nút sử dụng InfiniBand (đạt đỉnh ở mức 400Gbps), điều này làm phát sinh thêm những thách thức về độ trễ.

Phóng to

Trong bài viết này, chúng tôi khám phá bài toán về tính song song chuyên gia, mô tả việc triển khai các GPU kernel thực hiện việc điều phối token đến các thiết bị nắm giữ các chuyên gia được chỉ định một cách hiệu quả và thu thập lại chúng bằng cách tính tổng trọng số của các đầu ra chuyên gia. Mặc dù các bản triển khai được tối ưu hóa mạnh mẽ nhắm đến phần cứng mạng cụ thể đã tồn tại, chúng tôi giới thiệu một bản triển khai có tính di động cao, dựa trên một tập hợp tối thiểu các nguyên hàm NVSHMEM nhưng vẫn mang lại hiệu suất nhanh hơn 10 lần so với giao tiếp all-to-all tiêu chuẩn. Quá trình triển khai của chúng tôi đạt được hiệu suất tiên tiến trên các cấu hình đơn nút đồng thời duy trì tính di động tuyệt vời trên nhiều môi trường mạng khác nhau.

Phân đoạn và Tính song song

Để suy luận hiệu quả, trọng số của một mô hình phải được lưu giữ trong bộ nhớ thiết bị, đồng thời để lại đủ Space cho các hoạt động kích hoạt, bộ nhớ đệm KV và các bộ đệm khác cần thiết cho quá trình truyền xuôi (forward pass) qua mô hình. Các mô hình mạnh mẽ nhất vượt quá dung lượng của cả những GPU mạnh mẽ nhất, do đó quá trình suy luận phải được trải rộng trên nhiều thiết bị, nơi có thể cùng nhau lưu trữ trọng số của mô hình. Dựa trên các sơ đồ phân đoạn trọng số, các sơ đồ giao tiếp và tính toán khác nhau phải được sử dụng để giảm tải tính toán và đồng bộ hóa các thiết bị.

Phóng to

Tính song song chuyên gia (Expert Parallelism - EP), được minh họa trong hình đầu tiên, chỉ thực hiện song song hóa quá trình tính toán của chuyên gia. Các chuyên gia khác nhau được gán cho các thiết bị khác nhau, nơi lưu giữ trọng số của họ. Sau khi định tuyến, các token được gửi đến thiết bị tương ứng, với các kết quả được thu thập và tích lũy sau đó. Độ phức tạp của việc định tuyến phụ thuộc vào mức độ song song trong các phần khác của mô hình: việc nhân bản các lớp khác hoàn toàn có thể loại bỏ nhu cầu định tuyến, vì mỗi thứ bậc (rank) có thể chọn các token từ bảng định tuyến được nhân bản cục bộ. Tuy nhiên, nếu chỉ có một thứ bậc chạy định tuyến, việc phát quảng bá (broadcast) là bắt buộc để điều phối các token, chỉ số và trọng số đến các chuyên gia tương ứng của chúng. Cuối cùng, thao tác all-gather hoặc phát quảng bá all-to-all sẽ đồng bộ hóa các token đầu ra với bất kỳ thứ bậc nào tiếp tục thực thi mô hình. Việc triển khai như vậy tương đối đơn giản vì torch đã cung cấp các nguyên hàm cần thiết, mặc dù một số lợi ích có thể đạt được bằng cách kết hợp (fusion) quá trình tích lũy với thao tác thu thập (gather) nhằm gom các token từ các chuyên gia.

Phóng to

Tính song song chỉ dành cho chuyên gia không mở rộng quy mô một cách tối ưu, vì các nút trong một cụm có thể ở trạng thái nhàn rỗi trong khi mô hình đang chạy các lớp không phải chuyên gia, chẳng hạn như attention, norm và sampling. Tuy nhiên, việc tính toán các lớp này, chủ yếu là attention, cũng có thể hưởng lợi từ Tính song song Ten-xơ (Tensor Parallelism - TP). Hầu hết các mô hình đều dựa vào multi-head attention, có nghĩa là các attention head và các phép chiếu Q, K, V tương ứng của chúng cũng có thể được phân đoạn trên các thiết bị, bằng cách nhân bản hoặc thu thập các lát cắt giữa các lớp khác nhau. Nếu attention được trải đều trên tất cả các thiết bị, thao tác all-gather có thể đồng bộ hóa các hoạt động kích hoạt, cho phép quá trình định tuyến được nhân bản, từ đó đòi hỏi các nguyên hàm đồng bộ hóa tương tự như trường hợp tính song song chuyên gia để triển khai hiệu quả. Tuy nhiên, có những giới hạn đối với tính song song ở cấp độ này, vì việc giảm số lượng attention head xuống dưới một ngưỡng nhất định sẽ mang lại hiệu quả giảm dần.

Phóng to

Để tận dụng tối đa tất cả các thiết bị và hỗ trợ mức độ song song chuyên gia rất cao, lên đến 128 hoặc 256 GPU, Tính song song dữ liệu (Data Parallelism - DP) là cần thiết. Theo sơ đồ này, các thiết bị mà mô hình được chia cắt sẽ được nhóm lại để xử lý các yêu cầu đồng thời, tính toán attention và duy trì các bộ nhớ đệm KV được phân đoạn trên nhóm cục bộ của chúng. Nhiều thể hiện (instance) của các nhóm song song này cộng tác trong việc đánh giá chuyên gia, với mỗi thể hiện lưu trữ một tập hợp con trọng số chuyên gia khác nhau. Dựa trên số lượng attention head, một nhóm thường có thể kéo dài tới kích thước của toàn bộ nút, vì giao tiếp nội bộ nút (intra-node) diễn ra nhanh hơn. Ví dụ, trong hình ở trên, một thứ bậc DP phục vụ độc lập hai yêu cầu A và B, xử lý attention, norm và bất kỳ sổ sách kế toán nào khác cho các yêu cầu đó. Thứ bậc DP còn lại xử lý một yêu cầu C riêng biệt. Tuy nhiên, nút đầu tiên lưu trữ một nửa số chuyên gia, trong khi nút kia lưu trữ nửa còn lại, do đó sau khi định tuyến, các token từ A và B có thể được gửi đến nút thứ hai và ngược lại. Điều này dẫn đến một bài toán giao tiếp thưa thớt (sparse communication): mỗi thiết bị có thể gửi một số lượng token khác nhau đến bất kỳ thứ bậc đích nào khác. Các nguyên hàm hiện có từ torch, chủ yếu là all_to_all, không thực sự phù hợp, vì chúng có thể yêu cầu một số hình thức đệm (padding) hoặc đồng bộ hóa GPU với CPU và phát quảng bá siêu dữ liệu. Để thực hiện giao tiếp hiệu quả, các kernel tùy chỉnh là cần thiết để xử lý thông tin định tuyến và khởi tạo giao tiếp từ các thiết bị mà không cần sự can thiệp của CPU. Sau khi định tuyến, một dispatch kernel phải gửi các token đến các thứ bậc mà chúng được định tuyến tới, trong khi ở phía kết hợp (combine), các hoạt động kích hoạt thuộc về các yêu cầu trong nhóm DP hiện tại phải được thu thập. Ngoài ra, công việc phải được cân bằng trong một nhóm DP có thể có nhiều thiết bị, nhằm giảm thiểu việc phát quảng bá và nhân bản chính xác các token trên tất cả các thứ bậc cho các lớp tiếp theo.

NVSHMEM

NVSHMEM là một triển khai OpenSHMEM dành riêng cho NVIDIA, cung cấp các tiện ích giao tiếp liên thiết bị có tính di động giúp trừu tượng hóa độ phức tạp của phần cứng bên dưới. API có thể diễn tả các lệnh đọc và ghi từ thiết bị đến thiết bị, được ánh xạ tới các nguyên hàm của từng lớp truyền tải riêng lẻ. Cả NVLink và RDMA đều được hỗ trợ, mang lại một mức độ di động và cho phép các NVSHMEM kernel hoạt động bên trong một nút hoặc trên một cụm gồm nhiều nút. Trong các kernel của mình, chúng tôi sử dụng các thao tác ghi bộ nhớ từ xa và báo hiệu (signaling) để truyền dữ liệu và đồng bộ hóa các thiết bị tham gia.

Các hoạt động của NVSHMEM được xây dựng xung quanh khái niệm bộ nhớ đối xứng (symmetric memory): chúng hoạt động trên các bộ đệm đã được phân bổ trên tất cả các thiết bị tham gia vào giao tiếp liên thiết bị. Một thiết bị cục bộ có thể sử dụng một địa chỉ có nguồn gốc từ bên trong bộ đệm cục bộ của chính nó cùng với chỉ số của thứ bậc đích để chỉ định đích đến của một thao tác. Hình dưới đây minh họa khái niệm này: cả hai GPU đều phân bổ các bộ đệm đối xứng có cùng kích thước, giữ lại các con trỏ src trỏ đến chúng. Thứ bậc đầu tiên muốn gửi 3 số nguyên đến thứ bậc thứ hai, đặt chúng ở đầu bộ đệm: nvshmem_int_put_nbi lấy ra địa chỉ bắt đầu từ bộ đệm cục bộ, chỉ định thiết bị đích. Thứ bậc thứ hai tạo ra một độ lệch (offset) từ bộ đệm của chính nó, gửi một phần tử đến thiết bị đầu tiên, với độ lệch là một. Mặc dù các địa chỉ đích phải luôn là các bộ đệm đối xứng được phân bổ bằng cách sử dụng nvshmem_alloc, các bộ đệm nguồn có thể là các vùng tùy ý của bộ nhớ thiết bị, miễn là chúng đã được đăng ký trước với NVSHMEM.

Mặc dù NVSHMEM cung cấp một loạt các nguyên hàm, các kernel của chúng tôi chỉ dựa vào 3 hàm, và xây dựng tất cả các cơ chế đồng bộ hóa và rào cản (fencing) dựa trên chúng.

Phóng to

nvshmemx_putmem_signal_nbi_warp: Truyền một khối dữ liệu từ thiết bị này sang thiết bị khác, đồng thời thiết lập một cờ (flag) trên thiết bị từ xa. Thao tác này hoặc thiết lập (NVSHMEM_SIGNAL_SET) hoặc tăng (NVSHMEM_SIGNAL_ADD) một vị trí 64-bit. Cờ được cập nhật sau khi toàn bộ khối bộ nhớ được truyền đi. Nếu thiết bị từ xa quan sát thấy sự thay đổi trong cờ, nó có thể truy cập an toàn vào bộ đệm trong bộ nhớ của chính mình. Hàm này hữu ích cho việc kết hợp truyền dữ liệu và đồng bộ hóa.

nvshmemx_signal_op hoạt động trên một vị trí bộ nhớ đơn lẻ, thường là một cờ 64-bit, tự động thiết lập hoặc tăng giá trị của nó. Nó rất hữu ích trong việc gửi siêu dữ liệu và đồng bộ hóa các thiết bị.

nvshmem_uint64_wait_until được sử dụng ở đầu nhận của một tín hiệu, để thăm dò (poll) một cờ cho đến khi thiết bị từ xa cập nhật nó.

Việc đảm bảo bất kỳ hình thức sắp xếp thứ tự nào giữa các thao tác thông qua nvshmem_quiet hoặc nvshmem_fence là rất tốn kém, do đó việc triển khai các kernel của chúng tôi tránh sử dụng bất kỳ nguyên hàm NVSHMEM nào khác và giảm thiểu mọi sự phụ thuộc ngoài các cặp gửi-nhận riêng lẻ. Khi gửi dữ liệu, chúng tôi luôn sử dụng phiên bản hàm không chặn (non-blocking), mà không chờ dữ liệu thậm chí được gửi ra khỏi thứ bậc cục bộ. Quá trình đồng bộ hóa ngầm khác trên nhóm kernel đảm bảo rằng các bộ đệm này không bị ghi đè một cách phá hoại trong khi các NIC cần quyền truy cập vào chúng.

NVSHMEM sử dụng GPUDirect RDMA để NIC có thể đọc và ghi trực tiếp bộ nhớ GPU nhằm trao đổi dữ liệu mà không cần sự tham gia của CPU. Hơn nữa, trên các NIC ConnectX, NVSHMEM hỗ trợ giao tiếp do GPU khởi tạo (còn được gọi là Infiniband GPUDirect Async, hoặc IBGDA), cho phép GPU gửi các thao tác RDMA trực tiếp đến NIC và thăm dò hàng đợi hoàn thành của NIC trực tiếp. Trên các nền tảng không hỗ trợ giao tiếp do GPU khởi tạo, NVSHMEM tạo ra một chuỗi (thread) "proxy" trên CPU máy chủ để khởi tạo giao tiếp và thăm dò trạng thái hoàn thành thay mặt cho GPU. Chương trình NVSHMEM có tính di động bất kể giao tiếp do GPU khởi tạo có được hỗ trợ hay không. Tuy nhiên, giao tiếp do GPU khởi tạo giúp cắt giảm đáng kể độ trễ vì nó hoàn toàn bỏ qua đường vòng đến CPU.

Kernel có tính di động

Chúng tôi triển khai giao tiếp MoE thông qua một cặp kernel điều phối (dispatch) và kết hợp (combine). Các dispatch kernel chịu trách nhiệm đọc các token và thông tin định tuyến trên mỗi thứ bậc, điều phối chúng đến các chuyên gia phù hợp. Các combine kernel thu thập các hoạt động kích hoạt do các chuyên gia tạo ra và gửi chúng trở lại các thứ bậc nguồn của chúng, đồng thời tính toán trọng số trung bình từ các chuyên gia đã chọn dựa trên các trọng số được tính toán bởi bộ định tuyến. Các kernel được chia nhỏ hơn nữa thành một thành phần gửi và nhận, nhằm cho phép việc truyền dữ liệu được chồng lấp với quá trình tính toán. Các kernel gửi đều là dạng không chặn và không đồng bộ hóa: chúng đơn giản là điều phối tất cả các thao tác ghi đến các thiết bị từ xa. Ở đầu bên kia, các combine kernel chỉ đọc từ bộ nhớ, chờ đợi cho đến khi tất cả dữ liệu cần thiết được truyền xong. Sau khi điều phối công việc đến các NIC, trong khi dữ liệu được truyền không đồng bộ qua đường truyền mạng, GPU có thể thực hiện các công việc hữu ích khác tại cục bộ, chẳng hạn như áp dụng các chuyên gia chia sẻ hoặc tính toán attention.

Mỗi cặp kernel có các bộ đệm bộ nhớ đối xứng riêng, với mỗi thứ bậc phân bổ bộ nhớ riêng tư để nhận dữ liệu từ tất cả các thứ bậc khác. Số lượng token tối đa có thể được điều phối từ một nhóm DP là cố định trên các thứ bậc, điều này cũng thiết lập giới hạn trên mà mỗi thứ bậc có thể nhận cho mỗi chuyên gia cục bộ từ mỗi thứ bậc DP. Điều này cho phép các thứ bậc gửi tạo ra một địa chỉ duy nhất trên thứ bậc đích để ghi vào, mà không yêu cầu bất kỳ sự đồng bộ hóa nào giữa chúng. Sau khi nhận được dữ liệu, các kernel sẽ xáo trộn các token để sắp xếp chúng trong bộ nhớ theo định dạng tối ưu được yêu cầu bởi các kernel nhận. Mặc dù các bộ đệm có kích thước khá lớn, chúng được tái sử dụng trên tất cả các lớp tuần tự của một mô hình.

Phóng to

Hình thức đồng bộ hóa toàn cục duy nhất, như được minh họa trong hình ở trên, được triển khai trong kernel combine-receive. Khi dữ liệu được điều phối từ phía gửi, một thứ bậc sẽ trả về và tiếp tục thực thi khi nó nhận được tất cả các token cần thiết từ các thứ bậc khác và bước vào các combine send kernel ngay sau khi chúng được xử lý. Rào cản trong combine-receive đảm bảo rằng không có thứ bậc nào có thể chạy vượt trước và bắt đầu dispatch-send trong khi bất kỳ thứ bậc nào khác vẫn đang chờ nhận dữ liệu, điều này có thể gây ra hiện tượng chồng lấp mang tính phá hoại. Quá trình đồng bộ hóa được thực hiện trong combine kernel để đơn giản hóa, như một sự tối ưu hóa, rào cản tương tự có thể được triển khai giữa dispatch-receive và combine-send kernel, giúp giảm hơn nữa độ trễ bằng cách cho NIC có nhiều thời gian hơn để ổn định các lượt truyền trong khi GPU đang chạy các kernel khác.

Cả hai kernel đều được chia nhỏ trên tất cả các SM khả dụng của thiết bị: trong khi dispatch-send và combine-receive kernel phải thực hiện song song trên số lượng token tối đa trên mỗi thứ bậc (max_m), dispatch-receive và combine-send kernel phải xử lý chừng đó token từ mỗi thứ bậc và chuyên gia cục bộ, với mức tối đa là max_m * (num_experts // EP) * (EP // TP).

Điều phối (Dispatch)

Dispatch kernel chịu trách nhiệm điều phối các token và hệ số tỷ lệ khối (block scaling factor) cho các thiết bị ngang hàng (peer) lưu trữ các chuyên gia mà các token được định tuyến tới. Phía gửi dựa vào chuyên môn hóa warp (warp specialization) để thực hiện song song hai tác vụ: tổng hợp thông tin định tuyến để xác định có bao nhiêu token được gửi từ mỗi thứ bậc và đóng gói các token thành các thông điệp để điều phối đến thiết bị từ xa. Phía nhận trước tiên chờ tất cả số lượng token được nhận, sau đó nó đọc và mở gói các dữ liệu tải trọng (payload) từ bộ nhớ chia sẻ thành các tensor sẽ được chuyển tiếp đến các kernel Grouped GEMM hạ nguồn. Nó cũng lưu trữ thông tin vào các bộ đệm được chia sẻ với các combine kernel, để chỉ ra nơi mỗi token nên được gửi trả lại. Thông tin này là bắt buộc vì kernel nhận thực hiện xáo trộn các token xung quanh trong các bộ đệm liền kề theo thứ tự phi định định (non-deterministic).

Phóng to

Trong phần gửi, một warp đơn gồm 32 chuỗi (thread) trên mỗi khối chịu trách nhiệm đọc thông tin định tuyến và đếm số lượng token mà mỗi chuyên gia đích dự kiến nhận được từ thứ bậc này, gửi số lượng cộng thêm một bằng cách sử dụng nvshmem_signal_op. Số lượng được tăng lên một, vì sự chuyển đổi từ số không sang khác không ở phía từ xa biểu thị việc đã nhận được số lượng từ một thứ bậc nguồn. Song song đó, các warp còn lại hợp tác để sao chép các token vào bộ nhớ đối xứng trên tất cả các khối một cách song song, đóng gói các hoạt động kích hoạt, hệ số tỷ lệ và chỉ số của chúng trên thứ bậc cục bộ thành một khối liền kề. Chỉ số này được yêu cầu bởi combine sender để xác định địa chỉ nơi token sẽ được ghi vào. Tiếp theo, sau khi đảm bảo tất cả dữ liệu đã được sao chép thông qua một rào cản trên các nhóm warp, các warp một lần nữa hoạt động độc lập, mỗi warp gửi cùng một bộ đệm đến một chuyên gia khác nhau một cách song song. Các token được gửi bằng cách sử dụng nvshmemx_putmem_signal_nbi_warp, thao tác này cũng tự động làm tăng số lượng token đã gửi từ thứ bậc cục bộ trên thiết bị từ xa. Trong một nhóm DP, vì mỗi thứ bậc sở hữu một bản sao của token cần gửi, việc điều phối được cân bằng đều, với mỗi thiết bị gửi ra một tập hợp con các token.

Ở phía nhận, tất cả các kernel trước tiên đồng bộ hóa với các thứ bậc gửi bằng cách chờ tổng số lượng token được gửi sang. Sau đó, tất cả chúng chờ số lượng token đã gửi được tăng tự động đạt đến tổng số lượng, biểu thị rằng tất cả các payload từ các thứ bậc nguồn cũng đã được gửi sang, nhờ vào ngữ nghĩa của lệnh putmem. Các kernel thăm dò số lượng bằng cách sử dụng nvshmem_uint64_wait_until, thực hiện song song thao tác trên tất cả các khối và chuỗi (thread). Sau đó, một rào cản xuyên khối đảm bảo rằng không có khối nào đọc các bộ đệm trừ khi tất cả dữ liệu đã được nhận chính xác. Được phân bố trên các khối và được đồng bộ hóa thông qua một bộ đếm tăng tự động, các token được sao chép từ bộ nhớ đối xứng vào các bộ đệm liền kề thông thường, trở thành đầu vào cho các kernel hạ nguồn thực hiện các chuyên gia. Thứ bậc nguồn, chỉ số chuyên gia và chỉ số token được lưu trữ riêng biệt, xác định chính xác vị trí mà combine kernel phải gửi các hoạt động kích hoạt. Mặc dù các token từ bên trong một nhóm DP được gửi từ các thiết bị khác nhau, tất cả chúng đều được nhóm lại với nhau để chuyển tiếp đến chuyên gia tương ứng.

Kết hợp (Combine)

Các combine kernel một lần nữa được chia thành nửa gửi và nhận: các bên gửi sao chép các hoạt động kích hoạt 16-bit chưa được lượng tử hóa sang các thiết bị từ xa, với các kernel nhận chờ dữ liệu được gửi sang và tính toán trọng số trung bình của tất cả các đóng góp của chuyên gia tại cục bộ. Ngoài ra, chúng cũng đóng vai trò như một rào cản để đồng bộ hóa chuỗi dispatch-combine: mỗi thứ bậc thiết lập một cờ trên mỗi thiết bị ngang hàng khi bắt đầu các kernel gửi, và các kernel nhận không được phép trả về trừ khi chúng quan sát thấy cờ được thiết lập. Độ trễ của việc đồng bộ hóa là tối thiểu, vì nó chồng lấp với quá trình giao tiếp và tính toán thực tế.

Phóng to

Ở phía gửi, các kernel duyệt qua danh sách các token được gán cho tất cả các chuyên gia cục bộ một cách song song, ghi chúng vào một bộ đệm trên thiết bị đích. Thứ bậc đích, chỉ số chuyên gia và chỉ số token được đọc từ các bộ đệm trên mỗi token được điền bởi các kernel scatter khi nhận được token. Mỗi bên gửi có vùng nhớ riêng tư của riêng mình cho mỗi chuyên gia để ghi vào, như được chỉ ra trong hình ở trên, giúp tránh nhu cầu đồng bộ hóa. Tương tự như điều phối, combine tự động tăng các bộ đếm trên mỗi token trên thiết bị đích để biểu thị việc nhận dữ liệu: khi bộ đếm khớp với số lượng chuyên gia mà một token được điều phối tới, nội dung token có thể được truy cập.

Trong kernel nhận, danh sách các token được duyệt qua một cách song song trên nhiều khối, chờ nội dung của chúng đến bằng cách thăm dò cờ được thiết lập bởi thao tác báo hiệu. Khi đến nơi, các payload được đọc từ các bộ đệm riêng tư, với bảng định tuyến cho biết nên đọc từ bộ đệm nào và gán trọng số nào cho mỗi chuyên gia. Các kết quả sau đó được ghi vào các tensor được phân bổ bên ngoài, với kernel kết thúc quá trình thực thi sau khi tất cả các thiết bị vượt qua rào cản. Trên một nhóm DP, tất cả các thứ bậc nhận được một bản sao của mỗi hoạt động kích hoạt chuyên gia để tính toán các bản sao của chính chúng.

Điểm chuẩn (Benchmarks)

Chúng tôi đánh giá các kernel của mình trên một cụm gồm 128 GPU trải rộng trên 16 nút được kết nối thông qua InfiniBand và các NIC CX-7. Chúng tôi thử nghiệm cả GPUDirect Async (IBGDA) và Reliable Connection (RC) với một CPU proxy. Chúng tôi so sánh chúng với bản triển khai DeepSeek đã được tối ưu hóa cao, cũng như các nguyên hàm dày đặc do PyTorch cung cấp (thông qua NCCL) hoặc NVSHMEM.

Phóng to

Mặc dù trên thao tác all-to-all NVSHMEM dày đặc, hiệu suất của IBRC và IBGDA là tương tự nhau (6378 µs so với 6180 µs), nhưng IBGDA nhanh hơn đáng kể với các kernel thưa thớt. Cộng lại Dispatch và Combine, IBGDA sử dụng 902 µs trong khi IBRC mất 3223 µs - cải thiện độ trễ gấp 3,6 lần. Trong khi all-to-all bị ràng buộc bởi băng thông, các kernel thưa thớt phát quảng bá lượng dữ liệu ít hơn nhiều bậc độ lớn, và bị ràng buộc bởi độ trễ. Bằng cách kích hoạt truyền tải mạng trực tiếp từ GPU mà không cần CPU proxy điều phối GPU và NIC, độ trễ từ đầu đến cuối được giảm đáng kể.

Mặc dù các kernel có tính di động của chúng tôi chậm hơn khoảng 2 lần so với các kernel DeepSeek được tối ưu hóa cao, chúng cải thiện độ trễ gấp 10 lần so với các kernel dày đặc (902 µs so với 9944 µs).

Ngoài ra, việc chia nhỏ thành các thành phần gửi-nhận cũng cho phép che giấu một phần độ trễ, không giống như các nguyên hàm của thư viện.

Phóng to

Trên cấu hình đơn nút (EP8), NVSHMEM sử dụng NVLINK để truyền tải, mang lại độ trễ thấp hơn và thông lượng cao hơn so với mạng liên nút. Các kernel có tính di động của chúng tôi nhanh hơn khoảng 2,5 lần so với DeepEP trên đơn nút (186 µs so với 481 µs).

Cải tiến tiếp theo

Các kernel được mô tả ở đây vượt trội hơn các nguyên hàm tích hợp sẵn của các khung (framework) ML và cung cấp hiệu suất khá tốt mà không cần chuyên Biệt hóa quá mức cho các phương tiện truyền tải liên thiết bị cụ thể, chẳng hạn như NVLink hoặc InfiniBand. Bên cạnh những cơ hội đã đề cập, có thể đạt được hiệu suất cao hơn nữa bằng cách thay thế các nguyên hàm giao tiếp bằng các phiên bản chuyên biệt hơn. Ví dụ, qua NVLink, việc sử dụng bộ nhớ đối xứng có thể được thay thế bằng các bộ đệm được chia sẻ giữa các thiết bị, loại bỏ một số thao tác sao chép và cho phép các sơ đồ đồng bộ hóa chi tiết hơn, trên từng token riêng lẻ thay vì các lô token (token batch). Qua InfiniBand, một bản triển khai có thể truy cập trực tiếp vào các cặp hàng đợi (queue pair) bên dưới, loại bỏ nhu cầu truyền cả số lượng token dự kiến lẫn đã gửi và xả các token từ hàng đợi vào các bộ đệm đầu ra khi chúng đến nơi. Tuy nhiên, các bản triển khai như vậy đi kèm với cái giá là tính di động, trong khi các bản triển khai của chúng tôi cho phép đánh giá các cụm có cấu hình cụ thể trước khi cam kết nỗ lực phát triển để tối ưu hóa cho phần cứng cụ thể.

Kết luận

Chúng tôi đã trình bày một thư viện có hiệu suất cao, có tính di động cho giao tiếp MoE, đạt hiệu suất nhanh gấp 10 lần so với giao tiếp all-to-all tiêu chuẩn đồng thời duy trì khả năng tương thích trên các cấu hình phần cứng đa dạng. Trên các triển khai đơn nút với NVLink, giải pháp của chúng tôi thể hiện độ trễ thấp hơn 2,5 lần so với các bản triển khai trước đó.

Phương pháp của chúng tôi cân bằng giữa hiệu suất và tính di động thông qua các cải tiến chính bao gồm hỗ trợ giao tiếp do GPU khởi tạo, kiến trúc kernel phân chia cho phép chồng lấp tính toán-giao tiếp và điều phối token hiệu quả sử dụng các nguyên hàm NVSHMEM tối thiểu. Mặc dù chậm hơn khoảng 2 lần so với các bản triển khai mang tính chuyên biệt cao trên các thiết lập đa nút, thư viện của chúng tôi mang lại tính linh hoạt vượt trội trên nhiều môi trường mạng khác nhau (NVLink, CX-7 và EFA).

Khi các mô hình MoE tiếp tục mở rộng quy mô, các chiến lược giao tiếp hiệu quả như của chúng tôi sẽ ngày càng trở nên quan trọng đối với việc triển khai thực tế. Bản triển khai hoàn toàn mã nguồn mở của chúng tôi có sẵn tại https://github.com/ppl-ai/pplx-kernels.

Tài liệu tham khảo

DeepEP

Báo cáo kỹ thuật DeepSeek-V3

NVIDIA NVSHMEM

Trích dẫn

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}