DeepGEMM: kernel FP8 GEMM yang bersih dan efisien melalui penskalaan granular

 (github.com/deepseek-ai)
2 poin oleh GN⁺ 2025-02-27 | 1 komentar | Bagikan ke WhatsApp
  • DeepGEMM adalah library kernel tensor core berperforma tinggi yang menyatukan primitive komputasi inti LLM modern seperti GEMM, fused MoE, MQA scoring, dan HyperConnection dalam satu codebase CUDA
  • Semua kernel dikompilasi saat runtime sebagai modul JIT yang ringan, sehingga tidak memerlukan kompilasi CUDA saat instalasi, dan membutuhkan C++20, CUDA Toolkit, PyTorch, serta CUTLASS 4.0 atau lebih baru
  • Library ini memanfaatkan sebagian konsep dari CUTLASS dan CuTe, tetapi tidak terlalu bergantung pada template dan struktur aljabar, serta dirancang agar pembelajaran optimisasi kernel GPU NVIDIA lebih mudah diakses melalui jumlah fungsi kernel inti yang terbatas
  • Cakupan dukungannya meliputi FP8, FP4, BF16 GEMM, grouped GEMM, kernel MQA logits untuk DeepSeek v3.2, hingga Mega MoE yang menumpangtindihkan komunikasi dan komputasi, dengan batasan layout memori yang berbeda pada SM90 dan SM100
  • Meski dirancang ringan, library ini menargetkan performa yang setara atau lebih tinggi daripada library yang dituning oleh pakar pada berbagai shape matriks, dan mencakup pembaruan yang mencapai hingga 1550 TFLOPS pada H800

Tujuan dan desain DeepGEMM

  • DeepGEMM adalah library kernel tensor core yang mengintegrasikan primitive komputasi utama yang digunakan pada model bahasa besar terbaru ke dalam satu codebase CUDA
    • GEMM: FP8, FP4, BF16
    • fused MoE dengan komunikasi yang ditumpangtindihkan: Mega MoE
    • MQA scoring untuk lightning indexer
    • HyperConnection(HC)
  • Semua kernel dikompilasi saat runtime sebagai modul Just-In-Time(JIT) yang ringan
    • Tidak memerlukan kompilasi CUDA selama proses instalasi
  • Memanfaatkan sebagian konsep dari CUTLASS dan CuTe
    • Namun tidak terlalu bergantung pada template berat atau struktur aljabar
    • Menjaga codebase tetap sederhana dengan membatasi jumlah fungsi kernel inti
  • Meski desainnya ringan, disebut mampu memberikan performa setara atau lebih baik daripada library yang dituning pakar pada berbagai shape matriks

Pembaruan utama

  • Pembaruan 16 April 2026 mencakup Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, kompilasi JIT yang lebih cepat, dan lainnya
    • Detail ada di #304
    • Benchmark Mega MoE ada di #316
  • Pembaruan 28 September 2025 menambahkan kernel scoring weighted ReLU MQA logits untuk lightning indexer DeepSeek v3.2
    • Detail ada di #200
  • Pembaruan 20 Juli 2025 mendukung SM90 dan SM100 sekaligus, serta direfaktor penuh menjadi modul JIT CPP dengan overhead CPU rendah
    • NVRTC dan optimisasi SASS pasca-kompilasi dinonaktifkan
    • NVRTC ditandai akan didukung nanti
    • Karena NVCC 12.9 melakukan FFMA interleaving secara otomatis, optimisasi pasca-kompilasi tidak lagi didukung
    • Detail ada di #112
  • Pembaruan 14 Mei 2025 menambahkan kernel weight gradient untuk dense dan backward MoE
    • Detail ada di #95
  • Pembaruan 7 Mei 2025 menghadirkan dukungan NVRTC dengan peningkatan kecepatan kompilasi hingga 10x
    • Dapat diaktifkan dengan DG_JIT_USE_NVRTC=1
    • Dalam beberapa kasus bisa ada penurunan performa
    • Detail ada di #94
  • Pembaruan 18 April 2025 mencapai hingga 1550 TFLOPS pada H800

