Transformer Internals

import math
import torch
import torch.nn.functional as F

def attention(q, k, v, mask=None):
    # q, k, v: [batch, heads, time, d_head]
    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.size(-1))
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights

Implementation Drill: Prove SDPA Matches The Reference

Before swapping a hand-written attention block for torch.nn.functional.scaled_dot_product_attention, build a tiny deterministic test that compares both paths. PyTorch's SDPA API can dispatch to optimized kernels, supports causal masking and optional GQA, and applies dropout according to dropout_p, so the safest first test fixes dropout at zero and checks mask semantics on small tensors.

def reference_sdpa(q, k, v, *, causal=False, mask=None):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    if causal:
        t_q, t_k = q.size(-2), k.size(-2)
        causal_mask = torch.ones(t_q, t_k, dtype=torch.bool).triu(1)
        scores = scores.masked_fill(causal_mask.to(q.device), float("-inf"))
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v

torch.manual_seed(0)
q = torch.randn(2, 4, 5, 16)
k = torch.randn(2, 4, 5, 16)
v = torch.randn(2, 4, 5, 16)

expected = reference_sdpa(q, k, v, causal=True)
actual = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)

1. Start small: use short sequences and fixed seeds so mask, transpose, and scale bugs are easy to isolate.
2. Freeze dropout: pass dropout_p=0.0 in correctness tests and evaluation code instead of relying on module mode.
3. Test both masks: verify causal-only, padding-mask-only, and combined-mask cases before benchmarking a fused backend.

Debugging Drill: Convert Boolean Mask Polarity At API Boundaries

PyTorch's boolean mask conventions are easy to mix up. In torch.nn.functional.scaled_dot_product_attention, a boolean attn_mask uses True for positions that are allowed to participate in attention. In torch.nn.Transformer masks and key-padding masks, True means a position is not allowed or should be ignored. Convert the mask at the wrapper boundary instead of letting both meanings leak through the model code.

def blocked_to_sdpa_keep_mask(blocked):
    # blocked: True means "do not attend" in many nn.Transformer-style masks.
    # SDPA bool attn_mask: True means "keep this attention score."
    return ~blocked

blocked = torch.tensor([
    [False, True, True],
    [False, False, True],
    [False, False, False],
])
keep = blocked_to_sdpa_keep_mask(blocked)
assert keep[0, 0] and not keep[0, 1]

out = F.scaled_dot_product_attention(
    q, k, v,
    attn_mask=keep[None, None, :, :],
    dropout_p=0.0,
)

1. Name the mask by meaning: use names such as blocked, padding_ignored, or sdpa_keep instead of generic mask.
2. Keep one conversion point: invert the boolean mask when calling SDPA, then test the wrapper rather than sprinkling ~mask across attention code.
3. Do not mix hints casually: SDPA documents is_causal and attn_mask as separate contracts, so use an explicit keep mask when you need padding, packing, or document-boundary rules.

Implementation Drill: Expand KV Heads Only When The Kernel Needs It

In GQA, the query tensor has more heads than the key/value tensors. The attention math still needs every query head to attend to a compatible key/value head, so a simple reference implementation repeats each KV head across its query-head group. Production kernels may avoid physically materializing that repeat, but your shape test should make the grouping explicit.

def repeat_kv_for_gqa(x, num_query_heads):
    # x: [batch, kv_heads, time, head_dim]
    b, kv_heads, t, d = x.shape
    assert num_query_heads % kv_heads == 0
    group_size = num_query_heads // kv_heads
    x = x[:, :, None, :, :].expand(b, kv_heads, group_size, t, d)
    return x.reshape(b, num_query_heads, t, d)

k_expanded = repeat_kv_for_gqa(k, q.size(1))
v_expanded = repeat_kv_for_gqa(v, q.size(1))
out = torch.nn.functional.scaled_dot_product_attention(
    q, k_expanded, v_expanded, is_causal=True, dropout_p=0.0
)

1. Invariant: query heads must be divisible by KV heads; MHA is q_heads == kv_heads, and MQA is kv_heads == 1.
2. Cache impact: KV cache memory scales with KV heads, not query heads, so reducing KV heads directly lowers long-context decode memory.
3. Correctness check: compare the repeated-KV reference path with a fused GQA path such as PyTorch SDPA enable_gqa=True for small deterministic inputs.

Reading Note: Treat enable_gqa As A Kernel Contract

