Mengaktifkan Model Berparameter Triliun pada AWS EFA

Di Perplexity, kami menggunakan model terbaik untuk produk, API, dan tim riset kami. Model Mixture-of-Experts (MoE) sumber terbuka berskala besar, seperti Kimi-K2, menghadirkan tantangan tersendiri, karena node inferensi terbesar dengan 8x GPU NVIDIA H200 tidak dapat menampungnya secara efisien, sehingga memerlukan penerapan multi-node. Kami menghadirkan serangkaian kernel untuk paralelisme ahli yang mencapai latensi mutakhir pada ConnectX-7, melampaui performa DeepEP. Kernel yang sama juga merupakan yang pertama mencapai latensi yang layak pada AWS Elastic Fabric Adapter (EFA), yang memungkinkan penerapan model dengan parameter triliunan.

Coba kernel kami di GitHub dan baca makalah riset lengkap di arXiv.

Pendahuluan

Mixture-of-Experts (MoE) telah muncul sebagai arsitektur standar untuk menskalakan hingga ratusan miliar atau triliunan parameter sekaligus menawarkan latensi yang wajar pada saat inferensi. Model MoE menggantikan lapisan padat transformer dengan sekumpulan ahli (expert) dan lapisan perutean linier kecil yang menentukan dengan ahli mana token harus dikalikan. Dengan mengelompokkan token yang terkait dengan ahli yang sama, masalah menjadi jarang (sparse) dan dapat diparalelkan dengan mudah, karena ahli yang berbeda sekarang dapat di-host pada GPU berbeda yang tersebar di beberapa node.

Tidak seperti bentuk paralelisme lainnya, seperti Tensor Parallelism (TP) dan Data Parallelism (DP) yang dapat diimplementasikan dengan mudah menggunakan pustaka komunikasi kolektif seperti torch.distributed dan NCCL, perutean MoE melibatkan komunikasi peer-to-peer yang jarang. Kernel pengiriman (dispatch) harus membagi kumpulan token yang masuk dan mengirimkannya ke rank yang menghosting para ahli, sementara masalah gandanya, yaitu kernel penggabung (combine), harus menerima padanan yang telah diproses dan menghitung rata-rata tertimbangnya. Transfer ini tidak dipetakan ke proses kolektif dan memerlukan kernel kustom untuk mencapai latensi rendah.

Di Perplexity, kami telah mengembangkan kernel antar-node yang portabel dan kernel khusus dalam-node, yang mempresentasikan karya kami secara lebih mendalam pada artikel sebelumnya tentang komunikasi MoE. Kami juga menunjukkan dalam posting blog lain bahwa penerapan multi-node dapat menghasilkan latensi yang lebih rendah dan throughput yang lebih tinggi untuk model MoE besar.

Dalam artikel ini, kami menyelami lebih dalam implementasi kernel antar-node baru kami yang sangat meningkatkan performa dari pekerjaan kami sebelumnya yang berbasis NVSHMEM. Selain menampilkan hasil mutakhir pada ConnectX-7, kami juga menunjukkan bahwa dalam runtime inferensi kami, kernel tersebut dapat digunakan secara efektif untuk melayani model sumber terbuka terbesar.

Inferensi melalui InfiniBand

Paralelisme tensor dan paralelisme ahli dapat diskalakan secara efektif dalam satu node, karena NVLink menawarkan throughput tinggi (900 GB/s) pada latensi mikrodetik. Meskipun sistem NVIDIA GB200 dan GB300 dilengkapi dengan memori yang cukup untuk menampung bobot dan cache KV dari model berparameter triliun, penerapan pada instans AWS p5en dengan hingga 8 GPU H200 terbatas pada 1120 GB HBM, yang harus dibagi antara bobot model dan cache KV. Akibatnya, model tertentu memerlukan penerapan antar-node melalui InfiniBand.

