Implementing the Attention Methods Used in LLMs: From Multi-Head Attention to Gated DeltaNet
MHA, MQA, GQA, sliding window, attention sinks, MLA with weight absorption, DSA-style sparse attention and Gated DeltaNet, each as a short PyTorch function with a check that it does what it claims. The MLA check runs on DeepSeek-V2-Lite's real layer-0 weights; the Qwen3.5-9B state size from the code matches llama.cpp's 50.25 MiB.
Implementing the Attention Methods Used in LLMs: From Multi-Head Attention to Gated DeltaNet
DeepSeek-V2-Lite stores 576 numbers per token per layer. Its 16 attention heads need, between their keys and values, 5,120. When we wrote multi-head latent attention the way its equations read, the first thing the code did was rebuild those 5,120 numbers for every cached token, at every step. That is correct, and it throws away the point: the memory is small, but each step still reads and builds the large version.
The fix is three lines that move one weight matrix from the key side to the query side. We ran both versions on DeepSeek-V2-Lite's real layer-0 weights and got the same output, to within 4.5 × 10⁻¹⁵. Only the second one never touches a per-head key.
That is the kind of thing this post is for. The first post in this series described seven attention designs in prose: what each stores, which models use it, what it gives up. Here each one is a short PyTorch function, followed by a check that the function does what the prose says: that a sliding cache of W tokens gives the same answer as masking, that grouped-query attention is multi-head attention with shared weights, that linear attention's running state equals its attention-shaped form. All checks run on a CPU in under a second, in float64, and every one passes. The code is written for reading, not speed; real kernels look nothing like it.
One helper for everything
Every design below ends in the same operation, so it gets one function. q, k and v are per head: (heads, tokens, head size). The mask says which tokens each query may look at; causal builds the usual "only the past" mask, optionally cut to a window. Ignore sink until section 3.
def attend(q, k, v, mask=None, sink=None):
"""q: (H, Tq, d), k and v: (H, Tk, d). mask: (Tq, Tk) bool, True = may look. sink: (H,) logits or None."""
s = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
if mask is not None:
s = s.masked_fill(~mask, float("-inf"))
if sink is not None: # one extra logit per head with no value behind it
s = torch.cat([s, sink.view(-1, 1, 1).expand(-1, s.shape[1], 1)], dim=-1)
return s.softmax(-1)[..., :-1] @ v
return s.softmax(-1) @ v
def causal(Tq, Tk, window=None):
i = torch.arange(Tk - Tq, Tk)[:, None] # positions of the new queries
j = torch.arange(Tk)[None, :]
m = j <= i
if window is not None:
m &= j > i - window
return m1. MHA, MQA and GQA: one class, three settings
Multi-head attention gives every head its own key and value. Multi-query attention gives all heads one shared key and value. Grouped-query attention sits between: heads in a group share one. In code they differ in one number, how many key-value heads are projected and stored:
class GQA:
"""n_kv == n_heads is MHA, n_kv == 1 is MQA, anything in between is GQA."""
def __init__(self, dim, n_heads, n_kv, head_dim):
assert n_heads % n_kv == 0
self.h, self.kv, self.d = n_heads, n_kv, head_dim
self.wq = torch.randn(dim, n_heads * head_dim) / dim**0.5
self.wk = torch.randn(dim, n_kv * head_dim) / dim**0.5
self.wv = torch.randn(dim, n_kv * head_dim) / dim**0.5
def __call__(self, x, cache): # x: (T_new, dim); cache: dict with "k", "v" of shape (n_kv, T_past, d)
T = x.shape[0]
q = (x @ self.wq).view(T, self.h, self.d).transpose(0, 1)
k = (x @ self.wk).view(T, self.kv, self.d).transpose(0, 1)
v = (x @ self.wv).view(T, self.kv, self.d).transpose(0, 1)
cache["k"] = torch.cat([cache["k"], k], dim=1) if "k" in cache else k # only n_kv heads are stored
cache["v"] = torch.cat([cache["v"], v], dim=1) if "v" in cache else v
g = self.h // self.kv # query heads per stored head
K = cache["k"].repeat_interleave(g, dim=0) # (n_heads, T_all, d), built on the fly, never stored
V = cache["v"].repeat_interleave(g, dim=0)
return attend(q, K, V, causal(T, K.shape[1]))The line that matters is repeat_interleave. The cache holds n_kv heads; the copies that line up with the query heads are made on the fly and never stored.
Two checks. First, for all three settings, feeding a 40-token sequence one token at a time with the cache gives the same output as feeding it all at once (largest difference 1.1 × 10⁻¹⁵). Second, GQA with 2 key-value heads gives the same output as MHA with 8, if the MHA copies each group's key and value weights to its 4 heads (6.7 × 10⁻¹⁶). That is what GQA is: MHA whose heads have been made to agree, in groups, on what to store. The cache shrinks by the group size. With 8 heads of 16, one token costs 256 numbers per layer in MHA, 64 with 2 groups, 32 with one.
For a real model, multiply by bytes and layers. Llama 3.1 8B: 2 × 8 KV heads × 128 × 2 bytes × 32 layers = 128 KiB per token.
2. Sliding window: a ring buffer
A sliding-window layer only looks at the last W tokens, so its cache only needs W slots. When a new token arrives, it overwrites the oldest:
class SlidingCache:
"""Keeps only the last `window` keys and values. Memory stops growing at the window."""
def __init__(self, n_kv, head_dim, window):
self.k = torch.zeros(n_kv, window, head_dim)
self.v = torch.zeros(n_kv, window, head_dim)
self.w, self.n = window, 0 # n = tokens seen so far
def add(self, k, v): # k, v: (n_kv, 1, d) for one new token
slot = self.n % self.w # overwrite the oldest token
self.k[:, slot], self.v[:, slot] = k[:, 0], v[:, 0]
self.n += 1
filled = min(self.n, self.w)
return self.k[:, :filled], self.v[:, :filled] # order does not matter to softmaxThe slots are out of order after the first wrap, and that is fine: softmax over a set of keys does not care about their order. (Position is already inside the keys, through RoPE, before they reach the cache.)
The check: with a window of 8 and 50 tokens, the ring buffer gives the same output at every step as full attention with a window mask (8.9 × 10⁻¹⁶), while holding 8 tokens instead of 50. A second check shows what the layer gives up. We changed the value of token 0 by 100. From step 8 on, the output did not move at all: difference exactly 0. That token is gone from this layer.
Models that use sliding windows mix them with full layers so that something can still see token 0: Gemma 3 has 5 sliding layers for each full one, gpt-oss alternates them one to one. In code that is just a list of layer kinds, each layer with its own cache class. Runtimes do not always keep exactly W slots; llama.cpp and vLLM both hold extra room next to the window, by different amounts (the deployment post has the measurements).
3. Attention sinks: a place to put attention
Softmax always hands out a full 100% of attention. A head that has nothing useful to look at still has to spread it over something. gpt-oss gives each head one extra learned logit, a "sink", that takes part in the softmax but has no value behind it. That is the sink branch in attend: append the logit as one more column, take the softmax, drop that column before multiplying by the values.
Two checks. With the sink at minus infinity, the output is exactly plain attention. With the sink logit at 20, far above the scores, the output's size falls to 3.8 × 10⁻⁸ of what it was: the head gives almost all its attention to nothing and contributes almost nothing. That is the option a sink buys.
The real weights have them. gpt-oss-20b's layer 0 stores a tensor self_attn.sinks of 64 numbers, one per query head, ranging from −2.45 to 4.06. The sinks cost no cache; they are weights, not per-token state. StreamingLLM's attention sinks are a different, earlier idea for the same problem: keep the first few real tokens in the cache forever.
4. MLA: store the latent, fold the weights
Multi-head latent attention stores, per token, one compressed vector c_kv (512 numbers in DeepSeek-V2-Lite) and one small key k_rope (64 numbers) shared by all heads. Each head's key and value are linear functions of c_kv: c_kv @ w_uk[h] and c_kv @ w_uv[h]. Here is the version that does what that sentence says, and the version that does the same thing without building them:
def mla_naive(q_nope, q_rope, c_kv, k_rope, w_uk, w_uv):
"""Rebuild every head's key and value from the latent, then attend as usual.
q_nope: (H, T, dn) q_rope: (H, T, dr) c_kv: (L, r) k_rope: (L, dr), one per token shared by all heads
w_uk: (H, r, dn) w_uv: (H, r, dv)"""
H, L = q_nope.shape[0], c_kv.shape[0]
k = torch.cat([c_kv @ w_uk, k_rope.expand(H, L, -1)], dim=-1) # (H, L, dn + dr)
v = c_kv @ w_uv # (H, L, dv)
q = torch.cat([q_nope, q_rope], dim=-1)
return attend(q, k, v, causal(q.shape[1], L))
def mla_absorbed(q_nope, q_rope, c_kv, k_rope, w_uk, w_uv):
"""Same result without ever building per-head keys or values: fold w_uk into the query
and w_uv into the output. Only c_kv and k_rope are read from the cache."""
scale = 1 / math.sqrt(q_nope.shape[-1] + q_rope.shape[-1])
q_lat = q_nope @ w_uk.transpose(-1, -2) # (H, T, r): query moved into latent space
s = (q_lat @ c_kv.T + q_rope @ k_rope.T) * scale # (H, T, L)
s = s.masked_fill(~causal(q_nope.shape[1], c_kv.shape[0]), float("-inf"))
return (s.softmax(-1) @ c_kv) @ w_uv # attend over latents, then expand once per headThe algebra is one step. A head's score for a cached token is q_nope · (c_kv @ w_uk), which is the same number as (q_nope @ w_ukᵀ) · c_kv. So multiply the query by w_ukᵀ once per step, and compare it with the stored latents directly. On the value side, attention is a weighted average, and averaging then multiplying by w_uv equals multiplying then averaging. The absorbed version reads 576 numbers per cached token and builds nothing per token.
The 64-number k_rope is why MLA has two parts. RoPE rotates a key by an angle that depends on its position, so a rotated key is not a fixed linear function of c_kv and the fold above would not go through. DeepSeek keeps a small separate key for position, shared by every head, and stores it next to the latent. That is the "+ 64" in 512 + 64.
Checks: the two functions agree on random weights (1.1 × 10⁻¹⁵), and on DeepSeek-V2-Lite's real kv_b_proj from layer 0, a 4,096 × 512 matrix that we split into the 16 heads' w_uk and w_uv the way the model code does (4.5 × 10⁻¹⁵). The equality is algebra, so it would hold for any weights; what the real-weight run confirms is the layout: 16 heads, a 128-wide key part and a 128-wide value part per head, all read out of 512 latent numbers. Stored as rebuilt heads, the same token would take 16 × (128 + 64 + 128) = 5,120 numbers per layer instead of 576.
5. Sparse attention: keep everything, read a little
DeepSeek Sparse Attention, in DeepSeek-V3.2-Exp, does not shrink the cache. It cuts how much of it each step reads. A small scorer, the lightning indexer, rates every earlier token for the current query, and attention runs only over the top k (2,048 in the model's config). The indexer here follows the shape of the one in the paper: a few small query heads, one small key per token, a ReLU and a weighted sum.
def indexer_scores(qi, ki, wi):
"""Lightning-indexer style: qi (Hi, T, di) small query heads, ki (L, di) one small key per token,
wi (Hi, T) per-head weights. score[t, s] = sum_j wi[j, t] * relu(qi[j, t] . ki[s])"""
return (wi[..., None] * torch.relu(qi @ ki.T)).sum(0) # (T, L)
def sparse_attend(q, k, v, scores, topk):
"""Each query attends only to its top-k tokens by indexer score (causal). All tokens stay in the cache."""
T, L = scores.shape
scores = scores.masked_fill(~causal(T, L), float("-inf"))
keep = torch.zeros(T, L, dtype=torch.bool)
idx = scores.topk(min(topk, L), dim=-1).indices
keep.scatter_(1, idx, True)
keep &= causal(T, L) # top-k may pick masked slots for early queries
return attend(q, k, v, keep), keep.sum(-1)Checks: with k at least the sequence length, the result is exactly full attention (difference 0). With k = 16 of 60, the last query reads 16 tokens. Our weights are random, so the gap from full attention (0.64 relative) says nothing about quality, only that the two are no longer the same computation. Whether the indexer picks the right tokens is a question about trained weights, and DeepSeek's report is where the answer lives.
Two things the code makes plain. The full cache stays, because any token may be in tomorrow's top k. And the indexer adds its own small key per token to the cache (ki; 128 numbers per token in V3.2-Exp's config), so sparse attention costs a little more memory, not less. What it saves is reading: the expensive attention touches k tokens per step instead of all of them, and the cheap indexer is the only part that scans everything.
6. Linear attention and Gated DeltaNet: a state instead of a cache
Drop the softmax, and attention can be rewritten so that nothing grows with the conversation. With a feature map φ, the attention-shaped form multiplies φ(q) φ(k)ᵀ by the values. Because there is no softmax in the way, the same sum can be built one token at a time into a matrix S = Σ v φ(k)ᵀ, and each query just reads S φ(q):
def linear_parallel(q, k, v):
"""Causal linear attention written like attention: no softmax, feature map phi = elu + 1."""
phi = lambda x: torch.nn.functional.elu(x) + 1
a = (phi(q) @ phi(k).transpose(-1, -2)) * causal(q.shape[1], k.shape[1])
return (a @ v) / a.sum(-1, keepdim=True)
def linear_recurrent(q, k, v):
"""The same thing as a running state: one (d_v x d_k) matrix and one vector per head, whatever the length."""
phi = lambda x: torch.nn.functional.elu(x) + 1
H, T, dk = k.shape
S, z, out = torch.zeros(H, v.shape[-1], dk), torch.zeros(H, dk), []
for t in range(T):
kt, qt = phi(k[:, t]), phi(q[:, t])
S = S + v[:, t, :, None] * kt[:, None, :] # add v k^T
z = z + kt
out.append((S @ qt[..., None])[..., 0] / (z * qt).sum(-1, keepdim=True))
return torch.stack(out, 1)
def gated_delta_step(S, k, v, q, alpha, beta):
"""Gated DeltaNet (Qwen3.5, Qwen3-Next) for one token, one head. S: (d_v, d_k).
alpha in (0, 1] decays the whole memory; beta in (0, 1] is how much to correct toward v at key k."""
S = alpha * S
S = S + beta * torch.outer(v - S @ k, k) # replace what is stored at k, instead of adding to it
return S, S @ qThe check: the running-state version and the attention-shaped version agree at every step (8.9 × 10⁻¹⁶), and the state is one 16 × 16 matrix per head whether the sequence is 40 tokens or 40,000.
The weakness is what that sum does to memory. Write two values at the same key and plain linear attention returns their sum: our check reads back exactly v1 + v2. The delta rule, which Gated DeltaNet builds on, first reads what is stored at the key and writes only the correction (v - S @ k). The same two writes then read back exactly v2. The gate alpha is the other half: it shrinks the whole memory a little at every token, so old entries fade. With alpha 0.9 and ten tokens that write nothing, 0.349 of v1 is left, which is 0.9¹⁰.
Real hybrids keep a few full-attention layers for exact recall and use this for the rest. Qwen3.5-9B has 24 Gated DeltaNet layers and 8 full ones. From its config.json: each layer holds one 128 × 128 state per value head (32 heads), plus a short convolution that keeps the last 3 inputs of q, k and v. In f32 that is 24 × (32 × 128 × 128 + 3 × (2 × 16 × 128 + 32 × 128)) × 4 bytes = 50.25 MiB per conversation, whatever its length. llama.cpp printed 50.25 MiB for this model (measured in the previous post).
The loop in linear_recurrent is how decoding works. For a long prompt, real implementations compute the same thing in chunks with matrix multiplications, because a Python loop over tokens would be far too slow.
What the code leaves out
- Speed. These functions build full score matrices and loop in Python. FlashAttention, paged caches, chunked DeltaNet kernels and FP8 indexers exist because of that.
- Training. Every check above is about what a design computes, not what a trained model learns to do with it.
- Positions, mostly. Only the MLA section applies RoPE, because that is the one design where it changes the structure.
The script with all checks is attention_code.py; it reads DeepSeek-V2-Lite's and gpt-oss-20b's weights from a local Hugging Face cache if they are there and skips those two checks if not.
Next in this series
How six models store their caches, measured in llama.cpp and vLLM, then what changes when you deploy on each runtime.