Komunikasi Mixture-of-Experts yang Efisien dan Portabel

Kami menghadirkan pustaka sumber terbuka berkinerja tinggi dan portabel untuk komunikasi Mixture-of-Experts (MoE) yang mencapai kinerja 10x lebih cepat dibandingkan dengan primitif komunikasi All-to-All standar.

Implementasi kami menampilkan beberapa inovasi teknis utama yang memberikan efisiensi komunikasi MoE yang superior:

Komunikasi yang diprakarsai GPU (IBGDA): Mendukung komunikasi langsung dari GPU ke NIC, secara signifikan mengurangi latensi dengan menghindari keterlibatan CPU

Tumpang tindih komunikasi dan komputasi: Arsitektur kernel terpisah dengan tahapan kirim dan terima yang terpisah memungkinkan komputasi berjalan secara bersamaan dengan transfer jaringan

Kinerja node tunggal tercepat: Latensi 2,5x lebih rendah daripada implementasi tercepat sebelumnya pada konfigurasi node tunggal

Kinerja multi-node yang efisien dan portabel: Implementasi kami mencapai kecepatan hingga 10x lebih cepat daripada komunikasi all-to-all standar. Meskipun sekitar 2x lebih lambat daripada implementasi yang sangat terspesialisasi, pendekatan kami menawarkan portabilitas yang lebih baik di berbagai versi NVSHMEM dan lingkungan jaringan (NVLink, CX-7, dan EFA)

Pustaka ini sepenuhnya bersumber terbuka dan tersedia di https://github.com/ppl-ai/pplx-kernels.

Dalam artikel ini, kami mengeksplorasi tantangan paralelisme pakar dalam model MoE berskala besar dan menjelaskan pendekatan kami untuk pengiriman dan penggabungan token yang efisien di lingkungan GPU terdistribusi.

Model Mixture-of-Experts (MoE), seperti DeepSeek R1 dan Mixtral 8x7B, menyempurnakan model padat dengan membatasi jumlah bobot yang diaktifkan untuk setiap token. Alih-alih mencapai jumlah parameter yang tinggi dengan meningkatkan jumlah lapisan atau memperbesar ukuran lapisan linear di Multi-Layer Perceptron (MLP) dari setiap lapisan dekoder, model MoE menggantikan MLP tradisional dengan beberapa pakar (expert) dan sebuah router. Sebagai contoh, dari 671B parameter DeepSeek R1, hanya 37B yang dikalikan dengan token tertentu selama inferensi. Hal ini sangat mengurangi jumlah komputasi yang diperlukan untuk mendekode token dan meningkatkan latensi dibandingkan dengan model padat.

Model MoE menghadirkan beberapa tantangan tambahan untuk inferensi, dibandingkan dengan model padat. Meskipun para pakar itu sendiri adalah lapisan MLP yang kecil, setiap lapisan dekoder menyertakan router yang menentukan ke pakar mana sebuah token dikirim, dengan setiap token dikirim ke beberapa pakar. Router biasanya berupa lapisan linear kecil yang menghasilkan distribusi probabilitas. Biasanya pakar dengan skor top-K dipilih dan aktivasi akhir dihitung sebagai rata-rata tertimbang, menjumlahkan keluaran pakar yang dikalikan dengan bobot yang berasal dari distribusi probabilitas router.

Untuk meminimalkan latensi pada sistem MoE terdistribusi, paralelisme dapat dieksploitasi di beberapa perangkat, namun hal ini memunculkan tantangan komunikasi. Model seperti Mixtral 8x7B atau Llama-70B pas dalam 8 perangkat di satu node tunggal, memanfaatkan interkoneksi NVLink yang cepat dan berlatensi rendah (hingga 900Gbps). Namun, model yang lebih besar memerlukan sharding pakar di beberapa node menggunakan InfiniBand (mencapai puncaknya pada 400Gbps), yang menghadirkan tantangan latensi tambahan.

Perutean token, komputasi pakar, dan pengumpulan keluaran dalam lapisan mixture-of-experts.