Meskipun InfiniBand masih sangat cepat, InfiniBand hanya menawarkan throughput 400 Gbps (50 GB/s), menambahkan latensi puluhan hingga ratusan mikrodetik untuk perutean MoE. Untungnya, tambahan latensi tersebut dapat disembunyikan melalui penggunaan ahli bersama (shared experts), batch mikro (micro-batching), dan tumpang tindih komputasi-komunikasi untuk menghasilkan throughput yang layak dengan biaya yang kompetitif. DeepSeek memelopori teknik-teknik ini melalui DeepSeek-R1, yang diterapkan pada kluster GPU H800 dan adaptor InfiniBand NVIDIA ConnectX-7.

Pada kluster AWS, situasinya berbeda karena penggunaan Elastic Fabric Adapter (EFA) kustom mereka. Meskipun adaptor ini menghasilkan throughput puncak 400 Gbps pada proses kolektif yang biasanya digunakan dalam beban kerja pelatihan, adaptor ini sedikit tertinggal dari ConnectX-7 pada ukuran pesan yang ditukar selama pengiriman dan penggabungan MoE. Selain itu, EFA tidak mendukung GPUDirect Async, yang memerlukan utas proksi CPU yang menjembatani GPU ke NIC untuk memulai transfer. Selain kompleksitas yang bertambah, transaksi PCIe tambahan menambah overhead senilai mikrodetik ke semua transfer. Kerangka kerja seperti NVSHMEM memungkinkan penulis kernel untuk menyediakan implementasi yang independen dari perangkat, namun API generik menimbulkan overhead yang membuat penerapan menjadi tidak efektif dalam hal biaya dan latensi.

pplx-kernels berbasis NVSHMEM

Kernel kami sebelumnya mengabstraksikan implementasi transpor mendasar melalui NVSHMEM. Kernel tersebut mengandalkan nvshmemx_putmem_signal_nbi_warp untuk mentransfer token satu per satu, menggunakan penghitung atomik untuk menyinkronkan peer. Meskipun ini menawarkan performa yang wajar pada NIC ConnectX-7 dengan IBGDA, implementasi IBRC ConnectX-7 berbasis proksi jauh lebih lambat, sementara implementasi EFA gagal memberikan latensi di bawah satu milidetik.

Kami menggunakan kernel tersebut sebagai uji coba untuk perangkat keras yang mendasarinya, sekaligus menentukan antarmuka yang paling sesuai untuk diekspos ke model kami. Kami berhipotesis bahwa seringnya transfer dan interaksi antara kernel dan NIC atau utas proksi menghambat performa. Selain itu, implementasi kami tidak mengandalkan jaminan apa pun dari perangkat keras selain keandalan, yang menyarankan bahwa kami dapat mempercepat berbagai hal dengan mengurangi overhead primitif sinkronisasi.

Mengikuti pengamatan ini, kami merancang kernel berbasis EFA yang mengekspos antarmuka yang sama, mengoptimalkan interaksi kernel perangkat dan proksi host di baliknya untuk performa puncak.

MoE all-to-all melalui EFA

Perbesar

Kernel baru mengikuti arsitektur hibrida CPU-GPU, di mana kernel GPU berinteraksi dengan model pada perangkat dan utas proksi yang berjalan pada CPU host menangani interaksi dengan NIC. Kernel baru ini menggunakan kembali TransferEngine yang sama yang kami bangun untuk transfer cache KV.

Perutean MoE dilakukan oleh kernel pengiriman (dispatch), yang menerima token dan ahli per token untuk merutenya dari model, mengirimkannya melalui jaringan ke rank yang menghosting para ahli. Setelah kernel GEMM yang dikelompokkan mengalikan token dengan bobot ahli, kernel penggabung (combine) mentransfer token ke rank sumbernya dan menghitung rata-rata tertimbangnya sebelum mengembalikannya ke model.

