Long-Context Attention Explained: Quadratic Cost, KV Cache and FlashAttention

The answer arrives late and costs much more than you expected. Long context attention, the step where each token looks back at all earlier tokens, grows with the square of the prompt length.

Picture a common setup. A team lets users paste a 300-page contract, about 128K tokens (1,31,072, roughly 1.3 lakh or 130,000), into a document chat, and asks one question. The setup is made up, but the maths below is real.

The same question on a 4K-token excerpt felt instant. On the full contract, the wait before the first word grows far more than 32 times, and each GPU can hold far fewer such chats at once. This guide explains both effects on one made-up model shape.

That shape is 36 layers, a model width of 4,608 numbers, 36 query heads of 128 numbers each, 6 key and value heads and 16-bit numbers. It matches no real model; it only keeps the arithmetic honest.

You need eight terms. Each one is used in every section below, so it is worth a minute here.

• Token: a piece of text, often part of a word, that the model reads as one unit. A context window is the most tokens a model accepts at once.
• Query, key and value: three short lists of numbers the model makes for every token. A query asks a question, keys say what each token offers, and values carry the content.
• Attention score: how well one token's query matches another token's key. Softmax then turns a row of scores into weights that add up to 1.
• Head: attention runs as several parallel copies called heads. Each head has its own smaller queries, keys and values, here 128 numbers long.
• FLOP: one floating-point operation, such as one multiply or one add. We count a multiply plus an add as 2 FLOPs, as the papers do. A matmul is a matrix multiplication.
• HBM (high bandwidth memory): the GPU's main memory, large but comparatively slow to read. SRAM is the small, very fast memory on the chip itself.
• KV cache: the keys and values of earlier tokens, stored so the model does not recompute them for every new token.
• Prefill and decode: prefill reads the whole prompt in one pass. Decode then makes the answer one token at a time.

Because attention compares every token with every other token. A prompt of N tokens makes an N by N table of scores in every head of every layer, so the work grows with N times N.

Attention is computed as softmax of Q times K-transposed, divided by the square root of the head size, then multiplied by V. The original transformer paper lists self-attention's cost per layer as order n squared times d, where d is the model width.

Here is an everyday picture. In a class, every student compares notes with every student, themselves included. Ten students make 100 comparisons; twenty students make 400, not 200.

Each student is a token, one comparison is one score, and the notes are the query and key. Where the picture breaks: a text-generating model lets each token look only backwards, so it needs about half the comparisons. Each score is also 128 multiply-adds, not one chat.

Pause and predict: what happens when you double the prompt?

Guess before reading on. Going from 32K to 64K tokens doubles the prompt, so does the work double?

No. On our made-up shape, attention work grows 4 times and the whole prefill grows 3.16 times. The weight matmuls, the parts that multiply each token by the model's fixed weights, only double.

A TFLOP is a trillion FLOPs. For short prompts, the weight matmuls dominate and attention is a small extra. Past about 23,800 tokens on this shape, attention does more work than all the weight matmuls together.

If the kernel skips the masked half, that crossover moves to about 47,600 tokens. The exact point depends on the model width, so treat it as a property of the shape, not a universal number.

The formula behind the attention column is 4 times N squared times width, per layer. The FlashAttention-2 paper uses the same count, 4 times sequence length squared times head size times number of heads.

Prefill is limited by raw compute, and decode is limited by memory reads. A 2023 study of inference serving describes the prompt phase as compute-intensive and token generation as memory-intensive.

In prefill, the GPU sees all N prompt tokens at once, so it does big matrix multiplications and keeps its maths units busy. This is where the N squared work from the table lands, and it sets the wait before the first answer token.

In decode, there is one new query per step. Its scores against the cached keys form one row of N numbers, not a square. But that row needs the whole KV cache read from HBM, on top of every model weight.

A GFLOP is a billion FLOPs. Per step, attention work grows in a straight line with the cache. Over a whole answer, though, every step is longer than the last.

Making a 1,000-token answer after a 4K prompt needs 3.0 TFLOP of attention in total. After a 128K prompt, the same answer needs 87.3 TFLOP, and each step also reads almost as many bytes of KV cache as weights.

It grows in a straight line: one fixed slice per token, per layer. For our shape that slice is 2 (a key and a value) times 36 layers times 6 heads times 128 numbers times 2 bytes.

That is 1,10,592 bytes, or 108 KiB per token (1 KiB is 1,024 bytes). A 4K chat needs 0.45 GB, a 32K chat 3.62 GB and a 128K chat 14.5 GB. At about 1,43,000 tokens, one chat's cache is as large as all the layer weights.

The 6 key and value heads matter. If all 36 heads kept their own keys and values, the cache would be 648 KiB per token and 87 GB for one 128K chat. The guide on grouped query attention explains that trade-off.

This is why long context limits how many users a GPU can serve. The cache for each chat normally sits in HBM next to the weights, so a few very long chats can fill a card that would hold dozens of short ones.

