DeepGEMM: İnce Ölçekleme ile Temiz ve Verimli FP8 GEMM Çekirdeği

DeepGEMM

DeepGEMM, DeepSeek-V3'te önerilen ince ölçeklemeyi destekleyen, FP8 genel matris çarpımı (GEMM) için bir kütüphanedir. Bu kütüphane genel ve Mix-of-Experts (MoE) gruplandırılmış GEMM'i destekler ve CUDA ile yazıldığı için kurulum sırasında derleme gerektirmez. NVIDIA Hopper tensor core'ları destekler ve FP8 tensor core birikiminin doğruluk sorununu çözmek için CUDA core tabanlı iki aşamalı birikim kullanır. CUTLASS ve CuTe kavramlarından kısmen yararlansa da, şablonlara ve cebirsel bağımlılıklara olan ihtiyacı en aza indirerek sadeliğini korur. Yaklaşık 300 satırlık koddan oluşan tek bir çekirdek kernel fonksiyonu sayesinde Hopper FP8 matris çarpımı ve optimizasyon tekniklerini öğrenmek için uygun bir kaynaktır. Hafif tasarımına rağmen, çeşitli matris biçimlerinde uzmanlarca ayarlanmış kütüphanelerle aynı ya da daha iyi performans gösterir.

Performans

DeepSeek-V3/R1 çıkarımında kullanılabilecek tüm biçimler, H800 SXM5 üzerinde NVCC 12.8 ile test edildi. Tüm hızlanma metrikleri, CUTLASS 3.6 tabanlı ve kurum içinde optimize edilmiş bir uygulamayla karşılaştırılarak hesaplandı. Bazı biçimlerde performans düşük olabilir; optimizasyon PR'ları memnuniyetle karşılanır.

Genel GEMM (yoğun modeller)

• Çeşitli matris boyutlarında yapılan DeepGEMM performans ölçümlerinde, belirli boyutlarda en fazla 2,7 kat hızlanma görüldü.

MoE modelleri için gruplandırılmış GEMM (ardışık yerleşim)

• Grup sayısına ve her grubun matris boyutuna bağlı olarak en fazla 1,2 kat hızlanma sağlar.

MoE modelleri için gruplandırılmış GEMM (maskeli yerleşim)

• Maskeli yerleşim kullanılarak en fazla 1,2 kat hızlanma elde edilir.

Hızlı başlangıç

Gereksinimler

• Hopper mimarili GPU, sm_90a desteği gerekli
• Python 3.8 veya üstü
• CUDA 12.3 veya üstü (en iyi performans için 12.8 veya üstü önerilir)
• PyTorch 2.1 veya üstü
• CUTLASS 3.6 veya üstü

Geliştirme

• Alt modüllerin klonlanması, sembolik bağlantı oluşturma, JIT derleme ve tüm GEMM uygulamalarının test edilmesini içeren geliştirme süreci açıklanır.

Kurulum

• deep_gemm, Python projelerine içe aktarılıp kullanılabilir.

Arayüz

Dikkat edilmesi gerekenler

• Bu kütüphane yalnızca GEMM çekirdeklerini içerir ve sadece NT biçimini destekler. Transpozisyon veya diğer FP8 cast işlemleri bağımsız olarak uygulanmalıdır.

Genel yoğun GEMM (gruplandırılmamış)

• Temel gruplandırılmamış FP8 GEMM çalıştırmak için bir işlev sunar.

Gruplandırılmış GEMM (ardışık yerleşim)

• MoE modellerinde uzmanların aynı biçimi paylaştığı senaryolar için tasarlanmıştır.

Gruplandırılmış GEMM (maskeli yerleşim)

• Çıkarımın decoding aşamasında mask tensor'ü verilerek yalnızca geçerli bölümler hesaplanır.

Yardımcı araçlar

• Performans optimizasyonuna yardımcı olmak için çeşitli yardımcı işlevler ve ortam değişkenleri sunar.

Optimizasyon

Kalıcı warp uzmanlaşması

• CUTLASS tasarımını izleyerek veri hareketini, tensor core MMA komutlarını ve CUDA core yükseltmesini üst üste bindirir.

Hopper TMA özelliği

• Veri hareketini hızlandırmak için TMA'dan yararlanır.

Ortak ayrıntı optimizasyonları

• Çeşitli optimizasyon teknikleriyle performansı artırır.

Birleşik ve optimize edilmiş blok zamanlayıcısı

• Tüm gruplandırılmamış ve gruplandırılmış çekirdekler için bir zamanlayıcı sunar.

Tam JIT tasarımı

• Kurulum sırasında derleme gerektirmeyen JIT tasarımıyla performansı artırır.

Hizalanmamış blok boyutları

• Belirli biçimlerde SM kullanımını en üst düzeye çıkarmak için hizalanmamış blok boyutlarını destekler.

FFMA SASS iç içe geçirme

• Performansı artırmak için FFMA komutlarını düzenleyerek warp düzeyinde paralelliği iyileştirir.

Teşekkür

• DeepGEMM, CUTLASS projesinden ilham almıştır; geliştiricilerine teşekkür ve saygı sunar.

Lisans

• MIT lisansı ile yayımlanmıştır.