Dalam artikel ini, kami mengeksplorasi masalah paralelisme pakar, menjelaskan implementasi kernel GPU kami yang secara efisien mengirimkan token ke perangkat yang menyimpan pakar yang ditetapkan dan mengumpulkannya kembali dengan menghitung jumlah tertimbang dari keluaran pakar. Meskipun ada implementasi yang dioptimalkan secara agresif yang menargetkan perangkat keras jaringan tertentu, kami menghadirkan implementasi yang sangat portabel yang mengandalkan sekumpulan kecil primitif NVSHMEM namun memberikan kinerja 10x lebih cepat daripada komunikasi all-to-all standar. Implementasi kami mencapai kinerja mutakhir pada konfigurasi node tunggal sambil mempertahankan portabilitas yang sangat baik di berbagai lingkungan jaringan.

Untuk inferensi yang efisien, bobot model harus disimpan dalam memori perangkat, sementara juga menyisakan ruang yang cukup untuk aktivasi, tembolok KV, dan penyangga (buffer) lain yang diperlukan oleh lolosan maju (forward pass) melalui model. Model yang paling mampu melampaui kapasitas GPU yang paling mumpuni sekalipun, sehingga inferensi harus disebar ke beberapa perangkat, yang secara kolektif dapat menyimpan bobot model. Berdasarkan skema sharding bobot, berbagai skema komunikasi dan komputasi harus digunakan untuk mengalihkan komputasi dan menyinkronkan perangkat.

Expert Parallelism (EP), yang diilustrasikan pada gambar pertama, hanya memparalelkan komputasi pakar. Pakar yang berbeda ditetapkan ke perangkat yang berbeda, yang menyimpan bobot mereka. Setelah perutean, token dikirim ke perangkat yang sesuai, dengan hasil yang dikumpulkan dan diakumulasikan setelahnya. Kompleksitas perutean bergantung pada tingkat paralelisme di bagian lain model: mereplikasi lapisan lain dapat sepenuhnya menghilangkan kebutuhan akan perutean, karena setiap peringkat (rank) dapat memilih token dari tabel perutean yang direplikasi secara lokal. Namun, jika hanya satu peringkat yang menjalankan perutean, siaran (broadcast) diperlukan untuk mengirimkan token, indeks, dan bobot ke pakar masing-masing. Terakhir, pengumpulan menyeluruh (all-gather) atau siaran all-to-all menyinkronkan token keluaran dengan peringkat mana pun yang melanjutkan eksekusi model. Implementasi semacam itu relatif sederhana karena torch telah mengekspos primitif yang diperlukan, meskipun beberapa keuntungan dapat diperoleh dengan menggabungkan akumulasi dengan operasi pengumpulan yang mengumpulkan token dari para pakar.

Menggabungkan paralelisme pakar dengan paralelisme tensor di dua GPU.

Paralelisme khusus pakar tidak berskala secara ideal, karena node dalam sebuah kluster mungkin menganggur saat model menjalankan lapisan non-pakar, seperti atensi (attention), norma, dan pengambilan sampel (sampling). Namun, komputasi lapisan-lapisan ini, terutama atensi, juga dapat diuntungkan oleh Tensor Parallelism (TP). Sebagian besar model mengandalkan atensi multi-kepala (multi-head attention), yang berarti kepala atensi dan proyeksi Q, K, dan V yang sesuai juga dapat di-shard di seluruh perangkat, mereplikasi atau mengumpulkan irisan di antara berbagai lapisan. Jika atensi disebar di semua perangkat, all-gather dapat menyinkronkan aktivasi, memungkinkan perutean direplikasi, yang memerlukan primitif sinkronisasi yang mirip dengan kasus paralelisme pakar untuk implementasi yang efisien. Namun, ada batas untuk paralelisme pada tingkat ini, karena mengurangi jumlah kepala atensi di bawah ambang batas tertentu akan menghasilkan hasil yang semakin berkurang.

Menggabungkan paralelisme pakar dengan paralelisme data untuk permintaan terpisah di dua GPU.