Kernel pengiriman dan penggabung dibagi menjadi bagian pengirim dan penerima untuk memfasilitasi tumpang tindih dan batch mikro. Pengirim menyusun ulang token ke dalam buffer kirim sehingga dapat ditulis ke peer melalui satu operasi tulis RDMA. Setelah mengeksekusi pengirim, GPU dapat melakukan pekerjaan lain (seperti perkalian ahli bersama) sementara transfer melalui jaringan tertunda di latar belakang. Kernel penerima memblokir hingga semua transfer diterima. Di latar belakang, utas pekerja sisi host menangani transfer: utas tersebut diberi sinyal oleh GPU ketika buffer kirim telah diisi, setelah itu utas tersebut mengirimkan penulisan RDMA. Begitu NIC memberi sinyal selesai, utas tersebut membuka blokir penerima, yang mengembalikan token ke model.

Kernel terhubung erat dengan host, berkomunikasi satu sama lain melalui memori terpadu dan GDRCopy. Memori terpadu CUDA memungkinkan memori perangkat atau memori host dipetakan ke dalam ruang alamat keduanya, menangani komunikasi secara transparan melalui PCIe.

Meskipun memori terpadu cocok untuk transfer massal, polling latensi rendah dilakukan melalui GDRCopy. GDRCopy menggunakan GPUDirect RDMA untuk membaca dari dan menulis ke memori GPU pada latensi mikrodetik.

Sebagai tambahan untuk RDMA, dalam satu node kernel juga memanfaatkan NVLink untuk bertukar token. Ini dapat mengalihkan sebagian besar trafik, hingga seperempat untuk EP32 dan hingga seperdelapan untuk EP64, sehingga menghasilkan latensi keseluruhan yang lebih rendah.

Buffer Kirim dan Terima

Desain buffer kirim untuk kernel pengiriman sangat penting untuk mencapai latensi yang layak yang melibatkan proksi host. Jika ada N rank yang masing-masing mengirimkan T token, masing-masing dirutekan ke R ahli dari E, sebuah rank dapat menerima hingga N * T * max(R, E / N) token dalam kasus terburuk, jika semuanya dirutekan ke rank tersebut. Selain itu, batas minimum ini hanya dapat dicapai ketika pengirim menata penulisannya secara kontigu, karena buffer terima pribadi akan menghabiskan terlalu banyak memori. Agar pengirim dapat menata penulisannya secara berturut-turut tanpa konflik, pengetahuan penuh tentang informasi perutean harus didistribusikan ke setiap rank. Akibatnya, kernel pengiriman kami terlebih dahulu bertukar jumlah token per ahli, setelah itu masing-masing dapat menentukan ke mana penulisan mereka pergi dan dari offset mana mereka menerima token dari setiap rank.

Hanya bertukar informasi perutean akan menghasilkan latensi yang tidak perlu tinggi, karena setiap rank harus menunggu penulisan tiba sebelum mengirimkan penulisan apa pun sendiri. Sebagai gantinya, kami memesan sejumlah kecil ruang privat per pengirim pada setiap rank penerima dan mengirimkan beberapa token bersamaan dengan penulisan informasi perutean. Ini memastikan pemanfaatan bandwidth penuh selama sebagian besar proses pengiriman. Setelah rute ditukar, token yang tersisa ditransfer melalui satu penulisan per peer jarak jauh. Pertukaran informasi perutean dan buffer berukuran tetap, beserta transfer yang mengisi buffer tersebut, diilustrasikan pada gambar di bawah ini.

Perbesar

Kernel penggabung jauh lebih sederhana, karena informasi perutean telah tersedia bagi mereka dari tahap pengiriman. Mereka dapat menyalin token dalam satu penulisan per peer, setelah itu mereka membaca dari buffer terima untuk menghitung rata-rata tertimbang token ke dalam tensor keluaran.

Untuk NVLink, buffer kirim dan terima dalam satu node dipetakan ke dalam peer agar mereka dapat langsung mendorong atau menarik token dari rank tetangga.

