My understanding is that FlashAttention's memory use is linear, or close to linear in practice, but computation is still O(n²). I'm unaware of anyone being able to apply FlashAttention on, say, a million tokens, because it must execute ~1/2 x 1,000,000^2 x n_head dot-products, each in a subspace with d_head dimensions. That's not exactly computationally cheap!
No you're right. I mistook you. Compute isn't linear yet.
It is only linear in terms of memory, not compute. Flash attention is a big advance, but not enough for 1 million tokens