Eine LLM-Inferenz-Engine von Grund auf mit C++ und CUDA bauen

• Wie man mit C++ und CUDA ohne Bibliotheken eine LLM-Inferenz-Engine aufbaut
• Dadurch lässt sich der gesamte Stack der LLM-Inferenz verstehen und konkret nachvollziehen, wie sich verschiedene Optimierungen auf die Inferenzgeschwindigkeit auswirken
• Ziel: ein Modell zu implementieren, das auf einem einzelnen CPU+GPU-Server mit Single-Batch schnell inferieren kann, und eine höhere Token-Verarbeitungsgeschwindigkeit als llama.cpp zu erreichen

1. Überblick über LLM-Architektur und Inferenz

• Die meisten wichtigen LLMs folgen derselben Architektur mit aufeinanderfolgenden Transformer-Blöcken.
• Das Laden des Modells besteht darin, eine anpassbare Transformer-Block-Klasse zu definieren, diese zu einer Sequenz zusammenzusetzen und sie mit safetensors-Gewichten zu initialisieren.
• Die Inferenz erfolgt hauptsächlich als Single-Batch, wobei die „Decode-Phase“ den Großteil der Ausführung ausmacht.

1.1 Überblick über die Inferenz

• Die Inferenz ist in eine Prefill-Phase unterteilt, in der die gegebenen Prompt-Tokens durch das Modell geschickt und damit der KV-Cache gefüllt werden, und in eine Decode-Phase, in der das Modell wiederholt ausgeführt wird, um Tokens zu erzeugen
  • Prefill-Phase: verarbeitet Prompt-Tokens und initialisiert den KV-Cache
  • Decode-Phase: erzeugt jeweils ein Token auf einmal
• KV-Cache: speichert frühere Schlüssel/Wert-Paare, damit Attention mit dem bisherigen Kontext schnell berechnet werden kann
• Der Forward-Pass des Modells verwendet eine Embedding-Tabelle, um Token-IDs auf Embedding-Vektoren abzubilden, und transformiert die Zustände durch eine Sequenz von Transformer-Blöcken

1.2 Engpässe und Benchmarks

• Engpass: Auf moderner Hardware ist die Speicherbandbreite der limitierende Faktor
  • Bei der Modellinferenz muss zur Erzeugung jedes Tokens das gesamte Modell gelesen werden; die Speicherbandbreite ist daher stärker limitierend als die Rechenleistung
• Modellquantisierung ist wirksam, um die Inferenzgeschwindigkeit zu verbessern
• Der theoretische maximale Token-Durchsatz variiert je nach Hardware, und die tatsächliche Leistung lässt sich über verschiedene Inferenz-Engines überprüfen
• Theoretische Geschwindigkeitsgrenze:
  • AMD EPYC 7702P: maximal 13.6 tok/s (FP16)
  • RTX 4090: maximal 67.1 tok/s (FP16)
• Benchmarks:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. CPU-basierte Inferenz

• Die anfängliche Implementierung auf der CPU ist Single-Threaded und unterstützt nur FP32-Gewichte
• Mit Multithreading kann die Parallelisierung des Codes beginnen, und mit SIMD lässt sich die Leistung weiter verbessern

2.1 Multithreading

• Mit OpenMP werden Matrix-Vektor-Multiplikation (matmul) und Multi-Head-Attention parallelisiert, um die Leistung zu verbessern
• Ergebnis der Optimierung: Geschwindigkeitssteigerung von 0.6 tok/s → 4.4 tok/s

2.2 Gewichtsquantisierung und SIMD-Optimierung

• Quantisierung: FP32-Gewichte werden auf FP16 quantisiert, wodurch sich der Speicherverbrauch halbiert und die Leistung verbessert
• SIMD: Mit AVX2 so optimiert, dass 8 FP32-Werte gleichzeitig verarbeitet werden
• Ergebnis: 8.4 tok/s erreicht

3. GPU-basierte Inferenz

• Durch Quantisierung des Modells auf FP16 und Laden auf eine RTX 4090 kann mit der Implementierung der GPU-Inferenz begonnen werden
• Mit CUDA lassen sich C++-Funktionen (Kernel) parallel auf der GPU ausführen

3.1 Einfaches Porting mit CUDA

• Das GPU-Backend kann implementiert werden, indem CPU-Operationen 1:1 in CUDA-Kernel umgewandelt werden
• CUDA-Kernel laufen asynchron, werden aber innerhalb desselben Streams sequenziell ausgeführt
• Problem: Wegen ineffizienter Thread-Nutzung werden die GPU-Ressourcen nicht ausreichend ausgelastet → langsam mit 2.9 tok/s

3.2 Bessere Matrixmultiplikation (matmul)