Pengiriman dan Penggabungan

Gambar di bawah ini merinci waktu dan interaksi antara GPU, CPU, dan NIC, yang menunjukkan satu iterasi pengiriman dan penggabungan yang diselingi dengan kernel model.

Pekerjaan dimulai pada pengiriman (dispatch send), dengan kernel menggabungkan informasi perutean dan menghitung jumlah token per ahli melalui akumulasi atomik dalam memori bersama dalam satu blok. Setelah jumlah token tersedia, token tersebut disalin ke memori host dan proksi segera diberi tahu untuk mulai mengirimkan informasi perutean ke semua rank peer lainnya. Secara paralel, pengiriman terus menghitung jumlah kumulatif dari jumlah token, mengidentifikasi offset yang diperlukan untuk mengepak token ke dalam buffer kirim per peer yang kontigu. Setelah offset tersedia, semua token yang akan dikirim melalui EFA dikepak ke dalam buffer kirim. Proksi diberi tahu untuk memposting penulisan ke dalam buffer kirim peer antar-node. Di dalam node (intra-node), kernel juga segera mendorong token ke dalam buffer privat dalam node yang sama melalui NVLink.

Perbesar

Perbesar

Begitu sebuah rank menerima informasi perutean dari semua peer, rank tersebut memusentralisasikannya dan menghitung berbagai offset yang diperlukan untuk mengidentifikasi token untuk pengiriman dan penggabungan. Pekerjaan ini dilakukan pada CPU saat transfer sudah tertunda, sehingga tidak bersifat krusial terhadap latensi. Selanjutnya, setelah semua offset diketahui, penulisan untuk sisa token, yang melebihi ukuran buffer privat, diposting dan fase pengiriman selesai.

Pada sisi penerima, kernel pertama-tama menarik token dari node yang sama yang berada di luar kapasitas buffer privat melalui NVLink. Selanjutnya, kernel memblokir hingga proksi host memberi sinyal selesainya semua transfer. Setelah sinyal diterima, token diacak dari buffer terima ke dalam tensor keluaran yang diteruskan ke GEMM grup. Karena buffer digunakan kembali antara pengiriman dan penggabungan, penghalang NVLink dan RDMA memastikan bahwa buffer penerima telah dibaca sebelum mengizinkan kernel berikutnya menimpanya.

Kernel penggabung mengikuti pola yang hampir serupa, mengacak token, memberi sinyal pada proksi, dan menunggu penyelesaian. Namun, metadata untuk memposting penulisan sudah siap tersedia dan semua offset jarak jauh diketahui, sehingga semua token yang ditransfer di dalam node didorong melalui NVLink ke penerima dan token antar-node ditransfer melalui satu penulisan tunggal. Sebelum mengembalikan hasil, kernel mengumpulkan token dari buffer penerima, menuliskannya ke dalam tensor keluaran.

TransferEngine untuk perutean MoE

Untuk membuat prototipe implementasi kami, awalnya kami menggunakan kembali TransferEngine yang dibangun untuk prefill dan dekode terdisagregasi guna mengeluarkan dan melacak penulisan RDMA. Namun, tidak seperti dekode terdisagregasi di mana setiap node memposting sejumlah besar penulisan ke peer yang sama, pengiriman MoE memposting transfer kecil ke semua peer dalam satu operasi. Secara langsung (out-of-the-box), performanya tidak optimal dan implementasi kami tertinggal di belakang DeepEP pada ConnectX-7 sekitar 20us. Akibatnya, kami mengkhususkan TransferEngine untuk mendukung dua operasi, yaitu sebar (scatter) dan penghalang (barrier), untuk mengoptimalkan interaksi dengan sejumlah besar peer.