Persyaratan dan alur instalasi

  • Lingkungan eksekusi membutuhkan GPU arsitektur NVIDIA SM90 atau SM100
  • Persyaratan perangkat lunaknya adalah sebagai berikut
    • Python 3.8 atau lebih baru
    • Compiler dengan dukungan C++20
    • CUDA Toolkit
      • SM90: CUDA 12.3 atau lebih baru
      • Untuk performa terbaik, sangat disarankan CUDA 12.9 atau lebih baru
      • SM100: CUDA 12.9 atau lebih baru
    • PyTorch 2.1 atau lebih baru
    • CUTLASS 4.0 atau lebih baru
    • Library {fmt}
  • Di lingkungan pengembangan, repositori di-clone beserta submodule, lalu develop.sh dijalankan untuk menghubungkan include yang diperlukan dan membangun modul CPP JIT
  • Instalasi dilakukan dengan menjalankan install.sh, lalu mengimpor deep_gemm di proyek Python

Antarmuka GEMM dan batasan layout

  • Konvensi penamaan kernel GEMM di DeepGEMM adalah D = C + A @ B
  • Layout shape input menggunakan NT sebagai acuan
    • fp8_gemm_nt menjalankan D = C + A @ B.T
  • Implementasi SM90 hanya mendukung layout memori NT
    • Ini sesuai dengan kombinasi row-major dan col-major
  • Implementasi SM100 mendukung seluruh layout memori NT, TN, NN, TT
  • Pada kedua arsitektur, scaling factor LHS harus berupa layout yang selaras TMA dan sudah ditransposisikan
    • SM90 mensyaratkan scaling factor dalam format FP32
    • SM100 mensyaratkan format packed UE8M0, dengan 4 UE8M0 dipaketkan ke dalam satu torch.int
  • Operasi seperti transpose input atau casting FP8 harus ditangani terpisah oleh pengguna
    • Library menyediakan fungsi utilitas PyTorch sederhana, tetapi performanya bisa lambat
    • Fokus utamanya adalah optimisasi kernel GEMM

Dense dan Grouped GEMM

  • FP8 GEMM dasar non-grouped menggunakan fungsi fp8_gemm_{nt, nn, tn, tt}
  • Grouped GEMM dengan contiguous layout, berbeda dari grouped GEMM tradisional di CUTLASS, hanya melakukan pengelompokan pada sumbu M
    • N dan K harus tetap
    • Ini dirancang untuk situasi pada model MoE ketika para expert berbagi shape yang sama
  • Pada training forward pass atau inference prefilling, jumlah token yang diproses tiap expert bisa berbeda
    • Bentuk token yang digabungkan ke dalam satu tensor ini disebut contiguous layout
    • Setiap segmen expert harus disejajarkan dengan ukuran blok GEMM M
    • Aturan alignment dapat diperiksa dengan get_mk_alignment_for_contiguous_layout()
  • Untuk backward bobot MoE, juga disediakan API grouped pada sumbu K
    • M dan N harus tetap
    • Fungsi terkait adalah k_grouped_fp8_gemm_tn_contiguous
  • Pada tahap inference decoding saat CUDA graph aktif dan CPU tidak mengetahui jumlah token per expert, didukung masked grouped GEMM
    • Jika tensor mask diberikan, kernel hanya menghitung area yang valid
    • Fungsinya adalah m_grouped_fp8_gemm_nt_masked
    • Ada contoh yang menggunakan output kernel latensi rendah dari DeepEP sebagai input

