PyTorch Profiling Teil 3: Attention ist alles was du profilierst

Die Profiling-in-PyTorch-Serie

Hugging Face setzt seine Profiling-Serie fort mit Attention-Mechanismen – dem Herzstück moderner Transformer-Architekturen. Nach den Grundlagen in Teil 1 (Mathe-Operationen) und Teil 2 (Linear Layers und Fused MLPs) geht es nun um die quadratische Zeitkomplexität und wie man sie optimiert.

Serie-Übersicht:

1. Part 1: Beginner’s Guide zu torch.profiler
2. Part 2: Von nn.Linear zu Fused MLP
3. Part 3: Attention Profiling (aktuell)

Attention unter der Lupe

Attention arbeitet mit Queries (q), Keys (k) und Values (v). Die Schritte:

1. Attention Scores berechnen: matmul(q, k.T)
2. Skalieren: scores * scale
3. Causal Mask anwenden: scores.masked_fill(mask, "-inf")
4. Softmax-Normalisierung: softmax(scores)
5. Values neu gewichten: matmul(attn, v)

Das Ziel: Nicht jede Optimierung im Detail zu erklären, sondern zu zeigen, wie jede unter dem Profiler aussieht.

Implementierungs-Varianten

Der Blogpost deckt vier Implementierungen ab:

1. Naive Attention

Die direkte Implementation aller fünf Schritte in PyTorch. Zeigt die Baseline-Performance und wo die Bottlenecks liegen.

2. In-Place Operations

Speicheroptimierte Variante mit In-Place-Operationen. Reduziert Memory-Allocation, aber wie sieht das im Profiler aus?

3. SDPA (Scaled Dot-Product Attention)

PyTorch’s eingebaute F.scaled_dot_product_attention – optimiert und fused.

4. Custom Kernels

Handoptimierte CUDA-Kernel wie Flash Attention. Was zeigt der Profiler bei maximaler Optimierung?

Profiler-Insights

Die Profiler-Traces und Tabellen zeigen:

• Memory-Bandwidth vs. Compute – Wo liegt der Engpass?
• Kernel-Fusion – Welche Operationen werden zusammengefasst?
• Quadratic Scaling – Wie verhält sich Attention bei steigender Sequenzlänge?
• Flash Attention Patterns – Wie unterscheidet sich die optimierte Implementation?

Setup für Experimente

Die Scripts laufen auf NVIDIA A100-SXM4-80GB GPUs. Hugging Face bietet zwei Wege zum Reproduzieren:

1. Dev Mode with Spaces – Einfache GPU-Umgebung
2. Hugging Face Jobs Pipeline – Skalierbare Ausführung

python 04_a_naive_attention.py
python 04_b_inplace_ops_attention.py
python 04_c_sdpa_attention.py
python 04_d_kernels_attention.py

Key-Takeaways für ML-Engineers

1. Profiler lesen lernen – Traces zeigen mehr als Scores: Sie zeigen Memory-Patterns
2. Attention ist nicht gleich Attention – Naive vs. SDPA vs. Flash Attention haben unterschiedliche Profiling-Signaturen
3. Quadratic Complexity verstehen – Der Profiler macht die Skalierungsproblematik sichtbar
4. Optimierung validieren – Behauptungen über Performance-Gains sollten im Profiler verifiziert werden

Original-Artikel: Profiling in PyTorch (Part 3): Attention is all you profile
Autoren: Aritra Roy Gosthipaty, Sergio Paniego, Sayak Paul, Rémi Ouazan Reboul
Datum: 10. Juli 2026
Scripts: GitHub Repository