Operasi sebar dan penghalang mengekspos grup peer yang dapat didaftarkan sebelumnya untuk mengisi terlebih dahulu struktur data yang digunakan oleh implementasi transpor yang mendasarinya (pembuatan templat Work Request (WR)). Sebar menyalin potongan berbeda dari buffer sumber yang sama ke dalam offset berbeda pada semua anggota grup peer. Penghalang memberi sinyal ke semua peer melalui pesan segera (immediate). Mengantrekan dan menangani operasi ini secara massal menghemat mikrodetik yang berharga pada jalur kritis. Pada EFA, tempat dua NIC 200Gbps digabungkan untuk menyediakan bandwidth 400Gbps, kami memecah (shard) peer grup di seluruh dua NIC alih-alih memecah bita dari setiap transfer.

Meskipun awalnya kami mengembangkan TransferEngine hanya pada EFA menggunakan libfabric, kernel MoE bekerja pada ConnectX-7 tanpa mengubah satu baris kode pun setelah menambahkan dukungan ConnectX-7 dalam TransferEngine menggunakan libibverbs. Dibandingkan dengan EFA berbasis Scalable Reliable Datagram (SRD), ConnectX-7 memerlukan pekerjaan tambahan untuk pengaturan koneksi dan manajemen peer. Kami menggunakan pasangan antrean Unreliable Datagram (UD) untuk melakukan bootstrap Reliable Connection (RC). Kami membuat dua pasangan antrean untuk setiap peer jarak jauh, satu untuk SEND/RECV dua sisi dan satu untuk WRITE_IMM satu sisi, karena RECV dan WRITE_IMM sama-sama menggunakan operasi RECV secara berurutan sementara kedua operasi tersebut memerlukan ukuran buffer yang berbeda. Selain itu, kami melonggarkan pengurutan transaksi PCIe antara ConnectX-7 dan memori GPU. Mirip dengan EFA, kami mengimplementasikan pengoptimalan templat WR. Selain itu, saat mengirim beberapa pesan kecil ke peer yang sama, kami menggunakan perantaian WR (WR chaining) untuk meningkatkan laju pesan.

Karena pekerjaan di sisi GPU identik di kedua adaptor, kami mendorong implementasi transpor yang mendasarinya hingga batasnya untuk melampaui performa DeepEP guna menyoroti bahwa kesenjangan antara perutean MoE EFA dan ConnectX-7 ditentukan oleh kemampuan perangkat keras yang mendasarinya.

Evaluasi

Latensi Pengiriman/Penggabungan

Kami mengevaluasi latensi pengiriman dan penggabungan pada sistem H200 yang dilengkapi dengan adaptor ConnectX-7 dan EFA pada konfigurasi standar DeepSeek-V3 dan Kimi-K2 dengan dimensi tersembunyi 7168, pengiriman fp8 berskala blok, dan penggabungan bf16. Kami mempertimbangkan ukuran batch 128 untuk pengiriman dan 4096 untuk penggabungan. Pada ConnectX-7, kami membandingkannya dengan DeepEP dan pplx-kernels berbasis NVSHMEM kami sebelumnya, menggunakan GPUDirect Async (IBGDA) maupun proksi host (IBRC). Kami mencapai latensi gabungan sebesar 459us, 582us, dan 692us masing-masing pada EP16, EP32, dan EP64. Kami unggul dari kernel berbasis EFA lainnya yang mengandalkan proksi generik, seperti UCCL-EP, yang mencapai 519us, 966us, and 1159us.

Perbesar

Pada dekode, meskipun kernel kami beberapa mikrodetik lebih lambat daripada DeepEP, kami mencapai latensi keseluruhan yang lebih rendah karena kernel penggabung kami jauh lebih cepat. Meskipun kami menggunakan proksi host, karena transfer massal kami, kami melampaui performa semua implementasi tanpa proksi berbasis IBGDA. Kami juga menunjukkan bahwa implementasi proksi itu penting, karena proksi kami yang sangat khusus memiliki kecepatan sepuluh kali lipat lebih cepat daripada implementasi NVSHMEM generik. Terlepas dari kenyataan bahwa EFA dua kali lebih lambat daripada ConnectX-7 saat mentransfer paket 256 KB, yang mendekati ukuran paket tipikal yang ditukar oleh perutean MoE, kesenjangan antara kedua adaptor secara keseluruhan tidak terlalu besar.

