Attention is O(n²) — and Mamba's linear escape

At 3:07 a.m., a document assistant has one job: read an 800,000-token support archive and answer which refund policy applies to a customer in Germany. The old policy appears near the beginning. A correction appears 600,000 tokens later. The answer needs both.

On a short prompt, a transformer handles this beautifully. On a huge one, the GPU starts charging rent by the millisecond. Double the context does not double the attention work. It can quadruple it. At enough tokens, the model spends more effort deciding which tokens may matter than actually answering the question.

Here is the correction I wish more architecture diagrams made explicit: Mamba is not a transformer killer. Attention is not obsolete. The useful idea is to stop paying for perfect, global random access at every layer. Use a cheap streaming state for most of the work, then spend attention where exact random access earns its keep.

A practical design is a hybrid: many linear state-space layers, a smaller number of attention layers, and retrieval or local attention where neither should process the entire archive. Some long-context workloads may favor this mix, but whether it wins depends on task quality, latency, memory, and cost on matched evaluations.

Where the quadratic bill comes from

A standard transformer layer uses self-attention, a mechanism where each token forms a query and compares it with permitted keys. In bidirectional attention, every token can compare with every other token.

In decoder-only causal attention, the query at position t can compare only with position t and the positions before it. Future tokens are masked because they have not been generated yet.

The model first creates three representations:

• a query, which describes what a token is looking for;
• a key, which describes what each token offers;
• a value, which contains the information returned when a key is selected.

They are usually written as Q, K, and V. The attention scores come from QKᵀ.

If the sequence contains n tokens, the raw score layout has n × n positions. In causal attention, only the triangular half is valid: position 17 can compare with positions 1 through 17, not positions 18 onward.

The number of valid causal query-key pairs is n(n + 1) / 2, roughly n² / 2. A naive implementation may still allocate the full square buffer and mask its future entries.

For fixed model width, the sequence-interaction part of the computation is therefore O(n²). Bidirectional attention has n² valid pairs; causal attention has roughly half as many. The constant changes, but the quadratic growth does not.

The model also performs projections whose cost grows roughly linearly with n. At long enough sequences, the global interaction dominates.

For one head at batch size one, these are the raw score slots in a naive full-matrix materialisation:

• A 2,000-token sequence has 4 million score slots.
• An 8,000-token sequence has 64 million.
• A 32,000-token sequence has 1.024 billion.
• A 1-million-token sequence has 1 trillion.

A causal mask makes only the triangular portion valid. For example, 8,000 causal queries have 8,000 × 8,001 / 2, or 32,004,000, valid pairs, even though a naive square buffer has 64 million slots.

Suppose a model has 32 heads and stores every score in 16-bit floating point. For batch one and naive full-matrix materialisation, at 8,000 tokens one layer would need space for 32 × 8,000², or 2.048 billion scores.

At 2 bytes per score, that is about 4.1 GB for the score tensor alone. At 32,000 tokens, the same tensor would be about 65.5 GB. At 1 million tokens, it would be 64 TB.

Those figures describe naive materialisation. Modern kernels such as FlashAttention tile the calculation so the entire n × n matrix does not sit in GPU memory at once. That is an enormous memory improvement.

It does not make the exact calculation linear, though. The model still has to account for the same allowed pairwise relationships; it simply does so in smaller blocks.

This distinction matters because “the context window is 1 million tokens” is a capacity statement, not a promise that a million-token prompt is cheap, fast, or equally reliable.

The deeper explanation is simple. Dense attention offers random access. In bidirectional attention, every token can reach directly to every other token. In causal attention, each token can reach every earlier token and itself.

That is still random access to the entire available history, rather like having an addressable record for every item in the archive. It is also expensive to scan at every layer.

For a gentle introduction to the underlying mechanism, self-attention is the useful foundation. The important systems question is what happens when that mechanism is repeated across a very long sequence.

One subtlety: prefill is not decode

People often say “attention is quadratic” as if every generated token requires a fresh square-shaped calculation. That is not quite right.

During prefill, the model reads the original prompt. In a typical decoder-only LLM, the query at position t attends to itself and the preceding prompt positions.

Across all n prompt queries, the valid set is triangular, with n(n + 1) / 2 pairs. A dense implementation may still use an n × n score buffer, but its future entries are masked.

During decode, the model generates one token at a time. A new token has one query, which compares against the keys cached for preceding tokens and its own newly computed key. For that single step, the work grows linearly with the current context length, not quadratically.

If the prompt has n tokens and the model generates m new tokens, the decode attention work is roughly proportional to m × n plus the interactions among the generated tokens, approximately m(m + 1) / 2.

The prompt still matters greatly, especially when n is large, but the two phases have different bottlenecks.

The KV cache, the stored keys and values used during decoding, also grows linearly with the number of tokens. It is not an n × n object.

This is why a long conversation can run out of memory even when the attention score matrix is handled efficiently: the model is keeping a key and value representation for every layer and every previous token.

That separation changes engineering decisions. If users send 4,000-token prompts but request 500 generated tokens, prefill may be the main concern.

