FlashAttention: Memory Traffic of Naive Attention
Beginner
The problem, from first principles
Standard attention writes the full n by n score matrix to GPU high-bandwidth memory and reads it back. For long sequences that traffic, not the arithmetic, is the bottleneck. FlashAttention avoids storing the matrix.
From theory to code
Implement naive_attention_hbm_bytes(n, bytes_per), the bytes needed to store the full score matrix.
Constraints
- Default precision is two bytes (fp16).
Hints
Hint 1
Count the n squared scores and multiply by the bytes per score.
Keywords
research-paperssystemsattentionkernels
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.