Attention, step by step

Code version: DeepInfer @ 9bd9cf2 · 2026-10-21

The idea#

The problem: a token alone has no context#

After the embedding, each token has its own vector. The vector of "bank" is the same in "river bank" and in "bank account". The model must mix in information from the other tokens to know which meaning is correct. Attention is the step that does this mix. It is the only place in a transformer where tokens see each other.

Query, key and value#

Each token makes three vectors from its hidden state, with three learned weight matrices (query, key, value):

• The query (q) says what this token looks for.
• The key (k) says what this token offers to others.
• The value (v) is the information that this token gives when another token selects it.

Four steps#

For one token with query q, attention does four steps:

1. Score. Compare q with the key of every token: the dot product q·k. A high score means a good match. Divide each score by √d_h, where d_h is the size of one head's vectors. The deep dive explains why.
2. Mask. A token may look only at itself and at earlier tokens. The causal mask sets the scores of later tokens to −∞.
3. Softmax. Softmax turns the scores into weights: exp(score) divided by the sum of all exp values. The weights are positive and add up to 1. A masked token gets weight 0, because exp(−∞) = 0.
4. Mix. The output is the weighted sum of the values: Σ weight_j · v_j.

A worked example#

Here is one head with vectors of 2 numbers (d_h = 2). The text is "The cat sat on", and we compute the row of "on". These are the same numbers as on the page The KV cache.

The output of "on" is 0.065·(2, 0) + 0.131·(0, 2) + 0.266·(1, 1) + 0.539·(0, 3) = (0.395, 2.144). The token "on" takes most of its new information from itself and from "sat".

Many heads#

One attention computation is one head. A layer runs many heads in parallel, often 32 to 128. Each head has its own weight matrices, so each head can look for a different relation. For example, one head can find the subject of a verb, and another the previous word. The outputs of all heads are joined and multiplied by one more matrix. This is multi-head attention.

An everyday comparison#

Think of a search in a library. Your question is the query. The label on each book is its key. The text in the book is its value. You compare your question with every label, and you read the books with the best labels.

The comparison stops here. Attention does not select one book. It reads all earlier books and mixes them, with more weight for the better matches. Also, nobody writes the labels: training learns how each token makes its query, key and value.

The picture#

Read panel 2 from left to right. Each row of the grid is one query token, and each column is one key token. The grey cells above the diagonal are later tokens, so the mask hides them. Panel 3 shows the KV heads that the KV cache must store.

Try it: the attention toy#

The toy below has 6 tokens and one head with d_h = 2. Tokens 1 to 4 have the keys and values of the worked example, and row "on" has its query q = (1, 2). The queries of the other rows are made up for the toy.

Things to try:

1. Click each token from left to right. Each row has one more non-zero weight than the row before. This is the causal mask.
2. Move Sharpness to 0. Every visible key gets the same weight. Move it to 4. The best key takes almost all the weight.
3. Select row "on". Edit the key of "cat" and set k[1] to 3. Watch the weight of "cat" grow and the output move toward the value of "cat", (0, 2).
4. Turn off the causal mask. Row "The" can now see "mat". A decoder LLM never does this; an encoder model, such as a model for text classification, can.

In DeepInfer#

Attention is in v4_modules/attn.py. The page MLA attention: one latent vector per token explains the full design. Here are three pointers.

Step 1: the scale is 1/√head_dim#

        self.qk_rope_head_dim = cfg.qk_rope_head_dim
        self.qk_head_dim = cfg.head_dim
        self.q_lora_rank = cfg.q_lora_rank
        self.num_heads = cfg.num_attention_heads
        self.n_wv_group = cfg.o_groups
        self.wv_group_dim = (self.num_heads // self.n_wv_group) * self.qk_head_dim
        self.o_head_dim = cfg.o_lora_rank
        self.max_model_len = cfg.runtime_max_model_len
        self.softmax_scale = self.qk_head_dim**-0.5

Step 2: prefill and decode use different kernels#

        if input_metadata.is_prompt:
            context = self._prefill_attention(
                q_b,
                kv,
                kv_ca_index,
                kv_swa_full_index if self.use_swa else None,
                total_k_swa_cached,
                total_k_swa_cur,
                input_metadata,
            )
            del kv
        else:
            context = self._decode_attention(
                q_b,
                kv_ca_index,
                kv_swa_index if self.use_swa else None,
                kv_cache,
                input_metadata,
            )

Step 3: an index list instead of a mask#

The attention backend calls the FlashMLA library. Its docstring gives the shapes (deepinfer/model_executor/layers/attention_backend.py:35-49):

• query: [total_q, num_heads_q, head_dim], BF16. One row per new token, all heads.
• kv_a_with_mqa: [total_k, 1, head_dim]. The "1" is the single shared KV head.
• index_topk: [total_q, topk_tokens], int32. For each query, the key rows that it may read.

The plain decode call turns off the causal flag, because the index list already contains only allowed keys. On GPU with 64 or 128 heads, sparse layers and layers without a compressed cache use a fused FlashMLA kernel instead (deepinfer/model_executor/models/v4_modules/attn.py:87-92). That call has no causal flag at all, and it also reads only the listed rows (deepinfer/model_executor/layers/attention_backend.py:107-121).

    output, lse = flash_mla_with_kvcache(
        q=query,
        k_cache=key_cache,
        block_table=None,  # FlashMLA 保留形参但当前忽略
        cache_seqlens=None,  # 同上
        head_dim_v=head_dim_v,
        tile_scheduler_metadata=sched_meta,
        softmax_scale=softmax_scale,
        causal=False,  # sparse attention 只支持非 causal
        is_fp8_kvcache=True,  # V4.1 的 KV cache 为 paged FP8
        indices=indices_in_kvcache.unsqueeze(1),
        attn_sink=attn_sink,

Deep dive#

Q : [T, d_h]      K : [S, d_h]      V : [S, d_v]

S_raw = Q · Kᵀ                         [T, S]
S     = S_raw · scale                  scale = 1 / sqrt(d_h)
S     = S + M                          M[i, j] = 0 if key j is allowed for query i, else −∞
P     = softmax(S, along each row)     [T, S], each row sums to 1
O     = P · V                          [T, d_v]

m      = max_j z_j
p_j    = exp(z_j − m) / Σ_k exp(z_k − m)

bytes per token per layer = 2 × n_kv_heads × d_h × bytes_per_number

What can go wrong#

• No scale. Without the division by √d_h, softmax becomes very sharp for large heads (deep dive "Why divide by √d_h?"). DeepInfer passes the scale to every attention call (deepinfer/model_executor/models/v4_modules/attn.py:621, deepinfer/model_executor/models/v4_modules/attn.py:664).
• A wrong mask. If a query can read a later token, the model sees text that it must predict. In DeepInfer, a wrong index list would do the same, because the kernel does not add its own causal mask.
• Overflow. Softmax without the maximum subtracted gives infinity or NaN for large scores.
• Results that depend on the batch. The same request can give slightly different numbers in different batches, because the kernels sum in a different order. DeepInfer has a batch-invariant mode, DEEPINFER_BATCH_INVARIANT, that it passes to the decode kernel (deepinfer/model_executor/models/v4_modules/attn.py:666). See Determinism and batch invariance.

Why it matters. Attention is the only part of the model whose cost grows with the length of the text. It is also the reason for the KV cache. Most of DeepInfer's model pages (MLA, sliding window, sparse indexer, KV compression) are ways to make attention cheaper in time or in memory.