Ist ein ganztägiger Kontext besser als ein zehnminütiger Kontext? Flash Attention und die Frage nach der Sequenzlänge

Hier ist die Frage, die in diesem Artikel beantwortet werden soll: Wenn ein Transformator einen ganzen Handelstag mit einer Auflösung von einer Sekunde anstelle eines Zehn-Minuten-Fensters abdecken könnte, würde er dann bessere Vorhersagen treffen?

Bis vor Kurzem konnte man nicht einmal fragen. Standard-Aufmerksamkeitsbedürfnisse $\mathcal{O}(N^2)$ Speicher, also benötigt ein Tag mit 23.400 Schritten und 12 Köpfen in float16 etwa 12,9 GB allein für die Score-Matrix – mehr als die Modellparameter und mehr, als die meisten Karten Ihnen bieten. Die Frage wurde durch Arithmetik gelöst, bevor jemand sie testen konnte.

Flash Attention (Dao et al., 2022) öffnet es. Nicht durch Annäherung der Aufmerksamkeit – es berechnet das genau gleiche Ergebnis – sondern durch Umstrukturierung der Berechnung, um IO-bewusst zu sein und den Datenverkehr zwischen GPU-Speicherebenen zu minimieren. Das ist der wirklich interessante Inhalt hier, und der Großteil dieses Artikels beschäftigt sich damit, wie es funktioniert: Kacheln, die Online-Softmax-Wiederholung, das $\Theta(N^2 d^2 / M)$ IO-gebunden und Rückwärtsdurchlauf-Neuberechnung.

Aber der Mechanismus ist der Wegbereiter, nicht der Anspruch. „Längerer Kontext ist besser“ ist eine empirische Aussage über Märkte und die Position dieses Blogs – von Temporal Fusion Transformers, die feststellte, dass Vanilla-Transformatoren auf wiederkehrenden Finanzserien-Overfit- und Short-Lookback-Modellen bei hoher Frequenz wettbewerbsfähig bleiben – Kürzungen in die andere Richtung. Der Artikel schließt also mit der Messung ab, nicht mit dem Mechanismus.

Warum Aufmerksamkeit gedächtnisgebunden ist

Aufmerksamkeit berechnet $\text{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d_k})\mathbf{V}$ – Das Grundelement selbst wird im Handelskontext in Temporal Fusion Transformers for Multi-Horizon Portfolio Forecasting. Das ganze Problem ist eine Zeile davon: die Zwischenpunktmatrix $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top/\sqrt{d_k}$ Ist $N \times N$, wird es in den Speicher geschrieben, für den Softmax zurückgelesen, erneut geschrieben und für den endgültigen Matmul erneut gelesen – und es muss für die Backpropagation aufbewahrt werden.

Die arithmetische Intensität der Aufmerksamkeit beträgt $\approx \min(d, N)$, also etwa 64 FLOP/Byte bei $d = 64$ — deutlich links vom A100-Kammpunkt. Es liegt an der schrägen Bandbreitenobergrenze, nicht an der flachen Rechenobergrenze: Die GPU verbringt mehr Zeit mit der Bewegung $\mathbf{S}$ herum, als irgendetwas zu multiplizieren. Das hier verwendete Dachlinien-Framework – Firstpunkt, geneigte oder flache Decke und warum die gleiche Argumentation darüber entscheidet, ob sich der Kauf einer GPU überhaupt lohnt – basiert auf gemessenen Zahlen in When the GPU Pays Off.

Die Speicherhierarchie, die der Algorithmus ausnutzt

SRAM bietet ungefähr die 10-fache Bandbreite und eine 100-mal geringere Latenz bei einem Tausendstel der Kapazität. Alles, was Flash Attention tut, ergibt sich aus diesem Handel: Kapazität aufgeben, Bandbreite und Latenz kaufen. Der gleiche Schritt „den Algorithmus umstrukturieren statt Hardware zu kaufen“, gemessen an einem CPU-Backtest, ist die Backtest-Geschwindigkeitsleiter.

Der Flash Attention-Algorithmus

Flash Attention verarbeitet die Aufmerksamkeit in Kacheln, die so dimensioniert sind, dass sie in den SRAM passen, und materialisiert nie die volle Aufmerksamkeit $N \times N$ Matrix in HBM überhaupt.