Untuk memanfaatkan semua perangkat sebaik-baiknya dan mendukung tingkat paralelisme pakar yang sangat tinggi, hingga 128 atau 256 GPU, Data Parallelism (DP) diperlukan. Di bawah skema ini, perangkat tempat model dipecah dikelompokkan untuk menangani permintaan secara bersamaan, menghitung atensi, dan memelihara tembolok KV yang di-shard di seluruh grup lokal mereka. Beberapa instans grup paralel ini berkolaborasi dalam evaluasi pakar, dengan masing-masing menghosting subset bobot pakar yang berbeda. Berdasarkan jumlah kepala atensi, sebuah grup biasanya dapat mencakup ukuran seluruh node, karena komunikasi antarnode lebih cepat. Sebagai contoh, pada gambar di atas, satu peringkat DP melayani dua permintaan secara mandiri, A dan B, menangani atensi, norma, dan pembukuan lainnya untuk permintaan tersebut. Peringkat DP lainnya memproses permintaan C yang berbeda. Namun, node pertama menghosting separuh pakar, sementara node lainnya menghosting separuh sisanya, sehingga setelah perutean, token dari A dan B dapat dikirim ke node kedua dan sebaliknya. Hal ini mengarah pada masalah komunikasi jarang (sparse communication): setiap perangkat mungkin mengirimkan jumlah token yang berbeda ke tujuan peringkat mana pun. Primitif yang ada dari torch, terutama all_to_all, tidak terlalu cocok, karena mungkin memerlukan beberapa bentuk padding atau sinkronisasi GPU-ke-CPU dan siaran metadata. Untuk menerapkan komunikasi secara efektif, kernel kustom diperlukan untuk memproses informasi perutean dan memulai komunikasi dari perangkat tanpa intervensi CPU. Setelah perutean, kernel pengiriman (dispatch) harus mengirim token ke peringkat tempat mereka dirutekan, sementara di sisi penggabungan (combine), aktivasi milik permintaan dalam grup DP saat ini harus dikumpulkan. Selain itu, pekerjaan harus diseimbangkan dalam grup DP yang mungkin memiliki beberapa perangkat, untuk meminimalkan siaran dan mereplikasi token dengan benar pada semua peringkat untuk lapisan berikutnya.

NVSHMEM adalah implementasi OpenSHMEM khusus NVIDIA, yang menyediakan fasilitas komunikasi antarperangkat portabel yang mengabstraksi kompleksitas perangkat keras mendasar. API dapat mengekspresikan pembacaan dan penulisan antarperangkat, yang dipetakan ke primitif lapisan transpor individual. Baik NVLink maupun RDMA didukung, memberikan tingkat portabilitas dan memungkinkan kernel NVSHMEM beroperasi di dalam satu node atau di seluruh kluster yang terdiri dari beberapa node. Dalam kernel kami, kami menggunakan penulisan memori jarak jauh dan operasi penyiyalan (signaling) untuk mentransfer data dan menyinkronkan perangkat yang berpartisipasi.

Operasi NVSHMEM dibangun di sekitar konsep memori simetris: operasi tersebut berjalan pada buffer yang telah dialokasikan pada semua perangkat yang berpartisipasi dalam komunikasi antarperangkat. Perangkat lokal dapat menggunakan alamat yang diturunkan dalam buffer lokalnya sendiri beserta indeks peringkat target untuk menentukan tujuan suatu operasi. Gambar di bawah ini mengilustrasikan konsep ini: kedua GPU mengalokasikan buffer simetris dengan ukuran yang sama, mempertahankan pointer src ke buffer tersebut. Peringkat pertama ingin mengirim 3 bilangan bulat ke peringkat kedua, menempatkannya di awal buffer: nvshmem_int_put_nbi menurunkan alamat awal dari buffer lokal, menentukan perangkat target. Peringkat kedua menurunkan offset dari buffernya sendiri, mengirim satu elemen ke perangkat pertama, dengan offset sebesar satu. Meskipun alamat tujuan harus selalu berupa buffer simetris yang dialokasikan menggunakan nvshmem_alloc, buffer sumber dapat berupa wilayah memori perangkat apa pun, asalkan telah didaftarkan sebelumnya dengan NVSHMEM.

Meskipun NVSHMEM menyediakan berbagai macam primitif, kernel kami hanya mengandalkan 3 fungsi, dan membangun semua sinkronisasi dan pemagaran (fencing) di atasnya.

Memori simetris dan penulisan jarak jauh antara dua GPU menggunakan NVSHMEM.

nvshmemx_putmem_signal_nbi_warp: Mentransfer blok data dari satu perangkat ke perangkat lain, sekaligus menyetel tanda (flag) pada perangkat jarak jauh. Operasi ini menyetel (NVSHMEM_SIGNAL_SET) atau menambah (NVSHMEM_SIGNAL_ADD) lokasi 64-bit. Tanda diperbarui setelah seluruh blok memori ditransfer. Jika perangkat jarak jauh mengamati perubahan pada tanda tersebut, perangkat tersebut dapat dengan aman mengakses buffer di memorinya sendiri. Fungsi ini berguna untuk menggabungkan transfer data dan sinkronisasi.