PyTorch's SDPA API can repeat key and value heads internally when enable_gqa=True, but the flag is not a license to pass any head layout. The public contract is the same one the GQA paper makes visible: query heads are grouped over fewer key/value heads. Keep the slow repeated-KV path as a reference test, then use the fused path only after the shape and backend constraints are explicit.

def assert_gqa_contract(q, k, v):
    q_heads = q.size(-3)
    k_heads = k.size(-3)
    v_heads = v.size(-3)
    assert k_heads == v_heads
    assert q_heads % k_heads == 0

ref = F.scaled_dot_product_attention(
    q,
    repeat_kv_for_gqa(k, q.size(-3)),
    repeat_kv_for_gqa(v, q.size(-3)),
    dropout_p=0.0,
    is_causal=True,
)

assert_gqa_contract(q, k, v)
fused = F.scaled_dot_product_attention(
    q, k, v, dropout_p=0.0, is_causal=True, enable_gqa=True
)
torch.testing.assert_close(fused, ref, rtol=1e-5, atol=1e-5)

1. Use the flag for production shape: pass fewer KV heads to the kernel so the cache layout matches the model, not an expanded debugging tensor.
2. Keep the reference path: explicit repetition is easier to inspect in unit tests and catches accidental head-order changes.
3. Check backend support: PyTorch documents GQA support as backend-limited, so benchmark the exact dtype, device, head dimension, and sequence shape you plan to serve.

Implementation Drill: Build The Head-Layout Memory Ledger

MQA was proposed for faster incremental decoding because decode steps repeatedly load old keys and values from memory. GQA generalizes that idea by choosing an intermediate number of KV heads. When interviewing or designing a model variant, make the memory ratio explicit before arguing about quality: for the same layers, sequence length, head dimension, and dtype, the KV-cache difference is mostly the KV-head count.

def kv_cache_bytes(layers, batch, seq_len, kv_heads, head_dim, dtype_bytes):
    # Factor 2 counts K and V.
    return 2 * layers * batch * seq_len * kv_heads * head_dim * dtype_bytes

def head_layout_report(query_heads, kv_heads):
    assert query_heads % kv_heads == 0
    if kv_heads == query_heads:
        layout = "MHA"
    elif kv_heads == 1:
        layout = "MQA"
    else:
        layout = "GQA"
    return {
        "layout": layout,
        "queries_per_kv_head": query_heads // kv_heads,
        "kv_cache_ratio_vs_mha": kv_heads / query_heads,
    }

1. Estimate first: compare kv_heads / query_heads before changing kernels; a 32-query-head model with 8 KV heads uses one quarter of the MHA KV cache.
2. Keep the contract visible: the attention layer still returns query-head outputs, even when fewer KV heads are stored.
3. Test the API path: PyTorch SDPA's GQA path expects query heads to divide cleanly over key/value heads and key/value to have matching head counts.
4. Separate model and serving claims: lower cache traffic can improve decode throughput, but quality depends on training or uptraining, not only on the inference kernel.

Reading Drill: MLA Caches Latents, Not Full K/V Heads

Multi-head latent attention is a model-architecture change, not just a serving flag. DeepSeek-V2 introduces MLA as low-rank joint compression for keys and values, so the decode cache stores a compact latent state that can reconstruct per-head K/V when attention runs. DeepSeek-V3 keeps the same MLA direction after validating it in V2. Compare it with GQA carefully: GQA reduces the number of stored KV heads, while MLA changes what the cache stores.

def cache_items_per_token_mha(query_heads, head_dim):
    # K and V for every query head.
    return 2 * query_heads * head_dim

def cache_items_per_token_gqa(kv_heads, head_dim):
    # K and V for fewer shared KV heads.
    return 2 * kv_heads * head_dim

def cache_items_per_token_mla(latent_dim, rope_dim):
    # Simplified ledger: compressed KV latent plus decoupled positional part.
    return latent_dim + rope_dim

ledger = {
    "mha": cache_items_per_token_mha(query_heads=32, head_dim=128),
    "gqa": cache_items_per_token_gqa(kv_heads=8, head_dim=128),
    "mla": cache_items_per_token_mla(latent_dim=512, rope_dim=64),
}

1. Name the cached tensor: in MHA/GQA, cache entries are keys and values; in MLA-style designs, cache entries are compressed latent states plus any positional components the architecture keeps separate.
2. Do not retrofit blindly: changing a trained MHA model into MLA changes projections and usually requires architecture-aware training, conversion, or fine-tuning evidence.
3. Track decode cost: a smaller cache can reduce memory bandwidth, but reconstructing K/V adds compute, so benchmark tokens per second and quality together.
4. Keep the contract testable: record latent dimension, positional dimension, reconstruction projections, model revision, and cache dtype in the serving profile.