Partition $\mathbf{Q}$ hinein $T_r = \lceil N/B_r \rceil$ Reihenblöcke und $\mathbf{K}, \mathbf{V}$ hinein $T_c = \lceil N/B_c \rceil$ Säulenblöcke, mit $B_r, B_c$ so gewählt, dass eine Kachel samt Akkumulatoren auf den Chip passt. Iterieren Sie für jeden Abfrageblock alle Schlüsselwertblöcke:

For each query block Q_i:
    Initialize: O_i = 0, m_i = -inf, l_i = 0    # output, running max, running sum
    For each KV block (K_j, V_j):
        1. Load Q_i, K_j, V_j from HBM to SRAM
        2. Compute S_ij = Q_i @ K_j^T / sqrt(d)  # in SRAM
        3. Compute local max: m_ij = rowmax(S_ij)
        4. Compute P_ij = exp(S_ij - m_ij)        # in SRAM
        5. Compute local sum: l_ij = rowsum(P_ij)
        6. Update running statistics:
           m_new = max(m_i, m_ij)
           l_new = l_i * exp(m_i - m_new) + l_ij * exp(m_ij - m_new)
           O_i = O_i * (l_i * exp(m_i - m_new) / l_new)
                + P_ij @ V_j * (exp(m_ij - m_new) / l_new)
           m_i = m_new, l_i = l_new
    Write O_i to HBM

Die Online-Softmax-Wiederholung

Der Trick, der das Kacheln ermöglicht, ist Online-Softmax. Ein naiver Softmax benötigt zwei Durchläufe über die Zeile: einen zum Ermitteln des Maximums (für numerische Stabilität), einen zum Potenzieren und Normalisieren. Zwei Durchgänge über eine Zeile, die Sie nicht speichern möchten, sind ein Widerspruch – daher führt Flash Attention weiterhin Statistiken aus und skaliert sie dabei neu.

Nach Blöcken $1, \ldots, j$:

und der Ausgangsakkumulator wird um denselben Faktor korrigiert:

Jedes Mal, wenn ein neuer Block das laufende Maximum erhöht, wird die zuvor akkumulierte Ausgabe rückwirkend um neu skaliert $e^{m^{(j-1)} - m^{(j)}}$ – als ob das neue Maximum von Anfang an bekannt gewesen wäre. Das Ergebnis ist algebraisch identisch mit dem Softmax mit zwei Durchgängen. In der exakten Arithmetik ist dies keine Näherung; es ist eine Neuvereinigung. (Bei endlicher Präzision kommt es auf einen anderen Rundungspfad an – siehe die Genauigkeitsprüfung unten.)

IO-Komplexität

Dies ist die offizielle Erklärung des Sieges. Flash Attention führt durch

HBM-Zugriffe, wo $M$ ist die SRAM-Größe dagegen $\Theta(Nd + N^2)$ für die Standardimplementierung. Beachten Sie, dass $M$ erscheint im Nenner: Je größer der On-Chip-Scratchpad, desto weniger Roundtrips, weshalb der Algorithmus in Bezug auf die Speicherhierarchie und nicht in Bezug auf die FLOP-Anzahl angegeben wird. Für typisch $d = 64$ Und $M \approx 100$ KB, das Verhältnis begünstigt Flash Attention durch etwa 5-10x weniger Zugriffe.

Rückwärtsdurchlauf: Neu berechnen statt speichern

Die Rückausbreitung durch Aufmerksamkeit erfordert normalerweise die $\mathbf{P}$ Matrix, die der Vorwärtspass einfach nicht behalten wollte. Flash Attention berechnet die Kacheln neu $\mathbf{Q}, \mathbf{K}$ Während des Rückwärtsdurchlaufs wird nur die Ausgabe gespeichert $\mathbf{O}$ und die Softmax-Statistiken $(m, \ell)$ – beides $\mathcal{O}(N)$, nicht $\mathcal{O}(N^2)$. Es tauscht eine bescheidene Menge redundanter Arithmetik gegen den Speicherterm ein, der das gesamte Problem darstellte. Dies ist derselbe Vorteil wie das Gradienten-Checkpointing, das auf Kachelgranularität innerhalb eines einzelnen Operators angewendet wird.

FA2: Parallelität

Flash Attention 2 (Dao, 2023) behielt den Algorithmus bei und korrigierte die Planung:

