Scaled dot-product attention with masking

Yesterday I built the autoregressive decoder loop and called into the model as a black box. Today I'm opening that box and implementing the core operation inside it: scaled dot-product attention. This is the function that every transformer layer calls, the one that dominates runtime for long sequences, and the one that FlashAttention, PagedAttention, and every other attention optimization is trying to speed up.

The math

Scaled dot-product attention takes three inputs: queries (Q), keys (K), and values (V), all matrices. The formula from "Attention Is All You Need":

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

Where d_k is the dimension of each key vector (typically 128 for modern LLMs). The scaling by 1/sqrt(d_k) prevents the dot products from growing large and pushing softmax into regions where its gradients vanish.

Here's the naive implementation:

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V):
    """
    Q: (batch, num_heads, seq_q, head_dim)
    K: (batch, num_heads, seq_kv, head_dim)
    V: (batch, num_heads, seq_kv, head_dim)
    """
    d_k = Q.size(-1)
    # (batch, num_heads, seq_q, seq_kv)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    weights = F.softmax(scores, dim=-1)
    # (batch, num_heads, seq_q, head_dim)
    output = torch.matmul(weights, V)
    return output

Four lines of math. But those four lines are where most of inference time and memory go for long contexts.

Adding the causal mask

In autoregressive language models, each token should only attend to tokens at earlier positions (and itself). This is enforced with a causal mask: a lower-triangular matrix of ones, applied before softmax by setting masked positions to negative infinity.

def causal_attention(Q, K, V):
    d_k = Q.size(-1)
    seq_q = Q.size(-2)
    seq_kv = K.size(-2)

    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)

    # Create causal mask: True where attention is NOT allowed
    # For prefill: full lower-triangular mask
    # For decode with KV cache: seq_q=1, seq_kv=full, no masking needed
    if seq_q > 1:
        mask = torch.triu(torch.ones(seq_q, seq_kv, device=Q.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))

    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, V)
    return output

A subtle point: during decode with KV cache, the query has length 1 (just the new token) and the keys/values have the full sequence length. Since the new token is at the last position, it should attend to all previous tokens. No masking is needed. The causal mask only matters during prefill, when you process multiple tokens at once.

Multi-head attention

Transformers don't run one attention computation. They split the model dimension into multiple heads, run attention independently on each, and concatenate the results. For a model with dimension 4096 and 32 heads, each head works with dimension 128.

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model=4096, num_heads=32):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads  # 128

        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, past_kv=None):
        batch, seq_len, _ = x.shape

        # Project to Q, K, V
        Q = self.W_q(x)  # (batch, seq, d_model)
        K = self.W_k(x)
        V = self.W_v(x)

        # Reshape to (batch, num_heads, seq, head_dim)
        Q = Q.view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Append to KV cache if present
        if past_kv is not None:
            past_k, past_v = past_kv
            K = torch.cat([past_k, K], dim=2)
            V = torch.cat([past_v, V], dim=2)

        # Attention
        output = causal_attention(Q, K, V)

        # Reshape back and project
        output = output.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.W_o(output), (K, V)

Grouped-query attention (GQA)

Modern LLMs like Llama 3 use grouped-query attention, where multiple query heads share a single key-value head. Llama 3.1 70B has 64 query heads but only 8 KV heads: a group size of 8. This reduces the KV cache size by 8x compared to full multi-head attention.

def grouped_query_attention(Q, K, V, num_kv_heads):
    """
    Q: (batch, num_q_heads, seq_q, head_dim)   e.g. 64 heads
    K: (batch, num_kv_heads, seq_kv, head_dim)  e.g. 8 heads
    V: (batch, num_kv_heads, seq_kv, head_dim)  e.g. 8 heads
    """
    num_q_heads = Q.size(1)
    group_size = num_q_heads // num_kv_heads  # 8

    # Expand K,V to match Q's head count by repeating
    K = K.repeat_interleave(group_size, dim=1)  # (batch, 64, seq_kv, head_dim)
    V = V.repeat_interleave(group_size, dim=1)

    return causal_attention(Q, K, V)

The repeat_interleave doesn't actually copy data in optimized implementations. FlashAttention and PyTorch's F.scaled_dot_product_attention handle GQA natively by broadcasting rather than expanding.

The memory problem

The naive attention implementation has a memory problem that becomes severe at long sequence lengths. The scores matrix Q @ K^T has shape (batch, heads, seq_q, seq_kv). For a single head at sequence length 8192, that's 8192 x 8192 = 67 million elements. In FP32, that's 256 MB per head. With 64 heads, the scores matrix alone consumes 16 GB.

This is why FlashAttention was invented. FlashAttention never materializes the full scores matrix. Instead, it tiles the computation: it processes Q, K, V in blocks that fit in GPU SRAM (shared memory), computing attention one tile at a time and accumulating results. The peak memory is proportional to the tile size, not the sequence length.

PyTorch's built-in path

Since PyTorch 2.0, F.scaled_dot_product_attention(Q, K, V, is_causal=True) automatically dispatches to FlashAttention (if available), memory-efficient attention, or the math fallback. In practice, always use this function rather than the manual implementation above. But writing it by hand first helps you understand what it's doing and why it's necessary.

What I learned building this

Implementing attention from scratch made a few things click for me:

The causal mask is only needed during prefill. During decode with KV cache, the single query token naturally attends to all cached positions.
GQA reduces KV cache memory but doesn't change the compute. The Q @ K^T matmul is the same size regardless of how many KV heads there are (because the KV heads get expanded to match Q).
The N^2 memory scaling of naive attention is real and painful. At sequence length 32k, the scores matrix is untenable without FlashAttention.
The scaling factor 1/sqrt(d_k) isn't optional. Without it, the dot products grow with dimension, softmax saturates, and gradients die. For d_k=128, the scale factor is about 0.088.

Next: profiling attention memory growth, where I'll instrument this implementation to measure exactly how memory scales with sequence length and batch size.