TopolearnTransformers dominate language, vision and audio; attention is their core and its kernel choice dominates speed and memory.
Attention lets each position mix information from all others. For queries Q, keys K and values V of shape (batch, heads, seq, head_dim):
softmax(Q K^T / sqrt(d)) V
A causal mask stops position i from attending to later positions, as needed for language-model decoding.
torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=False, dropout_p=0.0) computes exactly this but dispatches to fused kernels such as FlashAttention and memory-efficient attention when the inputs allow. These never materialise the full seq x seq score matrix, so memory grows linearly with sequence length instead of quadratically. Always prefer it over writing the softmax by hand.
To implement multi-head attention yourself, project the input with nn.Linear to Q, K, V, reshape (B, T, H*D) to (B, H, T, D), call SDPA, then transpose back and project the output.
PyTorch also provides nn.MultiheadAttention, nn.TransformerEncoderLayer and nn.TransformerDecoderLayer (pass batch_first=True; norm_first=True gives the pre-norm variant used by modern models). Most production LLM code writes its own blocks around SDPA, adding rotary position embeddings and KV caching.
Example
import torch
import torch.nn.functional as F
from torch import nn
B, H, T, D = 2, 4, 16, 32 # batch, heads, sequence, head dim
q, k, v = (torch.randn(B, H, T, D) for _ in range(3))
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape) # torch.Size([2, 4, 16, 32])
# The same computation written out (materialises a T x T matrix):
scores = q @ k.transpose(-2, -1) / D ** 0.5
causal = torch.ones(T, T, dtype=torch.bool).tril()
scores = scores.masked_fill(~causal, float("-inf"))
manual = scores.softmax(dim=-1) @ v
print(torch.allclose(out, manual, atol=1e-5)) # True
layer = nn.TransformerEncoderLayer(d_model=128, nhead=4, batch_first=True, norm_first=True)
print(layer(torch.randn(2, 16, 128)).shape) # torch.Size([2, 16, 128])Remember
- Attention = softmax(QK^T / sqrt(d)) V, per head
scaled_dot_product_attentionpicks fused Flash/efficient kernels- Heads live in a separate dim: (B, H, T, D)
Watch out for
- Passing tensors as (B, T, H, D) to SDPA, which then treats T as the head dimension.