1. Weniger Nicht-Matmul-FLOPs. FA1 verbrachte Echtzeit mit Neuskalierung, Max-Ermittlung und Potenzierung – Operationen, die auf CUDA-Kernen und nicht auf Tensorkernen ausgeführt werden. FA2 verschiebt die Neuskalierung bis zum Ende der inneren Schleife.
2. Parallelität über die Sequenzlänge. FA1 parallelisiert nur über Batch und Köpfe. FA2 parallelisiert auch über Abfrageblöcke. Dies ist insbesondere für den Handelsfall wichtig, bei dem es oft eine sehr lange Sequenz pro Asset und eine Batch-Größe von 1–4 gibt – genau das Regime, bei dem Batch-and-Head-Parallelität die GPU aushungert.
3. Warp-Arbeitspartitionierung. Jeder Warp verwendet eine andere Teilmenge von Abfrageblöcken, anstatt eine Score-Berechnung aufzuteilen und über Warps hinweg zu reduzieren, wodurch eine Cross-Warp-Reduktion entfällt.

Berichtetes Ergebnis: ~70 % der theoretischen Spitzen-FLOPs auf A100 gegenüber ~35 % für FA1.

FA3: Hopper-Mechanik

Flash Attention 3 (Dao, Shah, 2024) ist architekturspezifisch für H100:

1. Asynchrone Warp-Spezialisierung. Hoppers Tensor Memory Accelerator (TMA) bewegt HBM→SRAM asynchron. FA3 teilt Warps in Erzeuger auf, die TMA-Lasten für den nächsten KV-Block ausgeben, und Verbraucher, die für den aktuellen Block rechnen, sodass die Datenbewegung hinter der Arithmetik verborgen bleibt.
2. Verschachteltes Matmul und Softmax. $\mathbf{Q}_i\mathbf{K}_j^\top$ für einen Block läuft auf Tensorkernen, während der Softmax für den vorherigen Block auf CUDA-Kernen läuft – zwei verschiedene Hardwareeinheiten, wirklich gleichzeitig und nicht in Zeitscheiben unterteilt.
3. FP8 mit inkohärenter Verarbeitung. H100 führt FP8 mit 2x FP16-Durchsatz aus. Naive FP8-Aufmerksamkeit wird durch Ausreißer zerstört; FA3 dreht Vektoren vor der blockweisen Quantisierung zufällig, um die Größe der Ausreißer über die Koordinaten zu verteilen, was mit einem 2,6-fach geringeren numerischen Fehler als bei naivem FP8 gemeldet wird.

Durch kausale Maskierung erhält der Handel einen Rabatt

Kausale Maskierung ist für Zeitreihen zwingend erforderlich – das Modell darf sich nicht auf die Zukunft konzentrieren – und bei der Kachelung handelt es sich nicht um zusätzliche, sondern um eingesparte Kosten. Jede Kachel, deren Schlüssel im Verhältnis zu ihren Abfragen vollständig in der Zukunft liegen, wird komplett übersprungen, nie geladen und nie berechnet, wodurch etwa die Hälfte der Arbeit eingespart wird. In PyTorch ist dies der Fall is_causal=True; nichts anderes ist erforderlich.

Die Integration umfasst acht Zeilen

Bei fast keinem Code, den Sie benötigen, geht es um Flash Attention. Tauschen Sie den expliziten Score-Matrix-Pfad für den fusionierten Kernel aus:

scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_head)
scores = scores.masked_fill(causal_mask, float("-inf"))
out = torch.matmul(torch.softmax(scores, dim=-1), v)

out = torch.nn.functional.scaled_dot_product_attention(
    q, k, v,
    dropout_p=self.dropout if self.training else 0.0,
    is_causal=True,
)

Das ist die ganze Veränderung. q, k, v geformt sind (batch, heads, seq, head_dim); Die Kausalmaske ist verschwunden, weil der Kernel sie erstellt. Für ein vollständig annotiertes PyTorch-Handelsmodell, in das dies eingefügt werden soll – Eingabeprojektion, Blöcke und ein 3-Klassen-Up/Flat/Down-Kopf – verwenden Sie DeepLOB, und eine vollständige Trainingspipeline finden Sie unter Temporal Fusion Transformers. Der Bau einer vierten Kopie dieses Gerüsts hier würde nichts lehren.

Anforderungen: Rechenleistung >= 8,0 (A100, H100, RTX 3090+), Eingaben mit halber Genauigkeit, PyTorch >= 2,0. Überprüfen Sie den tatsächlich genutzten Fast Path torch.backends.cuda.sdp_kernel Diagnostik und torch.cuda.max_memory_allocated() – SDPA greift stillschweigend auf den Mathematikkern zurück, wenn eine Vorbedingung fehlschlägt, und ein stiller Fallback sieht genauso aus wie ein funktionierendes Modell, das nur langsam ist.

