Créer un moteur d’inférence LLM from scratch avec C++ et CUDA

• Comment construire un moteur d’inférence LLM sans bibliothèque avec C++ et CUDA
• Cela permet de comprendre toute la pile de l’inférence LLM et de mesurer concrètement l’impact des différentes optimisations sur la vitesse d’inférence
• Objectif : implémenter un modèle capable d’une inférence rapide en batch unique sur un serveur CPU + GPU unique, et atteindre une vitesse de traitement des tokens supérieure à llama.cpp

1. Vue d’ensemble de l’architecture et de l’inférence des LLM

• La plupart des principaux LLM suivent la même architecture, basée sur une succession de blocs transformer.
• Le chargement du modèle consiste à définir une classe de bloc transformer personnalisable, à l’assembler en séquence, puis à l’initialiser avec des poids safetensors.
• L’inférence se fait principalement en batch unique, et l’« étape de décodage » représente l’essentiel de l’exécution.

1.1 Vue d’ensemble de l’inférence

• L’inférence se divise entre une phase de prefill, qui envoie les tokens du prompt au modèle pour remplir le cache KV, et une phase de decode, qui repasse ensuite le modèle de façon répétée pour générer des tokens
  • Phase de prefill : traitement des tokens du prompt et initialisation du cache KV
  • Phase de decode : génération d’un token à la fois
• Cache KV : stocke les paires clé/valeur précédentes afin de calculer rapidement l’attention sur le contexte passé
• La passe forward du modèle utilise une table d’embeddings pour mapper les ID de tokens vers des vecteurs d’embedding, puis transforme l’état via une séquence de blocs transformer

1.2 Goulots d’étranglement et benchmarks

• Goulot d’étranglement : sur le matériel moderne, la bande passante mémoire est le facteur limitant
  • Lors de l’inférence, la génération de chaque token nécessite de relire l’ensemble du modèle ; la bande passante mémoire est donc une contrainte plus forte que le calcul
• La quantification du modèle est efficace pour améliorer la vitesse d’inférence
• Le débit théorique maximal en tokens dépend du matériel, et les performances réelles peuvent être vérifiées via plusieurs moteurs d’inférence
• Limite théorique de vitesse :
  • AMD EPYC 7702P : jusqu’à 13.6 tok/s (FP16)
  • RTX 4090 : jusqu’à 67.1 tok/s (FP16)
• Benchmarks :
  • llama.cpp : CPU 8.7 tok/s, GPU 61 tok/s
  • calm : GPU 66 tok/s

2. Inférence sur CPU

• L’implémentation initiale sur CPU est mono-thread et ne prend en charge que des poids FP32
• On peut commencer à paralléliser le code avec le multithreading, puis améliorer les performances avec le SIMD

2.1 Multithreading

• Utilisation d’OpenMP pour paralléliser la multiplication matrice-vecteur (matmul) et l’attention multi-têtes afin d’améliorer les performances
• Résultat de l’optimisation : vitesse améliorée de 0.6 tok/s → 4.4 tok/s

2.2 Quantification des poids et optimisation SIMD

• Quantification : quantification des poids FP32 en FP16 pour réduire de moitié l’usage mémoire et améliorer les performances
• SIMD : optimisation avec AVX2 pour traiter 8 valeurs FP32 simultanément
• Résultat : 8.4 tok/s atteints

3. Inférence sur GPU

• En quantifiant le modèle en FP16 et en le chargeant sur une RTX 4090, on peut commencer l’implémentation de l’inférence GPU
• CUDA permet d’exécuter en parallèle sur le GPU des fonctions C++ (kernels)

3.1 Portage simple vers CUDA

• Il est possible d’implémenter un backend GPU en convertissant les opérations CPU en kernels CUDA un à un
• Les kernels CUDA s’exécutent de manière asynchrone, mais de façon séquentielle dans un même stream
• Problème : l’inefficacité des threads empêche d’exploiter pleinement les ressources du GPU → seulement 2.9 tok/s

3.2 Une meilleure multiplication matricielle (matmul)

