ACE Journal

Long-Context Attention Beyond a Million Tokens

Who this is for: ML engineers and systems-minded practitioners who already understand the standard transformer attention block and the KV cache, and who now have to make a long-context model actually run. The assumption is that you can read attention pseudocode and reason about memory on a single accelerator, but you have not yet had to choose between sparse, linear, and hybrid attention for a real workload. If you are staring at a one-million-token requirement and wondering why you cannot just turn up the context length, this is the layer that explains what breaks and what to reach for instead.

The wall you hit at one million tokens

Full self-attention compares every token to every other token. For a sequence of length N, that is N squared pairwise scores per attention head, per layer. The number is abstract until you put a real length into it. At N equal to one million, a single head computes on the order of a trillion score entries before any value projection is applied, and the score matrix itself, at float16 with a 128-dimensional head, would occupy hundreds of gigabytes for one layer alone. No single accelerator in 2026, not an H100 with 80 GB nor an H200 with 141 GB, can hold that.

This is not a small-constant problem you optimize away with a faster kernel. FlashAttention and its successors removed the need to ever write the full score matrix to high-bandwidth memory, which was a genuine and large win, but they did not change the quadratic count of multiply-accumulate operations [1]. The work is still there; it is just streamed through on-chip SRAM instead of spilled. Double the context and you quadruple the compute. That asymptote is the wall, and everything in this article is a different way of refusing to climb it.

The wall matters now because the workloads people want have grown into it: whole-codebase understanding, multi-hour agent transcripts, book-length retrieval-augmented generation, and long video token streams all want context lengths where N squared is not a budget you can pay. The field’s response has been to give up exact full attention in one of three directions, illustrated in Figure 1: keep the softmax but score only a sparse subset of the pairs, replace the softmax with a kernel that factors into a recurrence, or interleave a little exact attention with a lot of cheap attention.

flowchart TB
    Full[full self-attention<br/>O of N squared] --> Wall{N approaches<br/>one million?}
    Wall -->|cost too high| Branch[give up exact full attention]
    Branch --> Sparse[sparse attention<br/>compute a subset of pairs]
    Branch --> Linear[linear / state-space<br/>factor into a recurrence]
    Branch --> Hybrid[hybrid stack<br/>interleave full + cheap]
    Sparse --> Cost[O of N times window]
    Linear --> Cost2[O of N, fixed state]
    Hybrid --> Cost3[O of N plus periodic N squared]

Figure 1. The three escape routes from quadratic attention and the asymptotic cost each buys. Note that the hybrid branch does not eliminate the quadratic term, it rations it to a few layers, which is why it keeps a closer recall profile to full attention than the other two.

How the three families actually work

Sparse attention: attend to a chosen subset

Sparse attention keeps the softmax but restricts which token pairs are scored at all. The simplest and most widely deployed pattern is sliding-window attention, used in Mistral and its derivatives [2]: each token attends only to a fixed window of preceding tokens, which makes the per-layer cost linear in N for a fixed window size. The cost of locality is that a single layer can no longer form a direct edge between two far-apart tokens. Stacked layers recover long-range information indirectly, because a token at the top of the stack can see, through intermediate layers, a receptive field that grows with depth, but the effective reach is bounded by window size times the number of layers.

Pure locality loses too much, so practical sparse schemes add a small set of global positions. Longformer combined a sliding window with a handful of global tokens that attend to, and are attended by, the entire sequence [3]. Those global tokens, typically anchored on document starts or query positions, act as low-bandwidth long-range channels: the rest of the sequence operates locally and cheaply, while the global budget carries the information that has to travel far. For retrieval-heavy work where the needle is sparse in a long haystack, the size of that global budget is a critical hyperparameter, because it is the only direct path long-range information has.

Linear and state-space attention: factor into a recurrence

Linear attention attacks the cost from the other side. Standard attention computes softmax over query-key dot products, and the softmax is what forces the full pairwise matrix. Replace it with a kernel feature map and the computation re-associates: instead of scoring every query against every key, you accumulate a running key-value state and read each query against that state [4]. The result is linear in N and, crucially, requires a fixed-size state rather than a growing cache. RWKV, Mamba, and Griffin each implement a variant of this recurrence, with state-space models in particular framing the recurrence as a learned linear dynamical system [5][6][7].