Perbesar

Meskipun kernel kami yang dioptimalkan untuk dekode dapat diskalakan ke ukuran batch prefill tanpa penyesuaian tambahan, performanya tertinggal di belakang DeepEP. Ini terutama disebabkan oleh fakta bahwa kami mengumpulkan token hanya pada sisi penerima penggabungan, sedangkan DeepEP mengumpulkan token di pengirim untuk mengurangi jumlah data yang ditransfer. Namun, pengoptimalan ini dapat berdampak pada akurasi. Selain itu, karena kernel mentransfer data dalam satu penulisan, kami menggunakan lebih banyak memori untuk buffer kirim dan terima.

Throughput Kimi-K2 / DeepSeek-V3

Untuk membuktikan kelayakan kernel kami, kami menguji performa ujung ke ujung (end-to-end) dari mesin inferensi internal kami yang melayani DeepSeek-V3 (671B parameter) dan Kimi-K2 (1T parameter) pada beberapa instans p5en yang dilengkapi dengan Elastic Fabric Adapter (EFA) dan GPU H200. Kami menjalankan model dalam mode verifikasi draf dan melaporkan latensi dengan asumsi tingkat penerimaan 80% menggunakan satu lapisan MTP.

Pada DeepSeek-V3, kami membandingkannya terhadap baseline dalam node (intranode) EP=8 yang menjalankan kernel pengiriman dan penggabungan khusus NVLink yang efisien. Penerapan antar-node secara keseluruhan dapat menyamai atau melampaui performa penerapan node tunggal, terutama pada batch berukuran sedang, yang menunjukkan bahwa penskalaan di seluruh node memang membawa peningkatan pada throughput secara keseluruhan.

Pada Kimi-K2, kernel memungkinkan peny layanan. Karena ukurannya, model tidak dapat muat dalam node yang sama menggunakan Data Parallelism (DP) pada node H200, sehingga penerapan antar-node diperlukan untuk melayani model tersebut. Karena model ini memiliki lebih sedikit kepala atensi (attention heads), kami mencapai latensi yang lebih baik daripada DeepSeek-V3 dan kami dapat melayani model pada EFA pada latensi yang layak.

Perbesar

Perbesar

Kami juga telah mempertimbangkan batch mikro, meskipun hal itu hanya meningkatkan latensi pada ukuran batch yang sangat tinggi atau saat menggunakan lebih dari satu token draf. Untuk mendapatkan manfaat dari pemecahan masalah GEMM yang dikelompokkan menjadi dua, ukuran batch token harus cukup padat. Sebagian besar waktu masih dihabiskan untuk pengiriman dan penggabungan MoE, yang menunjukkan bahwa perbaikan lebih lanjut pada kernel akan menghasilkan peningkatan latensi dan memfasilitasi batch mikro.

Pekerjaan Mendatang

Kami bekerja sama erat dengan para insinyur AWS untuk meningkatkan performa pada EFA. Di pihak kami, kami akan mengikuti pembaruan pada libfabric yang akan mengurangi overhead bidang data. Selain itu, kami akan bereksperimen dengan efa-direct untuk lebih mengurangi overhead ruang pengguna (userspace) dan berinteraksi langsung dengan perangkat keras.

Referensi

Laporan Teknis DeepSeek R1

Laporan Teknis Kimi K2

DeepEP

UCCL-EP

Meningkatkan Performa Jaringan Sistem HPC Menggunakan NVIDIA Magnum IO NVSHMEM dan GPUDirect Async

Kutipan

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}