DeepGEMM : des kernels GEMM FP8 propres et efficaces grâce à un scaling fin

 (github.com/deepseek-ai)
2 points par GN⁺ 2025-02-27 | 1 commentaires | Partager sur WhatsApp
  • DeepGEMM est une bibliothèque de kernels tensor core haute performance qui regroupe dans une même base de code CUDA des primitives de calcul essentielles aux LLM modernes, comme GEMM, fused MoE, MQA scoring et HyperConnection
  • Tous les kernels sont compilés à l’exécution via un module JIT léger ; aucune compilation CUDA n’est nécessaire lors de l’installation, et C++20, CUDA Toolkit, PyTorch et CUTLASS 4.0 ou supérieur sont requis
  • Elle s’appuie sur certains concepts de CUTLASS et CuTe, mais ne dépend pas lourdement des templates ni de structures algébriques ; elle est conçue pour rendre l’apprentissage de l’optimisation de kernels GPU NVIDIA plus accessible avec un nombre limité de fonctions de kernels clés
  • Le périmètre pris en charge couvre FP8, FP4, BF16 GEMM, grouped GEMM, les kernels MQA logits pour DeepSeek v3.2, ainsi que Mega MoE, qui chevauche communication et calcul ; les contraintes de layout mémoire diffèrent entre SM90 et SM100
  • Malgré sa conception légère, la bibliothèque vise des performances égales ou supérieures à celles de bibliothèques tunées par des experts sur diverses formes de matrices, avec une mise à jour atteignant jusqu’à 1550 TFLOPS sur H800

Objectifs et conception de DeepGEMM

  • DeepGEMM est une bibliothèque de kernels tensor core qui unifie dans une même base de code CUDA les principales primitives de calcul utilisées dans les grands modèles de langage récents
    • GEMM : FP8, FP4, BF16
    • fused MoE avec chevauchement des communications : Mega MoE
    • MQA scoring pour lightning indexer
    • HyperConnection (HC)
  • Tous les kernels sont compilés à l’exécution via un module Just-In-Time (JIT) léger
    • Aucune compilation CUDA n’est nécessaire pendant l’installation
  • Elle exploite certains concepts de CUTLASS et CuTe
    • Mais elle ne dépend pas fortement de templates lourds ni de structures algébriques
    • La base de code reste simple en limitant le nombre de fonctions de kernels essentielles
  • Malgré sa conception légère, le projet indique atteindre des performances équivalentes ou supérieures à celles de bibliothèques tunées par des experts sur plusieurs formes de matrices

Principales mises à jour

  • La mise à jour du 16 avril 2026 inclut Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, une compilation JIT plus rapide, etc.
    • Détails dans #304
    • Benchmarks de Mega MoE dans #316
  • La mise à jour du 28 septembre 2025 ajoute un kernel de scoring weighted ReLU MQA logits pour le lightning indexer de DeepSeek v3.2
    • Détails dans #200
  • La mise à jour du 20 juillet 2025 prend en charge à la fois SM90 et SM100, avec un refactoring complet autour d’un module JIT CPP à faible overhead CPU
    • NVRTC et l’optimisation SASS post-compilation sont désactivés
    • NVRTC est indiqué comme pris en charge ultérieurement
    • Comme NVCC 12.9 effectue automatiquement l’interleaving FFMA, l’optimisation post-compilation n’est plus prise en charge
    • Détails dans #112
  • La mise à jour du 14 mai 2025 ajoute des kernels weight gradient pour le backward dense et MoE
    • Détails dans #95
  • La mise à jour du 7 mai 2025 apporte jusqu’à 10× d’accélération de la compilation grâce au support NVRTC
    • Activable avec DG_JIT_USE_NVRTC=1
    • Dans certains cas, des pertes de performance peuvent survenir
    • Détails dans #94
  • La mise à jour du 18 avril 2025 atteint jusqu’à 1550 TFLOPS sur H800

Prérequis et installation

  • L’environnement d’exécution nécessite un GPU d’architecture NVIDIA SM90 ou SM100
  • Les prérequis logiciels sont les suivants
    • Python 3.8 ou supérieur
    • Compilateur compatible C++20
    • CUDA Toolkit
      • SM90 : CUDA 12.3 ou supérieur
      • CUDA 12.9 ou supérieur est fortement recommandé pour obtenir les meilleures performances
      • SM100 : CUDA 12.9 ou supérieur
    • PyTorch 2.1 ou supérieur
    • CUTLASS 4.0 ou supérieur
    • Bibliothèque {fmt}
  • En environnement de développement, il faut cloner le dépôt avec les submodules, puis exécuter develop.sh pour relier les includes nécessaires et construire le module JIT CPP
  • L’installation consiste à exécuter install.sh, puis à importer deep_gemm dans un projet Python

