Why standard time complexity is a poor measure for GPUs

The "Attention is logarithmic, actually" article argues that it doesn't make sense to classify attention as a quadratic operation due to the highly parallelized nature of GPUs. He argues that the work-depth model of complexity is more informative for parallelized hardware.

In a nutshell, number of operations performed ("work") hides the fact that many operations can happen simultaneously. So quadratic work $\ne$ quadratic latency. You need to account for "depth" -- the number of steps required even if you had unlimited hardware.

• element-wise multiplication is $O(n)$ work but $O(1)$ depth because every operation can happen simultanously
• sum reduction is $O(N)$ work but $O(\log n)$ depth since you can sum pairs in a reduction tree
• matrix multiplication is $O(n^3)$ work but $O(\log n)$ depth
• softmax is $O(n)$ work but $O(\log n)$ depth

The blog derives that vanilla attention is $O(d{n}^{2}d)$ but $O(\log n+\log d)$ depth ($n$ = seq len, $d$ = embedding dimension, $b$ = batch size). Since $n>>d$, you can simplify to $O(\log d)$.

However, this is just the theoretical floor. The $Q{K}^T$ attention matrix is too large for the fast caches, so it spills down the memory hierarchy, so computations have to be shared, adding more sequential steps. So given today's hardware, the depth is more like $O(n\log n)$.