Because a plain implementation writes the N by N score table to HBM and reads it back, several times. The FlashAttention paper notes that steps such as softmax are memory-bound: moving the data takes longer than the maths.

At 32K tokens, the score table for one layer, across all 36 heads at 16-bit, would be 77.3 GB. At 128K it would be 1,237 GB. Even one head at 128K is 34.4 GB.

The FlashAttention paper gives the speeds that make this hurt. Its example data-centre GPU, the A100, has 40 to 80 GB of HBM at 1.5 to 2.0 TB/s. Its on-chip SRAM is 192 KB in each of 108 streaming multiprocessors, the chip's compute units, at about 19 TB/s.

So the fast memory is around ten times quicker but holds only kilobytes per unit. A plain implementation writes scores out, reads them for softmax, writes the weights out and reads them again to multiply by V.

The paper's Theorem 2 puts numbers on it. Standard attention needs on the order of N times d plus N squared HBM accesses, where d is the head size here.

It computes the same softmax of Q times K-transposed times V, in small tiles that stay in on-chip SRAM. The full score table is never written to HBM. The 2022 paper calls this an IO-aware exact attention algorithm.

IO-aware means it is designed around memory reads and writes, not just FLOPs. Exact means it is not an approximation: its Theorem 1 states it returns the same softmax of Q times K-transposed times V.

Tiling: work on blocks that fit in fast memory

Tiling means cutting Q, K and V into blocks of rows. The paper picks block sizes from the SRAM size and the head size, so one query block and one key and value block fit at once.

For each query block, the kernel streams every key and value block past it. Scores for one block pair live only in SRAM, and are thrown away once used.

The diagram and our code use the loop order of FlashAttention-2, with queries in the outer loop. The first paper looped the other way round and rescaled the output after every tile.

Online softmax: normalise without seeing the whole row

The hard part is softmax, which needs the maximum and the total of a whole row before any weight is final. Online softmax, described in a 2018 paper, keeps a running maximum and a running sum instead.

1. 1

   Keep three totals per query row

   m is the largest score seen so far, l is the sum of e raised to each score minus m, and acc is the weighted sum of values.

2. 2

   Score the next tile

   Compute the block of scores for this key tile and find the new maximum, the larger of m and this tile's biggest score.

3. 3

   Rescale the old totals

   Multiply l and acc by e raised to the old m minus the new m. This makes them as if the new maximum had been used from the start.

4. 4

   Add this tile's part

   Add this tile's exponentials to l, and this tile's exponentials times its values to acc. Then set m to the new maximum.

5. 5

   Divide once at the end

   After the last tile, the output row is acc divided by l. That is exactly the softmax-weighted sum of all values.

Subtracting the maximum is the usual safe-softmax trick: e raised to a large score would overflow, but e raised to a score minus the maximum never exceeds 1.

What FlashAttention does not change

It does not remove the quadratic compute. The paper's Theorem 1 still counts on the order of N squared times d FLOPs. Its Figure 2, forward plus backward at 1,024 tokens on an A100, shows 75.2 GFLOPs against 66.6 for standard attention.

Yet HBM reads and writes fell from 40.3 GB to 4.4 GB, and runtime from 41.7 ms to 7.3 ms. The extra FLOPs come from recomputation. In training, the backward pass rebuilds the score blocks from Q, K and small saved row statistics instead of storing the N by N table.

What changes is memory traffic and memory size. HBM accesses drop to the order of N squared times d squared divided by M, where M is the SRAM size. Extra memory becomes linear in N instead of quadratic.

It also does not shrink the KV cache. It reads the same keys and values, so decode still pays for every cached token.

Run both versions on the same random inputs and compare outputs. The NumPy code below is a teaching version of the idea for one head, not the real GPU kernel.

import numpy as np

def naive_attention(Q, K, V):
    S = Q @ K.T / np.sqrt(Q.shape[1])           # N x N score matrix, all in memory
    S = S - S.max(axis=1, keepdims=True)         # safe softmax: subtract row max
    P = np.exp(S)
    P = P / P.sum(axis=1, keepdims=True)
    return P @ V

def tiled_attention(Q, K, V, block=64):
    N, d = Q.shape
    scale = 1 / np.sqrt(d)
    O = np.zeros_like(Q)
    for i in range(0, N, block):                 # one tile of queries at a time
        q = Q[i:i+block]
        m = np.full(len(q), -np.inf)             # running max per query row
        l = np.zeros(len(q))                     # running sum of exp per row
        acc = np.zeros((len(q), d))              # running weighted sum of V
        for j in range(0, N, block):             # stream key/value tiles
            s = q @ K[j:j+block].T * scale       # only block x block scores
            m_new = np.maximum(m, s.max(axis=1))
            fix = np.exp(m - m_new)              # rescale old totals to new max
            p = np.exp(s - m_new[:, None])
            l = l * fix + p.sum(axis=1)
            acc = acc * fix[:, None] + p @ V[j:j+block]
            m = m_new
        O[i:i+block] = acc / l[:, None]          # divide once at the end
    return O

