FlashMLA: Optimierung von MLA-Dekodierungskerneln für Hopper-GPUs (DeepSeek Open Source Week Day 1)

Allgemeine Einführung

FlashMLA ist ein effizienter MLA-Dekodierungskern (Multi-head Latent Attention), der von DeepSeek AI entwickelt und für NVIDIA Hopper Architecture GPUs optimiert wurde, um die Leistung der Verarbeitung von Sequenzen variabler Länge zu verbessern. Das Projekt wurde auf GitHub als Open Source veröffentlicht und steht Entwicklern kostenlos zur Verfügung. Es unterstützt BF16-Präzisionsberechnungen und ausgelagertes KV-Caching (Blockgröße 64) und zeigt auf der H800 SXM5 eine gute Leistung mit bis zu 3000 GB/s Bandbreite in speicherintensiven Konfigurationen und bis zu 580 TFLOPS in rechenintensiven Konfigurationen. FlashMLA wurde von FlashAttention 2&3 und dem Cutlass-Projekt inspiriert. DeepSeek AI hat mit diesem viel beachteten Open-Source-Projekt seine Innovationskraft in der KI-Technologie unter Beweis gestellt.

Funktionsliste

• Effiziente MLA-DekodierungOptimiert für Hopper GPUs, um die Verarbeitung von Sequenzen variabler Länge erheblich zu beschleunigen.
• Unterstützt BF16 GenauigkeitGleitkommaoperationen mit halber Genauigkeit zur Verbesserung der Berechnungseffizienz unter Beibehaltung der Genauigkeit.
• Paging KV CachingEin Paging-Mechanismus mit einer Blockgröße von 64 wird verwendet, um den Speicher effektiv zu verwalten und die Inferenzleistung zu verbessern.
• Hohe LeistungBietet eine Speicherbandbreite von bis zu 3000 GB/s und eine Rechenleistung von 580 TFLOPS auf der H800 GPU.
• offene QuelleVollständiger Quellcode wird zur Verfügung gestellt, um kundenspezifische Änderungen und Integration durch Entwickler zu unterstützen.

Hilfe verwenden

Ablauf der Installation

FlashMLA ist ein Open-Source-Projekt, das auf GitHub basiert. Sie müssen sicherstellen, dass die Umgebung den Anforderungen entspricht und die Installation vor der Verwendung abschließen. Hier sind die detaillierten Schritte:

1. ökologische Bereitschaft

• BetriebssystemUnterstützung für Linux-Systeme (Ubuntu 20.04 oder höher empfohlen).
• Hardware-VoraussetzungErfordert einen NVIDIA Hopper Architecture-Grafikprozessor (wie den H800 SXM5).
• Software-Abhängigkeit::
  • CUDA 12.6 oder höher (siehe NVIDIA-Website für Installationsanweisungen).
  • PyTorch 2.0 oder höher (empfohlen über pip install torch (Installation).
  • Python 3.8 oder höher.
• InspektionswerkzeugeStellen Sie sicher, dass Git installiert ist, um Code von GitHub herunterzuladen.

2. den Quellcode herunterladen

1. Öffnen Sie ein Terminal und geben Sie den folgenden Befehl ein, um das FlashMLA-Repository zu klonen:

   git clone https://github.com/deepseek-ai/FlashMLA.git

2. Rufen Sie den Projektkatalog auf:

   cd FlashMLA
   

3. die Installation von Abhängigkeiten

Das Projekt hängt von PyTorch und CUDA ab, die mit den folgenden Befehlen installiert werden können:

pip install -r anforderungen.txt

Falls nicht Anforderungen.txt Datei ist es einfach, sicherzustellen, dass PyTorch installiert ist:

pip install torch torchvision

Überprüfen Sie, ob CUDA verfügbar ist:

python -c "import torch; print(torch.cuda.is_available())"

Ausfuhren Wahr Zeigt die erfolgreiche Konfiguration der Umgebung an.

4. kompilieren und testen

FlashMLA bietet vorkompilierte CUDA-Plug-ins, aber stellen Sie sicher, dass Sie Ihre lokale CUDA-Version verwenden:

1. Wechseln Sie in das Quellverzeichnis und führen Sie das Kompilierungsskript (falls vorhanden) aus:

   python setup.py installieren
   

2. Testen Sie die Installation auf Erfolg und führen Sie den Beispielcode aus:

   python beispiel.py
   

Wenn keine Fehler gemeldet werden, ist die Installation abgeschlossen.

Wie zu verwenden

Das Hauptmerkmal von FlashMLA ist die effiziente Unterstützung der MLA-Dekodierung für AI-Modellinferenzaufgaben. Hier sind die Schritte:

Funktion 1: FlashMLA laden und ausführen

1. Importmodul::
   Einführung in die FlashMLA-Kernfunktionen in Python-Skripten:

   from flash_mla import get_mla_metadata, flash_mla_with_kvcache
   

2. Vorbereiten der Dateneingabe::
   • cache_seqlensLegt die Sequenzlänge des KV-Cache fest.
   • q_i: Abfrage-Tensor.
   • kvcache_iKV zwischengespeicherte Daten.
   • block_tableBlocktabelle für Paging-Cache.
3. Abrufen von Metadaten::

   tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)
   