Rotary position embeddings encode relative position by rotating pairs of query and key dimensions before the dot product. Values are not rotated.

def rotate_half(x):
    x1, x2 = x[..., ::2], x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

def apply_rope(x, cos, sin):
    return x * cos + rotate_half(x) * sin

The original Transformer paper defines scaled dot-product attention. FlashAttention keeps the result exact but changes the memory schedule: it tiles attention so the GPU moves less data between high-bandwidth memory and faster on-chip memory. PyTorch exposes a high-level scaled_dot_product_attention function that can select optimized backends for supported inputs.

import torch.nn.functional as F

y = F.scaled_dot_product_attention(
    query=q,
    key=k,
    value=v,
    attn_mask=None,
    dropout_p=0.0,
    is_causal=True,
)

Implementation Drill: Explain A Fused-Attention Fallback

A good systems answer connects the math to GPU memory movement. Standard attention materializes a large score matrix; FlashAttention computes the same result with a tiled schedule that reduces high-bandwidth-memory traffic. FlashAttention-2 keeps the exact-attention contract but improves parallel work partitioning, so production code should ask which backend ran, not only whether the formula is correct.

from torch.nn.attention import SDPBackend, sdpa_kernel

def require_flash_attention(q, k, v):
    # Use this as a deployment probe, not as the only correctness test.
    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        return F.scaled_dot_product_attention(
            q, k, v, is_causal=True, dropout_p=0.0
        )

try:
    y = require_flash_attention(q, k, v)
except RuntimeError as exc:
    raise RuntimeError(
        "FlashAttention did not run for this device, dtype, shape, or mask."
    ) from exc

1. Shape audit: confirm q, k, and v are laid out as batch, heads, time, head_dim, with mask broadcastable to batch, heads, query positions, key positions.
2. Dropout audit: pass dropout_p=0.0 during eval; PyTorch applies dropout according to the argument, independent of module training mode.
3. Backend audit: test the fused backend intentionally with sdpa_kernel; if it cannot run, record the exact device, dtype, shape, mask, and PyTorch warning instead of filing a vague "attention is slow" note.
4. GQA audit: when enabling grouped-query attention, verify query heads are divisible by key/value heads and that key and value have the same head count.

Benchmark Drill: Prove Which Attention Backend Ran

FlashAttention-style kernels keep attention exact while changing the memory schedule, but the best backend is conditional on hardware, dtype, head dimension, masks, and sequence shape. FlashAttention-3 adds Hopper-specific asynchronous execution and FP8-oriented techniques, while PyTorch's SDPA dispatcher can choose among several backends or fall back when the requested one is unsupported. A useful benchmark therefore records both latency and dispatch evidence.

from torch.nn.attention import SDPBackend, sdpa_kernel

def time_attention_backend(q, k, v, backend):
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    with sdpa_kernel(backend):
        start.record()
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=True, dropout_p=0.0
        )
        end.record()

    torch.cuda.synchronize()
    return out, start.elapsed_time(end)

out, ms = time_attention_backend(q, k, v, SDPBackend.FLASH_ATTENTION)

1. Record the shape: log batch, heads, query length, KV length, head dimension, dtype, mask type, device, and PyTorch version with every timing.
2. Force and catch: use sdpa_kernel in a probe so unsupported FlashAttention inputs fail loudly instead of silently measuring a different backend.
3. Compare numerics: check the selected backend against a math or reference path with tolerances appropriate for fp16, bf16, or FP8 experiments.
4. Do not generalize across hardware: a result from an H100, A100, consumer GPU, or CPU says little about another deployment target.

Implementation Checklist: Make nn.MultiheadAttention Hit The Fast Path

PyTorch's reference nn.MultiheadAttention can call scaled_dot_product_attention when possible, but the highest-speed inference path is a contract between module settings, tensor layout, masks, and autograd state. Treat it as something to assert in a benchmark harness, not something to assume from the layer name alone.

mha = torch.nn.MultiheadAttention(
    embed_dim=768,
    num_heads=12,
    batch_first=True,
).eval()

with torch.inference_mode():
    # Self-attention: query, key, and value are the same batched tensor.
    y, _ = mha(
        x, x, x,
        need_weights=False,  # allows the optimized SDPA path
    )

