DeepGEMM: Sauberer und effizienter FP8-GEMM-Kernel durch feingranulares Scaling

DeepGEMM

DeepGEMM ist eine Bibliothek für FP8 General Matrix Multiplication (GEMM), die das in DeepSeek-V3 vorgeschlagene feingranulare Scaling unterstützt. Die Bibliothek unterstützt allgemeine sowie für Mixture-of-Experts (MoE) gruppierte GEMMs, ist in CUDA geschrieben und erfordert bei der Installation keine Kompilierung. Sie unterstützt NVIDIA-Hopper-Tensor-Cores und verwendet eine zweistufige Akkumulation auf CUDA-Cores, um die Ungenauigkeit der FP8-Tensor-Core-Akkumulation zu beheben. Dabei nutzt sie teilweise Konzepte aus CUTLASS und CuTe, hält die Implementierung jedoch durch minimale Abhängigkeit von Templates oder Algebra bewusst einfach. Mit einer einzigen zentralen Kernelfunktion von etwa 300 Zeilen Code ist sie eine geeignete Ressource, um Hopper-FP8-Matrixmultiplikation und Optimierungstechniken zu erlernen. Trotz des leichtgewichtigen Designs erreicht oder übertrifft sie bei verschiedenen Matrixformen die Leistung von von Expert:innen abgestimmten Bibliotheken.

Leistung

Alle Formen, die in der DeepSeek-V3/R1-Inferenz verwendet werden können, wurden auf H800 SXM5 mit NVCC 12.8 getestet. Alle Speedup-Metriken wurden durch Vergleich mit einer intern optimierten Implementierung auf Basis von CUTLASS 3.6 berechnet. Bei einigen Formen kann die Leistung schlechter ausfallen; Optimierungs-PRs sind willkommen.

Allgemeines GEMM (dichtes Modell)

• Messungen der DeepGEMM-Leistung über verschiedene Matrixgrößen hinweg zeigen bei bestimmten Größen bis zu 2,7-fache Beschleunigung.

Gruppiertes GEMM für MoE-Modelle (kontinuierliches Layout)

• Je nach Anzahl der Gruppen und Matrixgröße pro Gruppe werden bis zu 1,2-fache Beschleunigungen erreicht.

Gruppiertes GEMM für MoE-Modelle (Masken-Layout)

• Mit Masken-Layout wird eine Beschleunigung von bis zu 1,2-fach erreicht.

Schnellstart

Anforderungen

• GPU mit Hopper-Architektur, Unterstützung für sm_90a erforderlich
• Python 3.8 oder höher
• CUDA 12.3 oder höher (für beste Leistung wird 12.8 oder höher empfohlen)
• PyTorch 2.1 oder höher
• CUTLASS 3.6 oder höher

Entwicklung

• Beschreibung des Entwicklungsablaufs einschließlich Klonen von Submodulen, Erstellen symbolischer Links, JIT-Kompilierung und Testen aller GEMM-Implementierungen.

Installation

• deep_gemm kann in Python-Projekten importiert und verwendet werden.

Schnittstellen

Hinweise

• Diese Bibliothek enthält nur GEMM-Kernel und unterstützt ausschließlich das NT-Format. Transponierung oder andere FP8-Casting-Aufgaben müssen separat implementiert werden.

Allgemeines dichtes GEMM (nicht gruppiert)

• Stellt Funktionen zur Ausführung grundlegender nicht gruppierter FP8-GEMMs bereit.

Gruppiertes GEMM (kontinuierliches Layout)

• Für Szenarien in MoE-Modellen entworfen, in denen die Expert:innen dieselbe Form teilen.

Gruppiertes GEMM (Masken-Layout)

• In der Decoding-Phase der Inferenz wird ein Masken-Tensor bereitgestellt, sodass nur gültige Teile berechnet werden.

Dienstprogramme

• Bietet verschiedene Dienstprogrammfunktionen und Umgebungsvariablen zur Leistungsoptimierung.

Optimierung

Persistente Warp-Spezialisierung

• Folgt dem CUTLASS-Design und überlappt Datenbewegung, Tensor-Core-MMA-Instruktionen und Promotion auf CUDA-Cores.

Hopper-TMA-Funktion

• Nutzt TMA zur Beschleunigung der Datenbewegung.

Allgemeine Detailoptimierungen

• Verbessert die Leistung durch verschiedene Optimierungstechniken.

Einheitlicher und optimierter Block-Scheduler

• Stellt einen Scheduler für alle nicht gruppierten und gruppierten Kernel bereit.

Vollständiges JIT-Design

• Verbessert die Leistung durch ein JIT-Design, das bei der Installation keine Kompilierung erfordert.

Nicht ausgerichtete Blockgrößen

• Unterstützt nicht ausgerichtete Blockgrößen, um bei bestimmten Formen die SM-Auslastung zu maximieren.

FFMA-SASS-Interleaving

• Modifiziert FFMA-Instruktionen zur Verbesserung der Warp-Level-Parallelität und damit der Leistung.

Danksagung

• DeepGEMM wurde vom CUTLASS-Projekt inspiriert und drückt den Entwickler:innen Dank und Respekt aus.

Lizenz

• Veröffentlicht unter der MIT-Lizenz.