Self-Attention Complexity, FlashAttention, PagedAttention, and RoPE in Code

Answer four questions on transformer attention: how self-attention works and what it costs, what FlashAttention is and what it improves, what PagedAttention is and how it saves memory, and how a transformer accounts for word order, including how RoPE is applied in code. ### Clarifying Questions - Should the complexity be given for training on full sequences, for token-by-token generation, or both? - Is the attention causal, as in a decoder, or bidirectional, as in an encoder? - Is a PyTorch implementation acceptable for the RoPE code? ### Part 1 — Self-attention and its cost How does self-attention work, and what is its asymptotic complexity in time and in memory? ```hint Write the shapes Write the shape of every matrix for a sequence of n tokens and model width d, and find the largest intermediate result. ``` #### What This Part Should Cover - Queries, keys and values, scaled dot-product attention, masking and multiple heads - Time and memory complexity in terms of n and d, and which term dominates when - How the costs differ between training on a full sequence and generating with a KV cache ### Part 2 — FlashAttention What is FlashAttention, how is it better than ordinary attention, and does it improve time, memory or both? ```hint Where the time goes Compare the arithmetic attention needs with the bytes it moves between GPU high-bandwidth memory and on-chip memory. ``` #### What This Part Should Cover - Tiling, the online softmax and recomputation in the backward pass - Whether the result is exact or approximate - Precisely what improves: asymptotic memory, wall-clock time, or the operation count ### Part 3 — PagedAttention What is PagedAttention, and how does it optimize memory usage? ```hint Think like an operating system Recall how an operating system avoids reserving contiguous physical memory for each process, and map that onto each request's KV cache. ``` #### What This Part Should Cover - Why a naively allocated KV cache wastes memory during serving - Fixed-size blocks, block tables and on-demand allocation - Sharing blocks between sequences, and the effect on batch size and throughput ### Part 4 — Word order and RoPE How does a transformer take word order into account, and how is RoPE applied in code? ```hint Shuffle the input Check what happens to self-attention's outputs if the input tokens are permuted, then decide where position information has to enter. ``` #### What This Part Should Cover - Why attention alone ignores order - Absolute, relative and rotary position schemes - A correct RoPE implementation: what is rotated, by how much, and where in the attention computation ### What a Strong Answer Covers - Correct attention equations with shapes, and complexity in n and d for both time and memory - An accurate account of FlashAttention as an exact, IO-aware algorithm, including what it does not reduce - The KV-cache memory problem and how paging and block sharing address it - Why position must be injected, the main schemes, and RoPE applied to queries and keys - A clear separation between training-time and inference-time costs ### Follow-up Questions - How large is the KV cache for one sequence, and how do multi-query and grouped-query attention change it? - Why does FlashAttention recompute attention in the backward pass instead of storing it? - How does RoPE behave beyond the training context length, and what methods extend it? - Why are the dot products divided by $\sqrt{d_k}$?