• La multiplication matricielle représente une grande partie du temps d’exécution sur CPU et peut être optimisée avec OpenMP
• Sur GPU, on peut améliorer l’utilisation des threads en faisant traiter une ligne par bloc
• Méthode d’optimisation :
  1. Un bloc traite une ligne, et les threads du bloc collaborent au calcul
  2. Application d’une réduction au niveau du warp (warp reduction)
• Résultat : vitesse portée à 51.7 tok/s

3.3 Fusion de kernels et optimisations supplémentaires

• La fusion de kernels permet d’améliorer les performances
  • Fusion de kernels : regrouper des opérations consécutives dans un seul kernel afin de minimiser les accès mémoire et le temps de calcul
• Grâce à l’optimisation des patterns d’accès mémoire et à la réutilisation de l’espace, 56.1 tok/s sont atteints

3.4 Optimisation de l’attention et gestion des contextes longs

• Problème : avec des contextes longs, le kernel d’attention devient un goulot d’étranglement
• Solution :
  1. Optimisation des accès mémoire : refonte pour lire des blocs mémoire contigus
  2. Utilisation de mémoire partagée à la place de atomicAdd pour corriger les problèmes de valeurs décimales perdues
• Résultat des optimisations :
  • Contexte court : 63.8 tok/s (plus rapide que les 61.0 tok/s de llama.cpp)
  • Contexte long : 58.8 tok/s

3.5 Quantification du cache KV et problèmes d’optimisation du compilateur

• La quantification du cache KV en FP16 entraîne une baisse de performances (optimisations insuffisantes du compilateur)
• Solution : dérouler manuellement les boucles et appliquer du prefetch mémoire
• Résultat : environ 2× plus rapide que le FP32 tout en conservant 58.8 tok/s sur les contextes longs

4. Pistes d’amélioration futures

• Optimisation du prefill du prompt : traiter plusieurs tokens à la fois pour réduire le temps avant la génération du premier token
• Fusion des kernels d’attention : appliquer des techniques d’optimisation comme FlashAttention
• Quantification plus poussée : FP8, INT8, INT4, ainsi que la quantification des activations et du cache
• Optimisation des kernels : introduire des techniques avancées pour maximiser la bande passante mémoire et l’efficacité de calcul
• Utilisation de bibliothèques : exploiter des bibliothèques comme cuDNN et cuBLAS pour réduire le temps nécessaire à l’optimisation

Résumé des résultats :

• Grâce à diverses optimisations sur CPU et GPU, une vitesse de 63.8 tok/s a été atteinte
• Des performances proches de, ou supérieures à, celles de llama.cpp et calm ont été obtenues
• Un moteur d’inférence LLM haute performance a été implémenté uniquement avec C++ et CUDA, sans bibliothèque

[GN⁺]

Commentaires Hacker News

• L’auteur se dit heureux que son billet de blog ait attiré l’attention et aimerait avoir des retours
• Un lecteur salue la qualité de l’article et se demande combien de temps il a fallu pour l’écrire
  • Travaillant dans le domaine du GPGPU, il aimerait rédiger un article similaire, mais hésite à cause de l’incertitude sur le temps nécessaire
• Un autre lecteur estime que le code n’utilise ni les tensor cores ni les instructions wgmma
  • Il explique que ce type de programmation est difficile car il faut gérer plusieurs tâches en parallèle
  • Il mentionne qu’en raison des limites de bande passante, des calculs supplémentaires ne sont peut-être pas nécessaires
  • Il juge que le code du blog a de bonnes chances de bien fonctionner lors d’un portage vers d’autres accélérateurs
  • Il craint que l’usage de wgmma réduise la portabilité entre générations de Nvidia
• Un autre lecteur cherche des ressources Python similaires à partager avec son équipe
  • Il souhaite un contenu concis, complet sur le plan conceptuel et au format tutoriel, plutôt qu’axé sur les performances
• Un utilisateur aimerait comparer sa version de Mistral et ses performances en tokens/seconde
  • Il est recommandé de consulter la section sur la quantification dans le README
• Certains estiment que __shfl_down n’est plus recommandé aujourd’hui à cause des problèmes modernes de synchronisation de warp