1. Prefer batch-first self-attention: the documented inference fastpath expects 3D batched inputs with batch_first=True and query, key, and value as the same tensor.
2. Disable training overhead: run under .eval() and torch.inference_mode() or otherwise ensure no tensor argument requires gradients.
3. Request no weights: pass need_weights=False when attention maps are not part of the product requirement, because returning weights can block the optimized SDPA route.
4. Keep projection shapes simple: leave add_bias_kv and add_zero_attn off, and keep kdim and vdim equal to embed_dim for the fast inference conditions.
5. Benchmark mask choices: padding and attention masks change eligibility and backend behavior; measure the exact shape, dtype, mask, and autocast setting used in serving.

Reading Note: Use FlexAttention For Custom Attention Contracts

Plain SDPA is the first production target for standard full, causal, or padding-masked attention. FlexAttention is useful when the attention rule itself is part of the model or serving contract: sliding windows, document boundaries, ALiBi-style score changes, packed examples, or paged-cache lookup. PyTorch's API lets a score_mod change each attention score after the query-key dot product, while a BlockMask describes block-sparse connectivity. Treat those functions as compiled kernel inputs: changing shapes, mask logic, or captured values can change compile behavior as well as correctness.

from torch.nn.attention.flex_attention import (
    create_block_mask,
    flex_attention,
)

def sliding_window_mask(b, h, q_idx, kv_idx):
    return (kv_idx <= q_idx) & (q_idx - kv_idx < 256)

block_mask = create_block_mask(
    sliding_window_mask,
    B=None,
    H=None,
    Q_LEN=max_seq_len,
    KV_LEN=max_seq_len,
    device="cuda",
)

compiled_flex_attention = torch.compile(flex_attention, dynamic=True)
out = compiled_flex_attention(q, k, v, block_mask=block_mask)

1. Start with SDPA: if the only requirement is standard causal attention, keep the simpler API and benchmark its backend first.
2. Make the rule explicit: write the mask or score modification as a small function with absolute query and KV indices, then compare it against a dense reference mask.
3. Cache the mask: precompute reusable block masks for fixed maximum shapes instead of rebuilding sparse structure inside the decode loop.
4. Watch recompiles: prefill and decode can need different compiled kernels, and per-token Python function changes can destroy the intended speedup.

Implementation Drill: Pack Documents Without Cross-Document Attention

Packed training and retrieval batches often concatenate several short documents into one sequence to reduce padding. The attention rule then needs two constraints at the same time: causal order and same-document visibility. PyTorch's FlexAttention docs define a mask_mod as returning True for positions that participate in attention, and the FlexAttention paper explains that BlockMask lets fully masked score blocks be skipped. Before trusting the fused path, compare the mask function against a dense boolean reference for the exact packed layout.

from torch.nn.attention.flex_attention import create_block_mask

doc_ids = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2], device="cuda")

def packed_causal_mask(b, h, q_idx, kv_idx):
    same_doc = doc_ids[q_idx] == doc_ids[kv_idx]
    causal = kv_idx <= q_idx
    return same_doc & causal

block_mask = create_block_mask(
    packed_causal_mask,
    B=None,
    H=None,
    Q_LEN=doc_ids.numel(),
    KV_LEN=doc_ids.numel(),
    device="cuda",
)

q_idx = torch.arange(doc_ids.numel(), device="cuda")[:, None]
kv_idx = torch.arange(doc_ids.numel(), device="cuda")[None, :]
dense = (doc_ids[:, None] == doc_ids[None, :]) & (kv_idx <= q_idx)
assert not dense[3, 2]  # first token of doc 1 cannot attend to doc 0

1. Define ownership: keep a token-to-document vector next to the packed sequence so the mask rule is data-driven, not inferred from padding.
2. Test boundaries: assert that the first token of every packed document cannot attend to the previous document's final token.
3. Cache by shape: build reusable block masks for repeated packed layouts or length buckets; creating sparse structure inside the hot path can erase the kernel win.
4. Compare semantics: verify a small dense mask before benchmarking, because a fast block-sparse mask that leaks documents trains on the wrong objective.

Sources: Attention Is All You Need, FlashAttention, Multi-Query Attention, Grouped-Query Attention, DeepSeek-V2 / Multi-head Latent Attention, DeepSeek-V3 Technical Report, FlashAttention-2, FlashAttention-3, PyTorch SDPA docs, PyTorch Transformer docs, PyTorch SDPA backend selector, PyTorch MultiheadAttention docs, PyTorch FlexAttention docs, PyTorch FlexAttention inference note, FlexAttention paper.