DeepGEMM: Kernel FP8 GEMM yang bersih dan efisien dengan penskalaan granular

DeepGEMM

DeepGEMM adalah library untuk perkalian matriks umum FP8 (GEMM) yang mendukung penskalaan granular yang diusulkan di DeepSeek-V3. Library ini mendukung GEMM terkelompok untuk skenario umum maupun Mix-of-Experts (MoE), ditulis dengan CUDA, dan tidak memerlukan kompilasi saat instalasi. Library ini mendukung NVIDIA Hopper Tensor Core, dan menggunakan akumulasi dua tahap di CUDA core untuk mengatasi ketidakakuratan akumulasi FP8 Tensor Core. Library ini memanfaatkan sebagian konsep dari CUTLASS dan CuTe, tetapi tetap sederhana dengan meminimalkan ketergantungan pada template maupun aljabar. Dengan satu fungsi kernel inti yang terdiri dari sekitar 300 baris kode, ini adalah sumber yang cocok untuk mempelajari perkalian matriks FP8 Hopper dan teknik optimasinya. Meski berdesain ringan, performanya setara atau lebih baik daripada library yang telah disetel oleh pakar pada berbagai bentuk matriks.

Kinerja

Semua bentuk yang dapat digunakan dalam inferensi DeepSeek-V3/R1 diuji pada H800 SXM5 dengan NVCC 12.8. Semua metrik peningkatan kecepatan dihitung dengan membandingkannya terhadap implementasi yang dioptimalkan secara internal berbasis CUTLASS 3.6. Beberapa bentuk mungkin menunjukkan kinerja yang kurang baik, dan PR optimasi sangat disambut.

GEMM umum (model padat)

• Hasil pengukuran kinerja DeepGEMM pada berbagai ukuran matriks menunjukkan peningkatan kecepatan hingga 2,7x pada ukuran tertentu.

GEMM terkelompok untuk model MoE (layout kontigu)

• Bergantung pada jumlah grup dan ukuran matriks tiap grup, peningkatan kecepatan dapat mencapai 1,2x.

GEMM terkelompok untuk model MoE (layout mask)

• Dengan menggunakan layout mask, peningkatan kecepatan dapat mencapai 1,2x.

Mulai cepat

Persyaratan

• GPU arsitektur Hopper, dukungan sm_90a diperlukan
• Python 3.8 atau lebih baru
• CUDA 12.3 atau lebih baru (12.8 atau lebih baru direkomendasikan untuk performa terbaik)
• PyTorch 2.1 atau lebih baru
• CUTLASS 3.6 atau lebih baru

Pengembangan

• Menjelaskan proses pengembangan termasuk clone submodule, pembuatan symbolic link, kompilasi JIT, dan pengujian semua implementasi GEMM.

Instalasi

• deep_gemm dapat diimpor dan digunakan dalam proyek Python.

Antarmuka

Perhatian

• Library ini hanya mencakup kernel GEMM dan hanya mendukung format NT. Transposisi atau pekerjaan casting FP8 lainnya harus diimplementasikan secara terpisah.

GEMM padat umum (tidak terkelompok)

• Menyediakan fungsi untuk menjalankan FP8 GEMM dasar yang tidak terkelompok.

GEMM terkelompok (layout kontigu)

• Dirancang untuk skenario model MoE di mana para expert berbagi bentuk yang sama.

GEMM terkelompok (layout mask)

• Pada tahap decoding inferensi, tensor mask disediakan agar hanya bagian yang valid yang dihitung.

Utilitas

• Menyediakan berbagai fungsi utilitas dan variabel lingkungan untuk membantu optimasi performa.

Optimasi

Spesialisasi warp persisten

• Mengikuti desain CUTLASS, dengan menumpangtindihkan perpindahan data, instruksi Tensor Core MMA, dan promosi CUDA core.

Fitur Hopper TMA

• Memanfaatkan TMA untuk mempercepat perpindahan data.

Optimasi detail umum

• Meningkatkan performa melalui berbagai teknik optimasi.

Scheduler blok terpadu dan teroptimasi

• Menyediakan scheduler untuk semua kernel yang tidak terkelompok maupun terkelompok.

Desain JIT penuh

• Meningkatkan performa melalui desain JIT yang tidak memerlukan kompilasi saat instalasi.

Ukuran blok yang tidak disejajarkan

• Mendukung ukuran blok yang tidak disejajarkan untuk memaksimalkan pemanfaatan SM pada bentuk tertentu.

Interleaving FFMA SASS

• Meningkatkan paralelisme tingkat warp dengan memodifikasi instruksi FFMA demi peningkatan performa.

Ucapan terima kasih

• DeepGEMM terinspirasi dari proyek CUTLASS, dan menyampaikan terima kasih serta rasa hormat kepada para pengembangnya.

Lisensi

• Dirilis di bawah lisensi MIT.