-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgenerate_docs.py
More file actions
234 lines (167 loc) · 7.75 KB
/
Copy pathgenerate_docs.py
File metadata and controls
234 lines (167 loc) · 7.75 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
from pathlib import Path
ROOT = Path(__file__).parent
LICENSE = """\
MIT License
Copyright (c) 2025 Joao Felipe De Souza
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
DESIGN = """\
# Design Document -- attention-kernel-profiler
## 1. Motivation
Attention is the central computation of transformer-based LLMs.
Its cost scales quadratically with sequence length in the naive implementation,
but modern kernels (FlashAttention, mem_efficient_attention) reduce this to
near-linear by avoiding materializing the full n x n attention matrix.
This project measures how much difference the kernel choice makes on real hardware.
---
## 2. Hardware Context
The RTX 2070 is a Turing architecture GPU (sm75).
FlashAttention requires sm80+ (Ampere, e.g., A100, RTX 3090).
PyTorch's SDPA automatically falls back to mem_efficient_attention on sm75.
This project documents the performance of mem_efficient_attention as the
practical FlashAttention substitute on consumer-grade hardware.
---
## 3. Attention Variants
### naive_math
Manual implementation:
scores = Q @ K^T / sqrt(d_head)
scores = causal_mask(scores)
attn = softmax(scores)
out = attn @ V
Complexity: O(n^2) time, O(n^2) memory (materializes full attention matrix).
### sdpa_default
torch.nn.functional.scaled_dot_product_attention with auto backend selection.
On sm75: uses mem_efficient_attention.
On sm80+: uses FlashAttention.
### sdpa_flash
Forces FlashAttention backend. Falls back to sdpa_default on sm75.
### sdpa_math
Forces the quadratic math backend (no tiled attention).
Used as a baseline to measure the benefit of kernel optimization.
### sdpa_memeff
Forces memory-efficient attention backend explicitly.
### sliding_w128
Sliding window attention: tokens outside a window of 128 are masked.
Current implementation allocates full n x n matrix and masks.
A true sliding window kernel would compute only the window.
### gqa_4kv
Grouped Query Attention: 12 query heads, 3 KV heads (group size = 4).
Expands K/V via repeat_interleave before SDPA.
---
## 4. Key Results
### Time (us) at batch_size=1
Variant seq=256 seq=512 seq=1024 seq=2048
naive 288 549 1264 4947
sdpa_default 96 120 255 647
sdpa_math 489 1186 3766 8483
sdpa_memeff 115 257 228 707
gqa_4kv 147 175 332 945
### Speedup vs naive at seq=2048
sdpa_default: 7.64x
sdpa_memeff: 7.00x
sdpa_flash: 7.02x
gqa_4kv: 5.24x
sdpa_math: 0.58x (worse than naive!)
### Memory scaling (batch_size=1, seq 64->2048, 32x longer)
naive: 8.7 -> 217 MB (25x growth, O(n^0.93))
sdpa_default: 8.5 -> 20 MB (2.4x growth, O(n^0.25))
sdpa_math: 9.7 -> 489 MB (50x growth, O(n^1.13))
The sublinear memory scaling of sdpa_default confirms that
mem_efficient_attention achieves near O(n) memory on this hardware,
closely matching FlashAttention's theoretical O(n) guarantee.
---
## 5. Connections to Prior Projects
Finding Connection
──────────────────────────────────────────────────────────────
O(n^2) memory = KV cache pressure kv-cache-compaction-lab
Attention cost scales with seq latency-breakdown-simulator
FlashAttention reduces memory kv-cache-disaggregation-sim
GQA reduces KV footprint real-model-profiler
sm75 vs sm80+ matters real-model-profiler calibration
"""
README = """\
# 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](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.
"""
(ROOT / "LICENSE").write_text(LICENSE)
(ROOT / "DESIGN.md").write_text(DESIGN)
(ROOT / "README.md").write_text(README)
print("Wrote LICENSE, DESIGN.md, README.md")