nvshmemx_signal_op beroperasi pada satu lokasi memori, biasanya tanda 64-bit, secara atomik menyetel atau menambahkannya. Ini berguna dalam mengirimkan metadata dan menyinkronkan perangkat.

nvshmem_uint64_wait_until digunakan pada ujung penerima dari suatu sinyal, untuk melakukan polling pada suatu tanda hingga perangkat jarak jauh memperbaruinya.

Memastikan segala bentuk pengurutan antaroperasi melalui nvshmem_quiet atau nvshmem_fence memerlukan biaya yang mahal, sehingga implementasi kernel kami menghindari penggunaan primitif NVSHMEM lainnya dan meminimalkan dependensi apa pun selain pasangan kirim-terima individual. Saat mengirim data, kami selalu menggunakan versi fungsi non-pemblokiran (non-blocking), tanpa menunggu data bahkan dikirim keluar dari peringkat lokal. Sinkronisasi implisit lainnya di seluruh grup kernel memastikan bahwa buffer ini tidak ditimpa secara destruktif saat NIC memerlukan akses ke buffer tersebut.

NVSHMEM menggunakan GPUDirect RDMA sehingga NIC dapat langsung membaca dan menulis memori GPU untuk pertukaran data tanpa melibatkan CPU. Selain itu, pada NIC ConnectX, NVSHMEM mendukung komunikasi yang diprakarsai GPU (juga dikenal sebagai Infiniband GPUDirect Async, atau IBGDA), yang memungkinkan GPU untuk memposkan operasi RDMA langsung ke NIC dan melakukan polling pada antrean penyelesaian (completion queue) NIC secara langsung. Pada platform yang tidak mendukung komunikasi yang diprakarsai GPU, NVSHMEM menelurkan utas "proksi" pada CPU host untuk memulai komunikasi dan melakukan polling penyelesaian atas nama GPU. Program NVSHMEM bersifat portabel terlepas dari apakah komunikasi yang diprakarsai GPU didukung atau tidak. Namun, komunikasi yang diprakarsai GPU secara signifikan memotong latensi karena sepenuhnya melewati jalan memutar ke CPU.

Kami mengimplementasikan komunikasi MoE melalui sepasang kernel pengiriman (dispatch) dan penggabungan (combine). Kernel pengiriman bertanggung jawab untuk membaca token dan informasi perutean pada setiap peringkat, mengirimkannya ke pakar yang sesuai. Kernel penggabungan mengumpulkan aktivasi yang dihasilkan oleh para pakar dan mengirimkannya kembali ke peringkat sumbernya, sekaligus menghitung rata-rata tertimbang dari pakar yang dipilih berdasarkan bobot yang dihitung oleh router. Kernel selanjutnya dibagi menjadi komponen kirim dan terima, untuk memungkinkan transfer data tumpang tindih dengan komputasi. Kernel kirim semuanya bersifat non-pemblokiran dan nir-sinkronisasi (non-synchronizing): kernel tersebut hanya mengirimkan semua penulisan ke perangkat jarak jauh. Di ujung lain, kernel penggabungan hanya membaca dari memori, menunggu hingga semua data yang diperlukan telah ditransfer. Setelah mengirimkan pekerjaan ke NIC, sementara data ditransfer secara asinkron melalui kabel, GPU dapat melakukan pekerjaan berguna lainnya secara lokal, seperti menerapkan pakar bersama atau menghitung atensi.

Setiappasang kernel memiliki buffer memori simetrisnya sendiri, di mana setiap peringkat mengalokasikan penyimpanan privat untuk menerima data dari semua peringkat lainnya. Jumlah maksimum token yang dapat dikirim dari grup DP tetap sama di semua peringkat, yang juga menetapkan batas atas yang dapat diterima oleh setiap peringkat untuk setiap pakar lokal dari setiap peringkat hal DP. Hal ini memungkinkan peringkat pengirim untuk menurunkan alamat unik pada peringkat tujuan untuk ditulisi, tanpa memerlukan sinkronisasi di antara mereka. Setelah data diterima, kernel mengacak token untuk menatanya di memori dalam format optimal yang diperlukan oleh kernel penerima. Meskipun buffer memiliki dimensi yang cukup besar, buffer tersebut digunakan kembali di semua lapisan sekuensial dari sebuah model.