4. laufende Dekodierung::

   o_i, lse_i = flash_mla_with_kvcache(q_i, kvcache_i, block_table, cache_seqlens, dv, tile_scheduler_metadata, num_splits, causal=True)
   

   Ausfuhren o_i ist das Dekodierungsergebnis.lse_i ist der Logsummenwert.

Funktion 2: Optimierung der Verarbeitung von Sequenzen variabler Länge

• Nehmen SieFlashMLA reduziert den Speicherbedarf bei Eingabesequenzen mit dynamischer Länge, indem es den KV-Cache auslagert.
• Rigg::
  1. Konfigurieren Sie die Paging-Parameter: Die Blockgröße ist auf 64 festgelegt und kann durch Anpassen der cache_seqlens Steuert die Länge der Sequenz.
  2. Zur Laufzeit festlegen kausal=True(c) Sicherstellen, dass der Mechanismus der kausalen Aufmerksamkeit wirksam ist.
• Wirkung3000 GB/s Speicherbandbreite auf dem H800 für umfangreiche Inferenzaufgaben.

Funktion 3: Leistungsprüfung

• Testmethoden::
  1. Bearbeiten Sie das Beispielskript (z. B. beispiel.py), wodurch sich der Umfang der Eingabedaten erhöht.
  2. Verwenden Sie den folgenden Code, um die Leistung aufzuzeichnen:

     Zeit importieren
     start = time.time()
     # Starten Sie den Dekodierungscode
     o_i, lse_i = flash_mla_mit_kvcache(...)
     print(f "Zeit: {time.time() - start} Sekunden")
     

• Erwartete ErgebnisseNahezu 3000 GB/s für speicherintensive Aufgaben und 580 TFLOPS für rechenintensive Aufgaben.

caveat

• Hardware-KompatibilitätNur Hopper GPUs werden unterstützt, H800 oder gleichwertig wird empfohlen.
• Tipps zur FehlersucheWenn Sie auf CUDA-Fehler stoßen, prüfen Sie, ob die Version übereinstimmt, oder suchen Sie die Unterstützung der Community in GitHub Issues.
• ProduktionsumgebungIntegration direkt in bestehende Modellinferenzprozesse, wobei sichergestellt wird, dass die Formate der Eingabedaten mit den FlashMLA-Anforderungen übereinstimmen.

Mit den oben genannten Schritten können Benutzer schnell mit FlashMLA beginnen und die Leistungsverbesserung durch die effiziente Dekodierung genießen. Der vollständige Code und die Dokumentation sind im GitHub-Repository zu finden, und es wird empfohlen, die Parameter entsprechend den tatsächlichen Projektanforderungen anzupassen.