Zyphra präsentiert ZYFA : FlashAttention FP8 für AMD MI355X, bis zu 1,76-mal schneller

Zyphra präsentiert ZYFA : FlashAttention FP8 für AMD MI355X, bis zu 1,76-mal schneller
3 Minuten Lesezeit

Zyphra hat ZYFA (Zyphra FlashAttention on AMD) veröffentlicht, eine Sammlung von FP8- und BF16-Attention-Computekernen für die AMD Instinct MI355X GPUs. Laut dem Unternehmen sind seine FP8-Kerne bei der Verarbeitung echter Modellaktivierungen zwischen 1,16- und 1,76-mal schneller als AITER, AMDs Standard-Kernelbibliothek. Der Code ist als Open Source unter der Apache-2.0-Lizenz veröffentlicht.

  • Zyphra veröffentlicht ZYFA, FP8- und BF16-Attention-Kerne für die AMD MI355X unter Apache-2.0-Lizenz.
  • Die FP8-Kerne sind laut Zyphra 1,16- bis 1,76-mal schneller als AITER.
  • Lean P und Ping-Pong-Scheduling überlappen Softmax und Matrixarbeit, wodurch die Matrix-Pipeline besser ausgelastet ist.
  • ZYFA ist auf kausales Prefill beschränkt und funktioniert nur auf der MI355X.

ZYFA ist eine FP8-Attention-Engine, die für die Prefill-Phase ausgelegt ist, also die Verarbeitung des vorgegebenen Kontexts vor der Antwortgenerierung, bei langen Kontexten. Sie treibt die Zyphra Optimization Engine (ZOE) an, die Zyphra Cloud betreibt. Die BF16-Kerne sind in HIP geschrieben, während die FP8-Kerne handoptimierter Assembler-Code sind, der von einem in Python geschriebenen Generator erzeugt wird.

Warum hat FP8 auf dem MI355X Schwierigkeiten?

FP8-Matrixinstruktionen erreichen eine doppelt so hohe Spitzenleistung wie BF16. Aber FP8 beschleunigt nur die beiden Matrixmultiplikationen, QKᵀ und PV, während das Softmax, das die Scores in Wahrscheinlichkeiten umwandelt, auf der Vektoreinheit mit gleicher Geschwindigkeit in beiden Präzisionen läuft. Bei 32 Anfragen gegenüber 128 Schlüsseln benötigt die FP8-Matrixarbeit 1.024 Zyklen, gegenüber 896 Vektorzyklen für das exakte Softmax: Der Kernel muss die beiden Schritte daher nahezu perfekt überlappen. Laut den von Zyphra veröffentlichten Ausführungsspuren bleibt AITERS Matrix-Pipeline 31,7 % der Zyklen inaktiv. Der MI355X verfolgt den Abschluss von Matrixinstruktionen nicht auf Hardwareebene, sodass die Befehlsreihenfolge diese Überlappung erzeugen muss.

Drei Techniken zur Umgehung des Engpasses

Die erste, genannt Lean P, ersetzt die Exponentialfunktion des Softmax durch eine einzige Multiplikations-Addition. Ein e4m3-Bitmuster, als Ganzzahl gelesen, erhöht sich um 8 bei jeder Verdoppelung des Werts, sodass der Code für 2^x etwa 8x plus einer Konstanten entspricht. Dieser Ansatz reduziert die Vektorzyklen für das Softmax pro Block von 896 auf 384.

Die zweite ist ein Ping-Pong-Zwei-Wellen-Scheduling: Zwei Waves teilen sich jeden SIMD, eine sendet Matrixinstruktionen, während die andere Wahrscheinlichkeiten berechnet, dann tauschen sie die Rollen. Mit Lean P ist die Matrix-Pipeline im Dauerbetrieb zu 86 % ausgelastet, gegenüber 60 % mit der exakten Berechnung. Die dritte besteht in einer Block-Skalierung: Jeder Block von 64 Schlüsseln wird vor der Rundung mit einer Zweierpotenz multipliziert, ein Faktor, den die Skalierungsinstruktion für Matrixmultiplikationen von CDNA4 ohne Kosten entfernt. Die Zeilensumme wird aus den gerundeten Wahrscheinlichkeiten mit einer zusätzlichen Matrixinstruktion pro Block berechnet.

Reale Gewinne und Genauigkeit

Bei echten Aktivierungen mit 16K Kontext erreichen die FP8-Gewinne das 1,16-fache für Mistral-7B, das 1,17-fache für Qwen3-8B und das 1,76-fache für DeepSeek-V2-Lite. Bei Gauß-Benchmarks mit 32K erreicht der FP8-Kernel mit einer Kopfgröße von 128 2.672 TFLOP/s, gegenüber 1.944 für AITER. Der Kernel führt auch bei der MLA-Form mit 2.464 TFLOP/s gegenüber 1.237, einem Verhältnis von 1,99.

Die Diskrepanz zwischen Benchmarks und echten Modellen liegt an der Taktfrequenz des MI355X, die von der Anzahl der in den Operanden umschaltenden Bits abhängt. AITER lässt mehr kleine Wahrscheinlichkeiten zu Nullen oder subnormalen Zahlen werden, was weniger Bits umschalten lässt und einen höheren Takt ermöglicht: 2.108 MHz gegenüber 1.628 MHz für ZYFA auf der Mistral-7B-Schicht. ZYFA benötigt insgesamt dennoch weniger Zyklen. Laut dem Zyphra-Team sollten die Leistungen bei echten Aktivierungen und nachhaltiger Taktfrequenz gemessen werden.

Im Vergleich zu einer fp64-Berechnung beträgt der FP8-Fehler von ZYFA das 0,22-fache desjenigen von AITER auf Mistral-7B und das 0,38-fache auf Qwen3-8B. Beim vollständigen Mistral-7B-Modell erreicht die KL-Divergenz, die den Unterschied zwischen zwei Wahrscheinlichkeitsverteilungen misst, 3,4×10⁻³ mit ZYFA und 5,9×10⁻³ mit AITER. Zyphra gibt auch an, dass ZYFA auf kausales Prefill im Vorwärtsdurchlauf beschränkt ist, ohne Decoding oder paginierten KV-Cache, und dass es nur auf MI355X funktioniert.

S’abonner
Notification pour
guest

0 Commentaires
Le plus ancien
Le plus récent Le plus populaire
0
Nous aimerions avoir votre avis, veuillez laisser un commentaire.x