Self-Attention: The Mechanism That Reads Every Token Against Every Other

A team I worked with spent three weeks debugging why their summarization pipeline started producing incoherent output when documents exceeded about 40,000 tokens. The model's advertised context was 128K. The first instinct was chunking strategy. Then retrieval logic. Eventually someone ran a needle-in-haystack test and found that recall for information placed in the middle of the document dropped below 60% past 32K tokens. The fix was architectural: switch to retrieval-augmented generation over the document rather than stuffing it into context.

That degradation is not a bug in the model they were using. It is a structural consequence of how attention works at scale. Understanding the mechanism explains the failure — and it explains most of the expensive decisions you will encounter when building LLM systems.

What attention actually computes

Start with a single token — say, the word "bank" in the sentence "She deposited cash at the river bank." To understand what "bank" means here, you need to look at "river." Self-attention is the mechanism that makes that lookup possible, and it does it for every token simultaneously.

For each token at position i, three vectors are computed by multiplying the token's embedding by three learned weight matrices:

• Query (Q): "What am I looking for?"
• Key (K): "What do I offer to others looking?"
• Value (V): "What do I contribute when someone attends to me?"

The attention score between position i (querying) and position j (being attended to) is:

score(i, j) = dot(Q_i, K_j) / sqrt(d_k)

where d_k is the dimension of the Key vectors. The square-root scaling prevents the dot products from growing too large as d_k increases, which would push the softmax into a saturated regime where gradients vanish during training.

Those scores for position i, computed against all positions j = 1...N, are then softmaxed to produce a probability distribution. The output at position i is the weighted sum of all Value vectors, using those probabilities as weights:

output_i = sum_j( softmax(score(i, j)) × V_j )

The result at each position is a blend of information from across the entire context, with the blend weights determined dynamically by the input — not fixed by position or proximity. That "river" tells "bank" how to mean is computed fresh for every new sentence. There are no hardcoded relationships.

One thing the formula does not contain: order. Attention by itself is permutation-invariant — shuffle the input tokens and you get the same scores, just shuffled. Word order enters through positional encodings; in modern models, RoPE rotates the Q and K vectors by a position-dependent angle before the dot product, so scores depend on relative distance as well as content. If you implement the MultiHeadAttention module below and skip that step, your model has no idea which token came first. The embeddings and positional encoding article works through how RoPE does this.

sequenceDiagram
    participant T1 as "She"
    participant T2 as "deposited"
    participant T3 as "river"
    participant T4 as "bank"

    Note over T4: bank queries all tokens
    T4->>T1: score(Q_bank, K_she)
    T4->>T2: score(Q_bank, K_deposited)
    T4->>T3: score(Q_bank, K_river) ← HIGH
    T4->>T4: score(Q_bank, K_bank)
    Note over T4: softmax → high weight on "river"\noutput blends in river's Value vector

This is the critical departure from recurrent models. An RNN processing "She deposited cash at the river bank" compresses earlier tokens into a hidden state before it gets to "bank." Information degrades with distance. Self-attention has no such decay: the score between "bank" and "river" is computed as directly as the score between "bank" and the word immediately adjacent to it.

The N² problem

For a sequence of N tokens, computing the full attention matrix requires computing N² dot products. Storing that matrix in float16 requires:

Naive attention memory cost:
  N = 4,096 tokens (GPT-3.5 era context)
  Matrix: 4,096 × 4,096 × 2 bytes = ~33 MB  — per head, per layer

  N = 128,000 tokens (modern long-context)
  Matrix: 128,000 × 128,000 × 2 bytes = ~32 GB — per head, per layer

  A 32-layer model with 32 heads:
  32 GB × 32 heads × 32 layers = ~32 TB in GPU HBM
  (at N=128K, MHA, fp16 — completely infeasible)

This is why Flash Attention is not optional. It does not change the math. It reorders the computation into tiles that fit in the GPU's fast on-chip SRAM, so the N×N matrix is never written to high-bandwidth memory (HBM). Peak HBM usage drops from O(N²) to O(N). The same softmax scores are computed; they just live in SRAM briefly rather than being materialized as a giant matrix. Flash Attention v2 (2023) extended this with better work partitioning across GPU warps; v3 (2024) further improved FP8 throughput. Every serious serving system — vLLM, SGLang, TGI — uses it.

The compute still scales as O(N²). You cannot tile your way out of the number of multiply-accumulate operations. This is why inference costs rise sharply with context length, and why the KV cache adds a separate, persistent memory bill on top.

Multi-head attention: parallel specialization