Die Messung: Lohnt sich der längere Kontext?

Alles oben Genannte besagt, dass ein 32K- oder 128K-Kontext jetzt erschwinglich ist. Es sagt nichts darüber aus, ob es nützlich ist. Das ehrliche Experiment:

Trainieren Sie die gleiche Architektur bei $N \in \{512, 4096, 32768\}$ mit SDPA-Flash auf der BTC-Serie, die an anderer Stelle in dieser Serie verwendet wird und Parameter, Optimierer und Ziel festhält, sodass die Sequenzlänge die einzige Variable ist. Melden Sie zwei Dinge:

1. Kosten. Gemessene Wanduhr pro Epoche und torch.cuda.max_memory_allocated() bei jedem $N$.
2. Vorteil. Vorhersageleistung außerhalb der Stichprobe im Vergleich zu $N$, bei einem Walk-Forward-Split.

Ein früherer Entwurf dieses Artikels enthielt eine Tabelle mit Speicherzahlen pro Sequenzlänge, die analytisch aus dem abgeleitet wurde $\mathcal{O}(N)$ Aktivierungsformel. Diese Zeilen wurden entfernt: Sie wurden nie gemessen und stimmten nicht mit der eigenen Speicherbudget-Arithmetik des Artikels überein. Eine abgeleitete Zahl, die in einer Ergebnistabelle angezeigt wird, ist ein erfundenes Ergebnis, und dieser Blog liefert diese nicht.

Die interessante Eigenschaft dieses Experiments ist, dass es in beide Richtungen veröffentlicht werden kann. Wenn die Out-of-Sample-Leistung monoton mit ansteigt $N$, das rechtfertigt das gesamte Langkontextprogramm. Wenn es bei ein paar tausend Schritten ein Plateau erreicht oder nachlässt, ist das ein stärkeres Stück – ein Begleiter zu dem ehrlichen Negativ – und das würde bedeuten, dass die Memory Wall nie die verbindliche Einschränkung für den Handel mit Transformatoren war.

Mehr Kontext bedeutet mehr Kapazität und daher mehr überpassende Oberfläche

Es gibt einen bestimmten Grund, ein flaches oder negatives Ergebnis zu erwarten. Temporale Fusionstransformatoren dokumentiert bereits, dass Vanilla-Transformatoren naiv auf die Überanpassung von Finanzreihen angewendet werden – ihnen fehlen zeitliche induktive Verzerrungen und wiederkehrende Short-Lookback-Modelle bleiben bei hoher Frequenz wettbewerbsfähig. Durch die Erweiterung des Kontexts von 512 auf 32.768 Schritte werden keine Informationen proportional zur Länge hinzugefügt. Die geringfügige 32.000ste Verzögerung einer nahezu effizienten Preisreihe hat nur sehr geringe Auswirkungen. Was es zuverlässig hinzufügt, ist der Parameterwert der anzupassenden Dinge.

Also der Sweep vorbei $N$ muss als das behandelt werden, was es ist: eine Modellauswahlsuche, mit der gleichen Maschinerie, die dieser Blog auf jede andere Suche anwendet. Drei Sequenzlängen mal was auch immer sonst variiert, ist eine Versuchszählung, und der Gewinner muss ein deflationiertes Sharpe-Verhältnis erreichen, das anhand dieser Versuchszählung und eines PBO-Gates berechnet wird, und nicht nur seine Nachbarn schlagen. Ansonsten sind „lange Kontextgewinne“ nicht von der Auswahl des besten aus drei lauten Läufen zu unterscheiden.

Eine Genauigkeitsprüfung, denn „exakt“ macht viel Arbeit

Flash Attention ist nach exakter Arithmetik genau. Die damit verbundene Empfehlung – in fp16 oder bf16 laufen zu lassen und bei H100 FP8 in Betracht zu ziehen – ist nicht der Fall. Dabei handelt es sich um getrennte Behauptungen, und die zweite Behauptung dominiert in der Praxis: Das Neuzuordnen einer Summe und das Absenken auf die halbe Genauigkeit sind beides Störungen, und der Artikel, der die Ordnungsgarantie eingeführt hat, sollte die Präzisionsgarantie dann nicht mit der Hand bewegen.

