Transformer internals: attention without the mysticism

Most explanations of transformer attention treat it as a philosophical breakthrough in machine cognition. In practice, multi-head attention is a sequence of matrix multiplications, a softmax normalisation along a single axis, and a weighted linear combination of row vectors. When you strip away the anthropomorphic vocabulary, what remains is an elegant, bandwidth-heavy routing mechanism with very predictable memory characteristics.

If you are implementing custom inference kernels or optimising serving latency on modern accelerators, understanding the exact memory layout of the key-value cache and why naive self-attention scales quadratically with sequence length matters far more than conceptual analogies.

The mechanics of scaled dot-product

Given an input sequence projected into queries \(Q\), keys \(K\), and values \(V\) with dimension \(d_k\), the core operation computes a similarity score between every query token and every key token. The scaling factor \(\frac{1}{\sqrt{d_k}}\) is not decorative; as dimensionality grows, the dot products grow large in magnitude, pushing the softmax function into regions with vanishingly small gradients.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(
    q: torch.Tensor,  # [batch, heads, seq_len_q, head_dim]
    k: torch.Tensor,  # [batch, heads, seq_len_k, head_dim]
    v: torch.Tensor,  # [batch, heads, seq_len_k, head_dim]
    mask: torch.Tensor = None
) -> torch.Tensor:
    head_dim = q.size(-1)
    # Compute raw compatibility scores
    scores = torch.matmul(q, k.transpose(-2, -1)) / (head_dim ** 0.5)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
        
    attn_weights = F.softmax(scores, dim=-1)
    return torch.matmul(attn_weights, v)

During training, this entire tensor computation happens in parallel across the sequence dimension using causal masking. During autoregressive generation, however, you emit one token at a time. Evaluating the entire historical sequence from scratch at step \(t\) requires recomputing \(K\) and \(V\) projections for all \(t-1\) prior tokens—a catastrophic waste of compute.

At batch size 32 with an 8k context window, your bottleneck is almost never arithmetic compute (FLOPs). It is High Bandwidth Memory (HBM) throughput, driven by streaming gigabytes of past key and value projections.

The KV cache and the memory wall

To avoid recalculating projections at each generation step, autoregressive decoders retain the computed keys and values in GPU memory: the KV cache. While this reduces computation to a single vector-matrix multiply per layer, the memory footprint scales linearly with batch size, context length, layer count, and hidden dimension.

For a standard 70B parameter model with 80 layers, 64 heads, and head dimension 128 in 16-bit precision:

• Bytes per token per layer: \(2 \times (\text{layers}) \times (\text{kv\_heads}) \times (\text{head\_dim}) \times 2\text{ bytes}\).
• Standard MHA footprint: Without Grouped Query Attention (GQA), a single 4,096-token sequence consumes roughly 2.6 GB just for its cache.
• Grouped Query Attention (GQA): Sharing key/value heads across query head groups (e.g., 8 KV heads for 64 query heads) slashes this footprint by an 8× factor, making multi-tenant serving feasible on an 80GB SXM5 node.

Rotary positional embeddings (RoPE)

Absolute positional embeddings assign a static vector to each index. This fails to capture relative distances cleanly and degrades when extending beyond the initial training context window. Rotary Positional Embeddings solve this by rotating pairs of features in the complex plane by an angle proportional to their position index \(m\).

The inner product \(\langle R_{\Theta, m}^d q, R_{\Theta, n}^d k \rangle\) encodes relative distance \(m - n\) directly into the attention score. When extrapolating to longer contexts via NTK-aware scaling or YaRN, we modify the base frequency \(\theta_i = b^{-2(i-1)/d}\), effectively stretching or interpolating the rotation angles across high-frequency and low-frequency components.

When you inspect production implementations like FlashAttention or vLLM's PagedAttention, the secret is not algorithmic magic. It is hardware-conscious memory management: tiling the softmax reduction inside SRAM to avoid round-trips to HBM, and paging KV blocks to eliminate fragmentation.

Have questions about kernel-level attention optimisations or RoPE scaling? Send me a note — always glad to discuss.