Kernel MQA untuk DeepSeek v3.2 Indexer

  • Keluarga kernel MQA V3.2 menyediakan versi non-paged dan versi paged
    • non-paged untuk prefilling
    • paged untuk decoding
  • fp8_mqa_logits menerima 6 input
    • q: tensor E4M3, shape [seq_len, num_heads, head_dim]
    • kv: tensor E4M3 dan float scaling factor
      • shape tensor adalah [seq_len_kv, head_dim]
      • shape scaling factor adalah [seq_len_kv]
    • weights: tensor float, shape [seq_len, num_heads]
    • cu_seq_len_k_start, cu_seq_len_k_end: tensor int, shape [seq_len]
    • clean_logits: apakah logits yang tidak terisi akan dibersihkan menjadi -inf
  • Shape tensor output adalah [seq_len, seq_len_kv] dan merepresentasikan logits token-ke-token
  • Setiap token q i melakukan iterasi atas token kv j dari cu_seq_len_k_start[i] hingga sebelum cu_seq_len_k_end[i]
    • scaling factor dikalikan ke kv_j
    • nilai per head dihitung dengan q[i, :, :] @ kv_j
    • Setelah ReLU diterapkan, hasilnya dikalikan dengan weights[i, :] lalu dijumlahkan untuk membuat logit skalar
  • Fungsi versi paged adalah fp8_paged_mqa_logits

Mega MoE

  • Mega MoE menggabungkan beberapa tahap MoE ke dalam satu mega-kernel
    • EP dispatch
    • linear 1, FP8xFP4
    • SwiGLU
    • linear 2, FP8xFP4
    • EP combine
  • Mega MoE menumpangtindihkan komunikasi NVLink dan komputasi tensor core
  • Eksekusinya memerlukan multi-process launch yang menggunakan symmetric memory
  • Alur penggunaannya adalah sebagai berikut
    • Alokasikan buffer symmetric memory dengan deep_gemm.get_symm_buffer_for_mega_moe
      • Membutuhkan PyTorch 2.9 atau lebih baru
    • Ubah bobot ke layout yang dibutuhkan, termasuk FP4 dan UE8M0 SF, dengan deep_gemm.transform_weights_for_mega_moe
    • Sebelum pemanggilan, salin input, scaling factor, top-k index, dan top-k weight ke buffer
    • Jalankan kernel fused mega MoE dengan deep_gemm.fp8_fp4_mega_moe
  • Contoh lengkap konfigurasi multi-process dan benchmarking ada di tests/test_mega_moe.py

Utilitas dan variabel lingkungan

  • Fungsi utilitas utama mengontrol resource eksekusi, alignment, kompilasi JIT, dan konversi scaling factor
    • deep_gemm.set_num_sms / get_num_sms: mengatur dan mengambil jumlah maksimum SM yang akan digunakan
    • deep_gemm.set_tc_util / get_tc_util: mengatur dan mengambil rasio utilisasi tensor core perkiraan
    • deep_gemm.set_pdl / get_pdl: mengaktifkan dan menonaktifkan Programmatic Dependent Launch(PDL)
    • deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: mengatur dan mengambil alignment M/K level grup untuk contiguous layout
    • deep_gemm.transform_sf_into_required_layout: mengubah scaling factor ke layout yang dibutuhkan
    • deep_gemm.get_tma_aligned_size: mengambil ukuran alignment TMA yang diperlukan
  • Variabel lingkungan terkait JIT mengontrol output debug, lokasi cache, pemilihan compiler, dan opsi profiling
    • DG_JIT_DEBUG: menampilkan informasi debug JIT
    • DG_PRINT_CONFIGS: menampilkan config terpilih per shape
    • DG_JIT_CACHE_DIR: direktori cache kernel hasil kompilasi, default-nya $HOME/.deep_gemm
    • DG_JIT_USE_NVRTC: menggunakan NVRTC alih-alih NVCC; kompilasi bisa lebih cepat tetapi dalam beberapa kasus performanya lebih rendah
    • DG_JIT_NVCC_COMPILER: path compiler NVCC
    • DG_JIT_CPP_STANDARD: versi standar C++, default 20
  • Variabel lingkungan untuk debug dan profiling juga disediakan
    • DG_JIT_DUMP_ASM, DG_JIT_DUMP_PTX, DG_JIT_DUMP_SASS: dump keluaran PTX dan SASS
    • DG_JIT_WITH_LINEINFO: menyertakan informasi source line untuk alat profiling
    • DG_COMM_KERNEL_DEBUG: menginisialisasi buffer simetris dengan 0 sebelum pemanggilan Mega MoE
    • DG_USE_NVIDIA_TOOLS: melewati profiling internal saat menjalankan alat NVIDIA eksternal
  • Opsi build mengontrol instalasi dan cara pemuatan kernel
    • DG_SKIP_CUDA_BUILD: melewati build CUDA extension saat instalasi
    • DG_FORCE_BUILD: memaksa build lokal alih-alih mengunduh pre-built wheel
    • DG_JIT_USE_RUNTIME_API: menggunakan CUDA Runtime API untuk memuat kernel, membutuhkan CUDA runtime 12.8 atau lebih baru