The fixed state is the whole trick and the whole catch. Because the state does not grow with N, the model never materializes the full context, which is exactly why it scales. But a fixed-size state is a lossy summary of an arbitrarily long past, so a linear model cannot reliably recall an exact token from far back the way full attention can by simply attending to it. For tasks dominated by aggregate or recent context, that compression is acceptable and often invisible. For tasks that demand exact long-range lookup, it is the failure mode.

Hybrid stacks: ration the expensive layers

Hybrid architectures refuse the binary. Jamba and similar designs interleave full-attention layers at fixed intervals with linear or state-space layers in between [8]. Most of the depth runs cheap, while a minority of full-attention layers act as periodic anchors that can form the exact long-range edges the linear layers cannot. The bet is that you do not need every layer to have global exact attention, you need enough of them, placed well, to stitch long-range dependencies together. In practice this captures much of full attention’s recall at a fraction of the cost, which is why hybrids have become the pragmatic default for many long-context open-weight models in 2026.

Worked example: where the memory actually goes

Pseudocode hides the cost, so it helps to compute it directly. The following estimates the peak memory of the attention score matrix for exact attention, the quantity that blows up, as a function of sequence length.

def attn_score_bytes(seq_len, n_heads, bytes_per_elem=2):
    # The score matrix is [n_heads, seq_len, seq_len], one float per pair.
    # bytes_per_elem = 2 for float16/bf16.
    return n_heads * seq_len * seq_len * bytes_per_elem

for n in [8_192, 131_072, 1_000_000]:
    gib = attn_score_bytes(n, n_heads=32) / 1024**3
    print(f"{n:>9,} tokens -> {gib:>12,.1f} GiB scores (32 heads, fp16)")

Run it and the quadratic term speaks for itself: 8k tokens is a few gigabytes, 128k is already into the hundreds, and one million tokens is far past what any single device holds. This is the number FlashAttention stops you paying in HBM by never storing the full matrix, but the compute that produced those entries does not go away. The lesson is that the score matrix, not the parameters or the KV cache, is what you cannot afford to materialize at length.

Sliding-window attention changes the shape of that cost. Instead of every query seeing every key, each query sees only a window, so the score tensor is bounded by window size rather than by N:

def window_score_bytes(seq_len, window, n_heads, bytes_per_elem=2):
    # Each query attends to at most `window` keys, not all seq_len keys.
    effective = min(seq_len, window)
    return n_heads * seq_len * effective * bytes_per_elem

for n in [131_072, 1_000_000]:
    full = attn_score_bytes(n, 32) / 1024**3
    win = window_score_bytes(n, window=4_096, n_heads=32) / 1024**3
    print(f"{n:>9,} tokens -> full {full:>10,.1f} GiB | window {win:>8,.1f} GiB")

With a 4,096-token window the cost grows linearly in N, so a million tokens becomes a workload you can actually schedule. The trade is encoded right there in min(seq_len, window): once the sequence is longer than the window, the model has, in any single layer, simply stopped looking at the tokens outside it.

The exact-attention path that survives at length is chunked attention, where the sequence is split into segments processed sequentially with an accumulating KV cache. The accumulation is the part worth seeing concretely:

import torch

def chunked_kv_accumulate(chunks, attn_fn):
    # chunks: list of (q, k, v) per segment. Keys/values accumulate;
    # each new chunk's queries attend over all keys seen so far.
    k_cache, v_cache, outputs = [], [], []
    for q, k, v in chunks:
        k_cache.append(k)
        v_cache.append(v)
        K = torch.cat(k_cache, dim=-2)   # grows every chunk
        V = torch.cat(v_cache, dim=-2)
        outputs.append(attn_fn(q, K, V)) # exact attention, bounded q
    return torch.cat(outputs, dim=-2)

This keeps attention exact: every query still attends over every prior key. What it bounds is the query dimension per step, not the key dimension, so the KV cache, and the work each new chunk does against it, still grows with N. Chunked exact attention buys a higher ceiling, not an escape from the asymptote. On current single-device hardware that ceiling sits in the low hundreds of thousands of tokens per layer before throughput degrades to where a linear or sparse alternative wins on latency alone.

The hard part: positional signal collapses before the math does

The non-obvious lesson, the one you only learn by pushing a model past its training length, is that the binding constraint at the extreme is often not raw compute but the positional encoding and the attention distribution itself.