Der Blog hat bereits das passende Instrument. Die GPU-Präzisionsfalle legt den Standard fest: Niedrige Präzision warnt Sie nicht, sie gibt plausiblen Müll zurück, und Sie beweisen die Richtigkeit mit einem Paritätorakel für eine nachgeschaltete diskrete Menge – Handelszählungen – und nicht durch das Beobachten von Kurven. Hier angewendet:

• Berechnen Sie die Aufmerksamkeit mit SDPA-Flash in bf16 und mit einer fp64-Referenzimplementierung auf identischen Eingaben; Melden Sie maximalen relativen Fehler für den Ausgabetensor.
• Setzen Sie es bis zur Entscheidung durch: Geben Sie für ein Modell, das eine Aufwärts-/Flach-/Abwärtsbezeichnung ausgibt, als Bruchteil der Gesamtentscheidungen an, wie viele Beschriftungen zwischen den beiden Pfaden wechseln.

Kleine, begrenzte, erklärbare Meinungsverschiedenheiten sind das Zeichen für einen korrekten schnellen Weg. Eine unbegrenzte Zahl bedeutet, dass die FP8-Empfehlung für dieses Modell nie sicher war. Keine der beiden Zahlen ist bekannt, bis sie ausgeführt wird.

Wann man danach greifen sollte

Komprimiert auf die Entscheidung, die die gleiche Form hat wie der GPU-Entscheidungsleitfaden:

• Über ~2K-Zeitschritte auf einer CUDA-GPU: Ja, bedingungslos. Es handelt sich um eine einzeilige Änderung, die eine exakte Ausgabe erzeugt, und der Gewinn wächst mit $N$. Es gibt kein Szenario, in dem Sie möchten, dass das Wirklichkeit wird $N \times N$ Pfad stattdessen.
• Unterhalb von ~512 Zeitschritten, auf der CPU oder mit Nicht-Aufmerksamkeits-Architekturen (CNNs, SSMs wie Mamba): irrelevant. Links vom Grat stellt der feste Overhead die gesamten Kosten dar und Aufmerksamkeit war nie Ihr Engpass.
• Die oben genannten Schwellenwerte sind Folklore, keine Messung – sie stammen aus der allgemeinen Literatur, und der Crossover auf Ihrem eigenen Modell und Ihrer Karte ist ein Zehn-Zeilen-Benchmark. Führen Sie es aus, anstatt den runden Zahlen zu vertrauen.

Fazit

Flash Attention ist ein sauberes und wirklich wichtiges Ergebnis: Durch Respektierung der Speicherhierarchie und Neuzuordnung des Softmax wird die genaue Aufmerksamkeit berechnet $\mathcal{O}(N)$ Erinnerung statt $\mathcal{O}(N^2)$, und die $\Theta(N^2 d^2 / M)$ IObound erklärt genau, warum. Die Übernahme in einen Handelstransformator ist eine einzeilige Änderung ohne Genauigkeitseinbußen bei der exakten Arithmetik und einem großen Speichergewinn.

Was es nicht tut, ist die Frage oben zu beantworten. Es wandelt „ein ganztägiger Kontext ist unmöglich“ in „ein ganztägiger Kontext ist billig“ um, was eine Änderung der Kosten des Experiments darstellt, nicht seines Ergebnisses. Die herabstürzende Erinnerungsmauer ist eine Einladung zum Messen, und die Messung ist es, die daraus eine Zusammenfassung auf Papier in ein Ergebnis verwandelt.

Referenzen

1. Dao, T., Fu, D.Y., Ermon, S., Rudra, A., Re, C. „FlashAttention: Schnelle und speichereffiziente exakte Aufmerksamkeit mit IO-Awareness.“ NeurIPS (2022). arXiv:2205.14135
2. Dao, T. „FlashAttention-2: Schnellere Aufmerksamkeit mit besserer Parallelität und Arbeitspartitionierung.“ ICLR (2024). arXiv:2307.08691
3. Dao, T., Shah, J. „FlashAttention-3: Schnelle und genaue Aufmerksamkeit mit Asynchronität und geringer Präzision.“ NeurIPS (2024). arXiv:2407.08608
4. Vaswani, A., et al. „Aufmerksamkeit ist alles, was Sie brauchen.“ NeurIPS (2017).
5. Milakov, M., Gimelshein, N. „Online-Normalisierungsberechnung für Softmax.“ arXiv:1805.02867 (2018).