Kernel kirim dan terima terpisah untuk operasi dispatch dan combine.

Satu-satunya bentuk sinkronisasi global, seperti yang diilustrasikan pada gambar di atas, diimplementasikan dalam kernel combine-receive. Setelah data dikirim dari kirim, sebuah peringkat kembali dan melanjutkan eksekusi ketika ia menerima semua token yang dibutuhkannya dari peringkat lain dan memasuki kernel combine-send segera setelah diproses. Penghalang (barrier) di combine-receive memastikan bahwa tidak ada peringkat yang dapat berjalan mendahului dan memulai dispatch-send sementara peringkat lain masih menunggu untuk menerima data, yang potensial menyebabkan tumpang tindih yang bersifat destruktif. Sinkronisasi dilakukan di kernel penggabungan demi kesederhanaan, sebagai pengoptimalan, hambatan yang sama dapat diimplementasikan antara kernel dispatch-receive dan combine-send, yang selanjutnya mengurangi latensi dengan memberikan lebih banyak waktu bagi NIC untuk menyelesaikan transfer sementara GPU menjalankan kernel lain.

Kedua kernel dibagi di semua SM yang tersedia dari perangkat: sementara kernel dispatch-send dan combine-receive harus memparalelkan di seluruh jumlah token maksimum per peringkat (max_m), kernel dispatch-receive dan combine-send harus menangani token sebanyak itu dari setiap peringkat dan pakar lokal, dengan maksimum max_m * (num_experts // EP) * (EP // TP).

Kernel pengiriman bertanggung jawab untuk mengirimkan token dan faktor penskalaan blok ke rekan (peer) yang menghosting pakar tempat token tersebut dirutekan. Sisi pengirim mengandalkan spesialisasi warp untuk memparalelkan dua tugas: menggabungkan informasi perutean untuk menentukan berapa banyak token yang dikirim dari setiap peringkat dan mengepak token ke dalam pesan yang akan dikirim ke perangkat jarak jauh. Sisi penerima pertama-tama menunggu semua penghitungan token diterima, kemudian membaca dan membongkar muatan (payload) dari memori bersama ke dalam tensor yang akan diteruskan ke kernel Grouped GEMM hilir. Kernel ini juga menyimpan informasi ke dalam buffer yang dibagikan dengan kernel penggabungan, untuk menunjukkan ke mana setiap token harus dikirim kembali. Informasi ini diperlukan karena kernel penerima mengacak token di sekitar buffer yang berurutan secara kontigu dalam urutan non-deterministik.

Mendistribusikan token dan memberi sinyal kedatangannya antara rank sumber dan tujuan.

Di bagian pengirim, satu warp yang terdiri dari 32 utas pada setiap blok bertanggung jawab untuk membaca informasi perutean dan menghitung jumlah token yang diharapkan akan diterima oleh setiap pakar tujuan dari peringkat ini, mengirimkan jumlah ditambah satu menggunakan nvshmem_signal_op. Jumlah tersebut ditambah satu, karena transisi dari nol ke bukan nol di ujung jarak jauh menandakan penerimaan jumlah dari peringkat sumber. Secara paralel, warp yang tersisa bekerja sama untuk menyalin token ke dalam memori simetris di semua blok secara paralel, mengepak aktivasi, faktor penskalaan, dan indeksnya pada peringkat lokal ke dalam potongan kontigu. Indeks tersebut diperlukan oleh pengirim penggabungan untuk menentukan alamat tempat token akan ditulis. Selanjutnya, setelah memastikan semua data telah disalin melalui penghalang di seluruh grup warp, warp sekali lagi beroperasi secara mandiri, masing-masing mengirimkan buffer yang sama ke pakar yang berbeda secara paralel. Token dikirim menggunakan nvshmemx_putmem_signal_nbi_warp, yang juga secara atomik meningkatkan jumlah token terkirim dari peringkat lokal pada perangkat jarak jauh. Dalam grup DP, karena setiap peringkat memiliki replika token yang akan dikirim, pengiriman diseimbangkan secara merata, dengan setiap perangkat mengirimkan sebagian kecil (subset) dari token tersebut.

Pada ujung penerima, semua kernel terlebih dahulu disinkronkan dengan peringkat pengirim dengan menunggu jumlah total token dikirimkan. Setelah itu, semuanya menunggu jumlah token terkirim yang ditingkatkan secara atomik untuk menetap ke jumlah total, yang menunjukkan bahwa semua muatan dari peringkat sumber juga telah dikirim, berkat semantik panggilan putmem. Kernel melakukan polling pada jumlah menggunakan nvshmem_uint64_wait_until, memparalelkan operasi di semua blok dan utas. Selanjutnya, penghalang lintas-blok memastikan bahwa tidak ada blok yang membaca buffer kecuali semua data telah diterima dengan benar. Tersebar di seluruh blok dan disinkronkan melalui penghitung yang ditambah secara atomik, token disalin dari memori simetris ke dalam buffer kontigu biasa yang menjadi masukan bagi kernel hilir yang mengimplementasikan para pakar. Peringkat sumber, indeks pakar, dan indeks token disimpan secara terpisah, secara tepat menentukan lokasi tempat kernel penggabungan harus mengirim aktivasi. Meskipun token dari dalam grup DP dikirim dari perangkat yang berbeda, semuanya dikelompokkan bersama untuk diteruskan ke pakar yang sesuai.

Kernel penggabungan sekali lagi dibagi menjadi paruh kirim dan terima: pengirim menyalin aktivasi 16-bit yang tidak dikuantisasi ke perangkat jarak jauh, dengan kernel penerima menunggu data dikirim dan menghitung rata-rata tertimbang dari semua kontribusi pakar secara lokal. Selain itu, kernel ini juga bertindak sebagai penghalang, untuk menyinkronkan urutan dispatch-combine: setiap peringkat menyetel tanda pada setiap rekan pada saat masuk ke kernel kirim, dengan kernel terima tidak diizinkan untuk kembali kecuali mereka mengamati tanda tersebut telah disetel. Latensi sinkronisasi minimal, karena tumpang tindih dengan komunikasi dan komputasi aktual.

Menggabungkan keluaran pakar ke dalam token pada setiap rank paralel data.

Di sisi pengirim, kernel melintasi daftar token yang ditetapkan ke semua pakar lokal secara paralel, menuliskannya ke buffer pada peringkat tujuan. Peringkat target, indeks pakar, dan indeks token dibaca dari buffer per-token yang diisi oleh kernel pencar (scatter) saat menerima token. Setiap pengirim memiliki wilayah memori privat sendiri per pakar untuk ditulisi, seperti yang ditunjukkan pada gambar di atas, menghindari perlunya sinkronisasi. Mirip dengan pengiriman, penggabungan secara atomik menaikkan penghitung per token pada peringkat tujuan untuk menunjukkan penerimaan data: ketika penghitung cocok dengan jumlah pakar tempat token dikirimkan, isi token dapat diakses.

Dalam kernel penerima, daftar token dilintasi secara paralel di beberapa blok, menunggu isinya tiba dengan melakukan polling pada tanda yang disetel oleh operasi penyiyalan. Saat tiba, muatan dibaca dari buffer privat, dengan tabel perutean yang menunjukkan buffer mana yang harus dibaca dan bobot apa yang harus ditetapkan ke setiap pakar. Hasilnya kemudian ditulis ke tensor yang dialokasikan secara eksternal, dengan kernel menyelesaikan eksekusi setelah semua perangkat melewati penghalang. Di seluruh grup DP, semua peringkat menerima salinan dari setiap aktivasi pakar untuk menghitung replika mereka sendiri.

Kami mengevaluasi kernel kami pada kluster yang terdiri dari 128 GPU yang tersebar di 16 node yang dihubungkan melalui InfiniBand dan NIC CX-7. Kami mencoba GPUDirect Async (IBGDA) dan Reliable Connection (RC) dengan proksi CPU. Kami membandingkannya dengan implementasi DeepSeek yang sangat dioptimalkan, serta primitif padat yang disediakan oleh PyTorch (melalui NCCL) atau NVSHMEM.

Latensi komunikasi untuk 128 token per GPU dalam konfigurasi EP128 multi-node.

Meskipun pada operasi all-to-all NVSHMEM yang padat, kinerja IBRC dan IBGDA serupa (6378 µs vs 6180 µs), IBGDA secara signifikan lebih cepat dengan kernel jarang. Menambahkan Dispatch dan Combine, IBGDA menggunakan 902 µs sedangkan IBRC membutuhkan 3223 µs - peningkatan latensi sebesar 3,6x. Meskipun all-to-all dibatasi oleh lebar pita (bandwidth bound), kernel jarang menyiarkan data dengan besaran urutan yang jauh lebih sedikit, sehingga dibatasi oleh latensi (latency bound). Dengan memicu transfer jaringan langsung dari GPU, tanpa memerlukan proksi CPU untuk mengoordinasikan GPU dan NIC, latensi ujung-ke-ujung (end-to-end) dikurangi secara signifikan.

Meskipun kernel portabel kami sekitar 2x lebih lambat daripada kernel DeepSeek yang sangat dioptimalkan, kernel kami meningkatkan latensi sebesar 10x dibandingkan dengan kernel padat (902 µs vs 9944 µs).

Selain itu, pemisahan ke dalam komponen pengirim-penerima juga memungkinkan sebagian latensi disembunyikan, tidak seperti primitif pustaka.

Latensi komunikasi untuk 128 token per GPU dalam konfigurasi EP8 satu node.

Pada node tunggal (EP8), NVSHMEM menggunakan NVLINK untuk transportasi, memberikan latensi yang lebih rendah dan throughput yang lebih tinggi daripada jaringan antarnode. Kernel portabel kami sekitar 2,5x lebih cepat daripada DeepEP pada node tunggal (186 µs vs 481 µs).

Kernel yang dijelaskan di sini mengungguli primitif bawaan dari kerangka kerja ML dan menawarkan kinerja yang layak tanpa terlalu berspesialisasi untuk transpor antarperangkat tertentu, seperti NVLink atau InfiniBand. Selain peluang yang telah disebutkan, peningkatan kinerja lebih lanjut dapat dicapai dengan mengganti primitif komunikasi dengan versi yang lebih khusus. Sebagai contoh, di NVLink, penggunaan memori simetris dapat diganti dengan buffer yang dibagikan di seluruh perangkat, menghilangkan beberapa penyalinan dan memungkinkan skema sinkronisasi yang lebih halus, di seluruh token individual alih-alih batch token. Di InfiniBand, sebuah implementasi dapat mengakses pasangan antrean dasar secara langsung, menghilangkan kebutuhan untuk mentransfer penghitungan token yang diharapkan dan terkirim serta membilas (flush) token dari antrean ke dalam buffer keluaran saat tiba. Namun, implementasi semacam itu harus mengorbankan portabilitas, sementara implementasi kami memungkinkan kluster dengan konfigurasi tertentu dievaluasi sebelum mendedikasikan upaya pengembangan untuk mengoptimalkan perangkat keras tertentu.

Kami telah menghadirkan pustaka berkinerja tinggi dan portabel untuk komunikasi MoE yang mencapai kinerja 10x lebih cepat dibandingkan dengan komunikasi all-to-all standar sambil mempertahankan kompatibilitas di berbagai konfigurasi perangkat keras. Pada penyebaran node tunggal dengan NVLink, solusi kami menunjukkan latensi 2,5x lebih rendah daripada implementasi sebelumnya.

Pendekatan kami menyeimbangkan kinerja dengan portabilitas melalui inovasi utama termasuk dukungan komunikasi yang diprakarsai GPU, arsitektur kernel terpisah yang memungkinkan tumpang tindih komputasi-komunikasi, dan pengiriman token yang efisien menggunakan primitif NVSHMEM minimal. Meskipun sekitar 2x lebih lambat daripada implementasi yang sangat terspesialisasi pada penyiapan multi-node, pustaka kami menawarkan fleksibilitas superior di berbagai lingkungan jaringan (NVLink, CX-7, dan EFA).

Seiring model MoE terus berskala, strategi komunikasi yang efisien seperti milik kami akan menjadi semakin penting untuk penerapan praktis. Implementasi sumber terbuka penuh kami tersedia di https://github.com/ppl-ai/pplx-kernels.

DeepEP

Laporan Teknis DeepSeek-V3

NVIDIA NVSHMEM

BibTeX

@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}
}

Pendahuluan

Sharding dan Paralelisme

NVSHMEM

Kernel Portabel

Pengiriman (Dispatch)

Penggabungan (Combine)

Tolok Ukur (Benchmarks)

Peningkatan Lebih Lanjut

Kesimpulan

Referensi

Kutipan