Rotary positional embeddings (RoPE) are the dominant scheme in transformer language models, encoding position as a rotation applied to query and key vectors [9]. Vanilla RoPE extrapolates poorly: a model trained at 32k tokens does not simply keep working at 128k, because the rotation frequencies it learned do not generalize to angles it never saw. YaRN (Yet another RoPE extensioN) addresses this by rescaling the rotation frequencies so the positional signal spreads across the extended range, and it has been used to push effective context from 32k toward 128k and beyond without full retraining [10]. It works well enough to be a standard tool.

But even YaRN does not save you at the million-token extreme. Needle-in-a-haystack evaluations consistently show retrieval accuracy degrading for information placed far from the end of the sequence, and the reason is structural rather than positional. As context grows, the softmax has to distribute a fixed amount of attention probability across more and more candidate positions, so attention-weight entropy rises and the signal for any one distant token thins out. Figure 2 traces the path from sequence length to that failure.

flowchart LR
    Long[sequence far beyond<br/>training length] --> RoPE[RoPE angles<br/>never seen in training]
    RoPE --> YaRN[YaRN rescales<br/>rotation frequencies]
    YaRN --> Extend[effective context extended<br/>32k toward 128k+]
    Extend --> Entropy[but attention entropy<br/>rises with N]
    Entropy --> Fade[distant-token signal<br/>thins out]
    Fade --> Miss[needle far from end<br/>missed]

Figure 2. Why extending context is necessary but not sufficient. Watch the split after YaRN: it fixes the positional-angle problem yet leaves the entropy problem, which is why an extended-context model can pass a short retrieval test and still fail to find a needle buried deep in a million tokens.

The practical consequence is that you cannot fix long-range recall purely by stretching positions. Real fixes are architectural, sparse or linear attention that change what competes for attention bandwidth, or come from training explicitly on very long sequences with tasks that reward accurate long-range recall so the model learns to allocate its limited attention to the right places. Treating context extension as a config flag is the most common way teams ship a model that benchmarks well and fails in production.

Integration: choosing for a real workload

In day-to-day model selection the choice is rarely “the best long-context architecture” in the abstract. It is “what does this workload demand of recall, and what can the hardware afford,” and those two questions usually decide it for you. Figure 3 maps the decision the way it tends to play out in practice.

flowchart TB
    Start[long-context workload] --> Recall{need exact recall<br/>of arbitrary tokens?}
    Recall -->|no, aggregate context| Lin[linear / state-space:<br/>cheapest, fixed state]
    Recall -->|yes, some| Len{how long is the<br/>context, really?}
    Len -->|under ~200k| Chunk[chunked exact attention<br/>on one device]
    Len -->|past 200k| Need{recall must be<br/>reliable far back?}
    Need -->|yes| Hyb[hybrid: cheap layers +<br/>periodic full attention]
    Need -->|locality is enough| Spar[sparse: window<br/>+ global tokens]

Figure 3. A pragmatic decision path for picking a long-context approach. The first split, exact recall or not, is the one that most often gets skipped, and skipping it is how teams end up with a linear model on a task that needs exact lookup.

A retrieval-augmented system that mostly needs to integrate a long but recency-weighted context can often run a linear or state-space model and never feel the lossy state. A code-understanding or legal-analysis workload that must pull an exact clause from deep in the document cannot, and pays either for chunked exact attention up to a few hundred thousand tokens or for a hybrid past that. Make this decision explicitly rather than defaulting to whatever the framework makes easy, because the failure modes are silent: a linear model on an exact-recall task does not error, it quietly returns the wrong answer for the inputs that needed the token it forgot.

The three families are less competitors than points on a recall-versus-cost frontier, and an honest comparison names what each gives up.

Sparse attention [2][3] keeps the softmax, which means its local behavior is identical to full attention and it inherits all the tooling and intuition of standard transformers. Its weakness is structural: long-range information must route through depth or through a small global-token budget, so tasks needing many independent long-range edges in a single layer are a poor fit. It is the right call when locality genuinely dominates and you want to stay close to a standard transformer.

Linear and state-space models [4][5][6][7] are the cheapest at the extreme, with a fixed state that makes million-token and streaming workloads tractable in a way the others are not. The price is exact recall: the compressed state cannot reproduce an arbitrary past token on demand. They shine for aggregate, streaming, and recency-weighted context and struggle on precise long-range lookup. The research direction narrowing this gap is selective state mechanisms that learn what to keep, but the fundamental compression remains.

