FlashAttention-2: Attention FLOPs
Intermediate
The problem, from first principles
Attention does two matrix multiplications per head, each costing about two n squared d operations. Knowing that count tells you whether a kernel is limited by compute or by memory.
From theory to code
Implement attention_flops(n, d, causal), the floating-point operations of one attention head.
Constraints
- Causal attention does about half the work.
Hints
Hint 1
Count four n squared d operations for the two products, then halve for causal masking.
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.