Interface GEMM et contraintes de layout

  • La convention de nommage des kernels GEMM de DeepGEMM est D = C + A @ B
  • Le layout des shapes d’entrée est basé sur NT
    • fp8_gemm_nt exécute D = C + A @ B.T
  • L’implémentation SM90 ne prend en charge que le layout mémoire NT
    • Cela correspond à la combinaison row-major, col-major
  • L’implémentation SM100 prend en charge l’ensemble des layouts mémoire NT, TN, NN, TT
  • Sur les deux architectures, le scaling factor LHS doit être aligné TMA et utiliser un layout transposé
    • SM90 exige le scaling factor au format FP32
    • SM100 exige le format packed UE8M0, avec 4 valeurs UE8M0 empaquetées dans un torch.int
  • Les opérations comme la transposition des entrées ou le cast FP8 doivent être gérées séparément par l’utilisateur
    • La bibliothèque fournit des fonctions utilitaires PyTorch simples, mais leurs performances peuvent être lentes
    • L’objectif principal est l’optimisation des kernels GEMM

Dense et Grouped GEMM

  • Le GEMM FP8 non groupé de base utilise les fonctions fp8_gemm_{nt, nn, tn, tt}
  • Le grouped GEMM en contiguous layout, contrairement au grouped GEMM traditionnel de CUTLASS, ne groupe que l’axe M
    • N et K doivent rester fixes
    • Cette conception vise les cas où les experts partagent la même shape dans les modèles MoE
  • Pendant le forward pass d’entraînement ou le prefilling en inférence, le nombre de tokens traités par chaque expert peut varier
    • La concaténation de ces tokens dans un même tensor est appelée contiguous layout
    • Chaque segment d’expert doit être aligné sur la taille de bloc M du GEMM
    • Le critère d’alignement se consulte avec get_mk_alignment_for_contiguous_layout()
  • Une API grouped sur l’axe K est également fournie pour le backward des poids MoE
    • M et N doivent être fixes
    • La fonction correspondante est k_grouped_fp8_gemm_tn_contiguous
  • Lors de la phase d’inference decoding, quand CUDA graph est activé et que le CPU ne peut pas connaître le nombre de tokens par expert, masked grouped GEMM est pris en charge
    • En fournissant un tensor de masque, le kernel ne calcule que les zones valides
    • La fonction est m_grouped_fp8_gemm_nt_masked
    • Un exemple utilise en entrée la sortie des kernels low-latency de DeepEP