Hybrid stacks [8] are the pragmatic middle and, in 2026, the default for general-purpose long-context models, because they recover most of full attention’s recall at a cost dominated by the cheap layers. The trade is complexity: you now tune the ratio and placement of full-attention layers, and a poorly placed full layer wastes the budget it was supposed to anchor. None of the three removes the underlying tension; they relocate it. The choice is which relocation your workload can tolerate.

Where it breaks

Long-context attention fails in recognizable ways, and most production disappointments are one of these.

Context extension treated as a config flag. Raising a model’s max position with YaRN or similar [10] makes it accept longer inputs, but acceptance is not comprehension. Without training that rewards long-range recall, the extended model passes short tests and then misses needles buried deep, because the entropy problem from Figure 2 is untouched. The fix is to validate at the real length with a real retrieval task, not to trust the configured maximum.

A linear model on an exact-recall task. The most expensive silent failure. A state-space or linear model [5][6] on a workload that needs precise long-range lookup does not error, it confidently returns answers that omit the token its fixed state compressed away. The fix is the first split in Figure 3: decide whether the task needs exact recall before you pick the architecture.

Sliding windows that quietly sever dependencies. A window sized for the common case [2] silently drops any dependency longer than window size times depth. If the workload occasionally needs a link across the whole document, a pure window will miss it without any signal that it did. The fix is global tokens [3] or a hybrid full-attention layer to carry the rare long edge.

Mistaking a kernel win for a complexity win. FlashAttention [1] makes exact attention dramatically more memory-efficient, which tempts teams to assume the quadratic problem is solved. It is not. The compute still grows as N squared, so a length that fits in memory thanks to FlashAttention can still be far too slow to serve. The fix is to budget compute, not just memory, when projecting cost to longer contexts.

Hybrid ratios tuned by guesswork. The full-attention-layer budget in a hybrid [8] is a real hyperparameter, and a placement that clusters the expensive layers poorly wastes the recall they were meant to provide. The fix is to treat the ratio and placement as something you measure against long-range tasks, not a number copied from another model with a different workload.

Benchmarking on the wrong haystack. Averaged perplexity over long sequences hides retrieval failure, because a model can predict the bulk of a long text well while being unable to find a specific fact in it. The fix is to evaluate with position-stratified needle-in-a-haystack tests that report accuracy as a function of where in the context the needle sits, not a single aggregate number.

Further reading

Start with the FlashAttention work [1] to understand precisely which part of the cost a fast kernel removes and which part it cannot, because that distinction underpins every cost argument in this article. For sparse attention, the Mistral report [2] and the Longformer paper [3] are the canonical sources for windowed attention and global tokens respectively. For the linear and state-space family, the linear-attention paper [4] establishes the kernel re-association, and the Mamba [5], RWKV [6], and Griffin [7] papers give three concrete and distinct realizations of the fixed-state recurrence. The Jamba report [8] is the clearest description of the interleaved-hybrid pattern. On positional encoding, the RoPE paper [9] and the YaRN paper [10] cover the rotary scheme and its extension, and together they explain both why context extension is possible and why it is not sufficient.

Written with AI assistance, editorially reviewed.

References

[1] T. Dao, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” https://arxiv.org/abs/2307.08691

[2] A. Q. Jiang et al., “Mistral 7B.” https://arxiv.org/abs/2310.06825

[3] I. Beltagy, M. E. Peters, and A. Cohan, “Longformer: The Long-Document Transformer.” https://arxiv.org/abs/2004.05150

[4] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret, “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.” https://arxiv.org/abs/2006.16236

[5] A. Gu and T. Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces.” https://arxiv.org/abs/2312.00752

[6] B. Peng et al., “RWKV: Reinventing RNNs for the Transformer Era.” https://arxiv.org/abs/2305.13048

[7] S. De et al., “Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models.” https://arxiv.org/abs/2402.19427

[8] O. Lieber et al., “Jamba: A Hybrid Transformer-Mamba Language Model.” https://arxiv.org/abs/2403.19887

[9] J. Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding.” https://arxiv.org/abs/2104.09864

[10] B. Peng, J. Quesnelle, H. Fan, and E. Shippole, “YaRN: Efficient Context Window Extension of Large Language Models.” https://arxiv.org/abs/2309.00071

[11] N. Kitaev, Ł. Kaiser, and A. Levskaya, “Reformer: The Efficient Transformer.” https://arxiv.org/abs/2001.04451

[12] G. Xiao et al., “Efficient Streaming Language Models with Attention Sinks.” https://arxiv.org/abs/2309.17453