Flash Attention

• Good resource: https://www.youtube.com/watch?v=zEuwuCTEf_0

The task

• Assume for the moment 1 head, no batch dimension
  • This is okay, because they are fully independent, they are “embarrassingly parallel”
• We have $Q,K,V$ of shape $(N,d)$, seq. length $N$, head dimension $d$
• $P=softmax(\frac{Q{K}^T}{\sqrt{d}})$, $O=PV$
• How to parallelize and do this in one go?
  • $P$ is an intermediate, can we avoid materializing it?

Thoughts

• Attention is like 2-layer network, but the head dimension is very small