Fast attention with scaled_dot_product_attention

Setup#

pip install torch

Correctness runs anywhere, CPU included. The memory comparison needs a GPU (captured with torch 2.5.1 on an NVIDIA RTX A6000).

Same answer as the textbook formula#

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(2, 4, 128, 64)     # batch, heads, sequence, head_dim
k = torch.randn(2, 4, 128, 64)
v = torch.randn(2, 4, 128, 64)

# attention written by hand, straight from the paper
scores = q @ k.transpose(-2, -1) / (64 ** 0.5)
by_hand = torch.softmax(scores, dim=-1) @ v

fused = F.scaled_dot_product_attention(q, k, v)

print("same answer:", torch.allclose(by_hand, fused, atol=1e-5))
print("output shape:", tuple(fused.shape))

same answer: True
output shape: (2, 4, 128, 64)

Drop-in: if your module contains the three-line hand version, replacing it with the one-liner changes nothing about the result. Causal masking is the is_causal=True argument; arbitrary masks go in attn_mask.

The memory difference, measured#

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

if not torch.cuda.is_available():
    raise SystemExit("needs a GPU to show the memory difference")

def peak_mb(backend):
    torch.manual_seed(0)
    q = torch.randn(1, 8, 4096, 64, device="cuda", dtype=torch.float16)
    k, v = torch.randn_like(q), torch.randn_like(q)
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    with sdpa_kernel([backend]):
        F.scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()
    return torch.cuda.max_memory_allocated() / 1024**2

print(f"math backend (stores the score matrix): {peak_mb(SDPBackend.MATH):8.1f} MB peak")

# FlashAttention first; some builds (Windows wheels) ship without that kernel,
# so fall back to the memory-efficient one, which skips the score matrix too.
for label, backend in [("flash", SDPBackend.FLASH_ATTENTION),
                       ("mem-efficient", SDPBackend.EFFICIENT_ATTENTION)]:
    try:
        print(f"{label:<13} backend (never stores it): {peak_mb(backend):8.1f} MB peak")
        break
    except RuntimeError:
        continue                            # kernel absent in this build

math backend (stores the score matrix):   1716.2 MB peak
mem-efficient backend (never stores it):     24.1 MB peak

Seventy times less memory, at sequence length 4096. The score matrix alone is 8 heads x 4096 x 4096 in half precision — 256 MB that the math path must write, read and keep for backward, while the fused kernels never build it at all.

The loop exists because the machine that captured this output runs a Windows build with no FlashAttention kernel compiled in, so SDPBackend.FLASH_ATTENTION raises RuntimeError: No available kernel. The memory-efficient kernel uses the same tiling idea and lands on the same number. On a Linux build the first branch wins and the line reads flash. Print rather than assume: which kernels exist depends on your operating system, GPU generation, dtype and head dimension. The FlashAttention lesson shows how to list the available backends directly.

The sdpa_kernel context restricts which backend may run — here used to stage the comparison. In normal code you write no context at all and PyTorch dispatches: FlashAttention when dtype and hardware allow, a memory-efficient variant otherwise, and the explicit math fallback last.

Backend requirements matter. FlashAttention needs float16 or bfloat16 on a reasonably modern NVIDIA card, and head dimensions within kernel limits (typically ≤ 128 covers common models). Feed float32 and you silently get a slower backend — pair this with mixed precision.

Common mistakes#

Keeping the hand-written version out of habit. Every hand-rolled softmax-attention forgoes the fused kernels and stores the score matrix for backward. One function call is the whole migration for most modules.

Wrong tensor layout. The function expects (batch, heads, seq, head_dim). Code holding (batch, seq, heads, head_dim) needs a transpose(1, 2) first — and shape errors here produce wrong attention, not always a crash. See reshape, view and contiguity.

Benchmarking it in float32. The float32 path cannot use FlashAttention, so the comparison shows "no speedup" and the wrong conclusion gets drawn. Benchmark in half precision, with honest timing.

Assuming the answer is bit-identical. Tiled accumulation reorders floating-point sums; results match to tolerance (allclose), not exactly. Tests using torch.equal will fail for no real reason.

Try it yourself#

In sdpa_mem.py, double the sequence length to 8192 and predict both numbers before running: the math backend should roughly quadruple, the flash backend roughly double. Then add is_causal=True to both and see whether the gap survives masking.

What to learn next#

• Mixed precision training — the dtype prerequisite for the fast backends.
• Transformers — the architecture this function accelerates.
• torch.compile — fusing everything around the attention call.