AMD MI355X Menguasai NVIDIA B200 dalam Pelaksanaan Kimi K3

Ini mungkin saat yang telah lama ditunggu oleh AMD.

Baru-baru ini, Wafer AI telah melaksanakan di AMD MI355X Kimi K3. Hasilnya, model yang sebelumnya memerlukan 16 unit NVIDIA B200 yang tersebar di dua pelayan kini dapat dideploy dalam satu pelayan AMD dengan 8 unit MI355X.

Lebih penting lagi, ia bukan sekadar memasukkan model.

Dalam ujian dengan input 1024 Token dan output 400 Token, MI355X mencapai jumlah throughput 952 Token/s, dengan kelajuan penghasilan pengguna tunggal sebanyak 118 Token/s.

Berdasarkan pengiraan satu nod, throughput-nya kira-kira 3.8 kali ganda berbanding penyelesaian 16 unit B200, dan nilai berbanding kosnya juga melebihi B200 dan B300.

Dan yang paling mengejutkan, ROCm kali ini tidak begitu merepotkan.

Kimi K3 memiliki 2.8 triliun parameter, dan hanya bobot model sahaja memerlukan lebih daripada 1.5 TB memori GPU, belum termasuk KV Cache yang diperlukan untuk konteks jutaan Token.

Sebuah pelayan 8-keping B200, dengan setiap keping mempunyai 192 GB memori video, jumlah kapasiti sekitar 1.5 TB. Dengan kata lain, berat model sukar untuk dimuatkan sepenuhnya, apalagi meninggalkan ruang untuk KV Cache. Oleh itu, B200 mesti menggunakan dua pelayan, 16 GPU.

B300 memiliki 288 GB memori video setiap kad, membolehkan model dipasang dalam satu nod. Secara kebetulan, AMD MI355X juga mempunyai 288 GB memori video, dengan 8 kad MI355X menjumlahkan kira-kira 2.3 TB, cukup dengan satu pelayan.

Ini bukan sekadar mengurangkan satu mesin. Selepas model berjalan merentasi nod, setiap Token yang dihasilkan mungkin memerlukan penyegerakan data melalui rangkaian. Walaupun menggunakan rangkaian RoCE v2 sekelajuan kira-kira 195 Gb/s, komunikasi merentasi nod masih akan memperlambatkan dekod.

MI355X menggunakan memori video yang lebih besar untuk menyimpan keseluruhan model dalam satu nod.

Dari hasil akhir, jumlah puncak throughput untuk 8 unit MI355X mencapai 952 Token/s, dengan kelajuan penghasilan satu saluran sebanyak 118 Token/s.

Sebagai perbandingan, jumlah throughput keseluruhan untuk penghantaran dua nod dengan 16 B200 ialah 498 Token/s, yang setara dengan sekitar 249 Token/s setiap nod.

Dengan kata lain, throughput per node MI355X kira-kira 3.8 kali ganda berbanding throughput per node purata penyebaran dua node B200. Dari segi kelajuan penghasilan pengguna tunggal, 118 Token/s MI355X juga lebih tinggi berbanding 90 Token/s B200.

B300 masih merupakan solusi dengan prestasi tertinggi. Jumlah throughput keseluruhan nod dengan 8 unit B300 mencapai 1568 Token/s, dengan kelajuan penghasilan satu laluan sebanyak 172 Token/s, menjadikan jumlah throughput kira-kira 1.65 kali ganda MI355X.

Namun, harga yang berubah telah mengubah kesimpulan tersebut. Wafer dihitung pada kadar $2.50 per kad per jam untuk MI355X, $4.25 untuk B200, dan $6 untuk B300.

Dengan anggaran harga ini, MI355X boleh menyediakan throughput puncak sekitar 48 Token/s per dolar; B200 sekitar 7 Token/s; B300 sekitar 33 Token/s.

B300 lebih pantas, tetapi kecekapan kos unit MI355X lebih tinggi. Untuk pusat data yang memerlukan pengendalian model terbuka dalam skala besar, ini mungkin lebih penting daripada sekadar bersaing untuk gelaran prestasi terbaik.

Selama ini, masalah terbesar GPU pusat data AMD sering bukan pada peranti keras, tetapi pada perisian.

Model yang sama boleh dijalankan secara langsung di CUDA, tetapi di ROCm, mungkin perlu mengubah kerangka kerja, menambah operator, atau bahkan menulis semula kernel bawahannya.

tetapi Kimi K3 berbeza situasinya.

AMD menyediakan sokongan yang hampir serentak pada pelancaran. Wafer menyatakan bahawa model boleh dijalankan secara langsung di MI355X, dan kerja seterusnya terutamanya berfokus pada beberapa isu kompatibiliti dan pengoptimuman prestasi.

Satu masalah muncul dalam fasa spekulasi dekod. Kimi K3 tidak menyediakan parameter model draf yang diperlukan untuk MTP atau EAGLE, oleh itu Wafer menggunakan model draf blok luaran.

Rancangan ini boleh dijalankan secara langsung pada CUDA, tetapi dalam persekitaran ROCm, permintaan sebenar pertama menyebabkan ralat penjadual. Sebabnya ialah fungsi bernama top_k_renorm_prob tidak ditakrifkan dalam cawangan ROCm.

