Long context

import math

blocks = [([1.0, 2.0], [10.0, 20.0]),   # (scores, values) for keys 0-1
          ([3.0, 1.0], [30.0, 40.0])]   # (scores, values) for keys 2-3

m, total, acc = float("-inf"), 0.0, 0.0
for scores, values in blocks:
    m_new = max(m, max(scores))
    rescale = math.exp(m - m_new)
    powers = [math.exp(s - m_new) for s in scores]
    total = total * rescale + sum(powers)
    acc = acc * rescale + sum(p * v for p, v in zip(powers, values))
    m = m_new
    print(f"max {m}, rescale {rescale:.3f}, total {total:.3f}, weighted sum {acc:.2f}")
print("tiled answer:", round(acc / total, 2))

scores, values = [1.0, 2.0, 3.0, 1.0], [10.0, 20.0, 30.0, 40.0]
powers = [math.exp(s) for s in scores]
print("full answer: ", round(sum(p * v for p, v in zip(powers, values)) / sum(powers), 2))

The same trick at full size

The code below compares the two ways for 1,000 tokens. Q, K and V are random queries, keys and values with 32 numbers each. attention_full builds all the scores at once. attention_tiled reads the keys 100 at a time and keeps the three running numbers. In the code they're called m, total and acc.

import numpy as np

rng = np.random.default_rng(0)
T, d, block = 1000, 32, 100
Q, K, V = (rng.normal(size=(T, d)) for _ in range(3))

def attention_full(Q, K, V):
    S = Q @ K.T / np.sqrt(d)                          # all T × T scores at once
    W = np.exp(S - S.max(axis=1, keepdims=True))
    return (W / W.sum(axis=1, keepdims=True)) @ V

def attention_tiled(Q, K, V, block):
    m = np.full((len(Q), 1), -np.inf)                 # running max score for each query
    total = np.zeros((len(Q), 1))                     # running sum of exp(score - m)
    acc = np.zeros((len(Q), V.shape[1]))              # running weighted sum of values
    for start in range(0, len(K), block):
        S = Q @ K[start:start + block].T / np.sqrt(d)    # only T × block scores
        m_new = np.maximum(m, S.max(axis=1, keepdims=True))
        rescale = np.exp(m - m_new)                       # correct the old sums to the new max
        P = np.exp(S - m_new)
        total = total * rescale + P.sum(axis=1, keepdims=True)
        acc = acc * rescale + P @ V[start:start + block]
        m = m_new
    return acc / total

full, tiled = attention_full(Q, K, V), attention_tiled(Q, K, V, block)
print("largest difference:", float(np.abs(full - tiled).max()))
print(f"scores held at once: {T * T:,} (full) vs {T * block:,} (tiled)")

The largest difference is a tiny number, below 0.000000000000001. That's rounding noise from the computer's arithmetic, not a real difference. The answer is the same attention. The tiled version never held more than 100,000 scores at once, a tenth of the full million.

FlashAttention (Dao et al., 2022) does this on the GPU. A GPU has a small, very fast memory on the chip and a large, slower main memory. FlashAttention splits the queries into tiles too, so each tile is worked out in the fast memory. The full grid of scores never touches the slow memory.

The result is exact and faster, and the memory it needs grows linearly with $T$. The amount of arithmetic is still quadratic, though. A prompt 10 times longer still takes roughly 100 times the attention work to prefill. (Prefill is the first pass that reads the whole prompt.)

Idea 2: don't look at everything (sliding windows)

Think of a long line of people passing news along. Each person only talks to the few people right next to them. Yet the news still travels down the whole line, a few places per round.

Sliding-window attention limits which keys each query sees. Each token looks only at itself and the tokens right before it, up to $w$ tokens in all. Here $w$ is the window size. The first picture below shows 8 tokens with a window of 3. Each row is a query and each column is a key. A ■ means "may look", and a · means "blocked".

import numpy as np

def sliding_window_mask(T, window):
    i, j = np.indices((T, T))
    return (j <= i) & (i - j < window)

for row in sliding_window_mask(8, 3):
    print(" ".join("■" if allowed else "·" for allowed in row))

for layers in [1, 2, 4, 32]:
    print(f"{layers:>2} layer{'s' if layers > 1 else ' '} of window 4,096: information can travel {layers * 4095:,} tokens back")

Each token now makes at most $w$ scores, not $T$. So the cost grows linearly with the length of the text.

Information still travels further than one window. Each layer reads what the layer before it already gathered. So news moves back one window per layer. A window of $w$ tokens includes the token itself, so one layer reaches $w - 1$ tokens back:

In plain words: each layer lets information hop back one window, minus one for the token itself. Multiply by the number of layers.

Take a window of 3 and 2 layers. In layer 1, token 7 sees tokens 5 and 6, and token 5 sees tokens 3 and 4. In layer 2, token 7 reads what token 5 gathered, so it reaches token 3. That's 2 × (3 − 1) = 4 tokens back. The code above does the same sum for a window of 4,096.

Mistral 7B used a 4,096-token window in 32 layers. In theory that reaches 131,040 tokens, about 131,000. In practice, information that hops through many layers fades, like a message in a game of telephone. So many recent models mix the two kinds: some layers with sliding windows, some with full attention.

Researchers have also built models that avoid quadratic attention entirely. One family is state-space models, and Mamba is the best known. Several recent models are hybrids of these layers and attention.

Idea 3: stretch the positions the model knows

From the positional encoding lesson, you know most models use rotary position embeddings (RoPE). They rotate each query and key by an angle that depends on its position. A model trained on 8,000-token documents has never seen position 100,000. It's like a ruler that only has marks up to 8,000.

Labs fix this by stretching the rotation frequencies. Long distances then land in the range the model already knows. The simplest version is called position interpolation. Say you want 32,000 tokens from a model trained on 8,000. Multiply every position by 8,000 ÷ 32,000, which is 0.25. Position 32,000 then looks like position 8,000. Refinements such as YaRN stretch some frequencies more than others.

After the stretch, the model trains briefly on long documents. Most long-context models are pretrained at a moderate length and extended this way at the end.

Long context in practice

A long window doesn't mean the model uses all of it equally well.

• Lost in the middle. A 2023 study (Liu et al.) tested facts placed at different points in a long prompt. Models used facts near the start and the end better than facts in the middle. Newer models do much better, but position still matters.
• Needles are easy; reasoning is hard. Finding one planted fact in a long document is now routine. This test is called "needle in a haystack". Tasks that combine many facts spread across the context are harder. Benchmarks such as NVIDIA's RULER (2024) found that many models' effective context was shorter than advertised.
• You pay per token. A 500,000-token prompt costs 500,000 input tokens every time, and it takes longer to prefill. Prompt caching helps if the start of the prompt repeats.

Long context or RAG?

RAG (retrieval-augmented generation) looks up the passages that matter and puts only those in the prompt. So which should you use?

Put everything in the prompt when the material fits in the window. It also suits work that uses the material as a whole, such as "review this contract" or "compare these ten reports". And it suits material that changes with each request.

Retrieve when the collection is much bigger than the window. Retrieve when only a few passages matter for each question. Retrieve when cost and speed matter because you send many requests. Many systems combine both. They retrieve generously, then let a long context hold everything they found.

For marketers: when you paste long material into a prompt, put the documents first. Put your question and instructions at the end. Ask the model to quote the passages it relied on. That layout tends to improve answers on long inputs, and the quotes make answers easy to check.