if __name__ == "__main__":
    rng = np.random.default_rng(0)
    N, d = 1024, 64
    Q, K, V = (rng.standard_normal((N, d)) for _ in range(3))
    a = naive_attention(Q, K, V)
    b = tiled_attention(Q, K, V, block=64)
    print("numpy", np.__version__)
    print("max abs difference:", np.abs(a - b).max())
    print("allclose:", np.allclose(a, b))
    print("largest score array: naive", (N, N), "tiled", (64, 64))

We ran it with numpy 2.3.4, pinned with pip install numpy==2.3.4. Expect output like this, with N of 1,024 tokens and a head size of 64. The last digits may differ on your machine.

numpy 2.3.4
max abs difference: 3.885780586188048e-16
allclose: True
largest score array: naive (1024, 1024) tiled (64, 64)

What just happened? Both functions returned the same 1,024 output rows, apart from rounding near 1e-16. The tiled one never held more than a 64 by 64 block of scores.

We also ran a few harder checks. Each one is in the test script that ships with this guide's claim ledger.

• On 1,000 tokens, block sizes 1, 7, 32, 64, 128 and 1,000 all matched within 1.3e-15, including sizes that do not divide 1,000. None was bit-identical, because the additions happen in a different order.
• In float32 (32-bit numbers), tiled and plain attention differed by 9.3e-09. Both differed from the float64 answer by about 1.4e-07, so rounding from the number type is far bigger than the tiling difference.
• With scores as large as 162 in float32, softmax without subtracting the maximum produced NaN (not a number). The online version stayed finite.
• PyTorch 2.12.1 scaled_dot_product_attention on a CPU matched the NumPy plain version within 4.4e-16 in float64, and within 8.9e-16 with is_causal=True.
• For N of 4,096 in float64, Python's memory tracer, tracemalloc, saw a peak of 402.8 MB for the plain version and 2.3 MB for the tiled one.

So exact means equal up to floating-point rounding, which depends on the order of operations. It does not mean every bit matches a plain implementation.

They kept the same exact maths and made the GPU kernels faster. The improvements are about splitting work across the chip, not about a new formula.

The FlashAttention-2 paper reports about 2 times the speed of the first version, reaching 50 to 73 percent of the A100's peak. It cuts non-matmul FLOPs, splits one head's work along the sequence across groups of GPU threads, and divides work inside each group differently.

Its causal version skips blocks where every key comes after every query. That is about half the blocks for long sequences, and the paper reports a 1.7 to 1.8 times speed-up from it.

The FlashAttention-3 paper targets the H100 and reports 1.5 to 2.0 times the speed of FlashAttention-2. It reaches up to 740 TFLOP/s in FP16 and close to 1.2 PFLOP/s in FP8, an 8-bit number format.

Using it often needs no new code. As of the PyTorch 2.12 documentation, scaled_dot_product_attention can pick a FlashAttention-2 kernel, a memory-efficient kernel or a plain C++ version. The docs say these optimised kernels may be used on CUDA GPUs.

One detail is undocumented as of 26 October 2026, so pin your version. On a CPU with PyTorch 2.12.1, the default call gave bit-identical output to forcing the flash backend, and differed from the plain math backend by up to 3.9e-16.

When only a small part of the text matters for the question. You pay prefill for every token and KV cache for every token, whether the model needs them or not.

On our made-up shape, prefill with causal skipping is 70.3 TFLOP for 4K tokens and 7,770.5 TFLOP for 128K. That is 32 times the tokens for about 110 times the work.

At an assumed sustained 400 TFLOP/s, a made-up round number, that is about 0.18 seconds against 19.4 seconds before the first answer token. Real numbers depend on the GPU, batch and kernels.

Longer is not always more accurate either. One study found models often did best when the relevant facts sat at the start or end of the input. Accuracy dropped when they sat in the middle, even for models built for long contexts.

Prefix reuse is real in open-source serving. The vLLM docs say automatic prefix caching reuses the KV cache of a shared start of the prompt. They also say it does not reduce decode time.

Separate compute from memory, and prefill from decode, then show one number for each. Interviewers want to see that you know which part grows with N squared and which grows with N.

1. 1

   State the problem

   Long prompts slow the first token, raise cost per request and reduce how many chats fit on one GPU.

2. 2

   Name the quadratic part

   Attention scores are N by N per head per layer. Prefill attention FLOPs are 4 times N squared times width per layer.

3. 3

   Name the linear part

   The KV cache is 2 times layers times key and value heads times head size times bytes, per token. Decode rereads it every step.

4. 4

   Explain the kernel fix

   FlashAttention tiles attention in on-chip SRAM with an online softmax. Same result, far fewer HBM reads and writes, memory linear in N.

5. 5

   Explain the system fixes

   Grouped key and value heads shrink the cache. Retrieval and prefix caching cut the tokens you prefill.

6. 6

   Close with a check

   Measure time to first token, time per output token and GPU memory per chat at 4K and at your longest prompt.