Kernel MQA pour l’Indexer de DeepSeek v3.2

  • La famille de kernels MQA V3.2 fournit une version non-paged et une version paged
    • La version non-paged est destinée au prefilling
    • La version paged est destinée au decoding
  • fp8_mqa_logits prend 6 entrées
    • q : tensor E4M3, shape [seq_len, num_heads, head_dim]
    • kv : tensor E4M3 et scaling factor float
      • La shape du tensor est [seq_len_kv, head_dim]
      • La shape du scaling factor est [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 : indique s’il faut nettoyer les logits non remplis en les mettant à -inf
  • La shape du tensor de sortie est [seq_len, seq_len_kv] et représente les logits token-to-token
  • Chaque token q i parcourt les tokens kv j de cu_seq_len_k_start[i] jusqu’à avant cu_seq_len_k_end[i]
    • Multiplie kv_j par le scaling factor
    • Calcule les valeurs par tête avec q[i, :, :] @ kv_j
    • Après application de ReLU, multiplie par weights[i, :] et somme pour produire un logit scalaire
  • La fonction de la version paged est fp8_paged_mqa_logits

Mega MoE

  • Mega MoE fusionne plusieurs étapes MoE dans un seul mega-kernel
    • EP dispatch
    • linear 1, FP8xFP4
    • SwiGLU
    • linear 2, FP8xFP4
    • EP combine
  • Mega MoE overlap les communications NVLink et les calculs tensor core
  • L’exécution nécessite un multi-process launch utilisant de la symmetric memory
  • Le flux d’utilisation est le suivant
    • Allouer un buffer de symmetric memory avec deep_gemm.get_symm_buffer_for_mega_moe
      • PyTorch 2.9 ou supérieur est nécessaire
    • Transformer les poids, incluant FP4 et UE8M0 SF, dans le layout requis avec deep_gemm.transform_weights_for_mega_moe
    • Avant l’appel, copier les entrées, les scaling factors, l’index top-k et le poids top-k dans le buffer
    • Exécuter le kernel fused mega MoE avec deep_gemm.fp8_fp4_mega_moe
  • La configuration multi-process complète et un exemple de benchmark se trouvent dans tests/test_mega_moe.py

Utilitaires et variables d’environnement

  • Les principales fonctions utilitaires contrôlent les ressources d’exécution, l’alignement, la compilation JIT et la transformation des scaling factors
    • deep_gemm.set_num_sms / get_num_sms : définir et consulter le nombre maximal de SM à utiliser
    • deep_gemm.set_tc_util / get_tc_util : définir et consulter le ratio approximatif d’utilisation des tensor cores
    • deep_gemm.set_pdl / get_pdl : activer et désactiver Programmatic Dependent Launch (PDL)
    • deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout : définir et consulter l’alignement M/K au niveau groupe pour le contiguous layout
    • deep_gemm.transform_sf_into_required_layout : transformer le scaling factor dans le layout requis
    • deep_gemm.get_tma_aligned_size : consulter la taille d’alignement TMA nécessaire
  • Les variables d’environnement liées au JIT contrôlent les sorties de debug, l’emplacement du cache, le choix du compilateur et les options de profiling
    • DG_JIT_DEBUG : afficher les informations de debug JIT
    • DG_PRINT_CONFIGS : afficher la config choisie par shape
    • DG_JIT_CACHE_DIR : répertoire de cache des kernels compilés, valeur par défaut $HOME/.deep_gemm
    • DG_JIT_USE_NVRTC : utiliser NVRTC au lieu de NVCC, permettant une compilation rapide mais pouvant réduire les performances dans certains cas
    • DG_JIT_NVCC_COMPILER : chemin du compilateur NVCC
    • DG_JIT_CPP_STANDARD : version du standard C++, valeur par défaut 20
  • Des variables d’environnement de debug et de profiling sont également fournies
    • DG_JIT_DUMP_ASM, DG_JIT_DUMP_PTX, DG_JIT_DUMP_SASS : dumper les sorties PTX et SASS
    • DG_JIT_WITH_LINEINFO : inclure les informations de ligne source pour les outils de profiling
    • DG_COMM_KERNEL_DEBUG : initialiser le symmetric buffer à 0 avant l’appel Mega MoE
    • DG_USE_NVIDIA_TOOLS : ignorer le profiling interne lors de l’exécution d’outils NVIDIA externes
  • Les options de build contrôlent l’installation et le mode de chargement des kernels
    • DG_SKIP_CUDA_BUILD : ignorer le build de l’extension CUDA pendant l’installation
    • DG_FORCE_BUILD : forcer un build local au lieu de télécharger une wheel précompilée
    • DG_JIT_USE_RUNTIME_API : utiliser la CUDA Runtime API pour charger les kernels, nécessite CUDA runtime 12.8 ou supérieur

Licence et citation

  • Le dépôt DeepGEMM est publié sous MIT License
  • Le projet indique s’inspirer de CUTLASS
  • Le titre de la citation est DeepGEMM: clean and efficient BLAS kernel library on GPU

À lire aussi

1 commentaires

 
GN⁺ 2025-02-27
Commentaires sur Hacker News

  • L’entrelacement FFMA SASS semble vraiment impressionnant
    En constatant que les performances des kernels CUTLASS FP8 s’étaient améliorées entre NVCC 12.2 et 12.3, quelqu’un a comparé le SASS compilé et a vu qu’un bit était inversé selon un motif d’entrelacement dans plusieurs instructions FADD ; en se référant à une implémentation open source d’un assembleur CUDA, il semble avoir compris que ce bit était le bit yield, qui fait céder le warp courant pour permettre à un autre warp de s’exécuter
    Il est impressionnant qu’ils aient ensuite créé un script modifiant les instructions FFMA du binaire compilé et qu’en plus du bit yield, ils aient aussi inversé le bit reuse, puisque la réutilisation des registres n’est plus possible lorsqu’un warp cède la main, afin de mieux faire se chevaucher les instructions MMA et les instructions FFMA de promotion dans le GEMM FP8 à mise à l’échelle fine, avec dans certains cas plus de 10 % de gain de performance

    • D’après ce que j’ai lu ailleurs, ce genre de méthode est assez typique dans l’optimisation d’opérations matricielles critiques pour les performances
      Cela dit, pour ce problème précis, les autres entreprises d’IA ne semblaient pas encore en voir la nécessité et ne l’avaient donc pas appliquée ; au final, il est probable que tout le monde finisse par arriver à des points similaires
    • Scott Gray avait déjà découvert exactement ce genre de choses, et même davantage, sur Maxwell en 2015, et plusieurs personnes ont beaucoup travaillé dessus depuis

  • Ce type d’exemple montre bien à quel point les compilateurs actuels sont encore loin d’extraire les performances du matériel à partir du seul code de haut niveau
    Je me demande ce qu’il faudrait pour que des techniques de compilation traditionnelles ou des agents d’optimisation fondés sur l’IA produisent de tels résultats

    • Il faudrait probablement une quantité énorme d’essais et d’erreurs dans une boucle de rétroaction par apprentissage par renforcement

  • Les chiffres d’accélération rapportés sont comparés à leur propre référence basée sur CUTLASS
    Je me demande si quelqu’un a fait une comparaison directe des performances avec cuBLAS
    Jusqu’ici, les résultats CUTLASS GEMM que j’ai vus étaient à environ 10 % de cuBLAS ; si les gains de 2× à 2,5× mentionnés dans l’article se maintiennent, ce serait vraiment impressionnant

    • D’habitude j’évite le FP8 et je préfère l’I8, mais cette question m’a rendu curieux de voir ce que donne cuBLAS
      Tout d’abord, cuBLAS nécessite l’API d’extension cuBLASLt pour gérer des opérations en précision mixte comme le FP8
      De plus, il ne prend pas en charge certaines combinaisons de types qui sembleraient appropriées, comme E5M2 x E5M2 pour A x B, mais il prend en charge E5M2 x E4M3 ; et les contraintes continuent, par exemple sur Ampere, Hopper et Blackwell, où la matrice A doit toujours être dans une disposition transposée
      J’ai intégré le benchmark FP8 cuBLASLt dans mon dépôt « Less Slow C++ » <https://github.com/ashvardanian/less_slow.cpp>, en l’ajoutant à la liste des benchmarks cuBLAS existants et CUDA/PTX écrits à la main
      Il tourne sur un GPU H200, qui devrait avoir les mêmes performances qu’un H100, et sur des entrées carrées le débit culmine à environ 1,35 peta-ops
      Pour 256 on est à 2,68 T/s, 512 à 20,49 T/s, 1024 à 144,23 T/s, 2048 à 665,68 T/s, 4096 à 1,26 P/s, 8192 à 1,34 P/s et 16384 à 1,23 P/s environ, soit à peu près 67 % du chiffre annoncé par NVIDIA pour le GEMM dense <https://resources.nvidia.com/en-us-data-center-overview-mc/e...>
    • J’ai entendu dire qu’avec CUTLASS, on pouvait obtenir de meilleures performances qu’avec cuBLAS
      Je pensais que la référence retenait le meilleur des deux, cuBLAS ou CUTLASS

  • Ce genre d’open source illustre vraiment bien l’objectif de l’industrie : atteindre plus d’efficacité
    Cela dit, les principaux bénéficiaires de ce logiciel seront probablement davantage les grands acteurs qui servent des modèles à grande échelle — c’est-à-dire des concurrents potentiels de DeepSeek — que la communauté open source ordinaire qui apprend, expérimente ou sert des modèles sur du matériel grand public

    • Une meilleure efficacité peut au bout du compte se traduire par du matériel moins cher pour tout le monde, y compris pour DeepSeek lui-même

  • Je ne sais pas trop si optimiser vers des précisions de plus en plus faibles est une bonne direction à long terme
    Cela signifie que les modèles sont en réalité assez parcimonieux, et même si c’est possible aujourd’hui, je pense qu’il est probable que ce ne soit pas parce qu’ils doivent intrinsèquement être aussi parcimonieux, mais plutôt parce que de mauvaises idées se sont glissées dans la façon dont on les entraîne

    • Tant que la parcimonie obtenue gratuitement fonctionne, autant en profiter
      Faire en sorte que les très bons modèles ne puissent être entraînés qu’à plus haute précision est un problème de recherche, tandis que l’entraînement et l’inférence en basse précision relèvent de l’ingénierie
      On fait ce genre de choses depuis l’époque des CNN, depuis au moins 9 ans, et je pense qu’il en reste encore pour quelques années
    • Comme les fonctions d’activation abandonnent une bonne partie de la plage dynamique des nombres à virgule flottante, il est assez clair qu’il n’est probablement pas utile de réserver une large plage à des zones d’activation déjà saturées

  • Cela pourrait devenir sans intérêt à cause de MXFP, la prise en charge native du microscaling dans Blackwell
    Sur Hopper, ils l’ont en quelque sorte implémenté manuellement à une granularité plus grossière, mais avec des coefficients de mise à l’échelle en FP32

    • Exact
      Ce genre de démonstration publique de grande qualité montre bien où se trouve le moat de $NVDA
      Les GPU généralistes sont très flexibles : ils permettent de programmer de nombreuses tâches tout à fait raisonnables auxquelles le fournisseur matériel n’avait pas pensé dès le départ
      Mais si l’on prédit que l’avenir convergera de plus en plus vers une prise en charge matérielle dédiée, faisant disparaître cette marge d’optimisation logicielle, alors le fameux moat CUDA s’effondre
      Pour rester dans la partie, NVIDIA est en quelque sorte en train de détruire son propre moat :p

  • Waouh, c’est sous licence MIT
    J’aimerais que les grands acteurs adoptent ce mode de collaboration open source

  • Je me demande toujours pourquoi il existe des instructions non documentées
    Même si elles ne sont pas parfaitement stables, j’ai l’impression qu’il vaudrait mieux les mettre à disposition des utilisateurs
    Ce genre de choses doit bien être documenté en interne, mais je ne comprends pas pourquoi elles ne sont pas publiées
    La sécurité par l’obscurité ne fonctionne pas, et les concurrents feront de toute façon tous du rétro‑ingénierie

    • C’est probablement assez proche des raisons pour lesquelles ce que nous développons comporte aussi des parties non documentées
      Peut-être par manque de temps, ou parce qu’on ne veut pas laisser entendre qu’on prendra en charge des fonctionnalités instables ou expérimentales
      Et si les dégâts se limitent plus ou moins à l’équipe d’à côté, c’est aussi beaucoup plus facile à modifier
    • L’hypothèse selon laquelle « ce genre de choses est documenté en interne » peut déjà être fausse
      Il y a de fortes chances que cela ne figure que dans des documents de conception d’architecture ou des spécifications, et ce sont évidemment des documents qu’ils n’ont pas envie de partager

  • Honnêtement, cela dépasse mon périmètre d’usage et ma compréhension
    Mais je trouve vraiment appréciable et rafraîchissant que ce type de découvertes et d’améliorations soit partagé pour que tout le monde puisse en bénéficier

    • FFMA signifie Fused Floating-point Multiply-Add ; c’est une instruction GPU de base qui exécute D = A*B + C en une seule fois
      Elle est très importante pour la multiplication de matrices et les charges de travail de deep learning
      Dans le SASS de NVIDIA, l’instruction FFMA est encodée comme une instruction 64 bits ou 128 bits, avec plusieurs bits de contrôle qui déterminent son comportement exact
      Lorsque le bit yield est défini, il indique au planificateur de warps que le warp courant peut céder l’exécution après cette instruction, et le matériel peut exécuter un autre warp pour masquer la latence
      Le GPU obtient un débit élevé grâce à un parallélisme massif, et lorsqu’un warp est bloqué, par exemple en attente mémoire, un autre warp peut avancer
      Le bit reuse indique si un registre source peut être réutilisé dans l’opération immédiatement suivante ; lorsque le bit yield est défini, il doit impérativement être désactivé
      En effet, si un warp cède la main, le prochain warp exécuté n’est pas forcément le même, et un autre warp peut modifier l’état du fichier de registres ; le matériel ne peut donc pas garantir que les valeurs des registres soient conservées au-delà d’un yield
      Définir le bit yield des instructions FFMA selon un motif alterné permet au compilateur de créer des points de planification explicites où d’autres warps peuvent avancer, et il faut aussi effacer le bit reuse de ces instructions pour préserver la correction
      Ce changement aide particulièrement à chevaucher les instructions MMA, au cœur de la multiplication de matrices, avec les instructions FFMA de promotion, qui convertissent les valeurs FP8 pour les accumuler avec une précision plus élevée
      Le GEMM FP8 nécessite généralement de convertir vers une précision plus élevée pour l’accumulation, puis de reconvertir ensuite, ce qui ajoute des FFMA supplémentaires ; cela réduit les besoins en bande passante mémoire, mais crée un schéma de calcul complexe mêlant opérations de promotion et de rétrogradation
      Le « scaling fin » semble désigner le fait de gérer soigneusement la précision à différents points du calcul
      La manipulation du bit yield permet de mieux entrelacer les opérations de calcul et les conversions de format, ce qui rend l’utilisation des unités d’exécution du GPU plus efficace ; sans cette optimisation, le planificateur de warps pourrait ne pas trouver d’occasions naturelles de basculer, et les ressources de calcul seraient alors moins bien exploitées