If an agent generates thousands of tool-call steps while carrying the whole transcript, decode and KV-cache growth may dominate. One architecture can look excellent in one phase and poor in the other.

What Mamba changes

The linear recurrence

A state-space model, or SSM, processes a sequence by updating a running hidden state rather than comparing every position with every other position. The state is a compact summary of what has arrived so far.

A simplified recurrence looks like this:

h_t = A_t h_(t-1) + B_t x_t

y_t = C_t h_t + D x_t

Here, x_t is the input at position t, h_t is the state after reading it, and y_t is the output. The matrices control what information enters the state, how previous information changes, and what information is read out.

Mamba uses a selective SSM. “Selective” means that parts of this update depend on the current input, allowing the layer to decide what to retain or discard instead of applying one fixed transition to every token.

A token can influence the state strongly when it looks important and weakly when it looks like routine filler.

The key computational fact is that the state has a size determined by the model width and its state dimension, not by the number of tokens. Reading token 800,001 requires one more state update. It does not require comparing that token with all 800,000 previous tokens.

For fixed width and state size, the sequence-dependent work is therefore O(n). Double the sequence and you roughly double the number of updates.

From 8,000 to 32,000 tokens, the quadratic term grows by 16 times while a linear term grows by 4 times. From 8,000 to 1 million tokens, those factors are 15,625 and 125 respectively.

That is the scaling advantage, not a feasibility guarantee. Throughput, kernel efficiency, memory movement, activation memory, the context length used during training, wall-clock latency, and task quality still determine whether a million-token workload is practical.

Exact-recall quality also needs to be tested rather than inferred from the asymptotic notation.

Parallel scans and memory shape

It is worth killing one misleading picture before it spreads: Mamba is not necessarily processing one token in a painfully serial Python loop.

During training and prompt prefill, parallel-scan algorithms can parallelise the recurrence across positions. The straightforward recurrence still has a dependency from one position to the next.

During incremental generation of a single online sequence, that dependency remains. The next state cannot be updated until the previous state exists.

Standard training normally retains O(n) activations for backpropagation or recomputes them. The compact recurrent state is primarily an inference-cache advantage.

The implementation still needs good kernels and sensible memory movement. Linear work and a compact inference state do not guarantee lower latency.

Mamba also has a different memory shape. Instead of retaining an addressable key and value representation for every earlier token, it retains the current state for each layer.

This can make the inference cache for a long stream much smaller, but it does not guarantee lower end-to-end latency. It also creates the central trade-off.

A summary is not an address book

Return to the support archive. Near token 40,000, an old policy says refunds are available for 30 days. Near token 640,000, a legal update says the policy is now 14 days for a particular region.

An attention layer can compare a question with both passages directly. It can preserve the distinction between the old wording and the new wording because the source positions remain available to the layer.

A state-space layer has to carry useful information from those passages through its running state. It may learn to preserve the latest policy, a date, a location, or a contradiction flag.

That is often enough for normal language patterns. It is not the same as having a perfectly addressable record of every sentence.

In a finite-precision deployed model, the state is a lossy compression channel. As more unrelated material passes through it, information competes for representation.

The model may remember the gist of a 10-page policy while losing the exact exception in paragraph 37.

That is why a pure SSM can be excellent at streaming text, local dependencies, repeated patterns, and long sequences where a useful summary is enough.

Depending on the training and task, it may be less reliable for arbitrary-position retrieval, exact copying, long-range symbolic references, and “find the one sentence that contradicts this other sentence” tasks.

This is not a flaw that better prompting fixes. If the information was compressed away, a prompt cannot point the model back to a token it no longer has in an addressable form.

You need an attention path, external retrieval, or another memory mechanism.

The hybrid answer is a budget, not magic

A hybrid model interleaves Mamba-style layers with transformer attention layers. The linear layers process the bulk of the sequence with roughly linear sequence work. The attention layers provide occasional global communication and precise lookup.

The mental model is a memory hierarchy:

• State-space layers are the large, cheap, continuously updated summary.
• Attention is the smaller, expensive random-access memory.
• Retrieval is a separate index that decides which outside material deserves to enter the prompt at all.

The hybrid can work because information does not need to be globally compared at every layer to be globally useful. A sequence can be swept by several SSM layers, building representations that capture local structure and running context.

An attention layer can then connect distant positions when the task needs it. Later linear layers can propagate the result without reopening every pair of positions.

Jamba is a well-known example of this design direction. It combines Mamba-style layers with transformer attention and uses mixture-of-experts layers for parameter efficiency.

The exact mixture is architecture-specific. There is no universal “three Mamba layers, then one attention layer” rule waiting to be discovered.

There is also an important mathematical footnote. If a hybrid contains even one full-sequence dense attention layer, its strict asymptotic complexity still contains an O(n²) term.

Calling the whole model “linear” would be incorrect. The practical gain comes from reducing how many layers pay that price.

If the attention layers use a fixed local window, chunking, or another restricted pattern, the scaling can become linear in sequence length. The model has then changed the access pattern rather than preserving unrestricted global attention.