• Matrixmultiplikation beansprucht auf der CPU einen großen Teil der Laufzeit und kann mit OpenMP optimiert werden
• Auf der GPU lässt sich die Thread-Auslastung erhöhen, indem pro Block eine Zeile verarbeitet wird
• Optimierungsmethode:
  1. Ein Block verarbeitet eine Zeile, und die Threads im Block arbeiten bei der Berechnung zusammen
  2. Warp-Reduction anwenden
• Ergebnis: Geschwindigkeitssteigerung auf 51.7 tok/s

3.3 Kernel-Fusion und weitere Optimierungen

• Durch Kernel-Fusion lässt sich die Leistung verbessern
  • Kernel-Fusion: aufeinanderfolgende Operationen in einem Kernel zusammenfassen, um Speicherzugriffe und Rechenzeit zu minimieren
• Durch Optimierung der Speicherzugriffsmuster und Wiederverwendung von Speicherbereichen werden 56.1 tok/s erreicht

3.4 Attention-Optimierung und Verarbeitung langer Kontexte

• Problem: Bei langen Kontexten wird der Attention-Kernel zum Leistungsengpass
• Lösung:
  1. Optimierung der Speicherzugriffe: so umgestaltet, dass zusammenhängende Speicherblöcke gelesen werden
  2. Verwendung von Shared Memory statt atomicAdd, um Probleme mit verlorenen Fließkommawerten zu beheben
• Ergebnis der Optimierung:
  • Kurzer Kontext: 63.8 tok/s (schneller als die 61.0 tok/s von llama.cpp)
  • Langer Kontext: 58.8 tok/s erreicht

3.5 KV-Cache-Quantisierung und Probleme mit Compiler-Optimierungen

• Die Quantisierung des KV-Cache auf FP16 führt zu Leistungseinbußen (unzureichende Compiler-Optimierung)
• Lösung: Schleifen manuell unrollen und Memory Prefetching anwenden
• Ergebnis: etwa doppelte Geschwindigkeit gegenüber FP32 bei gleichbleibender Leistung von 58.8 tok/s in langen Kontexten

4. Künftige Verbesserungsrichtungen

• Optimierung des Prompt-Prefill: mehrere Tokens gleichzeitig verarbeiten, um die Zeit bis zum ersten Token zu verkürzen
• Fusion von Attention-Kerneln: Optimierungstechniken wie FlashAttention anwenden
• Stärkere Quantisierung: FP8, INT8, INT4 sowie Quantisierung von Aktivierungen/Cache anwenden
• Kernel-Optimierung: fortgeschrittene Techniken einführen, die Speicherbandbreite und Recheneffizienz maximieren
• Einsatz von Bibliotheken: Bibliotheken wie cuDNN und cuBLAS nutzen, um die Optimierungszeit zu verkürzen

Zusammenfassung der Ergebnisse:

• Durch verschiedene Optimierungen auf CPU und GPU wurde eine Geschwindigkeit von 63.8 tok/s erreicht
• Leistung auf dem Niveau von oder besser als llama.cpp und calm
• Eine hochperformante LLM-Inferenz-Engine wurde nur mit C++ und CUDA ohne Bibliotheken implementiert

[GN⁺]

Hacker-News-Kommentare

• Der Autor freut sich, dass sein Blogbeitrag Aufmerksamkeit bekommen hat, und möchte gern Feedback hören.
• Ein Leser lobt den Beitrag als hervorragend und fragt, wie viel Zeit das Schreiben gekostet hat.
  • Als jemand, der im Bereich GPGPU arbeitet, würde er gern einen ähnlichen Beitrag schreiben, zögert aber wegen der Unsicherheit über den nötigen Zeitaufwand.
• Ein anderer Leser meint, dass der Code keine Tensor Cores oder wgmma-Befehle nutzt.
  • Er erklärt, dass diese Art der Programmierung schwierig ist, weil man viele Aufgaben gleichzeitig handhaben muss.
  • Er erwähnt, dass wegen Bandbreitenbeschränkungen zusätzliche Berechnungen womöglich gar nicht nötig sind.
  • Er bewertet den Code aus dem Blog so, dass er sich wahrscheinlich gut auf andere Beschleuniger portieren lässt.
  • Er äußert die Sorge, dass die Nutzung von wgmma die Portabilität zwischen Nvidia-Generationen verschlechtern könnte.
• Ein weiterer Leser sucht nach ähnlichem Python-Material und möchte es mit seinem Team teilen.
  • Er wünscht sich Material, das weniger auf Performance als auf konzeptionelle Vollständigkeit setzt und im Tutorial-Stil kompakt ist.
• Ein Nutzer möchte seine eigene Mistral-Version und die Leistung in Token pro Sekunde vergleichen.
  • Es wird empfohlen, den Abschnitt zur Quantisierung in der README zu lesen.
• Es gibt die Meinung, dass __shfl_down heutzutage wegen Problemen mit der Warp-Synchronisation nicht mehr empfohlen wird.