KI & Technologie

Selbstaufmerksamkeit, CUDA-Kernel und FlashAttention

2026-11-15

Illustrativer Speicherblock eines Mikrocontrollers, nicht die im Artikel behandelte GPU.

Moderne Transformer nutzen die Matrixberechnungen der Selbstaufmerksamkeit, GPU-Compiler auf höherer Abstraktionsebene wie Triton und speicherbewusste CUDA-Kernel wie FlashAttention, um riesige Token-Sequenzen mit maximaler Hardwaregeschwindigkeit zu verarbeiten.

1. Wie Selbstaufmerksamkeit in einem Vision Transformer mathematisch funktioniert

In einem Vision Transformer (ViT) wird ein Bild (z. B. 224×224 Pixel) in ein Raster aus 16×16-Ausschnitten zerlegt, wodurch N = 196 abgeflachte Ausschnitte entstehen. Jeder Ausschnitt wird linear in einen Vektor der Dimension D (z. B. 768) projiziert und mit einem Positions-Embedding kombiniert. So entsteht eine Eingabematrix X ∈ R^(N × D).

Erzeugung von Queries, Keys und Values. Das Modell multipliziert die Ausschnittsmatrix X mit drei erlernten Gewichtsmatrizen (W_Q, W_K, W_V ∈ R^(D × d_k)), um für jeden Bildausschnitt drei unterschiedliche Repräsentationen zu erzeugen:

Q=XW_(Q), K=XW_(K), V=XW_(V)

Query (Q): Stellt dar, wonach ein bestimmter Bildausschnitt sucht (z. B. "Ich bin eine Felltextur und suche Augen oder Ohren in der Nähe").
Key (K): Gibt an, welche Informationen ein Ausschnitt enthält (z. B. "Ich enthalte eine scharfe gekrümmte Kante").
Value (V): Enthält die eigentlichen visuellen Merkmalsdaten, die bei einer Übereinstimmung weitergegeben werden.

Berechnung der Aufmerksamkeitswerte. Um zu messen, wie stark jeder Ausschnitt mit jedem anderen zusammenhängt, berechnet das Modell das Skalarprodukt zwischen sämtlichen Queries und Keys:

Scores=(QK^(T))/(sqrt(d_(k)))

Die Division durch sqrt(d_(k)) verhindert, dass große Skalarprodukte die nachfolgende softmax-Funktion in Bereiche mit winzigen Gradienten verschieben. Daraus entsteht eine N × N-Aufmerksamkeitsmatrix, deren Eintrag (i, j) bewertet, wie viel Aufmerksamkeit Ausschnitt i dem Ausschnitt j widmen sollte.

Softmax-Gewichtung und Aggregation. Die zeilenweise Anwendung von Softmax wandelt rohe Werte in Wahrscheinlichkeiten um, deren Summe 1 ergibt. Die Multiplikation dieser Wahrscheinlichkeiten mit der Value-Matrix V erzeugt für jeden Ausschnitt eine gemischte Repräsentation:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

Bei der Multi-Head-Selbstaufmerksamkeit läuft genau dieser Prozess parallel über mehrere Köpfe (z. B. 12 Köpfe), sodass sich verschiedene Köpfe gleichzeitig auf unterschiedliche räumliche Abstände, Farben oder Objektteile konzentrieren können.

2. Praktischer Codevergleich: direktes CUDA C++ gegenüber OpenAI Triton

Um zu sehen, wie sich die GPU-Programmierung entwickelt hat, betrachten wir eine einfache elementweise Vektoraddition (C = A + B) für N Elemente.

Direktes CUDA C++ (Verwaltung auf Thread-Ebene)

In CUDA C++ programmieren Sie aus der Perspektive eines einzelnen skalaren Threads. Sie müssen Thread-IDs manuell berechnen, Array-Grenzen prüfen und wiederkehrende hostseitige Startkonfigurationen schreiben.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (Python-Programmierung auf Blockebene)

In Triton programmieren Sie direkt in Python aus der Perspektive eines Datenblocks (einer Kachel). Der Triton-Compiler übernimmt im Hintergrund automatisch die Thread-Synchronisierung, die Zusammenfassung von Speicherzugriffen und die Vektorisierung.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: den Speicherengpass von Transformern lösen

Gewöhnliche Selbstaufmerksamkeit stößt bei hochauflösenden Bildern oder langen Textsequenzen an eine harte Grenze, weil die Aufmerksamkeitsmatrix S = QK^T mit der Sequenzlänge N quadratisch wächst (N × N).

Der Engpass bei der Speicherbandbreite (HBM gegenüber SRAM). Moderne NVIDIA-GPUs besitzen zwei wesentliche Speicherebenen:

HBM (High Bandwidth Memory / VRAM): Groß (z. B. 80 GB), aber beim Lesen und Schreiben vergleichsweise langsam.
SRAM (gemeinsam genutzter Speicher auf dem Chip): Winzig (~20 MB insgesamt über alle Streaming Multiprocessors hinweg), aber ungefähr 10- bis 20-mal schneller als HBM.

Bei gewöhnlicher Aufmerksamkeit berechnet die GPU S = QK^T, schreibt die riesige N × N-Matrix in den langsamen HBM, liest sie für die Anwendung von Softmax wieder ein, schreibt die Wahrscheinlichkeiten zurück in den HBM und liest sie für die Multiplikation mit V erneut ein. Die Tensor Cores verbringen den Großteil ihrer Zeit untätig mit dem Warten auf Speicherübertragungen.

Kachelung (ohne die N × N-Matrix jemals zu materialisieren). FlashAttention wurde von Tri Dao und anderen (2022) eingeführt und nutzt Kernel-Fusion, um die gesamte Aufmerksamkeitsberechnung im schnellen SRAM auf dem Chip zu halten. Es zerlegt Q, K und V in kleine Blöcke, die in den SRAM passen, berechnet die Aufmerksamkeit Kachel für Kachel und schreibt nur die endgültige Ausgabematrix zurück in den HBM. Die riesige N × N-Zwischenmatrix wird nie im VRAM gespeichert.

Der Online-Softmax-Trick. Normalerweise muss Softmax den Höchstwert und die Summe der Exponentialwerte einer ganzen Zeile kennen, bevor es auch nur ein Element normalisiert. FlashAttention verwendet einen Online-Softmax-Algorithmus, der fortlaufende Höchstwerte und Normalisierungsfaktoren blockweise erfasst und die Teilausgaben bei der Verarbeitung jeder neuen Kachel mathematisch neu skaliert.

Neuberechnung bei der Rückpropagation. Da FlashAttention die N × N-Aufmerksamkeitsmatrix beim Vorwärtsdurchlauf nicht speichert, nutzt es beim Training einen geschickten Kompromiss: Es speichert nur die kleinen Normalisierungsstatistiken im HBM und berechnet die Aufmerksamkeitskacheln bei der Rückpropagation im schnellen SRAM laufend neu. Zusätzliche Berechnungen im SRAM erweisen sich als viel schneller als das Lesen großer Matrizen aus dem HBM. So wird die tatsächliche Laufzeit um den Faktor 2 bis 4 beschleunigt, während der Speicherbedarf von quadratisch O(N²) auf linear O(N) sinkt.

Weitere Nachrichten: KI & Technologie