Flash Attention: Speicher-Effiziente Attention

Das Problem: Standard Attention ist Memory-Bound

Standard Attention materialisiert die komplette n×n Attention Matrix. Für lange Sequenzen wird dies zum Speicher-Engpass, nicht zum Compute-Engpass – die Hardware-Bandbreite, nicht die reine Rechenleistung, begrenzt die Performance.

Die Lösung: IO-Aware Tiling & Online Softmax

Flash Attention nutzt zwei zentrale Innovationen:

1. Block-weise Berechnung (Tiling)

Statt die gesamte Matrix zu materialisieren, lädt Flash Attention Blöcke von Q, K, V in den schnellen SRAM-Cache, berechnet block-weise Attention und führt die Ergebnisse zusammen – ohne je die vollständige n×n Matrix zu speichern.

2. Online Softmax Computation

Der klassische Softmax benötigt das Maximum all aller Eingaben, um numerisch stabil zu sein. Flash Attention nutzt ein cleveres mathematisches Derivat, das Softmax inkrementell aktualisiert während neue Blöcke geladen werden – ohne je alle Werte halten zu müssen.

Algorithmic Details & Implementation

Wie Flash Attention in der Praxis arbeitet:

Flash Attention 2 Pseudo-Code

def flash_attention_2(Q, K, V):
    N = Q.shape[0]        # Sequenzlaenge
    d = Q.shape[1]        # Head dimension
    B_r = 256             # Row block size (RAM)
    B_c = 32              # Col block size (RAM)

    O = zeros(N, d)               # Output
    L = zeros(N)                  # Running normalization
    m = -inf                      # Running max

    for i in range(0, N, B_r):
        Q_block = Q[i:i+B_r]      # Load Q block
        O_block = zeros(B_r, d)
        L_block = zeros(B_r)
        m_block = -inf

        for j in range(0, N, B_c):
            K_block = K[j:j+B_c]  # Load K, V blocks
            V_block = V[j:j+B_c]

            S_block = Q_block @ K_block.T / sqrt(d)           # Scores
            m_new = max(m_block, max(S_block, axis=1))  # Online max
            P_block = exp(S_block - m_new[:, None])

            O_block = O_block * exp(m_block - m_new)[:, None] + P_block @ V_block
            L_block = L_block * exp(m_block - m_new) + sum(P_block, axis=1)
            m_block = m_new

        O[i:i+B_r] = O_block / L_block[:, None]

    return O
            

Performance Impact & Context Scaling

Die praktischen Auswirkungen von Flash Attention auf Training und Inference:

Training Speed

Flash Attention ist 2-4× schneller beim Training als Standard Attention. Das kommt durch:

• Weniger HBM-Bandbreiten-Stalls
• Bessere GPU-Auslastung (75% statt 40-50%)
• Recomputation in Backprop statt Materialisierung

Memory Efficiency

Die Speichereffizienz ermöglicht ramatisch längere Kontexte:

Flash Attention Evolution

Seit der Original-Veröffentlichung hat Flash Attention sich weiterentwickelt:

Hinweis: Flash Attention v2 ist der aktuelle Standard für Training und Inference. Die meisten modernen LLMs nutzen FA v2 als Built-in.

Real-World Impact & Modell-Ökosystem

Flash Attention ist infrastrukturell für modernes LLM-Design. Combined mit GQA, RoPE und PagedAttention, ermöglichte es die Evolution der Context-Längen:

Warum Flash Attention so kritisch ist

Kernerkenntnisse