Skip to content

Tutorial 06: Scaled Dot-Product Attention

Attention is the mechanism that lets each token "look at" every other token in the sequence and decide what information to gather. This tutorial derives the attention formula from first principles, builds the implementation step by step, and explains the causal mask.

1. What Attention Computes

Given a sequence of token representations, attention computes a weighted sum of all token representations, where the weights are determined by pairwise relevance.

Think of it this way: when you read the sentence "The cat sat on the mat because it was tired," the word "it" needs to attend to "cat" (not "mat") to resolve the pronoun. Attention makes this connection explicit and learnable.

The formula:

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

where: - Q (queries): "What am I looking for?" - K (keys): "What do I contain?" - V (values): "What information do I provide?" - d_k: the dimension of each key/query vector

2. Step-by-Step Derivation

Step 1: Compute Relevance Scores

For every pair of tokens (i, j), compute a score that measures how relevant token j is to token i:

score(i, j) = q_i · k_j

In matrix form, for the entire sequence:

scores = Q @ K^T    # shape: (seq_len, seq_len)

Row i of the scores matrix contains the relevance of every token to token i.

import mlx.core as mx

# Simple case: same number of query and key heads, no batching
# Q shape: (num_heads, seq_len, head_dim)
# K shape: (num_heads, context_len, head_dim)
scores = mx.matmul(q, k.swapaxes(-2, -1))  # (heads, seq, context)

Step 2: Scale

The dot product grows with head_dim. If d_k = 128 and the values are normally distributed with mean 0 and variance 1, the dot product has variance ~128. This pushes softmax into regions with tiny gradients.

Dividing by sqrt(d_k) normalizes the variance back to 1:

scores = scores / sqrt(d_k)

# In practice, we precompute the scale factor
head_dim = q.shape[-1]
scale = head_dim ** -0.5   # = 1 / sqrt(head_dim)

scores = mx.matmul(q, k.swapaxes(-2, -1)) * scale

Step 3: Apply the Causal Mask

During autoregressive generation, token i can only attend to tokens j ≤ i. It cannot "see into the future." We enforce this by adding -inf to all positions where j > i:

mask[i][j] = 0       if j ≤ i  (allowed to attend)
mask[i][j] = -inf    if j > i  (blocked — future token)

After adding the mask, softmax turns -inf positions into zero attention weights:

def causal_mask(L: int, S: int, dtype) -> mx.array:
    """
    Create a causal (lower-triangular) attention mask.

    Args:
        L: Query length (number of tokens generating attention).
        S: Key/value length (total context, including new tokens).
        dtype: Output data type.

    Returns:
        Mask tensor of shape (L, S). Zeros for allowed positions,
        -inf for blocked positions.
    """
    # mx.tril extracts the lower triangle of a matrix of ones.
    # k=(S-L) shifts the diagonal to handle cases where L != S
    # (e.g., during decode when L=1 and S=grows with each step).
    mask = mx.tril(mx.ones((L, S)), k=(S - L))

    # Convert: 1.0 → 0.0 (attend), 0.0 → -inf (don't attend)
    mask = mx.where(mask, mx.array(0), mx.array(-mx.inf)).astype(dtype)
    return mask

Example: For a sequence of length 4:

mask = [[  0,  -inf, -inf, -inf],    ← token 0 attends only to itself
        [  0,    0,  -inf, -inf],    ← token 1 attends to 0 and 1
        [  0,    0,    0,  -inf],    ← token 2 attends to 0, 1, 2
        [  0,    0,    0,    0]]     ← token 3 attends to all

Step 4: Softmax to Get Weights

Apply softmax across the context dimension (axis=-1) to convert scores to a probability distribution:

weights = softmax(scores + mask, axis=-1)

Each row of weights sums to 1.0, and the masked positions contribute zero.

Step 5: Weighted Sum of Values

output = weights @ V    # shape: (num_heads, seq_len, head_dim)

Each output vector is a weighted combination of all value vectors, where the weights reflect relevance.

3. The Simple Implementation

This handles the case where Q, K, V all have the same shape:

def scaled_dot_product_attention_simple(
    query: mx.array,
    key: mx.array,
    value: mx.array,
    scale: float | None = None,
    mask: mx.array | None = None,
) -> mx.array:
    """
    Simple attention for equal-head-count Q/K/V.

    All inputs have shape: (batch, heads, seq_len, head_dim)
    Output has the same shape.
    """
    # Step 1: Compute scale factor
    factor = mx.rsqrt(query.shape[-1]) if scale is None else scale

    # Step 2: Compute relevance scores
    # query: (B, H, L, D) @ key^T: (B, H, D, S) → (B, H, L, S)
    scores = mx.matmul(query, key.swapaxes(-2, -1)) * factor

    # Step 3: Apply mask (if provided — already contains -inf for blocked positions)
    if mask is not None:
        scores = scores + mask

    # Step 4: Softmax to get attention weights
    weights = mx.softmax(scores, axis=-1)

    # Step 5: Weighted sum of values
    # weights: (B, H, L, S) @ value: (B, H, S, D) → (B, H, L, D)
    return mx.matmul(weights, value)

