Transformer Attention Math: Q/K/V, Softmax Weights, Masks, and KV Cache

The core concept of a Transformer’s Self-Attention mechanism can be intuitively understood as follows: every token in a sequence uses its own Query vector to evaluate the Key vectors of all other tokens. This process determines how “attention” should be distributed across the sentence, and finally, this weight distribution is used to compute a weighted sum of the information vectors (Values). Its mathematical expression is remarkably concise—just one line of code—but it conceals a staggering amount of engineering depth and model training nuances.

The formula is easy to memorise; knowing what each step computes is the hard part. So this article works Scaled Dot-Product Attention by hand over 3 tokens, writing out every intermediate matrix so you can check the arithmetic yourself. With that done, the rest follows: why Q/K/V are separate projections, what dividing by the square root of d protects against, exactly which positions a causal mask removes, what each head sees once they are split, and which part of the computation a KV cache actually stores at inference time.

1. The Core Mathematical Formula

This is the foundational equation of the Large Language Model era:

Attention(Q, K, V) = softmax((Q @ K^T) / sqrt(d_k)) @ V

Breaking it down:

• Q @ K^T produces an attention score matrix of size `[seq_len, seq_len]`. Because it’s a dot product, it measures the “similarity” or “affinity” between pairs of tokens in a high-dimensional space.
• Why must we divide by sqrt(d_k)? Suppose Q and K have a dimension of `d_k = 4096`, and their elements follow an independent distribution with a mean of 0 and a variance of 1. The variance of their dot product will scale up to `4096`. A massive variance creates extreme score values (e.g., 100 vs -100). When these are passed through a Softmax function, it forces the gradients to near-zero (vanishing gradients), a problem known as “Softmax saturation.”

2. Architectural Diagram: Data Flow and Dimensions

graph TD
    Input[Input Sequence: B, L, d_model] --> WQ(W_q Linear)
    Input --> WK(W_k Linear)
    Input --> WV(W_v Linear)
    
    WQ --> Q[Q: B, h, L, d_k]
    WK --> K[K: B, h, L, d_k]
    WV --> V[V: B, h, L, d_v]
    
    Q --> Dot[Dot Product: Q @ K^T]
    K --> Dot
    Dot --> Scale[Scale by 1/sqrt(d_k)]
    Scale --> Mask[Apply Causal Mask]
    Mask --> Softmax[Softmax along dim L]
    Softmax --> AttentionWeights[Attention Weights: B, h, L, L]
    
    AttentionWeights --> MatMulV[MatMul with V]
    V --> MatMulV
    
    MatMulV --> Context[Context Output: B, h, L, d_v]
    Context --> Concat[Concat Heads: B, L, d_model]
    Concat --> Out[W_o Linear]

3. Practical Demonstration: Self-Attention in NumPy

Formulas can be abstract. Let’s run a highly educational, purely Pythonic NumPy implementation. Imagine our input sequence consists of just 3 tokens (e.g., “AI”, “needs”, “math”) with an embedding dimension of 4:

import numpy as np

# 1. Simulate Q, K, V Matrices (Seq_len=3, d_k=4)
# Representing tokens: "AI", "needs", "math"
Q = np.array([
    [ 1.0,  0.5, -0.2,  0.1],  # AI
    [-0.5,  1.2,  0.8, -0.4],  # needs
    [ 0.2, -0.1,  1.5,  0.9]   # math
])
K = np.array([
    [ 0.8,  0.4, -0.3,  0.0],
    [-0.2,  1.0,  0.5, -0.1],
    [ 0.1, -0.2,  1.1,  0.7]
])
V = np.array([
    [ 1.0,  0.0],
    [ 0.0,  1.0],
    [-1.0, -1.0]
])

d_k = Q.shape[1]

# 2. Calculate Dot-Product Scores and Scale
scores = (Q @ K.T) / np.sqrt(d_k)
print("Scaled Scores:\\n", scores)

# 3. Causal Mask
# Mask out future positions to prevent information leakage (cheating)
mask = np.triu(np.ones((3, 3)), k=1)
scores[mask == 1] = -np.inf

# 4. Softmax Normalization
def softmax(x):
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / e_x.sum(axis=-1, keepdims=True)

weights = softmax(scores)
print("Attention Weights:\\n", np.round(weights, 3))

# 5. Value Weighting (Context Output)
context = weights @ V
print("Context Output:\\n", context)

If you run this code, you’ll observe that the first row (the word “AI”) assigns attention weights exclusively to itself. The third row (“math”) distributes its attention across the preceding two tokens. This perfectly demonstrates the essence of autoregressive models: they must synthesize historical context without peaking into the future.

4. What Does the Causal Mask Actually Change?

As demonstrated in the code above, in autoregressive generation tasks, if the model is currently predicting the 3rd token, it absolutely cannot “see” tokens 4 or 5. Right before the Softmax operation, we forcefully overwrite the upper triangular matrix of the attention scores to negative infinity (-inf). After passing through the Softmax, these specific weights are mathematically crushed to exactly 0. Therefore, the Mask does not delete tokens; rather, it performs a probabilistic cutoff, ensuring that illegal attention allocation is impossible.

5. An Engineer’s Perspective: The VRAM Killer and KV Cache

Real-World Insight: In textbooks, you see elegant matrix multiplication. But in industrial LLM deployment, what you see is a relentless stream of OOM (Out of Memory) exceptions.

During inference, large language models operate token-by-token. When generating token $t+1$, the K and V matrices for the previous $t$ tokens remain completely identical! If we were to naively multiply the full L x d_model matrices over and over, the computational waste would be catastrophic.

The KV Cache is the ultimate space-for-time tradeoff.

• We allocate a massive contiguous block in the GPU VRAM to cache historically generated K and V tensors.
• For every new token generated, we only compute $Q_{new}, K_{new}, V_{new}$ for that single token, and strictly append $K_{new}$ into the VRAM cache block.
• **The Cost is Staggering:** For slightly longer contexts (even just 10K tokens), a single batch’s KV Cache footprint can easily exceed the memory required to load the model weights themselves! This is exactly why the industry invented PagedAttention (the core of vLLM), MQA (Multi-Query Attention), and GQA (Grouped-Query Attention)—they are all desperate engineering hacks designed to shrink the KV Cache memory footprint.

6. Shape and Mask Checks

The easiest attention bugs are the ones that still produce a tensor. First, verify that batched attention uses batch x heads x tokens x dim and transposes only the final two dimensions for Q @ K^T. Then verify that the causal mask describes query-token to key-token visibility and broadcasts over batch and head. Finally, confirm that softmax is applied over the key dimension, not over the query dimension.

For the three-token toy example, every attention row should sum to approximately 1, future positions should be zero after masking, and the context output should have the same final dimension as V. A heatmap is useful for debugging, but it is not causal proof of why a model predicted a token; it only shows how the current weighted read used Value vectors.

7. Attention Verification Matrix

Self-attention often fails with code that runs but has the wrong semantics. Use the matrix below to audit the NumPy toy example here and to transfer the checks to batched, multi-head, or cached inference implementations.

8. Visualizations and Data Flow Summary

This mechanism may seem like simple linear algebra, but it currently supports the absolute frontier of global AI research. The next time your Transformer script crashes, your first instinct should always be: Print the shape of every single tensor, and trace the matrix multiplication on a piece of scrap paper.