So “linear-ish” should mean one of two things:

1. Most layers scale linearly, making the quadratic coefficient much smaller in realistic ranges.
2. The remaining attention is restricted enough that its own cost scales with a fixed window or a bounded number of retrieved tokens.

That precision is not pedantry. It tells you what to inspect in a model card and a profiler.

A model advertised as hybrid may still contain a full global attention layer that determines its million-token limit.

The strongest objection: optimize the transformer instead

The fair objection is that attention has a huge head start. FlashAttention reduces memory pressure. Local and sparse attention reduce the number of comparisons. Retrieval can shrink an 800,000-token archive to 12,000 relevant tokens.

Transformers have mature training recipes, fine-tuning tools, inference servers, and a large ecosystem of tested weights.

For many workloads, that objection wins.

A plain transformer is usually the sensible choice for a 4,000-token support chat or a 16,000-token RAG prompt. The quality is familiar, the tooling is good, and the theoretical long-context problem may never appear in production.

If RAG reliably narrows the document set, changing the sequence layer may be solving the wrong problem.

But each optimisation gives something up or leaves something behind:

• FlashAttention changes how memory is used; it does not remove pairwise arithmetic.
• Local attention is cheaper because it prevents some positions from seeing others.
• Retrieval is cheaper because it bets that the index can identify the right evidence.
• Summarisation is cheaper because it discards detail.

Those are all good bets when their assumptions hold. They are not equivalent to unrestricted random access.

The case for testing a hybrid becomes stronger when the traffic distribution contains a costly tail: prompts of 100,000, 500,000, or 1 million tokens that are common enough to affect GPU capacity, queueing, and cost.

A context window that works in a demo but makes the p99 request five times slower is not a production feature. It is an unusually polite denial of service.

What to do on Monday morning

Start with the workload, not the architecture diagram.

First, export the token-length distribution for real requests.

• Record prompt tokens and generated tokens separately.
• Measure prefill latency, decode tokens per second, peak GPU memory, and cost per request at p50, p95, and p99.

Averages hide the request that arrives with a 700,000-token transcript and occupies the GPU while everyone waits.

Second, build a small evaluation set around the actual failure mode. Include exact retrieval with distractors, contradictions between early and late documents, cross-file code references, long lists that must be copied accurately, and ordinary short prompts.

A model that produces a good summary can still fail the one question your customer cares about. Include the 800,000-token policy example if that is your product’s shape.

Third, establish a transformer baseline with the serving optimisations you would actually deploy. Test the realistic lengths, not only a short benchmark.

If a 32,000-token prompt is already rare and affordable, a Mamba migration may create more training and operational risk than value. If the p99 tail is dominated by very long prefill, the case is different.

The broader long-context design matters as much as the layer type.

Fourth, reduce unnecessary context before changing the model.

• Remove duplicate conversation turns.
• Retrieve only relevant documents.
• Keep tool outputs structured.

A model should not need a million tokens because an agent forgot to delete yesterday’s 200-page log. Use LLM cost and latency measurements to price the whole request, not just the model’s advertised token rate.

Fifth, compare a candidate hybrid or SSM model on the same prompts and the same quality checks. Record:

• wall-clock prefill and decode latency;
• peak memory;
• throughput;
• cost;
• exact-recall quality.

Check the context length used during training; accepting a long input is not evidence that the model learned to use it.

Do not assume you can replace attention layers in an existing transformer after the fact. The model must be trained for its recurrence, state size, and information flow.

This is an architecture migration, not a configuration toggle.

Finally, treat state as request-scoped data. Reset it between users and independent conversations. Preserve it only when a request is a genuine continuation.

If you checkpoint or resume a generation, checkpoint the corresponding model state as well. A small hidden state is cheap; accidentally carrying one customer’s conversation into another customer’s request is not.

Where the clever plan breaks

The first failure mode is a fluent answer with the wrong old fact. The model sounds confident, average loss looks normal, and only a targeted long-range retrieval test reveals that the late policy update disappeared from the running state.

Add exact-recall and contradiction tests; do not rely on perplexity alone.

The second is a latency cliff. Requests look healthy at 32,000 tokens and begin timing out at 128,000.

A profiler shows a full attention layer or an unexpectedly large KV cache. The fix may be local attention, retrieval, or a different layer mix.

“It is a hybrid” is not enough information.

The third is a scar at every chunk boundary. Teams process a long stream in pieces, reset the SSM state for each piece, and see quality drop at regular intervals.

The model is not broken; the application erased its memory at every boundary. Carry state across chunks that belong to one sequence, while keeping different sequences strictly isolated.

Use a transformer when the context is modest, arbitrary recall is central, and the ecosystem advantage matters more than long-stream efficiency.

Use retrieval when the main problem is selecting a small evidence set. Use an SSM or hybrid when the sequence is genuinely long, streaming, and expensive enough that global access is the dominant bill.

The right architecture for this workload is not the one that abolishes attention. It is the one that refuses to pay for global random access on every layer, every token, and every request.