Fungsi ini tidak rumit: memilih k nilai tertinggi daripada taburan kebarangkalian, menetapkan kebarangkalian lain kepada sifar, kemudian menormalkan semula kebarangkalian yang dikekalkan.

Wafer akhirnya menggunakan fungsi PyTorch biasa untuk melengkapi logik ini, tanpa perlu menulis kernel GPU secara manual atau mereka semula sistem penghuraian spekulatif.

Selepas diperbaiki, perkiraan dekod menghasilkan peningkatan kinerja satu laluan sekitar 2.2 kali, kinerja aliran tunggal di bawah kepadatan sederhana sekitar 1.7 kali, dan peningkatan throughput puncak keseluruhan sekitar 18%.

Lebih penting lagi, sistem mampu mencapai throughput puncak di bawah beban konsisten yang lebih tinggi.

Tentu, throughput bukanlah satu-satunya aspek dalam perkhidmatan inferens. Bagi pengguna sebenar, indikator lain yang secara langsung mempengaruhi pengalaman ialah TTFT, iaitu masa tunggu dari penghantaran permintaan hingga menerima Token pertama.

Pada tugas ini, MI355X menunjukkan prestasi awal yang tidak baik. Menghadapi tugas pra-pengisian permulaan sebanyak kira-kira 172,000 Token, MI355X memerlukan kira-kira 51 saat, manakala B300 hanya memerlukan kira-kira 23 saat.

Dalam model yang menyokong konteks jutaan token, tugas pra-isian boleh menjadi sangat besar. Jika pengguna perlu menunggu beberapa puluh saat atau lebih lama setiap kali mengendalikan konteks panjang, kelajuan dekod yang tinggi sekalipun sukar memperbaiki pengalaman pengguna.

Wafer akhirnya menemukan bahawa kesenjangan prestasi hampir seluruhnya datang dari satu kernel perhatian. Kimi K3 dalam konfigurasi paralel tensor 8-laluan, setiap GPU akan dibahagikan 12 kepala perhatian. Manakala kernel pra-isian MLA yang lebih pantas dalam AMD AITER hanya menyokong bentuk yang merupakan gandaan 4, 8, atau 16.

12 kepala tidak dapat disesuaikan, maka sistem memulihkan implementasi Triton umum yang lebih perlahan.

Penyelesaianannya sangat sederhana: tambah nol kepada 12 kepala perhatian hingga mencapai 16, panggil kernel pantas yang sedia ada, kemudian ambil semula 12 kepala yang benar-benar diperlukan selepas pengiraan selesai. Tiada perubahan terhadap struktur model, tiada penulisan kernel asm baru, hanya menambah empat nol.

Selepas pengoptimuman, kelajuan pra-isian stabil inti AITER MLA mencapai sekitar 13,000 token/s, manakala laluan fallback Triton sebelum ini hanya sekitar 4,000 hingga 7,000 token/s, menjadikan masa pra-isian sejuk berkurang sebanyak dua hingga tiga kali ganda.

Pengoptimuman ini tidak mengubah throughput dekod akhir, tetapi secara signifikan mengurangkan masa tunggu pengguna untuk munculnya huruf pertama.

Ini juga menunjukkan bahawa jurang perisian yang kelihatan besar antara AMD dan NVIDIA kadang-kadang bukan disebabkan oleh kekurangan kemampuan asas, tetapi kerana kernel berkelajuan tinggi yang sedia ada belum lagi merangkumi bentuk model baru tertentu.

Satu ujian tentu tidak membuktikan bahawa AMD telah mengejar semula NVIDIA sepenuhnya.

B200 terpaksa berjalan merentas nod disebabkan kekurangan memori video; prestasi mutlak B300 masih memimpin; alat, kerangka kerja, dan ekosistem pembangun ROCm juga masih ketinggalan berbanding CUDA.

Namun, model terbuka sedang dengan cepat memasuki era parameter triliunan. Apabila model menjadi sebesar yang tidak muat dalam satu server, kapasiti memori video bukan lagi nombor semata-mata dalam jadual parameter, tetapi akan secara langsung mempengaruhi kos komunikasi, kompleksiti pelaksanaan, dan throughput akhir.

Strategi AMD untuk memberikan lebih banyak HBM kepada konfigurasi satu kad sedang menjadi kelebihan sistem yang sebenarnya.

Jika AMD dapat terus meningkatkan kestabilan ROCm, memperluas sokongan bentuk inti berkelajuan tinggi, dan menyediakan penyesuaian hari pertama yang lebih pantas untuk model-model baru, maka pusat data mesti mempertimbangkan GPU ini dengan serius. Harga lebih rendah, memori video lebih besar, prestasi mencukupi, dan perisian tidak lagi memerlukan usaha berbulan-bulan.

Apakah pendapat anda tentang ini?

Pautan rujukan:

https://x.com/wafer_ai/status/2083628389903315406

https://x.com/ChiragAsarpota/status/2083864019870634151

Artikel ini berasal dari akaun WeChat "Machine Heart" (ID: almosthuman2014), penulis: Peminat LLM