A single attention operation might learn to track one type of relationship well. The insight of multi-head attention is to run H independent attention operations in parallel, each with its own Q, K, V projections, then concatenate and project the outputs:

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        # Separate projections per head — learned independently
        self.W_q = torch.nn.Linear(d_model, d_model, bias=False)
        self.W_k = torch.nn.Linear(d_model, d_model, bias=False)
        self.W_v = torch.nn.Linear(d_model, d_model, bias=False)
        self.W_o = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, mask=None):
        B, N, D = x.shape
        # Project and reshape into [B, H, N, d_k]
        Q = self.W_q(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        out, _ = scaled_dot_product_attention(Q, K, V, mask)
        # Merge heads and project back to d_model
        out = out.transpose(1, 2).contiguous().view(B, N, D)
        return self.W_o(out)

Each head learns a distinct projection of Q, K, and V. In practice, different heads demonstrably specialize on different relationships. Mechanistic interpretability research has identified heads that track subject-verb agreement, heads that resolve coreference (finding what "it" or "they" refers back to), and heads that copy specific tokens from earlier in the sequence ("induction heads"). These patterns are emergent — no label said "this head should do coreference." The model discovered these decompositions as useful ways to minimize prediction loss.

This has a practical implication: heads are not interchangeable. Pruning studies show that removing specific heads causes specific capability regressions. If you're working with a fine-tuned model that has degraded at a particular linguistic task, targeted head analysis is more informative than random probing.

{ "type": "attention", "sentence": "bank-river", "title": "Attention head specialization: hover any token to see what it attends to" }

Grouped-query attention: cutting the KV cache

In standard multi-head attention, the KV cache stores one Key and Value tensor per head per layer. A 70B model with 64 heads, at 128K context, in float16:

MHA KV cache at 128K context:
  Per layer: 2 tensors × 64 heads × 128,000 tokens × head_dim × 2 bytes
  Typical head_dim = 128
  = 2 × 64 × 128,000 × 128 × 2 = ~4.3 GB per layer

  80 layers: ~344 GB  — on top of 140 GB model weights
  Total: ~484 GB. A single H100 has 80 GB HBM.

That math is why you cannot naively run a 70B model at 128K context on a single GPU. The KV cache is the bottleneck, not the model weights.

Grouped-query attention (GQA) addresses this directly: instead of one independent K and V projection per head, K and V projections are shared across groups of Query heads. With H query heads and G groups, the KV cache shrinks by a factor of H/G.

flowchart TD
    subgraph MHA ["Multi-Head Attention (MHA)"]
        Q1["Q₁"] --> H1["Head 1 → KV₁"]
        Q2["Q₂"] --> H2["Head 2 → KV₂"]
        Q3["Q₃"] --> H3["Head 3 → KV₃"]
        Q4["Q₄"] --> H4["Head 4 → KV₄"]
    end
    subgraph GQA ["Grouped-Query Attention (GQA, G=2)"]
        Q1G["Q₁"] --> G1["Group 1\nKV shared"]
        Q2G["Q₂"] --> G1
        Q3G["Q₃"] --> G2["Group 2\nKV shared"]
        Q4G["Q₄"] --> G2
    end
    style G1 fill:#0e7490,color:#fff
    style G2 fill:#0e7490,color:#fff

At the extreme, G=1 (one KV pair shared by all heads) is multi-query attention (MQA). GQA with a moderate group size — Llama 3 8B uses 32 Q heads and 8 KV heads (4 Q per KV group); Llama 3 70B uses 64 Q heads and 8 KV heads (8 Q per KV group) — is the 2025-2026 standard. Quality loss relative to MHA is minimal on benchmarks; the memory savings are real and large. Qwen 3, Gemma 2, Mistral, and virtually every new open-weight model ships with GQA.

The KV cache is also why the context window article exists as its own topic. The memory math above is the direct reason that "a model supports 128K context" and "you can profitably run it at 128K context on your hardware" are different claims.

Sliding-window attention: bounding the quadratic

Between full attention and the hybrid architectures covered below sits a simpler mitigation, and after GQA and Flash it is the most widely deployed one: don't let every token attend to every other token. Sliding-window attention restricts each position to the previous W tokens, cutting compute from O(N²) to O(N·W). At N=128K with W=4,096 that is a ~31× reduction in attention FLOPs, and the KV cache for those layers stops growing once the window is full — it is capped at W tokens no matter how long the sequence gets.

The trade is direct long-range access. A token can no longer attend straight back to something 100K positions earlier; information has to propagate window-by-window across layers, which is lossier. That is why shipped models mix rather than commit: Mistral 7B used a 4K sliding window across all layers, while Gemma 2 alternates local and global layers 1:1 and Gemma 3 runs five local layers (1K window) per global layer — which is most of the reason Gemma 3's long-context KV cache is small. If your workload is mostly local reasoning over long inputs (code, logs), interleaved local-global layers keep nearly full quality at a fraction of the cost. If it is genuinely global retrieval across the whole context, the global layers carry the load — and those still pay N².

Causal masking: why decoders can only look backwards

In an encoder (like BERT), attention is bidirectional: every token can attend to every other token. That is fine for classification tasks where you have the whole sentence before you need to produce an output.

Autoregressive text generation requires a different constraint: at step t, the model must predict token t+1 using only tokens 1...t. If the attention matrix could see future tokens during training, the model would learn to cheat — it would just copy from the future rather than predicting. A causal mask zeros out (or sets to −∞ before softmax) all positions j > i in the attention matrix:

Causal mask for N=4 (1=attend, 0=masked):
  [1, 0, 0, 0]   token 1 sees only itself
  [1, 1, 0, 0]   token 2 sees tokens 1-2
  [1, 1, 1, 0]   token 3 sees tokens 1-3
  [1, 1, 1, 1]   token 4 sees all tokens

The lower-triangular structure means earlier tokens never have access to later ones. This is what makes the model trainable as a next-token predictor: at every position, it only has access to what a real system would have seen.

The flip side: during inference, you generate tokens left to right, and at each step you already have the Keys and Values for all previous tokens cached. You only need to compute the new token's Q, K, V and attend over the existing cache — you do not recompute the full attention for all previous tokens. That is what the KV cache actually is: the stored K and V tensors that avoid redundant computation.

Attention sinks: a production-relevant quirk

The StreamingLLM work (Xiao et al., 2023) named and documented what practitioners had been noticing empirically: specific tokens — the BOS (beginning-of-sequence) token and early punctuation like periods and commas — accumulate disproportionate attention mass in deep layers regardless of their semantic relevance.

The intuition is that softmax must produce a probability distribution that sums to 1.0. When no token in the context is particularly relevant to the query, the model needs somewhere to put the probability mass — and it learns to dump it on structurally predictable tokens that appear early in every sequence. These "sinks" absorb attention weight that contributes nothing semantically to the output.

The practical consequence is twofold. First, tokens positioned in the middle of long contexts sometimes fail to receive the attention weight their content warrants — the mass has already been allocated to sinks. Second, the sink positions themselves are not well-suited to carry semantic information (BOS is a formatting artifact), so attention on them is wasted compute. This is one structural reason why long-context retrieval degrades before the advertised limit — separate from the "lost in the middle" positional effect.

This phenomenon also explains why you sometimes see models give inexplicably high weight to the first token in their output even when it is uninformative. The model has trained to treat certain early positions as weight dumps, and that pattern persists at inference time.

flowchart LR
    subgraph context ["128K context window"]
        BOS["[BOS]\n← SINK"] --> C1["...content..."]
        C1 --> MID["...middle content\n(lower recall)..."]
        MID --> END["...end content\n(higher recall)"]
    end
    subgraph attention_mass ["Attention mass distribution (schematic)"]
        BOS2["[BOS]: high mass\n(sink)"]
        MID2["middle tokens:\ncompete for remaining mass"]
        END2["recent tokens:\nrecency bias"]
    end
    style BOS fill:#ff2e88,color:#111
    style BOS2 fill:#ff2e88,color:#111

What breaks

Context length vs. effective recall. Attention computes relevance correctly for the tokens it processes, but recall in practice degrades well before advertised limits. A model claiming 128K context may show reliable retrieval only to 32-64K, depending on how the RoPE position encoding was scaled and how much long-context data appeared during training. Needle-in-haystack benchmarks expose this, and the results vary substantially by model. The embeddings and positional encoding article covers why position encoding choices drive this — the point here is that attention itself is not the whole story.

The quadratic wall. Long-context inference does not scale linearly. Going from 32K to 128K quadruples the context length and increases attention compute ~16×. Prefill (processing the input) is where this hits hardest — it is O(N²) compute in a single forward pass. A 128K-token prefill on a modern model is seconds of GPU time on a single device. Systems like vLLM and SGLang handle this with chunked prefill and distributed attention, but the wall is real. Do not design a system that sends 100K-token prompts at high request volume and expect the same throughput as a 4K system.

Flash Attention is a kernel, not a magic override of quadratic compute. Engineers sometimes conflate "Flash Attention enabled" with "long context is cheap." Flash Attention removes the O(N²) memory problem. The O(N²) compute is still there. At N=128K, the attention computation requires ~16× more FLOPs than at N=32K. Memory no longer runs out; GPU time still scales quadratically.

GQA head count is a quality dial, not a binary. Reducing KV heads too aggressively does cost quality on tasks that require fine-grained context discrimination. For most instruction-following and generation tasks the degradation is within noise. For tasks requiring the model to distinguish many similar entities across a long document, it can matter. Check your eval on the specific task before assuming GQA with aggressive grouping is free.

Causal masking and encoder-style tasks. To be clear about what does not break: passing a document and asking a question at the end works fine — the question tokens, and every token the model generates, attend over the entire document. What breaks is using a decoder-only model's per-token hidden states in an encoder-like way: extracting embeddings for semantic similarity or classification. Under a causal mask, the representation of a token near the start of the document was computed without any right context — it never saw what came after it. Mean-pooling those states gives you a lopsided embedding; last-token pooling helps, but hangs the whole document's meaning on one position. For tasks where bidirectional encoding matters, use an encoder model, or a decoder-based embedding model that was adapted with the causal mask removed.

Hybrid architectures: where pure attention is losing ground

The quadratic cost of attention is not just an engineering nuisance — it is causing a structural shift in model architecture. By mid-2026, a viable alternative had emerged: hybrid models that use attention only for a fraction of layers, replacing the rest with linear or recurrent operations.

Jamba (AI21 Labs) interleaves one attention layer per eight-layer block — roughly 7:1 Mamba to attention — and IBM's Granite 4.0 runs about 9:1 Mamba-2 to attention. The pattern across shipped hybrids is consistent: most layers are SSM, with attention appearing every 4-10 layers. Mamba-3 (March 2026) claims 7× lower per-token cost versus equivalent transformers at 64K context. The theoretical reason is that state-space models process tokens in O(N) rather than O(N²), at the cost of compressing the past into a fixed-size hidden state rather than keeping the full KV cache.

This is not a replacement — attention layers in hybrid models still handle the tasks that benefit most from global context access. But the "one attention layer per block" assumption from the original transformer paper is being actively reconsidered for inference-heavy applications. If you are making architecture decisions for a custom model or evaluating which serving engines to support, the hybrid path is worth watching. For the next 12-18 months, though, pure-attention transformers will remain the dominant production reality.

Reading attention weights in production code

Most engineers interact with self-attention through API calls and never see weights. But when debugging a model's behavior — whether on an evaluation framework, a fine-tuned model, or a research checkpoint — being able to inspect attention weights is useful.

import torch
from transformers import AutoTokenizer, AutoModel

model_name = "meta-llama/Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# output_attentions requires a full model load, not quantized inference
model = AutoModel.from_pretrained(
    model_name,
    output_attentions=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

text = "She deposited cash at the river bank."
inputs = tokenizer(text, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

# outputs.attentions: tuple of (n_layers,) tensors
# Each tensor: [batch, n_heads, seq_len, seq_len]
attn = outputs.attentions  # tuple length = num layers

# Inspect layer 15, head 3, from "bank" (last real token)
layer_idx, head_idx = 15, 3
bank_position = inputs["input_ids"].shape[1] - 2  # exclude EOS
weights = attn[layer_idx][0, head_idx, bank_position, :]
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

for tok, w in zip(tokens, weights.tolist()):
    print(f"{tok:15s} {w:.4f}")

The shape [batch, n_heads, seq_len, seq_len] encodes the full story: rows are querying tokens, columns are keys. The value at [0, h, i, j] is the softmax weight that position i places on position j in head h. High values in expected columns (like "river" when querying "bank") confirm that the head is doing what you expect. Systematically high values on the BOS token across most rows is the attention sink pattern.

Note that output_attentions=True materalizes the full attention matrix — the thing Flash Attention was designed to avoid. For inspection on short sequences this is fine; on production-length sequences it will OOM.

The engineering judgment call

Self-attention is not a single choice — it is a space of tradeoffs between expressiveness, memory, and compute that you navigate through model selection and serving configuration.

For most teams consuming frontier API models, the relevant decisions are downstream: how much context to send, whether to use retrieval or full-context stuffing, and how to interpret degraded recall at long context as an architectural constraint rather than a model bug. The decoding and sampling article covers the next stage: once attention has computed contextual representations, how does the model turn those into actual token choices.

For teams running self-hosted models, the decisions become concrete:

• GQA is non-negotiable for anything running at long context or high batch sizes — the KV cache savings are too large to ignore.
• Flash Attention v2/v3 should be on by default in your serving engine. Verify it is enabled; do not assume.
• Context length × batch size is your primary memory planning equation. The KV cache scales linearly in batch size and quadratically in context — budget accordingly.
• Hybrid models (Jamba, Granite 4.0) are worth evaluating if you are serving at very long contexts (64K+) with latency constraints, since their per-token cost curve is better than pure transformers in that regime.

The intuition to carry into every LLM architecture decision is this: the power of self-attention is that it reads everything against everything. That is exactly what makes it expensive. Every design choice in this space — GQA, Flash Attention, hybrid layers, chunked prefill — is an attempt to keep the "read everything" expressiveness while paying less than the full N² price. The price is real and it keeps coming up in production.

For the broader picture of how attention fits into the full transformer block — residual stream, MLP, and layer normalization — see the transformer block article. For why that O(N²) cost has specific implications for the KV cache and long-context pricing, the context windows and KV cache article works through the memory math in full.