Lisensi dan sitasi

  • Repositori DeepGEMM dirilis dengan MIT License
  • Proyek ini menyatakan terinspirasi dari CUTLASS
  • Judul entri sitasinya adalah DeepGEMM: clean and efficient BLAS kernel library on GPU

Bacaan terkait

1 komentar

 
GN⁺ 2025-02-27
Komentar Hacker News

  • Interleaving FFMA SASS benar-benar terlihat mencengangkan
    Setelah melihat peningkatan performa kernel CUTLASS FP8 antara NVCC 12.2 dan 12.3 lalu membandingkan SASS hasil kompilasinya, tampaknya mereka menemukan satu bit yang dibalik dalam pola interleaving pada beberapa instruksi FADD. Dengan merujuk pada implementasi assembler CUDA open source, mereka menyimpulkan bahwa bit tersebut adalah yield bit yang membuat warp saat ini mengalah agar warp lain dapat berjalan
    Yang mengesankan, mereka membuat skrip untuk memodifikasi instruksi FFMA pada biner hasil kompilasi dengan memanfaatkan ini. Karena reuse register tidak mungkin dilakukan ketika warp mengalah, mereka juga membalik reuse bit bersama yield bit, sehingga instruksi MMA dan instruksi FFMA promosi pada GEMM FP8 dengan scaling halus dapat saling tumpang tindih lebih baik, dan dalam beberapa kasus meningkatkan performa lebih dari 10%

    • Dari yang pernah saya baca di tempat lain, pendekatan semacam ini cukup lazim dalam optimasi operasi matriks yang performanya kritis
      Hanya saja untuk masalah khusus ini, tampaknya perusahaan AI lain belum merasa perlu menerapkannya, dan pada akhirnya kemungkinan semua akan mencapai titik yang mirip
    • Scott Gray sudah menemukan hal yang persis seperti ini, bahkan lebih, di Maxwell pada 2015, dan setelah itu banyak orang juga telah membahasnya

  • Contoh seperti ini menunjukkan betapa jauhnya compiler yang ada sekarang dari kemampuan mengekstrak performa hardware hanya dari kode tingkat tinggi
    Saya penasaran apa yang dibutuhkan agar teknik compiler tradisional atau agen optimasi berbasis AI bisa menghasilkan hasil seperti ini

    • Sepertinya dibutuhkan trial-and-error yang sangat besar dalam loop umpan balik reinforcement learning

  • Angka percepatan yang dilaporkan dibandingkan dengan baseline internal berbasis CUTLASS
    Saya penasaran apakah ada yang sudah membandingkan performanya langsung dengan cuBLAS
    Sejauh ini hasil CUTLASS GEMM yang saya lihat kira-kira masih dalam 10% dibanding cuBLAS, jadi kalau peningkatan 2x–2,5x yang disebut di makalah tetap bertahan, itu akan benar-benar mengesankan

    • Biasanya saya menghindari FP8 dan lebih memilih I8, tetapi pertanyaan ini membuat saya penasaran seberapa bagus cuBLAS
      Pertama, cuBLAS membutuhkan API ekstensi cuBLASLt untuk menangani pekerjaan presisi campuran seperti FP8
      Selain itu, kombinasi tipe yang tampak masuk akal seperti E5M2 x E5M2 pada A x B tidak didukung, tetapi E5M2 x E4M3 didukung; batasan lain juga terus ada, misalnya pada Ampere, Hopper, dan Blackwell, matriks A selalu harus berada dalam layout tertransposisi
      Saya mengintegrasikan benchmark FP8 cuBLASLt ke repositori "Less Slow C++" saya <https://github.com/ashvardanian/less_slow.cpp>, dan menambahkannya ke daftar benchmark cuBLAS yang sudah ada serta benchmark CUDA/PTX yang saya tulis sendiri
      Saya menjalankannya di GPU H200, yang semestinya memiliki performa sama dengan H100, dan pada input persegi throughput-nya mencapai puncak sekitar 1,35 Peta-ops
      Untuk 256 sekitar 2,68T/s, 512 20,49T/s, 1024 144,23T/s, 2048 665,68T/s, 4096 1,26P/s, 8192 1,34P/s, dan 16384 1,23P/s; ini sekitar 67% dari angka yang dipromosikan NVIDIA untuk dense GEMM <https://resources.nvidia.com/en-us-data-center-overview-mc/e...>
    • Saya dengar CUTLASS bisa menghasilkan performa yang lebih baik daripada cuBLAS
      Saya kira baseline-nya memilih yang lebih baik antara cuBLAS dan CUTLASS

  • Open source seperti ini benar-benar menunjukkan dengan baik tujuan industri untuk mencapai efisiensi
    Namun, manfaat software ini tampaknya akan lebih besar bagi perusahaan besar yang melayani model dalam skala besar—yakni calon pesaing DeepSeek—daripada bagi komunitas open source umum yang ingin belajar, bereksperimen, atau melayani model di hardware konsumen

    • Jika efisiensi meningkat, pada akhirnya itu bisa berujung pada hardware yang lebih murah bagi semua orang, termasuk DeepSeek sendiri

  • Saya tidak yakin apakah arah optimasi ke presisi yang semakin rendah ini baik dalam jangka panjang
    Ini berarti model sebenarnya cukup sparse, tetapi meski sekarang bisa begitu, saya kira kemungkinan besar bukan karena secara intrinsik harus sesparse itu, melainkan karena ada ide buruk yang tercampur dalam cara pelatihannya

    • Selama sparsity yang didapat gratis masih berhasil, nikmati saja
      Membuat model yang sangat bagus hanya bisa dilatih pada presisi lebih tinggi adalah masalah riset, sedangkan pelatihan dan inferensi presisi rendah adalah masalah engineering
      Sejak era CNN, setidaknya sejak 9 tahun lalu, pekerjaan seperti ini sudah dilakukan, dan menurut saya masih akan berlangsung beberapa tahun lagi
    • Karena fungsi aktivasi membuang cukup banyak rentang dinamis angka floating-point, cukup jelas bahwa memberi rentang luas pada wilayah aktivasi yang sudah jenuh kemungkinan tidak terlalu berguna

  • Ini mungkin menjadi tidak relevan karena MXFP, yaitu dukungan microscaling native di Blackwell
    Di Hopper, pada dasarnya hal itu diimplementasikan secara manual dengan granularitas yang lebih kasar, tetapi menggunakan koefisien scaling FP32

    • Benar
      Demo publik berkualitas tinggi seperti ini menunjukkan dengan baik di mana moat $NVDA berada
      GPU general-purpose sangat fleksibel, sehingga berbagai pekerjaan yang masuk akal tetapi tidak terpikirkan oleh vendor hardware sejak awal dapat dilakukan melalui pemrograman
      Namun jika memprediksi bahwa masa depan akan semakin mengerucut pada dukungan hardware khusus sehingga ruang optimasi software seperti ini menghilang, maka apa yang disebut moat CUDA akan runtuh
      Demi tetap bertahan dalam permainan ini, NVIDIA pada dasarnya sedang meruntuhkan moat-nya sendiri :p

  • Wah, ini memakai lisensi MIT
    Semoga perusahaan-perusahaan besar mengadopsi cara kolaborasi open source seperti ini

  • Saya terus penasaran mengapa ada instruksi yang tidak terdokumentasi
    Rasanya lebih tepat jika tetap disediakan kepada pengguna meskipun belum sepenuhnya stabil
    Hal seperti ini semestinya terdokumentasi secara internal, tetapi saya tidak mengerti mengapa tidak dipublikasikan
    Keamanan yang bergantung pada ketidakjelasan tidak akan efektif, dan pesaing pada akhirnya akan melakukan rekayasa balik semuanya

    • Mungkin alasannya mirip dengan mengapa pada hal yang kita buat juga ada bagian yang tidak terdokumentasi
      Bisa karena kekurangan waktu, atau karena tidak ingin mengisyaratkan dukungan untuk fitur yang belum stabil atau masih eksperimental
      Jika dampaknya hanya sebatas tim di sebelah, mengubahnya juga jauh lebih mudah
    • Asumsi bahwa “hal seperti ini pasti terdokumentasi secara internal” bisa saja keliru sejak awal
      Kemungkinan besar hanya tercantum di dokumen seperti dokumen desain arsitektur atau spesifikasi, dan dokumen semacam itu tentu saja tidak ingin mereka bagikan

  • Sejujurnya ini di luar cakupan penggunaan dan pemahaman saya
    Meski begitu, saya sangat berterima kasih dan merasa ini menyegarkan karena temuan dan perbaikan seperti ini dibagikan sehingga semua orang bisa mendapat manfaat

    • FFMA adalah singkatan dari Fused Floating-point Multiply-Add, instruksi GPU dasar yang menjalankan D = A*B + C sekaligus
      Ini sangat penting dalam perkalian matriks dan beban kerja deep learning
      Dalam SASS NVIDIA, instruksi FFMA dienkode sebagai instruksi 64-bit atau 128-bit, dan memiliki berbagai bit kontrol yang menentukan perilaku persisnya
      Jika bit yield disetel, ia memberi tahu warp scheduler bahwa setelah instruksi ini warp saat ini dapat menyerahkan eksekusi, dan hardware dapat menjalankan warp lain untuk menyembunyikan latensi
      GPU memperoleh throughput tinggi melalui paralelisme masif, dan ketika satu warp berhenti karena menunggu memori atau hal lain, warp lain dapat terus berjalan
      Bit reuse menunjukkan apakah register sumber dapat digunakan kembali pada operasi berikutnya, dan jika bit yield disetel, bit ini harus dimatikan
      Sebab jika sebuah warp menyerahkan eksekusi, warp itu mungkin bukan warp berikutnya yang dijalankan, dan warp lain dapat mengubah status register file, sehingga hardware tidak dapat menjamin bahwa nilai register tetap dipertahankan melampaui yield
      Jika bit yield pada instruksi-instruksi FFMA disetel dengan pola berselang-seling, compiler membuat titik penjadwalan eksplisit tempat warp lain dapat berjalan, dan untuk menjaga kebenaran, bit reuse pada instruksi tersebut juga harus dihapus
      Perubahan ini khususnya membantu menumpangtindihkan instruksi MMA, yang merupakan inti perkalian matriks, dengan instruksi FFMA promosi yang melakukan konversi untuk mengakumulasikan FP8 dengan presisi lebih tinggi
      FP8 GEMM biasanya memerlukan proses konversi ke presisi lebih tinggi untuk akumulasi lalu mengembalikannya lagi, sehingga muncul FFMA tambahan; ini mengurangi kebutuhan bandwidth memori, tetapi menciptakan pola komputasi kompleks yang mencampurkan operasi promosi/demosi
      “Scaling halus” tampaknya berarti pekerjaan mengelola presisi secara hati-hati di berbagai titik komputasi
      Manipulasi bit yield membuat operasi komputasi dan konversi format dapat diinterleaving dengan lebih baik, sehingga unit eksekusi GPU digunakan lebih efisien; tanpa optimasi ini, warp scheduler mungkin tidak menemukan peluang transisi yang alami, sehingga sumber daya komputasi bisa kurang termanfaatkan