attention-kernel-profiler

Profiles 6 attention variants on GPU, measuring time, memory, and TFLOPS across sequence lengths from 64 to 2048 tokens.

For full design and results see DESIGN.md.

---

Hardware Note

The RTX 2070 is sm75 (Turing). FlashAttention requires sm80+ (Ampere). PyTorch automatically selects mem_efficient_attention as the substitute. This project profiles mem_efficient_attention as the real-world FlashAttention alternative on consumer-grade hardware.

---

Variants Profiled

Variant	Description
naive	Manual O(n²) — materializes full n x n attention matrix
sdpa_default	PyTorch auto (mem_efficient on sm75, flash on sm80+)
sdpa_flash	Forced FlashAttention (falls back on sm75)
sdpa_math	Forced quadratic math backend
sdpa_memeff	Forced memory-efficient backend
gqa_4kv	Grouped Query Attention (12 Q, 3 KV heads)

---

Key Findings

1. sdpa_default is 7.64x faster than naive at seq=2048

naive:        4,947 us
sdpa_default:   647 us
speedup:        7.64x

2. Memory scaling is the most dramatic result

seq 64 -> 2048 (32x longer):
naive:         8.7 MB ->  217 MB  (25x growth, ~O(n^2))
sdpa_default:  8.5 MB ->   20 MB  (2.4x growth, ~O(n^0.25))
sdpa_math:     9.7 MB ->  489 MB  (50x growth, worse than naive)

mem_efficient_attention achieves near-O(n) memory on sm75.

3. sdpa_math is the worst in both time and memory

seq=2048 vs sdpa_default:
Time:   8,483 us vs 647 us  (13x slower)
Memory: 489 MB vs 20 MB     (24x more)

Kernel choice matters more than model architecture at long sequences.

4. GQA gives 5x speedup vs naive but not vs sdpa_default

gqa (seq=2048): 945 us vs sdpa_default 647 us (1.46x slower)
Reason: KV head expansion via repeat_interleave adds overhead
Benefit: fewer KV parameters -> less KV cache memory in practice

---

Quick Start

python3 -m venv venv
source venv/bin/activate
pip install torch transformers matplotlib pandas

python3 profile_attention.py       # v1: basic variants
python3 profile_attention_v2.py    # v2: longer seqs, more batch sizes
python3 plot_attention.py          # plots
python3 analyze_attention.py       # analysis

---

Results

results/attention_profile.csv      60 rows (v1)
results/attention_profile_v2.csv  108 rows (v2, seq up to 2048)
results/plots/                    visualization

---

Portfolio Context

Project 16 in a series on LLM inference infrastructure. Complements real-model-profiler (project 15) by profiling the attention sub-component specifically, revealing how kernel choice dominates performance at long sequence lengths.