Why architecture matters here
Attention compute and memory matter because they determine LLM cost, latency, and the feasibility of long context -- the quadratic scaling is the fundamental constraint. Every aspect of LLM economics and performance traces partly to attention's cost: the cost of processing a long input (quadratic in length -- long documents are expensive), the memory needed (the KV cache growing with context -- limiting how much context fits in memory), the latency (attention compute contributing to the time per token), and the feasibility of very long context (the quadratic cost making naive attention infeasible at extreme lengths -- motivating the long-context techniques). Understanding attention's compute and memory is thus understanding the fundamental cost structure of LLMs -- why long context is expensive, why memory limits context, why certain optimizations (FlashAttention, KV cache management) matter so much. For anyone reasoning about LLM cost, capacity, or performance -- estimating costs, planning memory, choosing context lengths -- the compute and memory of attention (and the transformer) is the foundation, and it's why efficiency work in LLMs so heavily targets attention (the quadratic bottleneck).
The quadratic-scaling insight is the crux, and it's why long context is fundamentally expensive. Attention computes, for each position, its interaction with every other position (the all-to-all attention -- each query attends to all keys). For a sequence of length N, this is N x N interactions (the QK^T matrix -- N queries times N keys) -- O(N^2). So the attention cost scales quadratically with sequence length: doubling N quadruples the attention compute (and, naively, memory). This quadratic scaling is the fundamental constraint: it means long context is expensive in a way that grows faster than linearly (a 2x longer context is 4x the attention cost, a 10x longer context is 100x) -- so very long context (hundreds of thousands of tokens) faces a steep cost wall with naive attention. This is why the context length is a major cost driver (quadratic attention), why extending context is hard (the quadratic wall), and why long-context techniques (reducing the quadratic to sub-quadratic -- sparse, linear, sliding-window attention) are an active area (breaking the quadratic barrier for very long context). Understanding that attention is O(N^2) -- and thus long context is fundamentally, quadratically expensive with naive attention -- is understanding the central cost constraint of transformers and why so much effort goes into making attention more efficient.
And the memory-vs-compute-bound distinction (prefill vs decode) is essential to understanding LLM inference performance, because the two phases have very different characteristics. Prefill (processing the input prompt): the whole input is processed at once (in parallel over the sequence positions), so it's compute-bound (lots of parallel computation -- the attention and FFN over all positions -- limited by compute throughput). Decode (generating output tokens, one at a time): each token requires a forward pass, but only for the one new position (attending to all previous via the KV cache), so there's little computation per step but the model weights (and KV cache) must be loaded -- so it's memory-bandwidth-bound (limited by loading the weights/cache from memory, not by computation). This distinction matters enormously for inference optimization: prefill and decode have different bottlenecks (compute vs memory bandwidth), so they're optimized differently (prefill -- maximize compute utilization; decode -- the memory-bound regime that speculative decoding and batching target), and techniques like prefill-decode disaggregation (running prefill and decode on separate optimized hardware) exploit the difference. Understanding that prefill is compute-bound and decode is memory-bound -- and thus the two phases have different performance characteristics and optimizations -- is understanding LLM inference performance, and it's why inference serving so carefully manages the two phases.
The architecture: every piece explained
Top row: the attention computation and its cost. QK^T: the core -- query-key dot products for every position pair, forming the N x N attention matrix (for sequence length N -- each entry the interaction of a query and key). O(N^2) compute: the N x N matrix is N^2 dot products (each of dimension d) -- quadratic in sequence length (plus the AV step, also O(N^2)). O(N^2) memory: naively, storing the N x N attention matrix -- quadratic memory (the naive bottleneck FlashAttention addresses). Softmax + AV: the attention matrix is softmaxed (normalizing the attention weights per query) and multiplied by the values (AV -- the weighted sum of values per query -- the attention output) -- completing attention.
Middle row: the regimes and structures. FLOPs breakdown: attention is O(N^2) (in sequence) while the FFN is O(N) (linear in sequence, large constant -- the FFN processes each position independently) -- so at short context, the FFN dominates the FLOPs (the large constant), while at long context, attention dominates (the quadratic) -- the crossover depending on N and model dimensions. Memory-bound vs compute-bound: prefill (parallel over the sequence -- compute-bound) vs decode (one token at a time -- memory-bandwidth-bound, loading weights and KV cache) -- the two inference phases' different bottlenecks. KV cache: during decode, the keys and values of previous tokens are cached (so each new token attends to them without recomputing) -- the KV cache grows with context (memory proportional to context length x layers x heads), dominating decode memory (and limiting how much context fits). FlashAttention: computes attention in tiles (blocks) without materializing the full N x N matrix (fusing the operations, keeping intermediate results in fast on-chip memory) -- making attention memory O(N) not O(N^2) (a crucial optimization enabling long context and faster attention).
Bottom rows: long context and estimation. Long-context techniques: methods reducing the quadratic cost for very long sequences -- sparse attention (attending to a subset of positions, not all -- sub-quadratic), linear attention (approximating attention with linear complexity), sliding-window attention (each position attends to a local window -- linear) -- breaking the quadratic barrier for extreme lengths (at some approximation cost). Cost estimation: budgeting the compute (FLOPs -- attention plus FFN, for prefill and decode) and memory (weights plus KV cache) for a given model, context length, and workload -- estimating cost and capacity. The ops strip: context budgets (the context length as a cost/memory driver -- budgeting it, since it's quadratically expensive in attention and drives KV cache memory), memory planning (planning the memory -- weights plus KV cache, the KV cache growing with context and batch -- for the workload's context and concurrency), and kernel choice (using efficient kernels -- FlashAttention avoiding the O(N^2) memory, optimized kernels for the hardware -- for attention performance).
End-to-end flow
Trace the compute and memory of processing a long input. A model processes a long prompt (say 32K tokens). Prefill: the attention is O(N^2) in the sequence length -- for 32K tokens, the QK^T is a 32K x 32K interaction per head per layer (quadratic -- expensive, 4x the cost of 16K, 100x the cost of 3.2K) -- so the long prompt's prefill is dominated by the quadratic attention (attention dominates the FLOPs at this long context, versus the linear FFN). Naively, the attention would also need O(N^2) memory (a 32K x 32K matrix per head -- enormous) -- but FlashAttention avoids materializing it (computing attention in tiles, keeping the memory O(N)) -- so the long prompt is feasible (the quadratic memory would otherwise be prohibitive). The prefill is compute-bound (parallel over the 32K positions -- lots of computation, limited by compute). Then decode (generating the response): each token attends to all 32K+ previous via the KV cache (which holds the 32K tokens' keys/values -- large memory), and decode is memory-bound (loading the weights and the large KV cache per token). The compute (quadratic prefill attention) and memory (KV cache, FlashAttention avoiding O(N^2)) determined the cost and feasibility of the long context.
The KV-cache-memory and FlashAttention vignettes show the memory reality. A KV-cache case: serving many concurrent long-context conversations, the KV cache memory dominates -- each conversation's KV cache (proportional to its context length x layers x heads) is large, and many concurrent conversations' caches sum to a lot of memory (often the limiting factor for concurrency -- how many long-context conversations fit in memory). The team plans the memory (the KV cache growing with context and concurrency -- budgeting it, using techniques like KV cache quantization or paged attention to fit more) -- the KV cache being a primary memory concern for long-context serving. A FlashAttention case: without FlashAttention, the O(N^2) attention memory would make long context infeasible (a 32K x 32K matrix per head is huge); FlashAttention's tiled computation (O(N) memory, no full matrix materialization) makes it feasible (and faster -- better memory access patterns) -- the crucial optimization enabling efficient long-context attention. The team uses FlashAttention (or similar efficient kernels) -- essential for long context.
The long-context-technique and estimation vignettes complete it. A long-context case: for extreme context (hundreds of thousands of tokens), even FlashAttention's O(N) memory helps but the O(N^2) compute is still a wall -- so the team considers long-context techniques (sparse attention -- attending to a subset; sliding-window -- local attention; or other sub-quadratic methods) that reduce the quadratic compute (at some approximation cost) -- breaking the quadratic barrier for the extreme context. An estimation case: to budget for a workload, the team estimates the compute (attention O(N^2) plus FFN O(N) for prefill; the memory-bound decode) and memory (weights plus KV cache for the context and concurrency) -- projecting the cost and capacity (how much the long context costs, how many concurrent conversations fit) -- the cost estimation informing context budgets and capacity planning. The consolidated discipline the team documents: understand attention is O(N^2) (quadratic in sequence -- the fundamental cost constraint, making long context expensive), the FLOPs breakdown (attention vs FFN, dominating at long vs short context), the prefill-compute-bound / decode-memory-bound distinction (different phases, bottlenecks, optimizations), the KV cache dominating decode memory (and limiting concurrency), use FlashAttention (O(N) attention memory -- essential for long context), consider long-context techniques for extreme lengths (sub-quadratic), budget context (quadratically expensive, drives KV cache), plan memory (weights plus KV cache), and estimate cost -- because attention's quadratic compute and memory are the fundamental cost structure of transformers, determining LLM cost, latency, and long-context feasibility, and understanding them (and the optimizations -- FlashAttention, KV cache management, long-context techniques) is essential to reasoning about and optimizing LLM efficiency.