4. Grouped-Query Attention (GQA)

Modern models like Llama 3 and Qwen3 use Grouped-Query Attention: the number of key/value heads is smaller than the number of query heads. For example, Qwen3-8B has 32 query heads but only 8 key/value heads. Each group of 4 query heads shares one key/value head.

This saves memory (fewer KV vectors to cache) and compute (fewer KV matrix multiplies), with minimal quality loss.

The Math

If H_q query heads share H_kv key/value heads (where H_q is a multiple of H_kv), each KV head serves n_repeats = H_q / H_kv query heads:

// Reshape to expose the sharing structure
Q: (B, H_q, L, D) → (B, n_repeats, H_kv, L, D)
K: (B, H_kv, S, D) → (B, 1, H_kv, S, D)
V: (B, H_kv, S, D) → (B, 1, H_kv, S, D)

// Scores after matmul: (B, n_repeats, H_kv, L, S)
// The key and value are broadcast across n_repeats
scores = Q @ K^T * scale

// Attention weights and output
weights = softmax(scores, axis=-1)
output = weights @ V

// Reshape back: (B, n_repeats, H_kv, L, D) → (B, H_q, L, D)

The Implementation

def scaled_dot_product_attention_grouped(
    query: mx.array,
    key: mx.array,
    value: mx.array,
    scale: float | None = None,
    mask: mx.array | str | None = None,
) -> mx.array:
    """
    Attention with grouped-query support.

    Query shape:  (B, H_q, L, D)     — more query heads
    Key shape:    (B, H_kv, S, D)     — fewer key/value heads
    Value shape:  (B, H_kv, S, D)     — same as key
    Output shape: (B, H_q, L, D)
    """
    factor = mx.rsqrt(query.shape[-1]) if scale is None else mx.array(scale)
    factor = factor.astype(query.dtype)

    H_q, L, D = query.shape[-3:]
    H_kv, S, _ = key.shape[-3:]

    # Compute how many query heads share each KV head
    n_repeats = H_q // H_kv

    # Reshape Q to expose the sharing structure
    # (B, H_q, L, D) → (B, n_repeats, H_kv, L, D)
    B = query.shape[:-3]
    query = query.reshape(*B, -1, H_kv, n_repeats, L, D)

    # Reshape K and V to broadcast across the repeats dimension
    # (B, H_kv, S, D) → (B, 1, H_kv, S, D)
    key = key.reshape(*B, -1, H_kv, 1, S, D)
    value = value.reshape(*B, -1, H_kv, 1, S, D)

    # Compute scores — K is broadcast across n_repeats
    scores = mx.matmul(query, key.swapaxes(-2, -1)) * factor

    # Apply mask
    if isinstance(mask, mx.array):
        # Reshape mask to match the group structure
        mask = mx.broadcast_to(mask, (*B, H_q, L, S))
        mask = mask.reshape(*B, 1, H_kv, n_repeats, L, S)
        scores = scores + mask
    elif isinstance(mask, str):
        if mask != "causal":
            raise ValueError(f"unsupported mask: {mask!r}")
        scores = scores + causal_mask(L, S, scores.dtype)

    # Softmax + value weighted sum
    result = mx.matmul(mx.softmax(scores, axis=-1), value)

    # Reshape back to the original head layout
    return result.reshape(query.shape)

Why GQA Works

In standard multi-head attention (MHA), each head learns different attention patterns. But empirical research shows that key/value heads within a group tend to learn very similar patterns. GQA exploits this redundancy — sharing KV heads saves ~75% of KV cache memory (with 4:1 ratio) while preserving quality.

5. The Causal Mask in Detail

During decode (generating one token at a time), the query has length 1 and the key/value length is the full context:

Q shape: (B, H, 1, D)      — one new token
K shape: (B, H, S, D)      — all previous tokens + current
S = prompt_len + steps_so_far

The causal mask ensures the new token can attend to all previous tokens but not to future ones (which do not exist yet). Since L=1, the mask is just a row of zeros — the single query token can attend to everything in the context.

During prefill (processing the entire prompt at once), L=S and the mask is a full lower-triangular matrix.

6. Summary

  • Attention = relevance-weighted sum of value vectors.
  • Scale by 1/sqrt(d_k) to keep softmax in a numerically stable range.
  • Causal mask prevents attending to future tokens (lower-triangular -inf mask).
  • GQA shares KV heads across groups of query heads, saving memory.
  • The implementation is two matrix multiplications, a softmax, and a mask add.

Next: Tutorial 07 — Multi-Head and Grouped-Query Attention (Full Layer)